<commit_before>// Copyright (c) 2017 Fabian Schuiki

//! This module implements the scoreboard building blocks. Scoreboards are the
//! driving mechanism behind moore. They keep track of the results of each
//! compilation step for every node in the graph. Each node can be accessed in a
//! type safe manner by its ID.

use std;
use std::collections::{HashMap, BTreeMap};
use std::hash::Hash;

/// A context which provides a language-agnostic scoreboard. This is used by
/// the language-specific scoreboards to communicate with the global scoreboard.
pub trait GenericContext {
}

/// The `NodeStorage` trait allows for references to nodes to be stored and
/// retrieved via a unique node ID.
///
/// Once a node is created for example in an arena, a reference to it can be
/// stored in a `NodeStorage` to associate it with an ID. If that ID is
/// presented to the `NodeStorage` again, that same reference will be produced.
/// Implementors of this trait are expected to implement it multiple times, once
/// for each different ID/node type pair that they support. This then allows for
/// nodes to be looked up in a type safe manner based on their ID.
///
/// The `NodeStorage` does not assume ownership over the nodes added to it.
/// Therefore all nodes are references of at least the lifetime `'tn`.
///
/// # Example
///
/// ```
/// use moore_common::score::NodeStorage;
/// use std::collections::HashMap;
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
/// struct FooId(usize);
/// #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
/// struct BarId(usize);
///
/// struct Table<'tn> {
///     foos: HashMap<FooId, &'tn Foo>,
///     bars: HashMap<BarId, &'tn Bar>,
/// }
///
/// impl<'tn> NodeStorage<FooId> for Table<'tn> {
///     type Node = &'tn Foo;
///     fn get(&self, id: &FooId) -> Option<&&'tn Foo> { self.foos.get(id) }
///     fn set(&mut self, id: FooId, node: &'tn Foo) -> Option<&'tn Foo> { self.foos.insert(id, node) }
/// }
///
/// impl<'tn> NodeStorage<BarId> for Table<'tn> {
///     type Node = &'tn Bar;
///     fn get(&self, id: &BarId) -> Option<&&'tn Bar> { self.bars.get(id) }
///     fn set(&mut self, id: BarId, node: &'tn Bar) -> Option<&'tn Bar> { self.bars.insert(id, node) }
/// }
///
/// // Store node refs in table:
/// let foo = Foo;
/// let bar = Bar;
/// let mut tbl = Table{ foos: HashMap::new(), bars: HashMap::new() };
/// tbl.set(FooId(1), &foo);
/// tbl.set(BarId(2), &bar);
///
/// // Retrieve node refs again:
/// assert_eq!(tbl.get(&FooId(1)), Some(&&foo));
/// assert_eq!(tbl.get(&BarId(2)), Some(&&bar));
/// assert_eq!(tbl.get(&BarId(1)), None);
/// assert_eq!(tbl.get(&FooId(2)), None);
///
/// // The following would produce a compiler error due to type mismatch:
/// // let _: &Foo = *tbl.get(&BarId(1)).unwrap();
/// // let _: &Bar = *tbl.get(&FooId(2)).unwrap();
/// ```
pub trait NodeStorage<I> {
	/// The type of the node that is returned when presented with an ID of type
	/// `I`.
	type Node;

	/// Obtains a reference to the node with the given ID.
	///
	/// Returns `None` when no node with the given ID exists.
	fn get(&self, id: &I) -> Option<&Self::Node>;

	/// Store a reference to a node under the given ID.
	///
	/// Later that reference can be retrieved again by presenting the same ID to
	/// the `get` function. Returns the previously stored entry, if any.
	fn set(&mut self, id: I, node: Self::Node) -> Option<Self::Node>;
}

// Implement the NodeStorage trait for HashMaps.
impl<K,V> NodeStorage<K> for HashMap<K,V> where K: Hash + Eq {
	type Node = V;

	fn get(&self, id: &K) -> Option<&V> {
		HashMap::get(self, id)
	}

	fn set(&mut self, id: K, node: V) -> Option<V> {
		HashMap::insert(self, id, node)
	}
}

// Implement the NodeStorage trait for BTreeMaps.
impl<K,V> NodeStorage<K> for BTreeMap<K,V> where K: Ord {
	type Node = V;

	fn get(&self, id: &K) -> Option<&V> {
		BTreeMap::get(self, id)
	}

	fn set(&mut self, id: K, node: V) -> Option<V> {
		BTreeMap::insert(self, id, node)
	}
}

/// The `NodeMaker` trait allows for nodes to be generated from an ID.
///
/// This is useful in conjunction with the `NodeStorage` and `Scoreboard`
/// traits. If a value is requested that the scoreboard cannot find, this trait
/// allows for the node to be generated. For example, if the AST for a node is
/// requested but does not exist, this trait can be implemented in such a way
/// that said AST is loaded. This allows for complex on-demand loading and
/// compilation to be implemented.
///
/// The nodes are expected to be owned either by the caller or some arena.
/// Therefore only a reference to the created node is returned.
///
/// # Example
///
/// ```
/// use moore_common::score::{self, NodeMaker};
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// struct FooId(usize);
/// struct BarId(usize);
///
/// struct Table;
///
/// impl<'tn> NodeMaker<FooId, &'tn Foo> for Table {
///     fn make(&self, id: FooId) -> score::Result<&'tn Foo> {
///         Ok(unsafe { &*(1 as *const Foo) }) // usually you would allocate this in an arena
///     }
/// }
///
/// impl<'tn> NodeMaker<BarId, &'tn Bar> for Table {
///     fn make(&self, id: BarId) -> score::Result<&'tn Bar> {
///         Ok(unsafe { &*(1 as *const Bar) }) // usually you would allocate this in an arena
///     }
/// }
///
/// let tbl = Table;
/// let foo = tbl.make(FooId(1)).unwrap();
/// let bar = tbl.make(BarId(2)).unwrap();
/// assert_eq!(foo, &Foo);
/// assert_eq!(bar, &Bar);
///
/// // The following would produce a compiler error due to type mismatch:
/// // assert_eq!(foo, &Bar);
/// // assert_eq!(bar, &Foo);
/// ```
pub trait NodeMaker<I, N> {
	/// Creates the node with the given ID.
	///
	/// Returns `Err(())` upon failure. Note that the generated node has
	/// lifetime `'tn` that outlives the `NodeMaker`. This is required to allow
	/// for the `NodeMaker` to generate multiple nodes at the same time. The
	/// generated nodes should be owned by an arena or the owner of the
	/// `NodeMaker` itself.
	fn make(&self, id: I) -> Result<N>;
}

/// The result of making a node. Errors that occur while making a node should be
/// reported via a separate channel, e.g. diagnostics, which provide more
/// information to the user.
pub type Result<T> = std::result::Result<T, ()>;

/// Create a new node reference.
///
/// This is merely a wrapper around `NodeId` to provide a type safe
/// representation of a node.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate moore_common;
/// # extern crate rustc_serialize;
///
/// # fn main() {
/// node_ref!(FooRef);
/// node_ref!(BarRef);
/// # }
/// ```
///
/// This creates two structs `FooRef` and `BarRef` that both wrap around a
/// `NodeId`.
#[macro_export]
macro_rules! node_ref {
	($name:ident) => {
		#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, RustcEncodable, RustcDecodable, Hash)]
		pub struct $name($crate::NodeId);

		impl $name {
			/// Create a new reference.
			pub fn new(id: $crate::NodeId) -> $name {
				$name(id)
			}
		}

		impl std::fmt::Debug for $name {
			fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
				write!(f, "{}({})", stringify!($name), self.0)
			}
		}

		impl Into<$crate::NodeId> for $name {
			fn into(self) -> $crate::NodeId {
				self.0
			}
		}
	}
}

/// Create a new group of node references.
///
/// This is a simple enum that contains variants for each of the references.
/// Implements `From` for the various references, and `Into<NodeId>`.
#[macro_export]
macro_rules! node_ref_group {
	($name:ident: $($var:ident($ty:ty),)+) => {
		#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, RustcEncodable, RustcDecodable, Hash)]
		pub enum $name {
			$($var($ty),)*
		}

		impl std::fmt::Debug for $name {
			fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
				node_ref_group!(MATCHES *self, $($name::$var(id) => write!(f, "{}({:?})", stringify!($var), id)),*)
			}
		}

		impl Into<$crate::NodeId> for $name {
			fn into(self) -> $crate::NodeId {
				node_ref_group!(MATCHES self, $($name::$var(id) => id.into()),*)
			}
		}

		$(
		impl From<$ty> for $name {
			fn from(id: $ty) -> $name {
				$name::$var(id)
			}
		}
		)*
	};

	(MATCHES $value:expr, $($lhs:pat => $rhs:expr),+) => {
		match $value {
			$($lhs => $rhs),+
		}
	};
}

