Add signed_root to tree_hash crate

eth2/utils/tree_hash/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod cached_tree_hash;
+pub mod signed_root;
 pub mod standard_tree_hash;
 
 pub const BYTES_PER_CHUNK: usize = 32;
@@ -6,6 +7,7 @@ pub const HASHSIZE: usize = 32;
 pub const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;
 
 pub use cached_tree_hash::{BTreeOverlay, CachedTreeHashSubTree, Error, TreeHashCache};
+pub use signed_root::SignedRoot;
 pub use standard_tree_hash::{efficient_merkleize, TreeHash};
 
 #[derive(Debug, PartialEq, Clone)]

eth2/utils/tree_hash/src/signed_root.rs

@@ -0,0 +1,5 @@
+use crate::TreeHash;
+
+pub trait SignedRoot: TreeHash {
+    fn signed_root(&self) -> Vec<u8>;
+}

eth2/utils/tree_hash_derive/src/lib.rs

@@ -2,7 +2,7 @@
 extern crate proc_macro;
 
 use proc_macro::TokenStream;
-use quote::quote;
+use quote::{quote, ToTokens};
 use syn::{parse_macro_input, DeriveInput};
 
 /// Returns a Vec of `syn::Ident` for each named field in the struct, whilst filtering out fields
@@ -155,3 +155,69 @@ pub fn tree_hash_derive(input: TokenStream) -> TokenStream {
     };
     output.into()
 }
+
+/// Implements `tree_hash::TreeHash` for some `struct`, whilst excluding any fields following and
+/// including a field that is of type "Signature" or "AggregateSignature".
+///
+/// See:
+/// https://github.com/ethereum/eth2.0-specs/blob/master/specs/simple-serialize.md#signed-roots
+///
+/// This is a rather horrendous macro, it will read the type of the object as a string and decide
+/// if it's a signature by matching that string against "Signature" or "AggregateSignature". So,
+/// it's important that you use those exact words as your type -- don't alias it to something else.
+///
+/// If you can think of a better way to do this, please make an issue!
+///
+/// Fields are processed in the order they are defined.
+#[proc_macro_derive(SignedRoot, attributes(signed_root))]
+pub fn tree_hash_signed_root_derive(input: TokenStream) -> TokenStream {
+    let item = parse_macro_input!(input as DeriveInput);
+
+    let name = &item.ident;
+
+    let struct_data = match &item.data {
+        syn::Data::Struct(s) => s,
+        _ => panic!("tree_hash_derive only supports structs."),
+    };
+
+    let idents = get_signed_root_named_field_idents(&struct_data);
+
+    let output = quote! {
+        impl tree_hash::SignedRoot for #name {
+            fn signed_root(&self) -> Vec<u8> {
+                let mut leaves = Vec::with_capacity(4 * tree_hash::HASHSIZE);
+
+                #(
+                    leaves.append(&mut self.#idents.tree_hash_root());
+                )*
+
+                tree_hash::efficient_merkleize(&leaves)[0..32].to_vec()
+            }
+        }
+    };
+    output.into()
+}
+
+fn get_signed_root_named_field_idents(struct_data: &syn::DataStruct) -> Vec<&syn::Ident> {
+    struct_data
+        .fields
+        .iter()
+        .filter_map(|f| {
+            if should_skip_signed_root(&f) {
+                None
+            } else {
+                Some(match &f.ident {
+                    Some(ref ident) => ident,
+                    _ => panic!("tree_hash_derive only supports named struct fields"),
+                })
+            }
+        })
+        .collect()
+}
+
+fn should_skip_signed_root(field: &syn::Field) -> bool {
+    field
+        .attrs
+        .iter()
+        .any(|attr| attr.into_token_stream().to_string() == "# [ signed_root ( skip_hashing ) ]")
+}

eth2/utils/tree_hash_derive/tests/tests.rs

@@ -1,5 +1,5 @@
-use tree_hash::CachedTreeHashSubTree;
-use tree_hash_derive::{CachedTreeHashSubTree, TreeHash};
+use tree_hash::{CachedTreeHashSubTree, SignedRoot, TreeHash};
+use tree_hash_derive::{CachedTreeHashSubTree, SignedRoot, TreeHash};
 
 #[derive(Clone, Debug, TreeHash, CachedTreeHashSubTree)]
 pub struct Inner {
@@ -69,3 +69,32 @@ fn uneven_standard_vs_cached() {
 
     test_standard_and_cached(&original, &modified);
 }
+
+#[derive(Clone, Debug, TreeHash, SignedRoot)]
+pub struct SignedInner {
+    pub a: u64,
+    pub b: u64,
+    pub c: u64,
+    pub d: u64,
+    #[signed_root(skip_hashing)]
+    pub e: u64,
+}
+
+#[test]
+fn signed_root() {
+    let unsigned = Inner {
+        a: 1,
+        b: 2,
+        c: 3,
+        d: 4,
+    };
+    let signed = SignedInner {
+        a: 1,
+        b: 2,
+        c: 3,
+        d: 4,
+        e: 5,
+    };
+
+    assert_eq!(unsigned.tree_hash_root(), signed.signed_root());
+}