From 96fe99f8fa81a84ba3990565ced09031a0ae0671 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 17 Jul 2018 13:20:38 +1000
Subject: [PATCH] Implement vote_key for part. xlinks and agg votes

---
 src/state/aggregate_vote.rs           |  8 ++++
 src/state/partial_crosslink_record.rs | 66 +++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/src/state/aggregate_vote.rs b/src/state/aggregate_vote.rs
index c855d6aed..119e5f4f7 100644
--- a/src/state/aggregate_vote.rs
+++ b/src/state/aggregate_vote.rs
@@ -1,6 +1,7 @@
 use super::utils::types::*;
 use super::utils::bls::AggregateSignature;
 use super::rlp::{ RlpStream, Encodable };
+use super::bytes::{ BytesMut, BufMut };
 
 pub struct AggregateVote {
     pub shard_id: u16,
@@ -18,6 +19,13 @@ impl AggregateVote {
             aggregate_sig: AggregateSignature::new()
         }
     }
+
+    pub fn vote_key(&self) -> Vec<u8> {
+        let mut buf = BytesMut::with_capacity(34);
+        buf.extend_from_slice(&self.shard_block_hash.to_vec());
+        buf.put_u16_be(self.shard_id);
+        buf.to_vec()
+    }
 }
 
 /*
diff --git a/src/state/partial_crosslink_record.rs b/src/state/partial_crosslink_record.rs
index 7a7b3a5c5..a0e782cb3 100644
--- a/src/state/partial_crosslink_record.rs
+++ b/src/state/partial_crosslink_record.rs
@@ -1,6 +1,9 @@
+use std::io::Cursor;
 use super::utils::types::{ Sha256Digest, Bitfield };
 use super::rlp::{ RlpStream, Encodable };
+use super::bytes::{ BytesMut, BufMut, Buf };
 
+#[derive(Eq)]
 pub struct PartialCrosslinkRecord {
     pub shard_id: u16,
     pub shard_block_hash: Sha256Digest,
@@ -15,6 +18,37 @@ impl PartialCrosslinkRecord {
             voter_bitfield: Bitfield::new()
         }
     }
+    
+    pub fn new_from_vote_key(vote_key: &Vec<u8>, voter_bitfield: Bitfield)
+        -> Self 
+    {
+        let mut buf = Cursor::new(vote_key);
+        let mut hash_bytes = [0_u8; 32];
+        buf.copy_to_slice(&mut hash_bytes);
+        let shard_id: u16 = buf.get_u16_be();
+        let shard_block_hash = Sha256Digest::from_slice(&hash_bytes);
+        Self {
+            shard_id,
+            shard_block_hash,
+            voter_bitfield
+        }
+    }
+
+    pub fn vote_key(&self) -> Vec<u8> {
+        let mut buf = BytesMut::with_capacity(34);
+        buf.extend_from_slice(&self.shard_block_hash.to_vec());
+        buf.put_u16_be(self.shard_id);
+        buf.to_vec()
+    }
+}
+
+impl PartialEq for PartialCrosslinkRecord {
+    fn eq(&self, other: &PartialCrosslinkRecord) 
+    -> bool
+    {
+        (self.shard_id == other.shard_id) & 
+            (self.shard_block_hash == other.shard_block_hash)
+    }
 }
 
 /*
@@ -41,6 +75,38 @@ mod tests {
         assert_eq!(p.shard_block_hash.is_zero(), true);
         assert_eq!(p.voter_bitfield.num_true_bits(), 0);
     }
+    
+    #[test]
+    fn test_new_from_vote_key() {
+        let mut p = PartialCrosslinkRecord::zero();
+        p.shard_id = 223;
+        p.shard_block_hash = Sha256Digest::random();
+
+        let mut bitfield = Bitfield::new();
+        bitfield.set_bit(&42, &true);
+
+        let vk = p.vote_key();
+        let np = PartialCrosslinkRecord::new_from_vote_key(
+            &vk, bitfield.clone());
+
+        assert_eq!(np.shard_id, p.shard_id);
+        assert_eq!(np.shard_block_hash, p.shard_block_hash);
+        assert!(np.voter_bitfield == bitfield);
+    }
+    
+    #[test]
+    fn test_vote_key_formatting() {
+        let mut p = PartialCrosslinkRecord::zero();
+        let vk = p.vote_key();
+        assert_eq!(vk.len(), 34);
+        assert_eq!(vk, vec![0; 34]);
+
+        p.shard_id = 1; 
+        let vk = p.vote_key();
+        assert_eq!(vk.len(), 34);
+        assert_eq!(vk[0..33].to_vec(), vec![0; 33]);
+        assert_eq!(vk[33], 1);
+    }
 
     #[test]
     fn test_rlp_serialization() {