/// Create a new table that implements the `NodeStorage` trait.
///
/// The resulting table can then be used to store nodes in a type safe manner.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate moore_common;
/// use moore_common::score::NodeStorage;
/// # extern crate rustc_serialize;
/// # use std::collections::HashMap;
///
/// #[derive(PartialEq, Eq, Hash, Debug)]
/// struct FooRef(usize);
/// #[derive(PartialEq, Eq, Hash, Debug)]
/// struct BarRef(usize);
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// node_storage!(NodeTable<'tn>:
///     foos: FooRef => &'tn Foo,
///     bars: BarRef => &'tn Bar,
/// );
///
/// # fn main() {
/// let foo = &Foo;
/// let bar = &Bar;
///
/// let mut tbl = NodeTable::new();
/// tbl.set(FooRef(0), foo);
/// tbl.set(BarRef(1), bar);
///
/// assert_eq!(tbl.get(&FooRef(0)), Some(&foo));
/// assert_eq!(tbl.get(&BarRef(1)), Some(&bar));
///
/// // The following would produce a compiler error due to the type mismatch:
/// // assert_eq!(tbl.get(&BarRef(0)), Some(&foo));
/// // assert_eq!(tbl.get(&FooRef(1)), Some(&bar));
/// # }
/// ```
#[macro_export]
macro_rules! node_storage {
	($name:ident<$($lt:tt),+>: $($node_name:ident : $node_ref:ty => $node:ty,)+) => {
		pub struct $name<$($lt),*> {
			$($node_name: std::collections::HashMap<$node_ref, $node>,)*
		}

		node_storage!(@struct_impl $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	($name:ident<$($lt:tt),+> where ($($wh:tt)+): $($node_name:ident : $node_ref:ty => $node:ty,)+) => {
		pub struct $name<$($lt),*> where $($wh)* {
			$($node_name: std::collections::HashMap<$node_ref, $node>,)*
		}

		node_storage!(@struct_impl $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	(@struct_impl $name:ident; $($lt:tt),+; $($node_name:ident, $node_ref:ty, $node:ty;)*) => {
		impl<$($lt),*> $name<$($lt),*> {
			/// Create a new empty table.
			pub fn new() -> $name<$($lt),*> {
				$name {
					$($node_name: std::collections::HashMap::new(),)*
				}
			}
		}

		node_storage!(@trait_impl $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	(@trait_impl $name:ident; $($lt:tt),+; $node_name:ident, $node_ref:ty, $node:ty; $($tail_name:ident, $tail_ref:ty, $tail:ty;)*) => {
		impl<$($lt),*> $crate::score::NodeStorage<$node_ref> for $name<$($lt),*> {
			type Node = $node;

			fn get(&self, id: &$node_ref) -> Option<&$node> {
				self.$node_name.get(id)
			}

			fn set(&mut self, id: $node_ref, node: $node) -> Option<$node> {
				self.$node_name.insert(id, node)
			}
		}

		node_storage!(@trait_impl $name; $($lt),*; $($tail_name, $tail_ref, $tail;)*);
	};

	(@trait_impl $name:ident; $($lt:tt),*;) => {}
}
<commit_msg>common: drop use of @ to maintain 1.19 compat<commit_after>// Copyright (c) 2017 Fabian Schuiki

//! This module implements the scoreboard building blocks. Scoreboards are the
//! driving mechanism behind moore. They keep track of the results of each
//! compilation step for every node in the graph. Each node can be accessed in a
//! type safe manner by its ID.

use std;
use std::collections::{HashMap, BTreeMap};
use std::hash::Hash;

/// A context which provides a language-agnostic scoreboard. This is used by
/// the language-specific scoreboards to communicate with the global scoreboard.
pub trait GenericContext {
}

/// The `NodeStorage` trait allows for references to nodes to be stored and
/// retrieved via a unique node ID.
///
/// Once a node is created for example in an arena, a reference to it can be
/// stored in a `NodeStorage` to associate it with an ID. If that ID is
/// presented to the `NodeStorage` again, that same reference will be produced.
/// Implementors of this trait are expected to implement it multiple times, once
/// for each different ID/node type pair that they support. This then allows for
/// nodes to be looked up in a type safe manner based on their ID.
///
/// The `NodeStorage` does not assume ownership over the nodes added to it.
/// Therefore all nodes are references of at least the lifetime `'tn`.
///
/// # Example
///
/// ```
/// use moore_common::score::NodeStorage;
/// use std::collections::HashMap;
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
/// struct FooId(usize);
/// #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
/// struct BarId(usize);
///
/// struct Table<'tn> {
///     foos: HashMap<FooId, &'tn Foo>,
///     bars: HashMap<BarId, &'tn Bar>,
/// }
///
/// impl<'tn> NodeStorage<FooId> for Table<'tn> {
///     type Node = &'tn Foo;
///     fn get(&self, id: &FooId) -> Option<&&'tn Foo> { self.foos.get(id) }
///     fn set(&mut self, id: FooId, node: &'tn Foo) -> Option<&'tn Foo> { self.foos.insert(id, node) }
/// }
///
/// impl<'tn> NodeStorage<BarId> for Table<'tn> {
///     type Node = &'tn Bar;
///     fn get(&self, id: &BarId) -> Option<&&'tn Bar> { self.bars.get(id) }
///     fn set(&mut self, id: BarId, node: &'tn Bar) -> Option<&'tn Bar> { self.bars.insert(id, node) }
/// }
///
/// // Store node refs in table:
/// let foo = Foo;
/// let bar = Bar;
/// let mut tbl = Table{ foos: HashMap::new(), bars: HashMap::new() };
/// tbl.set(FooId(1), &foo);
/// tbl.set(BarId(2), &bar);
///
/// // Retrieve node refs again:
/// assert_eq!(tbl.get(&FooId(1)), Some(&&foo));
/// assert_eq!(tbl.get(&BarId(2)), Some(&&bar));
/// assert_eq!(tbl.get(&BarId(1)), None);
/// assert_eq!(tbl.get(&FooId(2)), None);
///
/// // The following would produce a compiler error due to type mismatch:
/// // let _: &Foo = *tbl.get(&BarId(1)).unwrap();
/// // let _: &Bar = *tbl.get(&FooId(2)).unwrap();
/// ```
pub trait NodeStorage<I> {
	/// The type of the node that is returned when presented with an ID of type
	/// `I`.
	type Node;

	/// Obtains a reference to the node with the given ID.
	///
	/// Returns `None` when no node with the given ID exists.
	fn get(&self, id: &I) -> Option<&Self::Node>;

	/// Store a reference to a node under the given ID.
	///
	/// Later that reference can be retrieved again by presenting the same ID to
	/// the `get` function. Returns the previously stored entry, if any.
	fn set(&mut self, id: I, node: Self::Node) -> Option<Self::Node>;
}

// Implement the NodeStorage trait for HashMaps.
impl<K,V> NodeStorage<K> for HashMap<K,V> where K: Hash + Eq {
	type Node = V;

	fn get(&self, id: &K) -> Option<&V> {
		HashMap::get(self, id)
	}

	fn set(&mut self, id: K, node: V) -> Option<V> {
		HashMap::insert(self, id, node)
	}
}

// Implement the NodeStorage trait for BTreeMaps.
impl<K,V> NodeStorage<K> for BTreeMap<K,V> where K: Ord {
	type Node = V;

	fn get(&self, id: &K) -> Option<&V> {
		BTreeMap::get(self, id)
	}

	fn set(&mut self, id: K, node: V) -> Option<V> {
		BTreeMap::insert(self, id, node)
	}
}

/// The `NodeMaker` trait allows for nodes to be generated from an ID.
///
/// This is useful in conjunction with the `NodeStorage` and `Scoreboard`
/// traits. If a value is requested that the scoreboard cannot find, this trait
/// allows for the node to be generated. For example, if the AST for a node is
/// requested but does not exist, this trait can be implemented in such a way
/// that said AST is loaded. This allows for complex on-demand loading and
/// compilation to be implemented.
///
/// The nodes are expected to be owned either by the caller or some arena.
/// Therefore only a reference to the created node is returned.
///
/// # Example
///
/// ```
/// use moore_common::score::{self, NodeMaker};
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// struct FooId(usize);
/// struct BarId(usize);
///
/// struct Table;
///
/// impl<'tn> NodeMaker<FooId, &'tn Foo> for Table {
///     fn make(&self, id: FooId) -> score::Result<&'tn Foo> {
///         Ok(unsafe { &*(1 as *const Foo) }) // usually you would allocate this in an arena
///     }
/// }
///
/// impl<'tn> NodeMaker<BarId, &'tn Bar> for Table {
///     fn make(&self, id: BarId) -> score::Result<&'tn Bar> {
///         Ok(unsafe { &*(1 as *const Bar) }) // usually you would allocate this in an arena
///     }
/// }
///
/// let tbl = Table;
/// let foo = tbl.make(FooId(1)).unwrap();
/// let bar = tbl.make(BarId(2)).unwrap();
/// assert_eq!(foo, &Foo);
/// assert_eq!(bar, &Bar);
///
/// // The following would produce a compiler error due to type mismatch:
/// // assert_eq!(foo, &Bar);
/// // assert_eq!(bar, &Foo);
/// ```
pub trait NodeMaker<I, N> {
	/// Creates the node with the given ID.
	///
	/// Returns `Err(())` upon failure. Note that the generated node has
	/// lifetime `'tn` that outlives the `NodeMaker`. This is required to allow
	/// for the `NodeMaker` to generate multiple nodes at the same time. The
	/// generated nodes should be owned by an arena or the owner of the
	/// `NodeMaker` itself.
	fn make(&self, id: I) -> Result<N>;
}

/// The result of making a node. Errors that occur while making a node should be
/// reported via a separate channel, e.g. diagnostics, which provide more
/// information to the user.
pub type Result<T> = std::result::Result<T, ()>;

/// Create a new node reference.
///
/// This is merely a wrapper around `NodeId` to provide a type safe
/// representation of a node.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate moore_common;
/// # extern crate rustc_serialize;
///
/// # fn main() {
/// node_ref!(FooRef);
/// node_ref!(BarRef);
/// # }
/// ```
///
/// This creates two structs `FooRef` and `BarRef` that both wrap around a
/// `NodeId`.
#[macro_export]
macro_rules! node_ref {
	($name:ident) => {
		#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, RustcEncodable, RustcDecodable, Hash)]
		pub struct $name($crate::NodeId);

		impl $name {
			/// Create a new reference.
			pub fn new(id: $crate::NodeId) -> $name {
				$name(id)
			}
		}

		impl std::fmt::Debug for $name {
			fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
				write!(f, "{}({})", stringify!($name), self.0)
			}
		}

		impl Into<$crate::NodeId> for $name {
			fn into(self) -> $crate::NodeId {
				self.0
			}
		}
	}
}

/// Create a new group of node references.
///
/// This is a simple enum that contains variants for each of the references.
/// Implements `From` for the various references, and `Into<NodeId>`.
#[macro_export]
macro_rules! node_ref_group {
	($name:ident: $($var:ident($ty:ty),)+) => {
		#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, RustcEncodable, RustcDecodable, Hash)]
		pub enum $name {
			$($var($ty),)*
		}

		impl std::fmt::Debug for $name {
			fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
				node_ref_group!(MATCHES *self, $($name::$var(id) => write!(f, "{}({:?})", stringify!($var), id)),*)
			}
		}

		impl Into<$crate::NodeId> for $name {
			fn into(self) -> $crate::NodeId {
				node_ref_group!(MATCHES self, $($name::$var(id) => id.into()),*)
			}
		}

		$(
		impl From<$ty> for $name {
			fn from(id: $ty) -> $name {
				$name::$var(id)
			}
		}
		)*
	};

	(MATCHES $value:expr, $($lhs:pat => $rhs:expr),+) => {
		match $value {
			$($lhs => $rhs),+
		}
	};
}

/// Create a new table that implements the `NodeStorage` trait.
///
/// The resulting table can then be used to store nodes in a type safe manner.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate moore_common;
/// use moore_common::score::NodeStorage;
/// # extern crate rustc_serialize;
/// # use std::collections::HashMap;
///
/// #[derive(PartialEq, Eq, Hash, Debug)]
/// struct FooRef(usize);
/// #[derive(PartialEq, Eq, Hash, Debug)]
/// struct BarRef(usize);
///
/// #[derive(PartialEq, Eq, Debug)]
/// struct Foo;
/// #[derive(PartialEq, Eq, Debug)]
/// struct Bar;
///
/// node_storage!(NodeTable<'tn>:
///     foos: FooRef => &'tn Foo,
///     bars: BarRef => &'tn Bar,
/// );
///
/// # fn main() {
/// let foo = &Foo;
/// let bar = &Bar;
///
/// let mut tbl = NodeTable::new();
/// tbl.set(FooRef(0), foo);
/// tbl.set(BarRef(1), bar);
///
/// assert_eq!(tbl.get(&FooRef(0)), Some(&foo));
/// assert_eq!(tbl.get(&BarRef(1)), Some(&bar));
///
/// // The following would produce a compiler error due to the type mismatch:
/// // assert_eq!(tbl.get(&BarRef(0)), Some(&foo));
/// // assert_eq!(tbl.get(&FooRef(1)), Some(&bar));
/// # }
/// ```
#[macro_export]
macro_rules! node_storage {
	($name:ident<$($lt:tt),+>: $($node_name:ident : $node_ref:ty => $node:ty,)+) => {
		pub struct $name<$($lt),*> {
			$($node_name: std::collections::HashMap<$node_ref, $node>,)*
		}

		node_storage!(STRUCT_IMPL $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	($name:ident<$($lt:tt),+> where ($($wh:tt)+): $($node_name:ident : $node_ref:ty => $node:ty,)+) => {
		pub struct $name<$($lt),*> where $($wh)* {
			$($node_name: std::collections::HashMap<$node_ref, $node>,)*
		}

		node_storage!(STRUCT_IMPL $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	(STRUCT_IMPL $name:ident; $($lt:tt),+; $($node_name:ident, $node_ref:ty, $node:ty;)*) => {
		impl<$($lt),*> $name<$($lt),*> {
			/// Create a new empty table.
			pub fn new() -> $name<$($lt),*> {
				$name {
					$($node_name: std::collections::HashMap::new(),)*
				}
			}
		}

		node_storage!(TRAIT_IMPL $name; $($lt),*; $($node_name, $node_ref, $node;)*);
	};

	(TRAIT_IMPL $name:ident; $($lt:tt),+; $node_name:ident, $node_ref:ty, $node:ty; $($tail_name:ident, $tail_ref:ty, $tail:ty;)*) => {
		impl<$($lt),*> $crate::score::NodeStorage<$node_ref> for $name<$($lt),*> {
			type Node = $node;

			fn get(&self, id: &$node_ref) -> Option<&$node> {
				self.$node_name.get(id)
			}

			fn set(&mut self, id: $node_ref, node: $node) -> Option<$node> {
				self.$node_name.insert(id, node)
			}
		}

		node_storage!(TRAIT_IMPL $name; $($lt),*; $($tail_name, $tail_ref, $tail;)*);
	};

	(TRAIT_IMPL $name:ident; $($lt:tt),*;) => {}
}
<|endoftext|>


<commit_before>//! Linux DMABUF protocol
//!
//! This module provides helper to handle the linux-dmabuf protocol, which allows clients to submit their
//! contents as dmabuf file descriptors. These handlers automate the aggregation of the metadata associated
//! with a dma buffer, and do some basic checking of the sanity of what the client sends.
//!
//! ## How to use
//!
//! To setup the dmabuf global, you will need to provide 2 things:
//!
//! - a list of the dmabuf formats you wish to support
//! - a closure to test if a dmabuf buffer can be imported by your renderer
//!
//! The list of supported format is just a `Vec<Format>`, where you will enter all the (code, modifier)
//! couples you support. You can typically receive a list of supported formats for one renderer by calling
//! [`crate::backend::renderer::Renderer::dmabuf_formats`].
//!
//! ```
//! # extern crate wayland_server;
//! # extern crate smithay;
//! use smithay::{
//!     backend::allocator::dmabuf::Dmabuf,
//!     reexports::{wayland_server::protocol::wl_buffer::WlBuffer},
//!     wayland::dmabuf::init_dmabuf_global,
//! };
//!
//! # let mut display = wayland_server::Display::new();
//! // define your supported formats
//! let formats = vec![
//!     /* ... */
//! ];
//! let dmabuf_global = init_dmabuf_global(
//!     &mut display,
//!     formats,
//!     |buffer, _| {
//!         /* validate the dmabuf and import it into your renderer state */
//!         let dmabuf = buffer.as_ref().user_data().get::<Dmabuf>().expect("dmabuf global sets this for us");
//!         true
//!     },
//!     None // we don't provide a logger in this example
//! );
//! ```

use std::{
    cell::RefCell,
    convert::TryFrom,
    os::unix::io::{IntoRawFd, RawFd},
    rc::Rc,
};

pub use wayland_protocols::unstable::linux_dmabuf::v1::server::zwp_linux_buffer_params_v1::Flags;
use wayland_protocols::unstable::linux_dmabuf::v1::server::{
    zwp_linux_buffer_params_v1::{
        Error as ParamError, Request as ParamsRequest, ZwpLinuxBufferParamsV1 as BufferParams,
    },
    zwp_linux_dmabuf_v1,
};
use wayland_server::{protocol::wl_buffer, DispatchData, Display, Filter, Global, Main};

use crate::backend::allocator::{
    dmabuf::{Dmabuf, DmabufFlags, Plane},
    Format, Fourcc, Modifier,
};

/// Handler trait for dmabuf validation
///
/// You need to provide an implementation of this trait

/// Initialize a dmabuf global.
///
/// You need to provide a vector of the supported formats, as well as a closure,
/// that will validate the parameters provided by the client and tests the import as a dmabuf.
pub fn init_dmabuf_global<F, L>(
    display: &mut Display,
    formats: Vec<Format>,
    handler: F,
    logger: L,
) -> Global<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>
where
    L: Into<Option<::slog::Logger>>,
    F: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    let log = crate::slog_or_fallback(logger).new(o!("smithay_module" => "dmabuf_handler"));

    let formats = Rc::<[Format]>::from(formats);
    let handler = Rc::new(RefCell::new(handler));

    trace!(
        log,
        "Initializing DMABUF handler with {} supported formats",
        formats.len()
    );

    display.create_global(
        3,
        Filter::new(
            move |(dmabuf, version): (Main<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>, u32), _, _| {
                let dma_formats = formats.clone();
                let dma_handler = handler.clone();
                let dma_log = log.clone();
                dmabuf.quick_assign(move |_, req, _| {
                    if let zwp_linux_dmabuf_v1::Request::CreateParams { params_id } = req {
                        let mut handler = ParamsHandler {
                            pending_planes: Vec::new(),
                            max_planes: 4,
                            used: false,
                            formats: dma_formats.clone(),
                            handler: dma_handler.clone(),
                            log: dma_log.clone(),
                        };
                        params_id.quick_assign(move |params, req, ddata| match req {
                            ParamsRequest::Add {
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                modifier_hi,
                                modifier_lo,
                            } => handler.add(
                                &*params,
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                ((modifier_hi as u64) << 32) + (modifier_lo as u64),
                            ),
                            ParamsRequest::Create {
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create(&*params, width, height, format, flags, ddata),
                            ParamsRequest::CreateImmed {
                                buffer_id,
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create_immed(&*params, buffer_id, width, height, format, flags, ddata),
                            _ => {}
                        });
                    }
                });

                // send the supported formats
                for f in &*formats {
                    dmabuf.format(f.code as u32);
                    if version >= 3 {
                        dmabuf.modifier(
                            f.code as u32,
                            (Into::<u64>::into(f.modifier) >> 32) as u32,
                            Into::<u64>::into(f.modifier) as u32,
                        );
                    }
                }
            },
        ),
    )
}

struct ParamsHandler<H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static> {
    pending_planes: Vec<Plane>,
    max_planes: u32,
    used: bool,
    formats: Rc<[Format]>,
    handler: Rc<RefCell<H>>,
    log: ::slog::Logger,
}

impl<H> ParamsHandler<H>
where
  H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static
{
    fn add(
        &mut self,
        params: &BufferParams,
        fd: RawFd,
        plane_idx: u32,
        offset: u32,
        stride: u32,
        modifier: u64,
    ) {
        // protocol checks:
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        // plane_idx is not too large
        if plane_idx >= self.max_planes {
            // plane_idx starts at 0
            params.as_ref().post_error(
                ParamError::PlaneIdx as u32,
                format!("Plane index {} is out of bounds.", plane_idx),
            );
            return;
        }
        // plane_idx has already been set
        if self.pending_planes.iter().any(|d| d.plane_idx == plane_idx) {
            params.as_ref().post_error(
                ParamError::PlaneSet as u32,
                format!("Plane index {} is already set.", plane_idx),
            );
            return;
        }
        // all checks passed, store the plane
        self.pending_planes.push(Plane {
            fd: Some(fd),
            plane_idx,
            offset,
            stride,
            modifier: Modifier::from(modifier),
        });
    }

    fn create<'a>(&mut self, params: &BufferParams, width: i32, height: i32, format: u32, flags: u32, ddata: DispatchData<'a>) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::new(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    format!("Provided buffer is incomplete, it has zero planes"),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            if let Some(buffer) = params
                .as_ref()
                .client()
                .and_then(|c| c.create_resource::<wl_buffer::WlBuffer>(1))
            {
                buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
                buffer.quick_assign(|_, _, _| {});

                trace!(self.log, "Created a new validated dma wl_buffer.");
                params.created(&buffer);
            } else {
                trace!(self.log, "Failed to create a wl_buffer");
                params.failed();
            }
        } else {
            trace!(self.log, "Refusing creation of an invalid dma wl_buffer.");
            params.failed();
        }
    }

    fn create_immed<'a>(
        &mut self,
        params: &BufferParams,
        buffer: Main<wl_buffer::WlBuffer>,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::new(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    format!("Provided buffer is incomplete, it has zero planes"),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
            buffer.quick_assign(|_, _, _| {});
            trace!(self.log, "Created a new validated dma wl_buffer.");
        } else {
            trace!(
                self.log,
                "Refusing creation of an invalid immediate dma wl_buffer, killing client."
            );
            params.as_ref().post_error(
                ParamError::InvalidWlBuffer as u32,
                "create_immed resulted in an invalid buffer.".into(),
            );
        }

    }
}

fn buffer_basic_checks(
    formats: &[Format],
    pending_planes: &[Plane],
    params: &BufferParams,
    format: Fourcc,
    width: i32,
    height: i32,
) -> bool {
    // protocol_checks:
    // This must be a known format
    let _format = match formats.iter().find(|f| f.code == format) {
        Some(f) => f,
        None => {
            params.as_ref().post_error(
                ParamError::InvalidFormat as u32,
                format!("Format {:?}/{:x} is not supported.", format, format as u32),
            );
            return false;
        }
    };
    // Width and height must be positivie
    if width < 1 || height < 1 {
        params.as_ref().post_error(
            ParamError::InvalidDimensions as u32,
            format!("Dimensions ({},{}) are not valid.", width, height),
        );
        return false;
    }
    // check the size of each plane buffer
    for plane in pending_planes {
        // check size for overflow
        let end = match plane
            .stride
            .checked_mul(height as u32)
            .and_then(|o| o.checked_add(plane.offset))
        {
            None => {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Size overflow for plane {}.", plane.plane_idx),
                );
                return false;
            }
            Some(e) => e,
        };
        if let Ok(size) = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekEnd) {
            // reset the seek point
            let _ = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekSet);
            if plane.offset as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid offset {} for plane {}.", plane.offset, plane.plane_idx),
                );
                return false;
            }
            if (plane.offset + plane.stride) as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid stride {} for plane {}.", plane.stride, plane.plane_idx),
                );
                return false;
            }
            // Planes > 0 can be subsampled, in which case 'size' will be smaller
            // than expected.
            if plane.plane_idx == 0 && end as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!(
                        "Invalid stride ({}) or height ({}) for plane {}.",
                        plane.stride, height, plane.plane_idx
                    ),
                );
                return false;
            }
        }
    }
    true
}
<commit_msg>wayland.dmabuf: fix doctest<commit_after>//! Linux DMABUF protocol
//!
//! This module provides helper to handle the linux-dmabuf protocol, which allows clients to submit their
//! contents as dmabuf file descriptors. These handlers automate the aggregation of the metadata associated
//! with a dma buffer, and do some basic checking of the sanity of what the client sends.
//!
//! ## How to use
//!
//! To setup the dmabuf global, you will need to provide 2 things:
//!
//! - a list of the dmabuf formats you wish to support
//! - a closure to test if a dmabuf buffer can be imported by your renderer
//!
//! The list of supported format is just a `Vec<Format>`, where you will enter all the (code, modifier)
//! couples you support. You can typically receive a list of supported formats for one renderer by calling
//! [`crate::backend::renderer::Renderer::dmabuf_formats`].
//!
//! ```
//! # extern crate wayland_server;
//! # extern crate smithay;
//! use smithay::{
//!     backend::allocator::dmabuf::Dmabuf,
//!     reexports::{wayland_server::protocol::wl_buffer::WlBuffer},
//!     wayland::dmabuf::init_dmabuf_global,
//! };
//!
//! # let mut display = wayland_server::Display::new();
//! // define your supported formats
//! let formats = vec![
//!     /* ... */
//! ];
//! let dmabuf_global = init_dmabuf_global(
//!     &mut display,
//!     formats,
//!     |buffer, dispatch_data| {
//!         /* validate the dmabuf and import it into your renderer state */
//!         true
//!     },
//!     None // we don't provide a logger in this example
//! );
//! ```

