snarkos_node_bft/helpers/
dag.rs1use snarkvm::{
17 console::types::{Address, Field},
18 ledger::narwhal::BatchCertificate,
19 prelude::Network,
20};
21
22use indexmap::IndexSet;
23use std::collections::{BTreeMap, HashMap};
24
25#[derive(Debug)]
34pub struct DAG<N: Network> {
35 graph: BTreeMap<u64, HashMap<Address<N>, BatchCertificate<N>>>,
38
39 recent_committed_ids: BTreeMap<u64, IndexSet<Field<N>>>,
42
43 last_committed_round: u64,
46}
47
48impl<N: Network> Default for DAG<N> {
49 fn default() -> Self {
51 Self::new()
52 }
53}
54
55impl<N: Network> DAG<N> {
56 pub fn new() -> Self {
58 Self { graph: Default::default(), recent_committed_ids: Default::default(), last_committed_round: 0 }
59 }
60
61 pub const fn graph(&self) -> &BTreeMap<u64, HashMap<Address<N>, BatchCertificate<N>>> {
63 &self.graph
64 }
65
66 pub const fn last_committed_round(&self) -> u64 {
68 self.last_committed_round
69 }
70
71 pub fn is_recently_committed(&self, round: u64, certificate_id: Field<N>) -> bool {
73 self.recent_committed_ids.get(&round).is_some_and(|ids| ids.contains(&certificate_id))
74 }
75
76 pub fn contains_certificate_in_round(&self, round: u64, certificate_id: Field<N>) -> bool {
78 self.graph.get(&round).is_some_and(|map| map.values().any(|certificate| certificate.id() == certificate_id))
79 }
80
81 pub fn get_certificate_for_round_with_author(&self, round: u64, author: Address<N>) -> Option<BatchCertificate<N>> {
83 self.graph.get(&round).and_then(|certificates| certificates.get(&author)).cloned()
84 }
85
86 pub fn get_certificate_for_round_with_id(
88 &self,
89 round: u64,
90 certificate_id: Field<N>,
91 ) -> Option<BatchCertificate<N>> {
92 self.graph
93 .get(&round)
94 .and_then(|map| map.values().find(|certificate| certificate.id() == certificate_id))
95 .cloned()
96 }
97
98 pub fn get_certificates_for_round(&self, round: u64) -> Option<HashMap<Address<N>, BatchCertificate<N>>> {
100 self.graph.get(&round).cloned()
101 }
102
103 pub fn insert(&mut self, certificate: BatchCertificate<N>) {
105 let round = certificate.round();
106 let author = certificate.author();
107
108 if !self.is_recently_committed(round, certificate.id()) {
110 let previous = self.graph.entry(round).or_default().insert(author, certificate);
112 if previous.is_none() {
114 trace!("Added new certificate for round {round} by author {author} to the DAG");
115 } else {
116 #[cfg(debug_assertions)]
117 error!("A certificate for round {round} by author {author} already existed in the DAG");
118 }
119 }
120 }
121
122 pub fn commit(&mut self, certificate: &BatchCertificate<N>, max_gc_rounds: u64) {
124 let certificate_id = certificate.id();
125 let certificate_round = certificate.round();
126 let author = certificate.author();
127
128 let is_new = self.recent_committed_ids.entry(certificate_round).or_default().insert(certificate_id);
132 if !is_new {
133 trace!("Certificate {certificate_id} was already committed for round {certificate_round}");
135 }
136
137 if self.last_committed_round < certificate_round {
139 trace!(
140 "Last committed round updated from {last_committed_round} to {certificate_round}",
141 last_committed_round = self.last_committed_round
142 );
143 self.last_committed_round = certificate_round;
144 } else {
145 }
147
148 self.recent_committed_ids.retain(|round, _| round + max_gc_rounds > self.last_committed_round);
152
153 self.graph.retain(|round, _| round + max_gc_rounds > self.last_committed_round);
155
156 self.graph.retain(|round, map| match *round > certificate_round {
159 true => true,
160 false => {
161 map.remove(&author);
162 !map.is_empty()
163 }
164 });
165 }
166}
167
168#[cfg(test)]
169pub(crate) mod test_helpers {
170 use super::*;
171
172 pub(crate) fn mock_dag_with_modified_last_committed_round<N: Network>(last_committed_round: u64) -> DAG<N> {
175 let mut dag = DAG::<N>::new();
176 dag.last_committed_round = last_committed_round;
177 dag
178 }
179}
180
181#[cfg(test)]
182mod tests {
183 use super::*;
184 use snarkvm::{
185 prelude::{MainnetV0, narwhal::batch_certificate::test_helpers::sample_batch_certificate_for_round},
186 utilities::TestRng,
187 };
188
189 #[test]
190 fn test_dag_empty() {
191 let dag = DAG::<MainnetV0>::new();
192
193 assert_eq!(dag.get_certificates_for_round(0), None);
194 assert_eq!(dag.last_committed_round(), 0);
195 }
196
197 #[test]
198 fn test_dag_insert() {
199 let rng = &mut TestRng::default();
200 let mut dag = DAG::<MainnetV0>::new();
201
202 const ROUND: u64 = 2;
203
204 let certificate = sample_batch_certificate_for_round(ROUND, rng);
206 dag.insert(certificate.clone());
207 assert!(dag.contains_certificate_in_round(ROUND, certificate.id()));
208 assert_eq!(dag.get_certificate_for_round_with_author(ROUND, certificate.author()), Some(certificate.clone()));
209 assert_eq!(dag.get_certificate_for_round_with_id(ROUND, certificate.id()), Some(certificate.clone()));
210 assert_eq!(
211 dag.get_certificates_for_round(ROUND),
212 Some(vec![(certificate.author(), certificate)].into_iter().collect())
213 );
214 assert_eq!(dag.last_committed_round(), 0);
215 }
216
217 #[test]
218 fn test_dag_commit() {
219 let rng = &mut TestRng::default();
220 let mut dag = DAG::<MainnetV0>::new();
221
222 let certificate_2 = sample_batch_certificate_for_round(2, &mut TestRng::fixed(123456789));
224 let certificate_3 = sample_batch_certificate_for_round(3, &mut TestRng::fixed(123456789));
225
226 dag.insert(certificate_2.clone());
228 assert!(dag.contains_certificate_in_round(2, certificate_2.id()));
229 assert_eq!(dag.get_certificate_for_round_with_author(2, certificate_2.author()), Some(certificate_2.clone()));
230 assert_eq!(dag.get_certificate_for_round_with_id(2, certificate_2.id()), Some(certificate_2.clone()));
231 assert_eq!(
232 dag.get_certificates_for_round(2),
233 Some(vec![(certificate_2.author(), certificate_2.clone())].into_iter().collect())
234 );
235 assert_eq!(dag.last_committed_round(), 0);
236
237 dag.insert(certificate_3.clone());
239 assert!(dag.contains_certificate_in_round(3, certificate_3.id()));
240 assert_eq!(dag.get_certificate_for_round_with_author(3, certificate_3.author()), Some(certificate_3.clone()));
241 assert_eq!(dag.get_certificate_for_round_with_id(3, certificate_3.id()), Some(certificate_3.clone()));
242 assert_eq!(
243 dag.get_certificates_for_round(3),
244 Some(vec![(certificate_3.author(), certificate_3.clone())].into_iter().collect())
245 );
246 assert_eq!(dag.last_committed_round(), 0);
247
248 let lower = sample_batch_certificate_for_round(2, rng);
250 dag.insert(lower.clone());
251
252 let higher = sample_batch_certificate_for_round(4, rng);
254 dag.insert(higher.clone());
255
256 dag.commit(&certificate_3, 10);
258 assert!(!dag.contains_certificate_in_round(2, certificate_2.id()));
259 assert!(!dag.contains_certificate_in_round(3, certificate_3.id()));
260 assert!(dag.contains_certificate_in_round(2, lower.id()));
261 assert!(dag.contains_certificate_in_round(4, higher.id()));
262 assert_eq!(dag.last_committed_round(), 3);
263 }
264
265 #[test]
266 fn test_is_recently_committed() {
267 let mut dag = DAG::<MainnetV0>::new();
268
269 let certificate_2 = sample_batch_certificate_for_round(2, &mut TestRng::fixed(123456789));
271 let certificate_3 = sample_batch_certificate_for_round(3, &mut TestRng::fixed(123456789));
272 let certificate_4 = sample_batch_certificate_for_round(4, &mut TestRng::fixed(123456789));
273
274 dag.insert(certificate_2.clone());
278 assert!(!dag.is_recently_committed(2, certificate_2.id()));
279 assert!(!dag.is_recently_committed(3, certificate_3.id()));
280
281 dag.commit(&certificate_2, 2);
283 assert!(dag.is_recently_committed(2, certificate_2.id()));
284 assert!(!dag.is_recently_committed(3, certificate_3.id()));
285
286 dag.insert(certificate_3.clone());
288 assert!(dag.is_recently_committed(2, certificate_2.id()));
289 assert!(!dag.is_recently_committed(3, certificate_3.id()));
290
291 dag.commit(&certificate_3, 2);
293 assert!(dag.is_recently_committed(2, certificate_2.id()));
294 assert!(dag.is_recently_committed(3, certificate_3.id()));
295
296 dag.insert(certificate_4.clone());
300 assert!(dag.is_recently_committed(2, certificate_2.id()));
301 assert!(dag.is_recently_committed(3, certificate_3.id()));
302 assert!(!dag.is_recently_committed(4, certificate_4.id()));
303
304 dag.commit(&certificate_4, 2);
306 assert!(!dag.is_recently_committed(2, certificate_2.id()));
307 assert!(dag.is_recently_committed(3, certificate_3.id()));
308 assert!(dag.is_recently_committed(4, certificate_4.id()));
309 }
310}