diff --git a/src/honey_badger/epoch_state.rs b/src/honey_badger/epoch_state.rs index 353a5fd8..fb89235c 100644 --- a/src/honey_badger/epoch_state.rs +++ b/src/honey_badger/epoch_state.rs @@ -117,7 +117,7 @@ pub struct EpochState { /// The status of threshold decryption, by proposer. decryption: BTreeMap>, /// N seen so far - map: BTreeMap>, + nodes_found_so_far: BTreeSet, _phantom: PhantomData, } @@ -134,7 +134,7 @@ where netinfo, subset: SubsetState::Ongoing(cs), decryption: BTreeMap::default(), - map: Default::default(), + nodes_found_so_far: Default::default(), _phantom: PhantomData, }) } @@ -190,8 +190,11 @@ where let plaintexts: BTreeMap = self .decryption .iter() - .flat_map(|(id, dec_state)| dec_state.plaintext().map(|pt| (id.clone(), pt))) - .collect(); + .flat_map(|(id, dec_state)| { + dec_state + .plaintext() + .nodes_found_so_far(|pt| (id.clone(), pt)) + }).collect(); if !proposer_ids.iter().eq(plaintexts.keys()) { return None; // Not all accepted contributions are decrypted yet. } @@ -228,15 +231,16 @@ where for cs_output in cs_outputs { match cs_output { SubsetOutput::Contribution(k, v) => { - step.extend(self.send_decryption_share(k, &v)?); + step.extend(self.send_decryption_share(k.clone(), &v)?); + self.nodes_found_so_far.insert(k); } SubsetOutput::Done => { - self.subset = SubsetState::Complete(self.map.keys().cloned().collect()); + self.subset = SubsetState::Complete(self.nodes_found_so_far.cloned().collect()); let faulty_shares: Vec<_> = self .decryption .keys() - .filter(|id| !self.map.contains_key(id)) + .filter(|id| !self.nodes_found_so_far.contains_key(id)) .cloned() .collect(); for id in faulty_shares { diff --git a/tests/subset.rs b/tests/subset.rs index b2b7b509..55f80b63 100644 --- a/tests/subset.rs +++ b/tests/subset.rs @@ -43,7 +43,6 @@ fn test_subset>>( network.step(); } // Verify that all instances output the same set. - let mut expected = None; for node in network.nodes.values() { let outputs = node.outputs(); let mut actual = BTreeMap::default(); @@ -61,10 +60,9 @@ fn test_subset>>( // The Subset algorithm guarantees that more than two thirds of the proposed elements // are in the set. assert!(actual.len() * 3 > inputs.len() * 2); - match expected { - None => expected = Some(actual), - Some(ref set) => assert_eq!(&actual, set), - }; + for (id, value) in actual { + assert_eq!(&inputs[&id], value); + } } }