use std::{
    cell::RefCell,
    convert::TryFrom,
    os::unix::io::{IntoRawFd, RawFd},
    rc::Rc,
};

pub use wayland_protocols::unstable::linux_dmabuf::v1::server::zwp_linux_buffer_params_v1::Flags;
use wayland_protocols::unstable::linux_dmabuf::v1::server::{
    zwp_linux_buffer_params_v1::{
        Error as ParamError, Request as ParamsRequest, ZwpLinuxBufferParamsV1 as BufferParams,
    },
    zwp_linux_dmabuf_v1,
};
use wayland_server::{protocol::wl_buffer, DispatchData, Display, Filter, Global, Main};

use crate::backend::allocator::{
    dmabuf::{Dmabuf, DmabufFlags, Plane},
    Format, Fourcc, Modifier,
};

/// Handler trait for dmabuf validation
///
/// You need to provide an implementation of this trait

/// Initialize a dmabuf global.
///
/// You need to provide a vector of the supported formats, as well as a closure,
/// that will validate the parameters provided by the client and tests the import as a dmabuf.
pub fn init_dmabuf_global<F, L>(
    display: &mut Display,
    formats: Vec<Format>,
    handler: F,
    logger: L,
) -> Global<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>
where
    L: Into<Option<::slog::Logger>>,
    F: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    let log = crate::slog_or_fallback(logger).new(o!("smithay_module" => "dmabuf_handler"));

    let formats = Rc::<[Format]>::from(formats);
    let handler = Rc::new(RefCell::new(handler));

    trace!(
        log,
        "Initializing DMABUF handler with {} supported formats",
        formats.len()
    );

    display.create_global(
        3,
        Filter::new(
            move |(dmabuf, version): (Main<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>, u32), _, _| {
                let dma_formats = formats.clone();
                let dma_handler = handler.clone();
                let dma_log = log.clone();
                dmabuf.quick_assign(move |_, req, _| {
                    if let zwp_linux_dmabuf_v1::Request::CreateParams { params_id } = req {
                        let mut handler = ParamsHandler {
                            pending_planes: Vec::new(),
                            max_planes: 4,
                            used: false,
                            formats: dma_formats.clone(),
                            handler: dma_handler.clone(),
                            log: dma_log.clone(),
                        };
                        params_id.quick_assign(move |params, req, ddata| match req {
                            ParamsRequest::Add {
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                modifier_hi,
                                modifier_lo,
                            } => handler.add(
                                &*params,
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                ((modifier_hi as u64) << 32) + (modifier_lo as u64),
                            ),
                            ParamsRequest::Create {
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create(&*params, width, height, format, flags, ddata),
                            ParamsRequest::CreateImmed {
                                buffer_id,
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create_immed(&*params, buffer_id, width, height, format, flags, ddata),
                            _ => {}
                        });
                    }
                });

                // send the supported formats
                for f in &*formats {
                    dmabuf.format(f.code as u32);
                    if version >= 3 {
                        dmabuf.modifier(
                            f.code as u32,
                            (Into::<u64>::into(f.modifier) >> 32) as u32,
                            Into::<u64>::into(f.modifier) as u32,
                        );
                    }
                }
            },
        ),
    )
}

struct ParamsHandler<H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static> {
    pending_planes: Vec<Plane>,
    max_planes: u32,
    used: bool,
    formats: Rc<[Format]>,
    handler: Rc<RefCell<H>>,
    log: ::slog::Logger,
}

impl<H> ParamsHandler<H>
where
  H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static
{
    fn add(
        &mut self,
        params: &BufferParams,
        fd: RawFd,
        plane_idx: u32,
        offset: u32,
        stride: u32,
        modifier: u64,
    ) {
        // protocol checks:
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        // plane_idx is not too large
        if plane_idx >= self.max_planes {
            // plane_idx starts at 0
            params.as_ref().post_error(
                ParamError::PlaneIdx as u32,
                format!("Plane index {} is out of bounds.", plane_idx),
            );
            return;
        }
        // plane_idx has already been set
        if self.pending_planes.iter().any(|d| d.plane_idx == plane_idx) {
            params.as_ref().post_error(
                ParamError::PlaneSet as u32,
                format!("Plane index {} is already set.", plane_idx),
            );
            return;
        }
        // all checks passed, store the plane
        self.pending_planes.push(Plane {
            fd: Some(fd),
            plane_idx,
            offset,
            stride,
            modifier: Modifier::from(modifier),
        });
    }

    fn create<'a>(&mut self, params: &BufferParams, width: i32, height: i32, format: u32, flags: u32, ddata: DispatchData<'a>) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::new(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    format!("Provided buffer is incomplete, it has zero planes"),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            if let Some(buffer) = params
                .as_ref()
                .client()
                .and_then(|c| c.create_resource::<wl_buffer::WlBuffer>(1))
            {
                buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
                buffer.quick_assign(|_, _, _| {});

                trace!(self.log, "Created a new validated dma wl_buffer.");
                params.created(&buffer);
            } else {
                trace!(self.log, "Failed to create a wl_buffer");
                params.failed();
            }
        } else {
            trace!(self.log, "Refusing creation of an invalid dma wl_buffer.");
            params.failed();
        }
    }

    fn create_immed<'a>(
        &mut self,
        params: &BufferParams,
        buffer: Main<wl_buffer::WlBuffer>,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::new(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    format!("Provided buffer is incomplete, it has zero planes"),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
            buffer.quick_assign(|_, _, _| {});
            trace!(self.log, "Created a new validated dma wl_buffer.");
        } else {
            trace!(
                self.log,
                "Refusing creation of an invalid immediate dma wl_buffer, killing client."
            );
            params.as_ref().post_error(
                ParamError::InvalidWlBuffer as u32,
                "create_immed resulted in an invalid buffer.".into(),
            );
        }

    }
}

