Skip to main content

snarkos_node_bft/helpers/
dag.rs

1// Copyright (c) 2019-2026 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use 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/// Maintains a directed acyclic graph (DAG) of certificates, from which we build a totally-ordered blockchain.
26///
27/// The DAG is updated in rounds, where each validator adds at most one new batch.
28/// Certificates older than GC are removed, as they are not longer needed.
29///
30/// Invariants:
31///  - For each entry in `graph` there is a corresponding certificate ID in `recent_committed_ids`.
32///  - `last_committed_round`, if >0, is equal to the largest key in `recent_commit_ids`.
33#[derive(Debug)]
34pub struct DAG<N: Network> {
35    /// The in-memory collection of certificates that comprise the DAG.
36    /// For each round, there is a mapping from node address to batch.
37    graph: BTreeMap<u64, HashMap<Address<N>, BatchCertificate<N>>>,
38
39    /// The in-memory collection of recently committed certificate IDs (up to GC),
40    /// where each entry has a round number as its key and a set of certificate IDs as its value.
41    recent_committed_ids: BTreeMap<u64, IndexSet<Field<N>>>,
42
43    /// The latest round a certificate was committed too.
44    /// This corresponds to the latest leader certificate's round, because the leader certificate always has the greatest round number of its subDAG.
45    last_committed_round: u64,
46}
47
48impl<N: Network> Default for DAG<N> {
49    /// Initializes a new DAG.
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl<N: Network> DAG<N> {
56    /// Initializes a new DAG.
57    pub fn new() -> Self {
58        Self { graph: Default::default(), recent_committed_ids: Default::default(), last_committed_round: 0 }
59    }
60
61    /// Returns the DAG.
62    pub const fn graph(&self) -> &BTreeMap<u64, HashMap<Address<N>, BatchCertificate<N>>> {
63        &self.graph
64    }
65
66    /// Returns the latest round a certificate was committed to.
67    pub const fn last_committed_round(&self) -> u64 {
68        self.last_committed_round
69    }
70
71    /// Returns `true` if the given certificate ID was recently committed.
72    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    /// Returns `true` if the given certificate ID exists in the given round.
77    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    /// Returns the batch certificate for the given round and author.
82    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    /// Returns the batch certificate for the given round and certificate ID.
87    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    /// Returns the batch certificates for the given round.
99    pub fn get_certificates_for_round(&self, round: u64) -> Option<HashMap<Address<N>, BatchCertificate<N>>> {
100        self.graph.get(&round).cloned()
101    }
102
103    /// Inserts a certificate into the DAG.
104    pub fn insert(&mut self, certificate: BatchCertificate<N>) {
105        let round = certificate.round();
106        let author = certificate.author();
107
108        // If the certificate was not recently committed, insert it into the DAG.
109        if !self.is_recently_committed(round, certificate.id()) {
110            // Insert the certificate into the DAG.
111            let previous = self.graph.entry(round).or_default().insert(author, certificate);
112            // If a previous certificate existed for the author, log it.
113            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    /// Commits a certificate, removing all certificates for this author at or before this round from the DAG.
123    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        /* Updates */
129
130        // Update the recently committed IDs.
131        let is_new = self.recent_committed_ids.entry(certificate_round).or_default().insert(certificate_id);
132        if !is_new {
133            //TODO (kaimast): return early here?
134            trace!("Certificate {certificate_id} was already committed for round {certificate_round}");
135        }
136
137        // Update the last committed round.
138        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            // TODO(kaimast); only remove old certificates for specific author here?
146        }
147
148        /* Garbage collect old commit ids */
149
150        // Remove committed IDs that are below the GC round.
151        self.recent_committed_ids.retain(|round, _| round + max_gc_rounds > self.last_committed_round);
152
153        // Remove certificates that are below the GC round.
154        self.graph.retain(|round, _| round + max_gc_rounds > self.last_committed_round);
155
156        // Remove any certificates for this author that are at or below the certificate round.
157        // TODO (kaimast): is this extra retain needed? It might be less expensive to keep the certificate and skip this additional iteration.
158        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    /// Returns a (malformed) mock DAG, with a modified `last_committed_round`.
173    /// Note: This is useful **only** for testing purposes.
174    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        // Sample a certificate for round 2.
205        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        // Sample a certificate for round 2 and 3 with the same author.
223        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        // Insert the certificate for round 2.
227        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        // Insert the certificate for round 3.
238        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        // Add a lower certificate. As the author is random, it's probably going to be different.
249        let lower = sample_batch_certificate_for_round(2, rng);
250        dag.insert(lower.clone());
251
252        // Add a higher certificate. As the author is random, it's probably going to be different.
253        let higher = sample_batch_certificate_for_round(4, rng);
254        dag.insert(higher.clone());
255
256        // Now commit the certificate for round 3, this will trigger GC.
257        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        // Sample a certificate for round 2, 3, and 4 with the same author.
270        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        /* Test normal case */
275
276        // Insert the certificate for round 2.
277        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        // Now commit the certificate for round 2.
282        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        // Insert the certificate for round 3.
287        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        // Now commit the certificate for round 3.
292        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        /* Test GC case */
297
298        // Insert the certificate for round 4.
299        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        // Now commit the certificate for round 4.
305        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}