Skip to main content

vrf_pbft/pbft/
round.rs

1use std::collections::HashSet;
2
3use crate::error::Result;
4use crate::pbft::stage::validate_transition;
5use crate::types::{Block, NodeId, Role, Stage};
6
7#[derive(Debug)]
8pub struct Round {
9    pub num: u64,
10    stage: Stage,
11    role: Role,
12    threshold: u64,
13    prepare_voters: HashSet<NodeId>,
14    commit_voters: HashSet<NodeId>,
15    block: Option<Block>,
16}
17
18impl Round {
19    pub fn new(num: u64, role: Role, threshold: u64) -> Self {
20        let stage = match role {
21            Role::Proposer => Stage::PrePrepare,
22            _ => Stage::Idle,
23        };
24        Self {
25            num,
26            stage,
27            role,
28            threshold,
29            prepare_voters: HashSet::new(),
30            commit_voters: HashSet::new(),
31            block: None,
32        }
33    }
34
35    pub fn stage(&self) -> Stage {
36        self.stage
37    }
38
39    pub fn role(&self) -> Role {
40        self.role
41    }
42
43    pub fn block(&self) -> Option<&Block> {
44        self.block.as_ref()
45    }
46
47    pub fn set_block(&mut self, block: Block) {
48        self.block = Some(block);
49    }
50
51    /// Record a prepare vote. Returns false if this voter already voted.
52    pub fn add_prepare_vote(&mut self, voter: NodeId) -> bool {
53        self.prepare_voters.insert(voter)
54    }
55
56    /// Record a commit vote. Returns false if this voter already voted.
57    pub fn add_commit_vote(&mut self, voter: NodeId) -> bool {
58        self.commit_voters.insert(voter)
59    }
60
61    pub fn prepare_count(&self) -> u64 {
62        self.prepare_voters.len() as u64
63    }
64
65    pub fn commit_count(&self) -> u64 {
66        self.commit_voters.len() as u64
67    }
68
69    pub fn has_prepare_quorum(&self) -> bool {
70        self.prepare_count() >= self.threshold
71    }
72
73    pub fn has_commit_quorum(&self) -> bool {
74        self.commit_count() >= self.threshold
75    }
76
77    pub fn advance(&mut self, to: Stage) -> Result<()> {
78        validate_transition(self.stage, to)?;
79        self.stage = to;
80        Ok(())
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn vote_counting() {
90        let mut round = Round::new(1, Role::Proposer, 3);
91        assert!(!round.has_prepare_quorum());
92
93        round.add_prepare_vote(1);
94        round.add_prepare_vote(2);
95        assert!(!round.has_prepare_quorum());
96
97        round.add_prepare_vote(3);
98        assert!(round.has_prepare_quorum());
99    }
100
101    #[test]
102    fn stage_advances() {
103        let mut round = Round::new(1, Role::Proposer, 3);
104        assert_eq!(round.stage(), Stage::PrePrepare);
105
106        assert!(round.advance(Stage::Prepare).is_ok());
107        assert_eq!(round.stage(), Stage::Prepare);
108
109        // Can't skip to Done
110        assert!(round.advance(Stage::Done).is_err());
111    }
112}