fn buffer_basic_checks(
    formats: &[Format],
    pending_planes: &[Plane],
    params: &BufferParams,
    format: Fourcc,
    width: i32,
    height: i32,
) -> bool {
    // protocol_checks:
    // This must be a known format
    let _format = match formats.iter().find(|f| f.code == format) {
        Some(f) => f,
        None => {
            params.as_ref().post_error(
                ParamError::InvalidFormat as u32,
                format!("Format {:?}/{:x} is not supported.", format, format as u32),
            );
            return false;
        }
    };
    // Width and height must be positivie
    if width < 1 || height < 1 {
        params.as_ref().post_error(
            ParamError::InvalidDimensions as u32,
            format!("Dimensions ({},{}) are not valid.", width, height),
        );
        return false;
    }
    // check the size of each plane buffer
    for plane in pending_planes {
        // check size for overflow
        let end = match plane
            .stride
            .checked_mul(height as u32)
            .and_then(|o| o.checked_add(plane.offset))
        {
            None => {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Size overflow for plane {}.", plane.plane_idx),
                );
                return false;
            }
            Some(e) => e,
        };
        if let Ok(size) = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekEnd) {
            // reset the seek point
            let _ = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekSet);
            if plane.offset as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid offset {} for plane {}.", plane.offset, plane.plane_idx),
                );
                return false;
            }
            if (plane.offset + plane.stride) as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid stride {} for plane {}.", plane.stride, plane.plane_idx),
                );
                return false;
            }
            // Planes > 0 can be subsampled, in which case 'size' will be smaller
            // than expected.
            if plane.plane_idx == 0 && end as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!(
                        "Invalid stride ({}) or height ({}) for plane {}.",
                        plane.stride, height, plane.plane_idx
                    ),
                );
                return false;
            }
        }
    }
    true
}
<|endoftext|>


