WIP for tree_hash

beacon_chain/types/src/shard_and_committee.rs

@@ -1,3 +1,5 @@
+use super::ssz::{merkle_hash, TreeHash};
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct ShardAndCommittee {
     pub shard: u16,
@@ -15,6 +17,26 @@ impl ShardAndCommittee {
     }
 }
 
+impl TreeHash for ShardAndCommittee {
+    // python sample code:
+    //    def hash_shard_and_committee(val):
+    //        committee = merkle_hash([x.to_bytes(3, 'big') for x in val.committee])
+    //        return hash(val.shard_id.to_bytes(2, 'big') + committee)
+    fn tree_hash(&self) -> Vec<u8> {
+        let mut committee_ssz_items = Vec::new();
+        for c in &self.committee {
+            let mut h = (*c as u32).tree_hash();
+            h.resize(3, 0);
+            committee_ssz_items.push(h);
+        }
+        let mut result = Vec::new();
+        result.append(&mut self.shard.tree_hash());
+        result.append(&mut merkle_hash(&mut committee_ssz_items));
+
+        result.tree_hash()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -25,4 +47,15 @@ mod tests {
         assert_eq!(s.shard, 0);
         assert_eq!(s.committee.len(), 0);
     }
+
+    #[test]
+    fn test_shard_and_committee_tree_hash() {
+        let s = ShardAndCommittee {
+            shard: 1,
+            committee: vec![1, 2, 3],
+        };
+
+        // should test a known hash value
+        assert_eq!(s.tree_hash().len(), 32);
+    }
 }

beacon_chain/types/src/validator_record.rs

@@ -1,4 +1,5 @@
 use super::bls::{Keypair, PublicKey};
+use super::ssz::TreeHash;
 use super::{Address, Hash256};
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -44,6 +45,46 @@ impl ValidatorRecord {
     }
 }
 
+impl TreeHash for ValidatorRecord {
+    /* python sample code:
+        def hash_validator_record(val):
+            return hash(val.pubkey.to_bytes(32, 'big') + val.withdrawal_shard.to_bytes(2, 'big') + \
+                val.withdrawal_address + val.randao_commitment + val.balance.to_bytes(16, 'big') + \
+                val.start_dynasty.to_bytes(8, 'big') + val.end_dynasty.to_bytes(8, 'big'))
+    */
+    fn tree_hash(&self) -> Vec<u8> {
+        // the serialized fields, to be hashed, should add up to 118 bytes in length.
+        // allocating it once here
+        let mut ssz = Vec::with_capacity(118);
+
+        // "val.pubkey.to_bytes(32, 'big')" logic
+        // TODO:
+        // probably all kinds of wrong here. Not sure how to convert (szz)
+        // pubkey into a big-endian 32 byte array. Note: as_bytes(), the only method on
+        // PublicKey, returns a 192 byte array.
+        let pub_key_bytes = &mut self.pubkey.as_bytes();
+        pub_key_bytes.resize(32, 0);
+        ssz.append(pub_key_bytes);
+
+        ssz.append(&mut self.withdrawal_shard.tree_hash());
+        ssz.append(&mut self.withdrawal_address.tree_hash());
+        ssz.append(&mut self.randao_commitment.tree_hash());
+
+        // balance is a 64bit number that serializes to 8 bytes.
+        // Right padding here to resize to 16 bytes - not sure why
+        // a 16 byte array is implemented in the python code: "val.balance.to_bytes(16, 'big')"
+        let mut balance = self.balance.tree_hash();
+        balance.resize(16, 0);
+        ssz.append(&mut balance);
+
+        // TODO:
+        // ... + val.start_dynasty.to_bytes(8, 'big') + val.end_dynasty.to_bytes(8, 'big')
+        // Our ValidatorRecord seems to be missing the start_dynasty and end_dynasty fields
+
+        ssz.tree_hash()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -59,4 +100,13 @@ mod tests {
         assert_eq!(v.status, 0);
         assert_eq!(v.exit_slot, 0);
     }
+
+    #[test]
+    fn test_validator_record_ree_hash() {
+        let (v, _kp) = ValidatorRecord::zero_with_thread_rand_keypair();
+        let h = v.tree_hash();
+
+        // TODO: should check a known hash result value
+        assert_eq!(h.len(), 32);
+    }
 }