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 pub fn add_prepare_vote(&mut self, voter: NodeId) -> bool {
53 self.prepare_voters.insert(voter)
54 }
55
56 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 assert!(round.advance(Stage::Done).is_err());
111 }
112}