<commit_before>extern crate x11;
extern crate libc;

use std::ptr;
use std::process::Command;
use std::boxed::Box;

use x11::xlib;
use super::WindowManager;
use super::Workspaces;
use super::container;
use super::layout::{self, LayoutDirection, MoveDirection};
use super::super::libx::{self, Context};

#[derive(Hash, Eq, PartialEq, Debug)]
pub struct KeyBind {
    pub key: xlib::KeySym,
    pub mask: u32,
}

impl KeyBind {
    pub fn build(mod_key: u32, tokens: &[&str]) -> KeyBind {
        let mut mask = 0;
        let mut sym = 0;
        for key in tokens {
            match *key {
                "$mod" => mask = mask | mod_key,
                "Shift" => mask = mask | xlib::ShiftMask,
                "Ctrl" => mask = mask | xlib::ControlMask,
                "Mod1" => mask = mask | xlib::Mod1Mask,
                "Mod2" => mask = mask | xlib::Mod2Mask,
                "Mod3" => mask = mask | xlib::Mod3Mask,
                "Mod4" => mask = mask | xlib::Mod4Mask,
                "Mod5" => mask = mask | xlib::Mod5Mask,
                _ => {
                    sym = libx::string_to_keysym(key);
                }
            }
        }

        KeyBind {
            mask: mask,
            key: sym
        }
    }
}

pub trait Handler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context);
}

pub struct ExecHandler {
    pub cmd: Command
}

pub struct LayoutHandler {
    layout_type: layout::Type,
}

pub struct FullscreenHandler;
/// switch to another workspace
pub struct WorkspaceHandler {
    pub key: char,
}

/// move window to target workspace
pub struct WindowToWorkspaceHandler {
    pub key: char,
}

pub struct WindowFocusHandler {
    pub direction: MoveDirection,
}

pub struct WindowCloseHandler;
pub struct SplitHandler;

pub enum Resize {
    Shrink,
    Grow,
}

pub struct WindowResizeHandler {
    pub direction: LayoutDirection,
    pub resize: Resize
}

impl Handler for WindowResizeHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        let focused = workspaces.get_focus();
        if focused.is_some() {
            let c = focused.unwrap();
            let parent = c.get_parent();
            if parent.is_some() {
                let p = parent.unwrap();
                let index = p.contain(c.raw_id()).unwrap();
                if p.direction == self.direction {
                    let step:f32 = match self.resize {
                        Resize::Shrink => {
                            -0.05
                        }
                        Resize::Grow => {
                            0.05
                        }
                    };

                    if index > 0 && index < (p.size() - 1) {
                        p.resize_children(index, index-1, step);
                        p.resize_children(index, index+1, step);
                    }
                    else if index > 0 {
                        p.resize_children(index, index-1, step*2.0);
                    }
                    else if index < (p.size() - 1) {
                        p.resize_children(index, index+1, step*2.0);
                    }
                    p.update_layout();
                }

            }
        }
    }
}

impl Handler for SplitHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        workspaces.current().print_tree(0);
        if let Some(c) = workspaces.get_focus() {
            c.split();
        }
        workspaces.current().print_tree(0);
    }
}

impl Handler for WindowCloseHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle window close");
        let (window, _) = libx::get_input_focus(context);
        libx::kill_window(context, window);
        // xlib::XWithdrawWindow(context, window, screen_num);
        // xlib::XDestroyWindow(context, window);
        println!("try kill window {}", window);

    }
}
impl Handler for WindowFocusHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("change focus in current workspace");
        if let Some(c) = workspaces.get_focus() {
            if let Some(p) = c.get_parent() {
                let index = p.contain(c.raw_id()).unwrap();
                let size = p.size();
                let next = match p.direction {
                    LayoutDirection::Vertical => {
                        match self.direction {
                            MoveDirection::Up => {
                                (index-1) % size
                            }
                            MoveDirection::Down => {
                                (index+1) % size
                            }
                            _ => { index }
                        }
                    }
                    LayoutDirection::Horizontal => {
                        match self.direction {
                            MoveDirection::Left => {
                                (index-1) % size
                            }
                            MoveDirection::Right => {
                                (index+1) % size
                            }
                            _ => { index }
                        }
                    }
                };
                if next != index {
                    c.unfocus();
                    let next_c = p.get_child(next).unwrap().focus();
                }
            }
        }
    }
}

impl Handler for WindowToWorkspaceHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        let id = match workspaces.get_focus() {
            Some(container) => {
                container.raw_id()
            }
            None => { return }
        };
        let from = workspaces.current_name();
        let to = self.key;
        workspaces.move_window(id, from, to);
    }
}

impl Handler for FullscreenHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("fullscreen toggle");
        if let Some(c) = workspaces.get_focus(){
            c.mode_toggle();
        }
    }
}
impl Handler for WorkspaceHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle workspace form {}", workspaces.current_name());
        if !workspaces.contain(self.key) {
            workspaces.create(self.key);
        }
        workspaces.switch_workspace(self.key, context);
        println!("to {}", workspaces.current_name());
    }
}

impl Handler for ExecHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle exec");
        self.cmd.spawn();
    }
}

impl Handler for LayoutHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle layout");
        let t = self.layout_type.clone();
        let old = workspaces.mode;
        workspaces.mode = container::Mode::Layout;
        let mut changed = false;
        if let Some(container) = workspaces.get_focus() {
            let c = if container.get_parent().is_some() {
                container.get_parent().unwrap()
            }
            else {
                container
            };
            c.change_layout(t);
            c.update_layout();
            changed = true;
        }
        if !changed {
            workspaces.mode = old;
        }
    }
}

impl LayoutHandler {
    pub fn new(layout: layout::Type) -> LayoutHandler{
        LayoutHandler {
            layout_type: layout
        }
    }
}

impl ExecHandler {
    pub fn new(tokens: &[&str]) -> ExecHandler {
        let (name, args) = tokens.split_at(1);
        let mut cmd = Command::new(name[0]);
        cmd.args(args);

        let handler = ExecHandler {
            cmd: cmd,
        };
        handler
    }
}
<commit_msg>avoid change focus when resizing window<commit_after>extern crate x11;
extern crate libc;

use std::ptr;
use std::process::Command;
use std::boxed::Box;

use x11::xlib;
use super::WindowManager;
use super::Workspaces;
use super::container;
use super::layout::{self, LayoutDirection, MoveDirection};
use super::super::libx::{self, Context};

#[derive(Hash, Eq, PartialEq, Debug)]
pub struct KeyBind {
    pub key: xlib::KeySym,
    pub mask: u32,
}

impl KeyBind {
    pub fn build(mod_key: u32, tokens: &[&str]) -> KeyBind {
        let mut mask = 0;
        let mut sym = 0;
        for key in tokens {
            match *key {
                "$mod" => mask = mask | mod_key,
                "Shift" => mask = mask | xlib::ShiftMask,
                "Ctrl" => mask = mask | xlib::ControlMask,
                "Mod1" => mask = mask | xlib::Mod1Mask,
                "Mod2" => mask = mask | xlib::Mod2Mask,
                "Mod3" => mask = mask | xlib::Mod3Mask,
                "Mod4" => mask = mask | xlib::Mod4Mask,
                "Mod5" => mask = mask | xlib::Mod5Mask,
                _ => {
                    sym = libx::string_to_keysym(key);
                }
            }
        }

        KeyBind {
            mask: mask,
            key: sym
        }
    }
}

