Merge pull request #92 from poanetwork/afck-memoize-keys

Avoid redundant key computations.

src/common_coin.rs

@@ -136,8 +136,7 @@ where
     }
 
     fn handle_share(&mut self, sender_id: &NodeUid, share: Signature) -> Result<()> {
-        if let Some(i) = self.netinfo.node_index(sender_id) {
-            let pk_i = self.netinfo.public_key_set().public_key_share(*i as u64);
+        if let Some(pk_i) = self.netinfo.public_key_share(sender_id) {
             if !pk_i.verify(&share, &self.nonce) {
                 // Silently ignore the invalid share.
                 return Ok(());

src/crypto/mod.rs

@@ -223,7 +223,7 @@ impl PublicKeySet {
 
     /// Returns the public key.
     pub fn public_key(&self) -> PublicKey {
-        PublicKey(self.commit.evaluate(0))
+        PublicKey(self.commit.coeff[0])
     }
 
     /// Returns the `i`-th public key share.

src/crypto/poly.rs

@@ -29,7 +29,7 @@ use rand::Rng;
 pub struct Poly {
     /// The coefficients of a polynomial.
     #[serde(with = "super::serde_impl::field_vec")]
-    coeff: Vec<Fr>,
+    pub(super) coeff: Vec<Fr>,
 }
 
 impl<B: Borrow<Poly>> ops::AddAssign<B> for Poly {
@@ -246,7 +246,7 @@ impl Poly {
 pub struct Commitment {
     /// The coefficients of the polynomial.
     #[serde(with = "super::serde_impl::projective_vec")]
-    coeff: Vec<G1>,
+    pub(super) coeff: Vec<G1>,
 }
 
 impl Hash for Commitment {

src/dynamic_honey_badger.rs

@@ -265,7 +265,7 @@ where
     /// Starts Key Generation for the set of nodes implied by the `change`.
     fn start_key_gen(&mut self, change: Change<NodeUid>) -> Result<()> {
         // Use the existing key shares - with the change applied - as keys for DKG.
-        let mut pub_keys = self.netinfo.public_key_map();
+        let mut pub_keys = self.netinfo.public_key_map().clone();
         if match change {
             Change::Remove(id) => pub_keys.remove(&id).is_none(),
             Change::Add(id, pub_key) => pub_keys.insert(id, pub_key).is_some(),

src/honey_badger.rs

@@ -278,9 +278,11 @@ where
         share: &DecryptionShare,
         ciphertext: &Ciphertext,
     ) -> bool {
-        let sender: u64 = *self.netinfo.node_index(sender_id).unwrap() as u64;
-        let pk = self.netinfo.public_key_set().public_key_share(sender);
-        pk.verify_decryption_share(&share, ciphertext)
+        if let Some(pk) = self.netinfo.public_key_share(sender_id) {
+            pk.verify_decryption_share(&share, ciphertext)
+        } else {
+            false
+        }
     }
 
     /// When selections of transactions have been decrypted for all valid proposers in this epoch,

src/messaging.rs

@@ -145,6 +145,7 @@ pub struct NetworkInfo<NodeUid> {
     is_peer: bool,
     secret_key: ClearOnDrop<Box<SecretKey>>,
     public_key_set: PublicKeySet,
+    public_keys: BTreeMap<NodeUid, PublicKey>,
     node_indices: BTreeMap<NodeUid, usize>,
 }
 
@@ -157,11 +158,14 @@ impl<NodeUid: Clone + Ord> NetworkInfo<NodeUid> {
     ) -> Self {
         let num_nodes = all_uids.len();
         let is_peer = all_uids.contains(&our_uid);
-        let node_indices = all_uids
+        let node_indices: BTreeMap<NodeUid, usize> = all_uids
             .iter()
-            .cloned()
             .enumerate()
-            .map(|(n, id)| (id, n))
+            .map(|(n, id)| (id.clone(), n))
+            .collect();
+        let public_keys = node_indices
+            .iter()
+            .map(|(id, idx)| (id.clone(), public_key_set.public_key_share(*idx as u64)))
             .collect();
         NetworkInfo {
             our_uid,
@@ -171,6 +175,7 @@ impl<NodeUid: Clone + Ord> NetworkInfo<NodeUid> {
             is_peer,
             secret_key,
             public_key_set,
+            public_keys,
             node_indices,
         }
     }
@@ -205,18 +210,13 @@ impl<NodeUid: Clone + Ord> NetworkInfo<NodeUid> {
     }
 
     /// Returns the public key share if a node with that ID exists, otherwise `None`.
-    pub fn public_key_share(&self, id: &NodeUid) -> Option<PublicKey> {
-        self.node_index(id)
-            .map(|idx| self.public_key_set.public_key_share(*idx as u64))
+    pub fn public_key_share(&self, id: &NodeUid) -> Option<&PublicKey> {
+        self.public_keys.get(id)
     }
 
     /// Returns a map of all node IDs to their public key shares.
-    pub fn public_key_map(&self) -> BTreeMap<NodeUid, PublicKey> {
-        let to_pair = |(idx, id): (usize, &NodeUid)| {
-            let pub_key = self.public_key_set().public_key_share(idx as u64);
-            (id.clone(), pub_key)
-        };
-        self.all_uids().iter().enumerate().map(to_pair).collect()
+    pub fn public_key_map(&self) -> &BTreeMap<NodeUid, PublicKey> {
+        &self.public_keys
     }
 
     /// The index of a node in a canonical numbering of all nodes.