pub trait Handler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context);
}

pub struct ExecHandler {
    pub cmd: Command
}

pub struct LayoutHandler {
    layout_type: layout::Type,
}

pub struct FullscreenHandler;
/// switch to another workspace
pub struct WorkspaceHandler {
    pub key: char,
}

/// move window to target workspace
pub struct WindowToWorkspaceHandler {
    pub key: char,
}

pub struct WindowFocusHandler {
    pub direction: MoveDirection,
}

pub struct WindowCloseHandler;
pub struct SplitHandler;

pub enum Resize {
    Shrink,
    Grow,
}

pub struct WindowResizeHandler {
    pub direction: LayoutDirection,
    pub resize: Resize
}

impl Handler for WindowResizeHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        let old = workspaces.mode;
        workspaces.mode = container::Mode::Layout;
        let mut changed = false;
        if let Some(c) = workspaces.get_focus() {
            if let Some(p) = c.get_parent() {
                let index = p.contain(c.raw_id()).unwrap();
                if p.direction == self.direction {
                    let step:f32 = match self.resize {
                        Resize::Shrink => {
                            -0.05
                        }
                        Resize::Grow => {
                            0.05
                        }
                    };

                    if index > 0 && index < (p.size() - 1) {
                        p.resize_children(index, index-1, step);
                        p.resize_children(index, index+1, step);
                    }
                    else if index > 0 {
                        p.resize_children(index, index-1, step*2.0);
                    }
                    else if index < (p.size() - 1) {
                        p.resize_children(index, index+1, step*2.0);
                    }
                    p.update_layout();
                    changed = true;
                }
            }
        }
        if !changed {
            workspaces.mode = old;
        }
    }
}

impl Handler for SplitHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        workspaces.current().print_tree(0);
        if let Some(c) = workspaces.get_focus() {
            c.split();
        }
        workspaces.current().print_tree(0);
    }
}

impl Handler for WindowCloseHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle window close");
        let (window, _) = libx::get_input_focus(context);
        libx::kill_window(context, window);
        // xlib::XWithdrawWindow(context, window, screen_num);
        // xlib::XDestroyWindow(context, window);
        println!("try kill window {}", window);

    }
}
impl Handler for WindowFocusHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("change focus in current workspace");
        if let Some(c) = workspaces.get_focus() {
            if let Some(p) = c.get_parent() {
                let index = p.contain(c.raw_id()).unwrap();
                let size = p.size();
                let next = match p.direction {
                    LayoutDirection::Vertical => {
                        match self.direction {
                            MoveDirection::Up => {
                                (index-1) % size
                            }
                            MoveDirection::Down => {
                                (index+1) % size
                            }
                            _ => { index }
                        }
                    }
                    LayoutDirection::Horizontal => {
                        match self.direction {
                            MoveDirection::Left => {
                                (index-1) % size
                            }
                            MoveDirection::Right => {
                                (index+1) % size
                            }
                            _ => { index }
                        }
                    }
                };
                if next != index {
                    c.unfocus();
                    let next_c = p.get_child(next).unwrap().focus();
                }
            }
        }
    }
}

impl Handler for WindowToWorkspaceHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        let id = match workspaces.get_focus() {
            Some(container) => {
                container.raw_id()
            }
            None => { return }
        };
        let from = workspaces.current_name();
        let to = self.key;
        workspaces.move_window(id, from, to);
    }
}

impl Handler for FullscreenHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("fullscreen toggle");
        if let Some(c) = workspaces.get_focus(){
            c.mode_toggle();
        }
    }
}
impl Handler for WorkspaceHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle workspace form {}", workspaces.current_name());
        if !workspaces.contain(self.key) {
            workspaces.create(self.key);
        }
        workspaces.switch_workspace(self.key, context);
        println!("to {}", workspaces.current_name());
    }
}

impl Handler for ExecHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle exec");
        self.cmd.spawn();
    }
}

impl Handler for LayoutHandler {
    fn handle(&mut self, workspaces: &mut Workspaces, context: Context) {
        debug!("handle layout");
        let t = self.layout_type.clone();
        let old = workspaces.mode;
        workspaces.mode = container::Mode::Layout;
        let mut changed = false;
        if let Some(container) = workspaces.get_focus() {
            let c = if container.get_parent().is_some() {
                container.get_parent().unwrap()
            }
            else {
                container
            };
            c.change_layout(t);
            c.update_layout();
            changed = true;
        }
        if !changed {
            workspaces.mode = old;
        }
    }
}

impl LayoutHandler {
    pub fn new(layout: layout::Type) -> LayoutHandler{
        LayoutHandler {
            layout_type: layout
        }
    }
}

impl ExecHandler {
    pub fn new(tokens: &[&str]) -> ExecHandler {
        let (name, args) = tokens.split_at(1);
        let mut cmd = Command::new(name[0]);
        cmd.args(args);

        let handler = ExecHandler {
            cmd: cmd,
        };
        handler
    }
}
<|endoftext|>


<commit_before>extern crate ssz_helpers;

use self::ssz_helpers::ssz_beacon_block::{
    SszBeaconBlock,
};
use std::sync::Arc;
use super::{
    ClientDB,
    DBError,
};
use super::BLOCKS_DB_COLUMN as DB_COLUMN;

#[derive(Clone, Debug, PartialEq)]
pub enum BeaconBlockAtSlotError {
    UnknownBeaconBlock,
    InvalidBeaconBlock,
    DBError(String),
}

pub struct BeaconBlockStore<T>
    where T: ClientDB
{
    db: Arc<T>,
}

impl<T: ClientDB> BeaconBlockStore<T> {
    pub fn new(db: Arc<T>) -> Self {
        Self {
            db,
        }
    }

    pub fn put_serialized_block(&self, hash: &[u8], ssz: &[u8])
        -> Result<(), DBError>
    {
        self.db.put(DB_COLUMN, hash, ssz)
    }

    pub fn get_serialized_block(&self, hash: &[u8])
        -> Result<Option<Vec<u8>>, DBError>
    {
        self.db.get(DB_COLUMN, hash)
    }

    pub fn block_exists(&self, hash: &[u8])
        -> Result<bool, DBError>
    {
        self.db.exists(DB_COLUMN, hash)
    }

    pub fn block_exists_in_canonical_chain(&self, hash: &[u8])
        -> Result<bool, DBError>
    {
        // TODO: implement logic for canonical chain
        self.db.exists(DB_COLUMN, hash)
    }

    /// Retrieve the block at a slot given a "head_hash" and a slot.
    ///
    /// A "head_hash" must be a block hash with a slot number greater than or equal to the desired
    /// slot.
    ///
    /// This function will read each block down the chain until it finds a block with the given
    /// slot number. If the slot is skipped, the function will return None.
    ///
    /// If a block is found, a tuple of (block_hash, serialized_block) is returned.
    pub fn block_at_slot(&self, head_hash: &[u8], slot: u64)
        -> Result<Option<(Vec<u8>, Vec<u8>)>, BeaconBlockAtSlotError>
    {
        match self.get_serialized_block(head_hash)? {
            None => Err(BeaconBlockAtSlotError::UnknownBeaconBlock),
            Some(ssz) => {
                let block = SszBeaconBlock::from_slice(&ssz)
                    .map_err(|_| BeaconBlockAtSlotError::InvalidBeaconBlock)?;
                match block.slot() {
                    s if s == slot => Ok(Some((head_hash.to_vec(), ssz.to_vec()))),
                    s if s < slot => Ok(None),
                    _ => {
                        match block.parent_hash() {
                            Some(parent_hash) => self.block_at_slot(parent_hash, slot),
                            None => Err(BeaconBlockAtSlotError::UnknownBeaconBlock)
                        }
                    }
                }
            }
        }
    }
}

impl From<DBError> for BeaconBlockAtSlotError {
    fn from(e: DBError) -> Self {
        BeaconBlockAtSlotError::DBError(e.message)
    }
}

#[cfg(test)]
mod tests {
    extern crate ssz;
    extern crate types;

    use self::types::beacon_block::BeaconBlock;
    use self::types::attestation_record::AttestationRecord;
    use self::types::Hash256;
    use self::ssz::SszStream;

    use super::*;
    use super::super::super::MemoryDB;
    use std::thread;
    use std::sync::Arc;

    #[test]
    fn test_block_store_on_memory_db() {
        let db = Arc::new(MemoryDB::open());
        let bs = Arc::new(BeaconBlockStore::new(db.clone()));

        let thread_count = 10;
        let write_count = 10;

        // We're expecting the product of these numbers to fit in one byte.
        assert!(thread_count * write_count <= 255);

        let mut handles = vec![];
        for t in 0..thread_count {
            let wc = write_count;
            let bs = bs.clone();
            let handle = thread::spawn(move || {
                for w in 0..wc {
                    let key = (t * w) as u8;
                    let val = 42;
                    bs.put_serialized_block(&vec![key], &vec![val]).unwrap();
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().unwrap();
        }

        for t in 0..thread_count {
            for w in 0..write_count {
                let key = (t * w) as u8;
                assert!(bs.block_exists(&vec![key]).unwrap());
                let val = bs.get_serialized_block(&vec![key]).unwrap().unwrap();
                assert_eq!(vec![42], val);
            }
        }
    }

    #[test]
    fn test_block_at_slot() {
        let db = Arc::new(MemoryDB::open());
        let bs = Arc::new(BeaconBlockStore::new(db.clone()));

        let blocks = (0..5).into_iter()
            .map(|_| {
                let mut block = BeaconBlock::zero();
                let ar = AttestationRecord::zero();
                block.attestations.push(ar);
                block
            });

        let hashes = [
            Hash256::from("zero".as_bytes()),
            Hash256::from("one".as_bytes()),
            Hash256::from("two".as_bytes()),
            Hash256::from("three".as_bytes()),
            Hash256::from("four".as_bytes()),
        ];

        let parent_hashes = [
            Hash256::from("genesis".as_bytes()),
            Hash256::from("zero".as_bytes()),
            Hash256::from("one".as_bytes()),
            Hash256::from("two".as_bytes()),
            Hash256::from("three".as_bytes()),
        ];

        let slots = [0, 1, 3, 4, 5];

        for (i, mut block) in blocks.enumerate() {
            block.slot = slots[i];
            let mut s = SszStream::new();
            s.append(&block);
            let ssz = s.drain();
            bs.put_serialized_block(&hashes[i].to_vec(), &ssz).unwrap();
        }

        let tuple = bs.block_at_slot(&hashes[4], 5).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 5);
        assert_eq!(tuple.0, hashes[4].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 4).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 4);
        assert_eq!(tuple.0, hashes[3].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 3).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 3);
        assert_eq!(tuple.0, hashes[2].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 0).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 0);
        assert_eq!(tuple.0, hashes[0].to_vec());

        let ssz = bs.block_at_slot(&hashes[4], 2).unwrap();
        assert_eq!(ssz, None);

        let ssz = bs.block_at_slot(&hashes[4], 6).unwrap();
        assert_eq!(ssz, None);

        let ssz = bs.block_at_slot(&Hash256::from("unknown".as_bytes()), 2);
        assert_eq!(ssz, Err(BeaconBlockAtSlotError::UnknownBeaconBlock));
    }
}
<commit_msg>Fix broken block_store test<commit_after>extern crate ssz_helpers;

use self::ssz_helpers::ssz_beacon_block::{
    SszBeaconBlock,
};
use std::sync::Arc;
use super::{
    ClientDB,
    DBError,
};
use super::BLOCKS_DB_COLUMN as DB_COLUMN;

#[derive(Clone, Debug, PartialEq)]
pub enum BeaconBlockAtSlotError {
    UnknownBeaconBlock,
    InvalidBeaconBlock,
    DBError(String),
}

pub struct BeaconBlockStore<T>
    where T: ClientDB
{
    db: Arc<T>,
}

impl<T: ClientDB> BeaconBlockStore<T> {
    pub fn new(db: Arc<T>) -> Self {
        Self {
            db,
        }
    }

    pub fn put_serialized_block(&self, hash: &[u8], ssz: &[u8])
        -> Result<(), DBError>
    {
        self.db.put(DB_COLUMN, hash, ssz)
    }

    pub fn get_serialized_block(&self, hash: &[u8])
        -> Result<Option<Vec<u8>>, DBError>
    {
        self.db.get(DB_COLUMN, hash)
    }

    pub fn block_exists(&self, hash: &[u8])
        -> Result<bool, DBError>
    {
        self.db.exists(DB_COLUMN, hash)
    }

    pub fn block_exists_in_canonical_chain(&self, hash: &[u8])
        -> Result<bool, DBError>
    {
        // TODO: implement logic for canonical chain
        self.db.exists(DB_COLUMN, hash)
    }

    /// Retrieve the block at a slot given a "head_hash" and a slot.
    ///
    /// A "head_hash" must be a block hash with a slot number greater than or equal to the desired
    /// slot.
    ///
    /// This function will read each block down the chain until it finds a block with the given
    /// slot number. If the slot is skipped, the function will return None.
    ///
    /// If a block is found, a tuple of (block_hash, serialized_block) is returned.
    pub fn block_at_slot(&self, head_hash: &[u8], slot: u64)
        -> Result<Option<(Vec<u8>, Vec<u8>)>, BeaconBlockAtSlotError>
    {
        match self.get_serialized_block(head_hash)? {
            None => Err(BeaconBlockAtSlotError::UnknownBeaconBlock),
            Some(ssz) => {
                let block = SszBeaconBlock::from_slice(&ssz)
                    .map_err(|_| BeaconBlockAtSlotError::InvalidBeaconBlock)?;
                match block.slot() {
                    s if s == slot => Ok(Some((head_hash.to_vec(), ssz.to_vec()))),
                    s if s < slot => Ok(None),
                    _ => {
                        match block.parent_hash() {
                            Some(parent_hash) => self.block_at_slot(parent_hash, slot),
                            None => Err(BeaconBlockAtSlotError::UnknownBeaconBlock)
                        }
                    }
                }
            }
        }
    }
}

impl From<DBError> for BeaconBlockAtSlotError {
    fn from(e: DBError) -> Self {
        BeaconBlockAtSlotError::DBError(e.message)
    }
}

#[cfg(test)]
mod tests {
    extern crate ssz;
    extern crate types;

    use self::types::beacon_block::BeaconBlock;
    use self::types::attestation_record::AttestationRecord;
    use self::types::Hash256;
    use self::ssz::SszStream;

    use super::*;
    use super::super::super::MemoryDB;
    use std::thread;
    use std::sync::Arc;

    #[test]
    fn test_block_store_on_memory_db() {
        let db = Arc::new(MemoryDB::open());
        let bs = Arc::new(BeaconBlockStore::new(db.clone()));

        let thread_count = 10;
        let write_count = 10;

        // We're expecting the product of these numbers to fit in one byte.
        assert!(thread_count * write_count <= 255);

        let mut handles = vec![];
        for t in 0..thread_count {
            let wc = write_count;
            let bs = bs.clone();
            let handle = thread::spawn(move || {
                for w in 0..wc {
                    let key = (t * w) as u8;
                    let val = 42;
                    bs.put_serialized_block(&vec![key], &vec![val]).unwrap();
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().unwrap();
        }

        for t in 0..thread_count {
            for w in 0..write_count {
                let key = (t * w) as u8;
                assert!(bs.block_exists(&vec![key]).unwrap());
                let val = bs.get_serialized_block(&vec![key]).unwrap().unwrap();
                assert_eq!(vec![42], val);
            }
        }
    }

    #[test]
    fn test_block_at_slot() {
        let db = Arc::new(MemoryDB::open());
        let bs = Arc::new(BeaconBlockStore::new(db.clone()));

        let blocks = (0..5).into_iter()
            .map(|_| {
                let mut block = BeaconBlock::zero();
                let ar = AttestationRecord::zero();
                block.attestations.push(ar);
                block
            });

        let hashes = [
            Hash256::from("zero".as_bytes()),
            Hash256::from("one".as_bytes()),
            Hash256::from("two".as_bytes()),
            Hash256::from("three".as_bytes()),
            Hash256::from("four".as_bytes()),
        ];

        let parent_hashes = [
            Hash256::from("genesis".as_bytes()),
            Hash256::from("zero".as_bytes()),
            Hash256::from("one".as_bytes()),
            Hash256::from("two".as_bytes()),
            Hash256::from("three".as_bytes()),
        ];

        let slots = [0, 1, 3, 4, 5];

        for (i, mut block) in blocks.enumerate() {
            block.ancestor_hashes.push(parent_hashes[i]);
            block.slot = slots[i];
            let mut s = SszStream::new();
            s.append(&block);
            let ssz = s.drain();
            bs.put_serialized_block(&hashes[i].to_vec(), &ssz).unwrap();
        }

        let tuple = bs.block_at_slot(&hashes[4], 5).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 5);
        assert_eq!(tuple.0, hashes[4].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 4).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 4);
        assert_eq!(tuple.0, hashes[3].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 3).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 3);
        assert_eq!(tuple.0, hashes[2].to_vec());

        let tuple = bs.block_at_slot(&hashes[4], 0).unwrap().unwrap();
        let block = SszBeaconBlock::from_slice(&tuple.1).unwrap();
        assert_eq!(block.slot(), 0);
        assert_eq!(tuple.0, hashes[0].to_vec());

        let ssz = bs.block_at_slot(&hashes[4], 2).unwrap();
        assert_eq!(ssz, None);

        let ssz = bs.block_at_slot(&hashes[4], 6).unwrap();
        assert_eq!(ssz, None);

        let ssz = bs.block_at_slot(&Hash256::from("unknown".as_bytes()), 2);
        assert_eq!(ssz, Err(BeaconBlockAtSlotError::UnknownBeaconBlock));
    }
}
<|endoftext|>


<commit_before>use rover;
use coords::Position2D;
use rover::Delta;
use rover::Orientation;

describe!(
    it "should create new rover" {
        let r = rover::new();
        assert_eq!(r.position, Position2D{x: 0, y: 0})
    }

    it "should move north" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(1));
        assert_eq!(r.position, Position2D{x: 0, y: 1})
    }

    it "should move north twice" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(2));
        assert_eq!(r.position, Position2D{x: 0, y: 2})
    }

    it "should move east" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(1));
        assert_eq!(r.position, Position2D{x: 1, y: 0})
    }

    it "should move east twice" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(2));
        assert_eq!(r.position, Position2D{x: 2, y: 0})
    }

    it "should move south" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(-1));
        assert_eq!(r.position, Position2D{x: 0, y: 100})
    }

    it "should move south twice" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(-2));
        assert_eq!(r.position, Position2D{x: 0, y: 99})
    }

    it "should move west" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(-1));
        assert_eq!(r.position, Position2D{x: 100, y: 0})
    }

    it "should move west twice" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(-2));
        assert_eq!(r.position, Position2D{x: 99, y: 0})
    }

    it "should move by vector" {
        let mut r = rover::new();
        r.move_(Delta::Vector(2, 3));
        assert_eq!(r.position, Position2D{x: 2, y: 3})
    }

    it "should advance north" {
        let mut r = rover::new();
        r.advance();
        assert_eq!(r.position, Position2D{x: 0, y: 1})
    }

    it "should advance east" {
        let mut r = rover::new();
        r.facing = Orientation::East;
        r.advance();
        assert_eq!(r.position, Position2D{x: 1, y: 0})
    }

    it "should advance south" {
        let mut r = rover::new();
        r.facing = Orientation::South;
        r.advance();
        assert_eq!(r.position, Position2D{x: 0, y: 100})
    }

    it "should advance west" {
        let mut r = rover::new();
        r.facing = Orientation::West;
        r.advance();
        assert_eq!(r.position, Position2D{x: 100, y: 0})
    }

    it "should not overflow default world while moving east" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(150));

        assert_eq!(r.position, Position2D{x: 49, y: 0})
    }

    it "should not overflow default world while moving north" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(150));

        assert_eq!(r.position, Position2D{x: 0, y: 49})
    }

    it "should not overflow default world while moving west" {
        let mut r = rover::new();
        r.facing = Orientation::West;
        r.advance();
        r.advance();

        assert_eq!(r.position, Position2D{x: 99, y: 0})
    }

    it "should not overflow both coordinates" {
        let mut r = rover::new();
        r.move_(Delta::Vector(1338, 143));

        assert_eq!(r.position, Position2D{x: 37, y: 42})
    }

    it "should not overflow while advancing" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(100));
        assert_eq!(r.position, Position2D{x: 0, y: 100})

        r.advance();

        assert_eq!(r.position, Position2D{x: 0, y: 0})
    }

    it "should abort transaction when colliding" {
        let mut r = rover::new();
        r.world.insert("moon monster", Position2D{x: 5, y: 5});
        r.move_(Delta::Vector(5, 5));
        assert_eq!(r.position, Position2D{x: 0, y: 0})
    }

    it "should abort last transaction when colliding in single steps" {
        let mut r = rover::new();
        r.world.insert("moon monster", Position2D{x: 5, y: 5});
        r.move_(Delta::Vector(5, 4));
        r.advance();

        assert_eq!(r.position, Position2D{x: 5, y: 4})
    }
  )
<commit_msg>fix syntax error<commit_after>use rover;
use coords::Position2D;
use rover::Delta;
use rover::Orientation;

describe!(
    it "should create new rover" {
        let r = rover::new();
        assert_eq!(r.position, Position2D{x: 0, y: 0});
    }

    it "should move north" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(1));
        assert_eq!(r.position, Position2D{x: 0, y: 1});
    }

    it "should move north twice" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(2));
        assert_eq!(r.position, Position2D{x: 0, y: 2});
    }

    it "should move east" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(1));
        assert_eq!(r.position, Position2D{x: 1, y: 0});
    }

    it "should move east twice" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(2));
        assert_eq!(r.position, Position2D{x: 2, y: 0});
    }

    it "should move south" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(-1));
        assert_eq!(r.position, Position2D{x: 0, y: 100});
    }

    it "should move south twice" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(-2));
        assert_eq!(r.position, Position2D{x: 0, y: 99});
    }

    it "should move west" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(-1));
        assert_eq!(r.position, Position2D{x: 100, y: 0});
    }

    it "should move west twice" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(-2));
        assert_eq!(r.position, Position2D{x: 99, y: 0});
    }

    it "should move by vector" {
        let mut r = rover::new();
        r.move_(Delta::Vector(2, 3));
        assert_eq!(r.position, Position2D{x: 2, y: 3});
    }

    it "should advance north" {
        let mut r = rover::new();
        r.advance();
        assert_eq!(r.position, Position2D{x: 0, y: 1});
    }

    it "should advance east" {
        let mut r = rover::new();
        r.facing = Orientation::East;
        r.advance();
        assert_eq!(r.position, Position2D{x: 1, y: 0});
    }

    it "should advance south" {
        let mut r = rover::new();
        r.facing = Orientation::South;
        r.advance();
        assert_eq!(r.position, Position2D{x: 0, y: 100});
    }

    it "should advance west" {
        let mut r = rover::new();
        r.facing = Orientation::West;
        r.advance();
        assert_eq!(r.position, Position2D{x: 100, y: 0});
    }

    it "should not overflow default world while moving east" {
        let mut r = rover::new();
        r.move_(Delta::XDelta(150));

        assert_eq!(r.position, Position2D{x: 49, y: 0});
    }

    it "should not overflow default world while moving north" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(150));

        assert_eq!(r.position, Position2D{x: 0, y: 49});
    }

    it "should not overflow default world while moving west" {
        let mut r = rover::new();
        r.facing = Orientation::West;
        r.advance();
        r.advance();

        assert_eq!(r.position, Position2D{x: 99, y: 0});
    }

    it "should not overflow both coordinates" {
        let mut r = rover::new();
        r.move_(Delta::Vector(1338, 143));

        assert_eq!(r.position, Position2D{x: 37, y: 42});
    }

    it "should not overflow while advancing" {
        let mut r = rover::new();
        r.move_(Delta::YDelta(100));
        assert_eq!(r.position, Position2D{x: 0, y: 100});

        r.advance();

        assert_eq!(r.position, Position2D{x: 0, y: 0});
    }

    it "should abort transaction when colliding" {
        let mut r = rover::new();
        r.world.insert("moon monster", Position2D{x: 5, y: 5});
        r.move_(Delta::Vector(5, 5));
        assert_eq!(r.position, Position2D{x: 0, y: 0});
    }

    it "should abort last transaction when colliding in single steps" {
        let mut r = rover::new();
        r.world.insert("moon monster", Position2D{x: 5, y: 5});
        r.move_(Delta::Vector(5, 4));
        r.advance();

        assert_eq!(r.position, Position2D{x: 5, y: 4});
    }
);
<|endoftext|>


<commit_before>#[repr(packed)]
pub struct ElfSectionTag {
    tag_type: u32,
    size: u32,
    num: u16,
    entsize: u16,
    shndx: u16,
    reserved: u16,
    first_section: ElfSectionHeader
}

impl ElfSectionTag {
    pub fn elf_sections(&'static self) -> ElfSectionIter {
        ElfSectionIter {
            current_section: &self.first_section,
            section_index: 0,
            total_sections: self.num as u32,
            entry_size: self.entsize as u64
        }
    }
}

#[repr(C)]
pub struct ElfSectionHeader {
    sh_name: u32,
    sh_type: u32,
    sh_flags: u64,
    sh_addr: u64,
    sh_offset: u64,
    sh_size: u64,
    sh_link: u32,
    sh_info: u32,
    sh_addralign: u64,
    sh_entsize: u64
}

impl ElfSectionHeader {
    pub fn name(&self) -> usize {
        self.sh_name as usize
    }

    pub fn section_type(&self) -> usize {
        self.sh_type as usize
    }

    pub fn start_address(&self) -> usize {
        self.sh_addr as usize
    }

    pub fn end_address(&self) -> usize {
        (self.sh_addr + self.sh_size) as usize
    }

    pub fn size(&self) -> usize {
        self.sh_size as usize
    }

    pub fn flags(&self) -> usize {
        self.sh_flags as usize
    }
}

pub struct ElfSectionIter {
    current_section: &'static ElfSectionHeader,
    section_index: u32,
    total_sections: u32,
    entry_size: u64
}

impl Iterator for ElfSectionIter {
    type Item = &'static ElfSectionHeader;

    fn next(&mut self) -> Option<Self::Item> {
        if self.section_index >= self.total_sections {
            return None;
        } else {
            let section = self.current_section;
            let next_section_addr =
                (self.current_section as *const ElfSectionHeader as u64) + self.entry_size;
            let next_section = unsafe {
                &*(next_section_addr as *const ElfSectionHeader)
            };

            self.current_section = next_section;
            self.section_index += 1;

            Some(section)
        }
    }
}
<commit_msg>Allow access to the number of ELF sections in a section tag as well as the total size of the Elf table in bytes.<commit_after>#[repr(packed)]
pub struct ElfSectionTag {
    tag_type: u32,
    size: u32,
    num: u16,
    entsize: u16,
    shndx: u16,
    reserved: u16,
    first_section: ElfSectionHeader
}

impl ElfSectionTag {
    pub fn elf_sections(&'static self) -> ElfSectionIter {
        ElfSectionIter {
            current_section: &self.first_section,
            section_index: 0,
            total_sections: self.num as u32,
            entry_size: self.entsize as u64
        }
    }

    pub fn section_count(&self) -> usize {
        self.num as usize
    }

    pub fn size(&self) -> usize {
        self.size as usize
    }
}

#[repr(C)]
pub struct ElfSectionHeader {
    sh_name: u32,
    sh_type: u32,
    sh_flags: u64,
    sh_addr: u64,
    sh_offset: u64,
    sh_size: u64,
    sh_link: u32,
    sh_info: u32,
    sh_addralign: u64,
    sh_entsize: u64
}

impl ElfSectionHeader {
    pub fn name(&self) -> usize {
        self.sh_name as usize
    }

    pub fn section_type(&self) -> usize {
        self.sh_type as usize
    }

    pub fn start_address(&self) -> usize {
        self.sh_addr as usize
    }

    pub fn end_address(&self) -> usize {
        (self.sh_addr + self.sh_size) as usize
    }

    pub fn size(&self) -> usize {
        self.sh_size as usize
    }

    pub fn flags(&self) -> usize {
        self.sh_flags as usize
    }
}

pub struct ElfSectionIter {
    current_section: &'static ElfSectionHeader,
    section_index: u32,
    total_sections: u32,
    entry_size: u64
}

impl Iterator for ElfSectionIter {
    type Item = &'static ElfSectionHeader;

    fn next(&mut self) -> Option<Self::Item> {
        if self.section_index >= self.total_sections {
            return None;
        } else {
            let section = self.current_section;
            let next_section_addr =
                (self.current_section as *const ElfSectionHeader as u64) + self.entry_size;
            let next_section = unsafe {
                &*(next_section_addr as *const ElfSectionHeader)
            };

            self.current_section = next_section;
            self.section_index += 1;

            Some(section)
        }
    }
}
<|endoftext|>