Skip to main content

yellowstone_block_machine/
forks.rs

1use {
2    rustc_hash::{FxHashMap, FxHashSet},
3    std::{
4        collections::{BTreeSet, HashSet, VecDeque},
5        hash::Hash,
6    },
7};
8
9///
10/// Efficient Ordered Set implementatoin that keeps track of the order of insertions and deletions.
11///
12/// Every operation is O(1).
13///
14#[derive(Default, Debug, Clone)]
15pub struct OrderedSet<K> {
16    versioned_keys: FxHashMap<K, u64>,
17    order: VecDeque<(K, u64)>,
18    deleted: FxHashSet<u64>,
19    version: u64,
20    len: usize,
21}
22
23impl<K> OrderedSet<K>
24where
25    K: Clone + Hash + Eq,
26{
27    pub fn new() -> Self {
28        Self {
29            versioned_keys: Default::default(),
30            order: Default::default(),
31            deleted: Default::default(),
32            version: 0,
33            len: 0,
34        }
35    }
36
37    fn next_version(&mut self) -> u64 {
38        let version = self.version;
39        self.version += 1;
40        version
41    }
42
43    pub fn first(&self) -> Option<&K> {
44        self.order
45            .iter()
46            .find(|(_, version)| !self.deleted.contains(version))
47            .map(|(key, _)| key)
48    }
49
50    ///
51    /// Insert a key into the set.
52    ///
53    /// Returns true if the key was already present in the set.
54    pub fn insert(&mut self, key: K) -> bool {
55        let is_present = match self.versioned_keys.get(&key) {
56            Some(version) => !self.deleted.contains(version),
57            _ => false,
58        };
59
60        if is_present {
61            return true;
62        }
63
64        let next_version = self.next_version();
65        self.versioned_keys.insert(key.clone(), next_version);
66        self.order.push_back((key, next_version));
67        self.len += 1;
68        false
69    }
70
71    ///
72    /// Pops the next element inserted into the set.
73    ///
74    /// Returns None if the set is empty.
75    pub fn pop_next(&mut self) -> Option<K> {
76        loop {
77            match self.order.pop_front() {
78                Some((key, version)) => {
79                    if !self.deleted.remove(&version) {
80                        self.versioned_keys.remove(&key);
81                        self.len -= 1;
82                        return Some(key);
83                    } else {
84                        continue;
85                    }
86                }
87                _ => {
88                    return None;
89                }
90            }
91        }
92    }
93
94    ///
95    /// Removes a key from the set.
96    ///
97    /// Returns true if the key was present in the set.
98    pub fn remove(&mut self, key: &K) -> bool {
99        match self.versioned_keys.remove(key) {
100            Some(version) => {
101                self.deleted.insert(version);
102                self.len -= 1;
103                true
104            }
105            _ => false,
106        }
107    }
108
109    ///
110    /// Returns true if the set is empty.
111    ///
112    pub fn is_empty(&self) -> bool {
113        self.versioned_keys.is_empty()
114    }
115
116    ///
117    /// Returns true if the set contains the key.
118    ///
119    pub fn contains(&self, key: &K) -> bool {
120        self.versioned_keys.contains_key(key)
121    }
122
123    ///
124    /// Returns an iterator over the keys in the set.
125    ///
126    pub fn iter(&self) -> OrderedSetIter<'_, K> {
127        OrderedSetIter {
128            ordered_set: self,
129            next_index: 0,
130        }
131    }
132
133    pub fn len(&self) -> usize {
134        self.len
135    }
136}
137
138pub struct OrderedSetIntoIterator<K> {
139    ordered_set: OrderedSet<K>,
140}
141
142impl<K> Iterator for OrderedSetIntoIterator<K>
143where
144    K: Clone + Hash + Eq,
145{
146    type Item = K;
147
148    fn next(&mut self) -> Option<Self::Item> {
149        self.ordered_set.pop_next()
150    }
151}
152
153impl<K> IntoIterator for OrderedSet<K>
154where
155    K: Clone + Hash + Eq,
156{
157    type Item = K;
158    type IntoIter = OrderedSetIntoIterator<K>;
159
160    fn into_iter(self) -> Self::IntoIter {
161        OrderedSetIntoIterator { ordered_set: self }
162    }
163}
164
165///
166/// Iterator for [`OrderedSet`].
167///
168/// Iterates over the keys in the order they were originally inserted.
169///  
170pub struct OrderedSetIter<'a, K> {
171    ordered_set: &'a OrderedSet<K>,
172    next_index: usize,
173}
174
175impl<'a, K> Iterator for OrderedSetIter<'a, K>
176where
177    K: Clone + Hash + Eq,
178{
179    type Item = &'a K;
180
181    fn next(&mut self) -> Option<Self::Item> {
182        loop {
183            match self.ordered_set.order.get(self.next_index) {
184                Some((key, version)) => {
185                    self.next_index += 1;
186                    if !self.ordered_set.deleted.contains(version) {
187                        return Some(key);
188                    } else {
189                        continue;
190                    }
191                }
192                _ => {
193                    return None;
194                }
195            }
196        }
197    }
198}
199
200///
201/// Forks is a utility to detect forks in a blockchain.
202/// Unlike the ForkBanks in Solana runtmie, this struct is more of a utility to detect forks in a blockchain incrementally.
203/// It retroactively detects forks in a blockchain.
204#[derive(Debug)]
205pub struct Forks<T>
206where
207    T: Ord,
208{
209    parent_children_map: FxHashMap<T /*parent */, BTreeSet<T> /*children */>,
210    rooted_slots: BTreeSet<T>,
211    // Map from child to parent
212    reverse_parent_children_map: FxHashMap<T /*child */, T /* parent */>,
213    forked_slots: FxHashSet<T>,
214    // Maximum number of rooted slots to keep track of.
215    max_rooted_depth_capacity: usize,
216}
217
218impl<T> Default for Forks<T>
219where
220    T: Clone + Eq + Hash + Ord,
221{
222    fn default() -> Self {
223        Self::with_max_capacity(1000)
224    }
225}
226
227#[derive(Debug, Clone, Copy, PartialEq)]
228pub enum ForksCapacityStatus {
229    BelowCapacity,
230    AtCapacity,
231    AboveCapacity,
232}
233
234pub struct ForksIterator<'forks, T>
235where
236    T: Ord,
237{
238    forks: &'forks Forks<T>,
239    to_visit: FxHashSet<T>,
240    visited: FxHashSet<T>,
241    queue: VecDeque<T>,
242}
243
244impl<T> Iterator for ForksIterator<'_, T>
245where
246    T: Clone + Eq + Hash + Ord,
247{
248    type Item = T;
249
250    fn next(&mut self) -> Option<Self::Item> {
251        while !self.to_visit.is_empty() {
252            while let Some(slot) = self.queue.pop_front() {
253                if self.visited.contains(&slot) {
254                    continue;
255                }
256                self.visited.insert(slot.clone());
257                self.to_visit.remove(&slot);
258
259                if let Some(children) = self.forks.parent_children_map.get(&slot) {
260                    for child in children.iter().cloned() {
261                        if !self.visited.contains(&child) {
262                            self.queue.push_back(child);
263                        }
264                    }
265                }
266                return Some(slot);
267            }
268
269            if let Some(slot) = self.to_visit.iter().next().cloned() {
270                self.queue.push_back(slot);
271            }
272        }
273        None
274    }
275}
276
277///
278/// Trait for tracing forks update during mutations of the Forks structure.
279///
280pub trait ForksMutationTracer<T> {
281    fn insert(&mut self, slot: T);
282
283    fn extend(&mut self, slots: impl IntoIterator<Item = T>) {
284        for slot in slots {
285            self.insert(slot);
286        }
287    }
288}
289
290impl<T> ForksMutationTracer<T> for FxHashSet<T>
291where
292    T: Eq + Hash,
293{
294    fn insert(&mut self, slot: T) {
295        self.insert(slot);
296    }
297
298    fn extend(&mut self, slots: impl IntoIterator<Item = T>) {
299        Extend::extend(self, slots);
300    }
301}
302
303impl<T> ForksMutationTracer<T> for HashSet<T>
304where
305    T: Eq + Hash,
306{
307    fn insert(&mut self, slot: T) {
308        self.insert(slot);
309    }
310
311    fn extend(&mut self, slots: impl IntoIterator<Item = T>) {
312        Extend::extend(self, slots);
313    }
314}
315
316impl<T> ForksMutationTracer<T> for Vec<T>
317where
318    T: PartialEq,
319{
320    fn insert(&mut self, slot: T) {
321        if self.contains(&slot) {
322            return;
323        }
324        self.push(slot);
325    }
326
327    fn extend(&mut self, slots: impl IntoIterator<Item = T>) {
328        for slot in slots {
329            if self.contains(&slot) {
330                continue;
331            }
332            self.push(slot);
333        }
334    }
335}
336
337impl<T> ForksMutationTracer<T> for BTreeSet<T>
338where
339    T: Ord,
340{
341    fn insert(&mut self, slot: T) {
342        self.insert(slot);
343    }
344
345    fn extend(&mut self, slots: impl IntoIterator<Item = T>) {
346        Extend::extend(self, slots);
347    }
348}
349
350///
351/// No-op implementation of ForksMutationTrace.
352///
353pub struct NoTrace;
354
355impl<T> ForksMutationTracer<T> for NoTrace {
356    fn insert(&mut self, _slot: T) {
357        // Do nothing
358    }
359}
360
361impl<T> Forks<T>
362where
363    T: Clone + Eq + Hash + Ord,
364{
365    pub fn with_max_capacity(capacity: usize) -> Self {
366        assert!(capacity > 0);
367        Self {
368            parent_children_map: Default::default(),
369            reverse_parent_children_map: Default::default(),
370            rooted_slots: Default::default(),
371            forked_slots: Default::default(),
372            max_rooted_depth_capacity: capacity,
373        }
374    }
375
376    pub fn oldest_rooted_slot(&self) -> Option<T> {
377        self.rooted_slots.first().cloned()
378    }
379
380    pub fn capacity_status(&self) -> ForksCapacityStatus {
381        match self.rooted_slots.len().cmp(&self.max_rooted_depth_capacity) {
382            std::cmp::Ordering::Less => ForksCapacityStatus::BelowCapacity,
383            std::cmp::Ordering::Equal => ForksCapacityStatus::AtCapacity,
384            std::cmp::Ordering::Greater => ForksCapacityStatus::AboveCapacity,
385        }
386    }
387
388    ///
389    /// Remove the old rooted slot and all forks derived from it.
390    fn pop_oldest_rooted_slot(&mut self, forks_removed: &mut FxHashSet<T>) {
391        if let Some(root) = self.rooted_slots.pop_first() {
392            let child = self.parent_children_map.remove(&root).unwrap_or_default();
393            let mut queue = VecDeque::from_iter(child);
394
395            while !queue.is_empty() {
396                let slot = queue.pop_front().unwrap();
397                if !self.forked_slots.contains(&slot) {
398                    continue;
399                }
400                forks_removed.insert(slot.clone());
401                self.reverse_parent_children_map.remove(&slot);
402                self.forked_slots.remove(&slot);
403
404                if let Some(children) = self.parent_children_map.remove(&slot) {
405                    queue.extend(children);
406                }
407            }
408
409            if let Some(new_root) = self.rooted_slots.first().cloned() {
410                self.reverse_parent_children_map.remove(&new_root);
411            }
412        }
413    }
414
415    pub fn truncate_excess_rooted_slots(&mut self, forks_removed: &mut FxHashSet<T>) {
416        while self.capacity_status() == ForksCapacityStatus::AboveCapacity {
417            self.pop_oldest_rooted_slot(forks_removed);
418        }
419    }
420
421    fn get_all_nodes(&self) -> FxHashSet<T> {
422        self.parent_children_map.keys().cloned().collect()
423    }
424
425    pub fn get_parent(&self, slot: &T) -> Option<T> {
426        self.reverse_parent_children_map.get(slot).cloned()
427    }
428
429    pub fn visit_slots(&self) -> ForksIterator<'_, T> {
430        ForksIterator {
431            forks: self,
432            to_visit: self.get_all_nodes(),
433            visited: Default::default(),
434            queue: VecDeque::from_iter(self.oldest_rooted_slot()),
435        }
436    }
437    #[allow(clippy::collapsible_else_if)]
438    fn mark_children_as_forks(&mut self, slot: T) -> FxHashSet<T> {
439        let mut queue = VecDeque::from([slot]);
440        let mut newly_forked_slots = FxHashSet::default();
441        let mut visited = FxHashSet::default();
442        while !queue.is_empty() {
443            let slot2 = queue.pop_front().unwrap();
444            if !visited.insert(slot2.clone()) {
445                continue;
446            }
447
448            if let Some(children) = self.parent_children_map.get(&slot2) {
449                for child in children.iter().cloned() {
450                    if self.rooted_slots.contains(&child) {
451                        continue;
452                    } else {
453                        if self.forked_slots.insert(child.clone()) {
454                            queue.push_back(child.clone());
455                            newly_forked_slots.insert(child);
456                        }
457                    }
458                }
459            }
460        }
461        newly_forked_slots
462    }
463
464    pub fn is_rooted_slot(&self, slot: &T) -> bool {
465        self.rooted_slots.contains(slot)
466    }
467
468    pub fn make_slot_rooted_with_rooted_trace<T1, T2>(
469        &mut self,
470        slot: T,
471        newly_forked_slot_out: &mut T1,
472        indirectly_rooted_slots: &mut T2,
473    ) where
474        T1: ForksMutationTracer<T>,
475        T2: ForksMutationTracer<T>,
476    {
477        if self.rooted_slots.contains(&slot) {
478            return;
479        }
480        self.rooted_slots.insert(slot.clone());
481
482        // The rest of the function retroactively detect forks of slot's sibblings, cousins and great* cousins.
483        let mut queue = VecDeque::new();
484        if let Some(parent) = self.reverse_parent_children_map.get(&slot) {
485            queue.push_back(parent.clone());
486        }
487        while !queue.is_empty() {
488            let slot2 = queue.pop_front().expect("empty");
489            // this line will only work the first iteration in the case
490            // we mark a slot i "rooted" and its parent is not rooted yet.
491            if self.rooted_slots.insert(slot2.clone()) {
492                indirectly_rooted_slots.insert(slot2.clone());
493            }
494
495            newly_forked_slot_out.extend(self.mark_children_as_forks(slot2.clone()));
496
497            if let Some(parent2) = self.reverse_parent_children_map.get(&slot2) {
498                // If the parent is already rooted, we don't need to mark its children as forks.
499                // because this process must have been done in the past.
500                if !self.rooted_slots.insert(parent2.clone()) {
501                    continue;
502                } else {
503                    indirectly_rooted_slots.insert(parent2.clone());
504                }
505                queue.push_back(parent2.clone());
506            }
507        }
508    }
509
510    pub fn mark_slot_as_forked<TTracer>(&mut self, slot: T, forks_detected: &mut TTracer)
511    where
512        TTracer: ForksMutationTracer<T>,
513    {
514        self.parent_children_map.entry(slot.clone()).or_default();
515        if self.forked_slots.insert(slot.clone()) {
516            forks_detected.insert(slot.clone());
517        }
518        forks_detected.extend(self.mark_children_as_forks(slot))
519    }
520
521    pub fn make_slot_rooted<TTracer>(&mut self, slot: T, newly_forked_slot_out: &mut TTracer)
522    where
523        TTracer: ForksMutationTracer<T>,
524    {
525        self.make_slot_rooted_with_rooted_trace(slot, newly_forked_slot_out, &mut NoTrace);
526    }
527
528    #[allow(clippy::collapsible_if)]
529    pub fn add_slot_with_parent_with_rooted_trace<T1, T2>(
530        &mut self,
531        slot: T,
532        parent: T,
533        newly_forked_slot_out: &mut T1,
534        indireclty_rooted: &mut T2,
535    ) -> bool
536    where
537        T1: ForksMutationTracer<T>,
538        T2: ForksMutationTracer<T>,
539    {
540        if self
541            .parent_children_map
542            .entry(parent.clone())
543            .or_default()
544            .contains(&slot)
545        {
546            // If the slot is already present, we don't need to do anything.
547            return false;
548        }
549
550        let mut parent_is_rooted = false;
551        if self.rooted_slots.contains(&parent) {
552            parent_is_rooted = true;
553            // If the parent is rooted, we must make sure that none of its children is rooted before inserting.
554            // Otherwise slot is a fork
555            let sibblings = self.parent_children_map.entry(parent.clone()).or_default();
556            for sibbling in sibblings.iter() {
557                if sibbling == &slot {
558                    break;
559                }
560                if self.rooted_slots.contains(sibbling) {
561                    // Parent is rooted so is one of its children.
562                    // This mean that `slot` is a fork
563                    if self.forked_slots.insert(slot.clone()) {
564                        newly_forked_slot_out.insert(slot.clone());
565                    }
566                }
567            }
568        }
569
570        self.parent_children_map
571            .entry(parent.clone())
572            .or_default()
573            .insert(slot.clone());
574        self.parent_children_map.entry(slot.clone()).or_default();
575        self.reverse_parent_children_map
576            .insert(slot.clone(), parent.clone());
577
578        let is_rooted = self.is_rooted_slot(&slot);
579
580        match (is_rooted, parent_is_rooted) {
581            (true, false) => {
582                indireclty_rooted.insert(parent.clone());
583                self.make_slot_rooted_with_rooted_trace(
584                    parent,
585                    newly_forked_slot_out,
586                    indireclty_rooted,
587                )
588            }
589            (false, false) => {
590                if self.forked_slots.contains(&parent) {
591                    if self.forked_slots.insert(slot.clone()) {
592                        newly_forked_slot_out.insert(slot);
593                    }
594                }
595            }
596            _ => {}
597        }
598        true
599    }
600
601    pub fn add_slot_with_parent<TTracer>(
602        &mut self,
603        slot: T,
604        parent: T,
605        newly_forked_slot_out: &mut TTracer,
606    ) -> bool
607    where
608        TTracer: ForksMutationTracer<T>,
609    {
610        self.add_slot_with_parent_with_rooted_trace(
611            slot,
612            parent,
613            newly_forked_slot_out,
614            &mut NoTrace,
615        )
616    }
617
618    pub fn len(&self) -> usize {
619        self.parent_children_map.len()
620    }
621
622    pub fn is_empty(&self) -> bool {
623        self.parent_children_map.is_empty()
624    }
625}
626
627#[cfg(test)]
628#[allow(dead_code)]
629fn module_path_for_test() -> &'static str {
630    module_path!()
631}
632
633#[cfg(test)]
634mod orderset_tests {
635    use super::*;
636
637    #[test]
638    fn pop_first_should_be_fifo_semantic() {
639        let mut ordered_set = OrderedSet::new();
640        assert!(ordered_set.is_empty());
641        assert!(!ordered_set.insert(1));
642        assert!(!ordered_set.insert(2));
643        assert!(!ordered_set.insert(3));
644        assert!(!ordered_set.insert(4));
645
646        assert_eq!(ordered_set.pop_next(), Some(1));
647        assert_eq!(ordered_set.pop_next(), Some(2));
648        assert_eq!(ordered_set.pop_next(), Some(3));
649        assert_eq!(ordered_set.pop_next(), Some(4));
650        assert_eq!(ordered_set.pop_next(), None);
651    }
652
653    #[test]
654    fn pop_first_should_handle_deleted_elements() {
655        let mut ordered_set = OrderedSet::new();
656        assert!(ordered_set.is_empty());
657        assert!(!ordered_set.insert(1));
658        assert!(!ordered_set.insert(2));
659        assert!(!ordered_set.insert(3));
660        assert!(!ordered_set.insert(4));
661
662        assert!(ordered_set.remove(&2));
663        assert!(!ordered_set.contains(&2));
664
665        assert_eq!(ordered_set.pop_next(), Some(1));
666        assert_eq!(ordered_set.pop_next(), Some(3));
667        assert_eq!(ordered_set.pop_next(), Some(4));
668        assert_eq!(ordered_set.pop_next(), None);
669    }
670
671    #[test]
672    fn pop_first_should_return_none_if_set_empty() {
673        let mut ordered_set = OrderedSet::<i32>::new();
674        assert!(ordered_set.is_empty());
675        assert_eq!(ordered_set.pop_next(), None);
676    }
677
678    #[test]
679    fn insert_should_return_true_if_key_already_present() {
680        let mut ordered_set = OrderedSet::new();
681        assert!(ordered_set.is_empty());
682        assert!(!ordered_set.insert(1));
683        assert!(ordered_set.insert(1));
684    }
685
686    #[test]
687    fn empty_ordereset_iter_should_return_none() {
688        let ordered_set = OrderedSet::<i32>::new();
689        let mut iter = ordered_set.iter();
690        assert_eq!(iter.next(), None);
691    }
692
693    #[test]
694    fn all_popped_orderedset_iter_should_return_none() {
695        let mut ordered_set = OrderedSet::new();
696        assert!(ordered_set.is_empty());
697        assert!(!ordered_set.insert(1));
698        assert!(!ordered_set.insert(2));
699        assert!(!ordered_set.insert(3));
700        assert!(!ordered_set.insert(4));
701
702        assert_eq!(ordered_set.pop_next(), Some(1));
703        assert_eq!(ordered_set.pop_next(), Some(2));
704        assert_eq!(ordered_set.pop_next(), Some(3));
705        assert_eq!(ordered_set.pop_next(), Some(4));
706        assert_eq!(ordered_set.pop_next(), None);
707
708        let mut iter = ordered_set.iter();
709        assert_eq!(iter.next(), None);
710    }
711
712    #[test]
713    fn ordereset_iter_should_ignore_removed_elements() {
714        let mut ordered_set = OrderedSet::new();
715        assert!(ordered_set.is_empty());
716        assert!(!ordered_set.insert(1));
717        assert!(!ordered_set.insert(2));
718        assert!(!ordered_set.insert(3));
719        assert!(!ordered_set.insert(4));
720
721        assert!(ordered_set.remove(&2));
722        assert!(!ordered_set.contains(&2));
723
724        let mut iter = ordered_set.iter();
725        assert_eq!(iter.next(), Some(&1));
726        assert_eq!(iter.next(), Some(&3));
727        assert_eq!(iter.next(), Some(&4));
728        assert_eq!(iter.next(), None);
729    }
730
731    #[test]
732    fn orderset_iter_should_go_through_all_elements() {
733        let mut ordered_set = OrderedSet::new();
734        assert!(ordered_set.is_empty());
735        assert!(!ordered_set.insert(1));
736        assert!(!ordered_set.insert(2));
737        assert!(!ordered_set.insert(3));
738        assert!(!ordered_set.insert(4));
739
740        let mut iter = ordered_set.iter();
741        assert_eq!(iter.next(), Some(&1));
742        assert_eq!(iter.next(), Some(&2));
743        assert_eq!(iter.next(), Some(&3));
744        assert_eq!(iter.next(), Some(&4));
745        assert_eq!(iter.next(), None);
746    }
747
748    #[test]
749    fn removing_then_reinsert_should_change_the_set_ordering() {
750        let mut ordered_set = OrderedSet::new();
751        assert!(ordered_set.is_empty());
752        assert!(!ordered_set.insert(1));
753        assert!(!ordered_set.insert(2));
754        assert!(!ordered_set.insert(3));
755        assert!(!ordered_set.insert(4));
756
757        assert!(ordered_set.remove(&2));
758        assert!(!ordered_set.contains(&2));
759
760        assert!(!ordered_set.insert(2));
761        assert!(ordered_set.contains(&2));
762
763        let mut iter = ordered_set.iter();
764        assert_eq!(iter.next(), Some(&1));
765        assert_eq!(iter.next(), Some(&3));
766        assert_eq!(iter.next(), Some(&4));
767        assert_eq!(iter.next(), Some(&2));
768        assert_eq!(iter.next(), None);
769
770        let x = ordered_set.pop_next();
771        assert_eq!(x, Some(1));
772        let x = ordered_set.pop_next();
773        assert_eq!(x, Some(3));
774        let x = ordered_set.pop_next();
775        assert_eq!(x, Some(4));
776        let x = ordered_set.pop_next();
777        assert_eq!(x, Some(2));
778        let x = ordered_set.pop_next();
779        assert_eq!(x, None);
780        assert!(ordered_set.is_empty());
781    }
782}
783
784#[cfg(test)]
785mod forks_tests {
786    use {crate::forks::Forks, rustc_hash::FxHashSet};
787
788    #[test]
789    fn adding_slot_with_parent_twice_should_shortcut() {
790        let mut fd = Forks::default();
791        let mut forked_detected = FxHashSet::default();
792        assert!(fd.add_slot_with_parent(2, 1, &mut forked_detected));
793        assert!(!fd.add_slot_with_parent(2, 1, &mut forked_detected));
794        assert!(forked_detected.is_empty());
795    }
796
797    #[test]
798    fn rooting_slot_should_be_retroactive() {
799        // test case 0
800        // 1 (rooted) -> 2 -> 3 -> 4
801        // make 4 rooted
802        // expected rooted: 1, 2, 3, 4
803        let mut fd = Forks::default();
804        let mut forked_detected = FxHashSet::default();
805        fd.add_slot_with_parent(2, 1, &mut forked_detected);
806        fd.make_slot_rooted(1, &mut forked_detected);
807        fd.add_slot_with_parent(3, 2, &mut forked_detected);
808        fd.add_slot_with_parent(4, 3, &mut forked_detected);
809
810        assert!(fd.rooted_slots.contains(&1));
811        assert!(fd.rooted_slots.len() == 1);
812        let mut indirectly_rooted = FxHashSet::default();
813        fd.make_slot_rooted_with_rooted_trace(4, &mut forked_detected, &mut indirectly_rooted);
814        assert_eq!(indirectly_rooted.len(), 2);
815        assert!(indirectly_rooted.contains(&2));
816        assert!(indirectly_rooted.contains(&3));
817
818        assert!(fd.rooted_slots.contains(&1));
819        assert!(fd.rooted_slots.contains(&2));
820        assert!(fd.rooted_slots.contains(&3));
821        assert!(fd.rooted_slots.contains(&4));
822        assert!(forked_detected.is_empty());
823    }
824
825    #[test]
826    fn forks_should_detect_sibblings_forks() {
827        // test case 0 (no sibblings)
828        // 1 -> 2
829        let mut fd = Forks::default();
830        let mut forked_detected = FxHashSet::default();
831        fd.add_slot_with_parent(2, 1, &mut forked_detected);
832        fd.make_slot_rooted(2, &mut forked_detected);
833        assert!(fd.forked_slots.is_empty());
834        assert!(fd.is_rooted_slot(&2));
835        assert!(fd.is_rooted_slot(&1));
836
837        // test case1 :
838        //  1 -> 2
839        //  1 -> 3
840        // make 2 rooted
841        // expected rooted 2, 1
842        // expected forks 3
843        let mut fd = Forks::default();
844        let mut forked_detected = FxHashSet::default();
845        fd.add_slot_with_parent(2, 1, &mut forked_detected);
846        fd.add_slot_with_parent(3, 1, &mut forked_detected);
847        fd.make_slot_rooted(2, &mut forked_detected);
848        assert!(fd.forked_slots.contains(&3));
849        assert_eq!(fd.forked_slots.len(), 1);
850        assert!(fd.is_rooted_slot(&2));
851        assert!(fd.is_rooted_slot(&1));
852
853        // test case 2
854        // 1 -> 2 -> 3
855        // 1 -> 4
856        // make 2 rooted
857        // expected rooted: 1, 2
858        // expected forks: 4
859        let mut fd = Forks::default();
860        let mut forked_detected = FxHashSet::default();
861        fd.add_slot_with_parent(2, 1, &mut forked_detected);
862        fd.add_slot_with_parent(3, 2, &mut forked_detected);
863        fd.add_slot_with_parent(4, 1, &mut forked_detected);
864        fd.make_slot_rooted(2, &mut forked_detected);
865        assert!(fd.forked_slots.contains(&4));
866        assert_eq!(fd.forked_slots.len(), 1);
867        assert!(fd.is_rooted_slot(&1));
868        assert!(fd.is_rooted_slot(&2));
869        assert!(!fd.is_rooted_slot(&3));
870
871        // test case 3
872        // 1 -> 2 -> 3
873        // 1 -> 4
874        // make 4 rooted
875        // expected rooted: 1, 4
876        // expected forks: 2, 3
877        let mut fd = Forks::default();
878        let mut forked_detected = FxHashSet::default();
879        fd.add_slot_with_parent(2, 1, &mut forked_detected);
880        fd.add_slot_with_parent(3, 2, &mut forked_detected);
881        fd.add_slot_with_parent(4, 1, &mut forked_detected);
882        fd.make_slot_rooted(4, &mut forked_detected);
883        assert!(fd.forked_slots.contains(&2));
884        assert!(fd.forked_slots.contains(&3));
885        assert_eq!(fd.forked_slots.len(), 2);
886        assert!(fd.is_rooted_slot(&1));
887        assert!(fd.is_rooted_slot(&4));
888        assert!(!fd.is_rooted_slot(&2));
889        assert!(!fd.is_rooted_slot(&3));
890    }
891
892    #[test]
893    fn forks_should_detect_cousins_forks() {
894        // test case 1:
895        // 1 (rooted) -> 2 -> 5
896        // 2 -> 3 -> 4
897        // expected rooted : 1,2,5
898        // expected forks : 3, 4
899
900        let mut fd = Forks::default();
901        let mut forked_detected = FxHashSet::default();
902        fd.add_slot_with_parent(2, 1, &mut forked_detected);
903        fd.make_slot_rooted(1, &mut forked_detected);
904
905        fd.add_slot_with_parent(3, 2, &mut forked_detected);
906        fd.add_slot_with_parent(4, 3, &mut forked_detected);
907
908        fd.add_slot_with_parent(5, 2, &mut forked_detected);
909
910        fd.make_slot_rooted(5, &mut forked_detected);
911
912        assert!(fd.is_rooted_slot(&1));
913        assert!(fd.is_rooted_slot(&2));
914        assert!(fd.is_rooted_slot(&5));
915        assert!(fd.forked_slots.contains(&3));
916        assert!(fd.forked_slots.contains(&4));
917        assert!(!fd.is_rooted_slot(&3));
918        assert!(!fd.is_rooted_slot(&4));
919    }
920
921    #[test]
922    fn forks_should_detect_retroactively_greater_cousins_forks() {
923        // Initial setup:
924        // 1 (rooted) -> 2 -> 5 -> ???
925        // 2 -> 3 -> 4
926        let mut fd = Forks::default();
927        let mut forked_detected = FxHashSet::default();
928
929        fd.add_slot_with_parent(2, 1, &mut forked_detected);
930        fd.make_slot_rooted(1, &mut forked_detected);
931
932        fd.add_slot_with_parent(3, 2, &mut forked_detected);
933        fd.add_slot_with_parent(4, 3, &mut forked_detected);
934
935        fd.add_slot_with_parent(5, 2, &mut forked_detected);
936
937        assert!(fd.is_rooted_slot(&1));
938        assert!(fd.forked_slots.is_empty());
939
940        // Then we define another rooted segment
941        // 9 -> 10
942        let mut retroactively_rooted_set = FxHashSet::default();
943        fd.add_slot_with_parent(10, 9, &mut forked_detected);
944        fd.make_slot_rooted_with_rooted_trace(
945            10,
946            &mut forked_detected,
947            &mut retroactively_rooted_set,
948        );
949        assert_eq!(retroactively_rooted_set.len(), 1);
950        assert!(retroactively_rooted_set.contains(&9));
951        assert!(fd.is_rooted_slot(&10));
952        assert!(fd.is_rooted_slot(&9));
953
954        assert!(fd.forked_slots.is_empty());
955
956        // then we connect both segment on 5 -> 9
957        // this should retroactively make 3,4 forks
958        // and make 2 and 5 rooted too
959        let mut retroactively_rooted_set = FxHashSet::default();
960        fd.add_slot_with_parent_with_rooted_trace(
961            9,
962            5,
963            &mut forked_detected,
964            &mut retroactively_rooted_set,
965        );
966        assert_eq!(retroactively_rooted_set.len(), 2);
967        assert!(retroactively_rooted_set.contains(&2));
968        assert!(retroactively_rooted_set.contains(&5));
969        assert!(fd.is_rooted_slot(&2));
970        assert!(fd.is_rooted_slot(&5));
971        assert!(fd.forked_slots.contains(&3));
972        assert!(fd.forked_slots.contains(&4));
973        assert_eq!(fd.forked_slots.len(), 2);
974    }
975
976    #[test]
977    fn forks_descendant_should_be_forks_aswell() {
978        // test case1 :
979        //  1 -> 2
980        //  1 -> 3
981        // make 2 rooted
982        // expected rooted 2, 1
983        // expected forks 3
984        let mut fd = Forks::default();
985        let mut forked_detected = FxHashSet::default();
986        fd.add_slot_with_parent(2, 1, &mut forked_detected);
987        fd.add_slot_with_parent(3, 1, &mut forked_detected);
988        fd.make_slot_rooted(2, &mut forked_detected);
989        assert!(fd.forked_slots.contains(&3));
990        assert_eq!(fd.forked_slots.len(), 1);
991        assert!(fd.is_rooted_slot(&2));
992        assert!(fd.is_rooted_slot(&1));
993
994        // add 4 as a descendant of 3
995        // expected rooted 2, 1
996        // expected forks 3, 4
997
998        fd.add_slot_with_parent(4, 3, &mut forked_detected);
999        assert!(fd.forked_slots.contains(&3));
1000        assert!(fd.forked_slots.contains(&4));
1001        assert_eq!(fd.forked_slots.len(), 2);
1002        assert!(fd.is_rooted_slot(&2));
1003        assert!(fd.is_rooted_slot(&1));
1004        assert!(forked_detected.contains(&4));
1005
1006        let all_slots = fd.visit_slots().collect::<FxHashSet<_>>();
1007        assert_eq!(all_slots.len(), 4);
1008    }
1009
1010    #[test]
1011    fn forks_iterator_should_work_on_disjoint_history() {
1012        //
1013        // Disjoint history comes from the fact that the local history may not contain the full history of the blockchain.
1014        // In this case, the local history is a subset of the full history.
1015        // Some slots may be missing from the local history making some slots unreachable from oldest known rooted slot.
1016        //
1017        // The iterator should be able to traverse the disjoint history.
1018
1019        //
1020        // 1 -> 2
1021        //
1022        // 4 -> 5
1023        // (missing 2 ~~> 4 path)
1024        let mut fd = Forks::default();
1025        let mut forked_detected = FxHashSet::default();
1026        fd.add_slot_with_parent(1, 2, &mut forked_detected);
1027        assert!(forked_detected.is_empty());
1028        fd.add_slot_with_parent(4, 5, &mut forked_detected);
1029        assert!(forked_detected.is_empty());
1030
1031        let slots = fd.visit_slots().collect::<FxHashSet<_>>();
1032        assert_eq!(fd.get_all_nodes().len(), 4);
1033        assert_eq!(slots.len(), 4);
1034        assert!(slots.contains(&1));
1035        assert!(slots.contains(&2));
1036        assert!(slots.contains(&4));
1037        assert!(slots.contains(&5));
1038    }
1039
1040    #[test]
1041    fn forks_should_truncate_old_rooted_slot_when_max_history_capacity_is_reached() {
1042        // test case
1043        //
1044        //                             14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20
1045        //                             ^
1046        //                             |
1047        // 1 (rooted) -> 2 (rooted) -> 3 (rooted) -> 4 (rooted) -> 5 -> 6 -> 7 -> 8 -> 9 -> 10
1048        // |
1049        // v
1050        // 11 -> 12 -> 13
1051        //
1052        // Popoing 1 should remove 11, 12, 13 which are forks
1053        //
1054
1055        let mut fd = Forks::with_max_capacity(1);
1056
1057        let mut forked_detected = FxHashSet::default();
1058        fd.add_slot_with_parent(2, 1, &mut forked_detected);
1059        fd.add_slot_with_parent(3, 2, &mut forked_detected);
1060        fd.add_slot_with_parent(4, 3, &mut forked_detected);
1061        fd.add_slot_with_parent(5, 4, &mut forked_detected);
1062        fd.add_slot_with_parent(6, 5, &mut forked_detected);
1063        fd.add_slot_with_parent(7, 6, &mut forked_detected);
1064        fd.add_slot_with_parent(8, 7, &mut forked_detected);
1065        fd.add_slot_with_parent(9, 8, &mut forked_detected);
1066        fd.add_slot_with_parent(10, 9, &mut forked_detected);
1067
1068        // 1 -> 11 -> 12 -> 13 (Forked path)
1069        fd.add_slot_with_parent(11, 1, &mut forked_detected);
1070        fd.add_slot_with_parent(12, 11, &mut forked_detected);
1071        fd.add_slot_with_parent(13, 12, &mut forked_detected);
1072
1073        // 3 -> 14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20 (Forked path)
1074        fd.add_slot_with_parent(14, 3, &mut forked_detected);
1075        fd.add_slot_with_parent(15, 14, &mut forked_detected);
1076        fd.add_slot_with_parent(16, 15, &mut forked_detected);
1077        fd.add_slot_with_parent(17, 16, &mut forked_detected);
1078        fd.add_slot_with_parent(18, 17, &mut forked_detected);
1079        fd.add_slot_with_parent(19, 18, &mut forked_detected);
1080        fd.add_slot_with_parent(20, 19, &mut forked_detected);
1081
1082        fd.make_slot_rooted(1, &mut forked_detected);
1083        assert_eq!(forked_detected.len(), 0);
1084
1085        fd.make_slot_rooted(2, &mut forked_detected);
1086
1087        let expected_forks_detected = FxHashSet::from_iter([11, 12, 13]);
1088        assert_eq!(forked_detected, expected_forks_detected);
1089
1090        forked_detected.clear();
1091
1092        fd.make_slot_rooted(4, &mut forked_detected);
1093        let expected_forks_detected = FxHashSet::from_iter([14, 15, 16, 17, 18, 19, 20]);
1094        assert_eq!(forked_detected, expected_forks_detected);
1095
1096        let mut forks_removed = FxHashSet::default();
1097
1098        fd.pop_oldest_rooted_slot(&mut forks_removed);
1099
1100        let expected_deleted_forks = FxHashSet::from_iter([11, 12, 13]);
1101        assert_eq!(forks_removed, expected_deleted_forks);
1102
1103        let all_slots = fd.visit_slots().collect::<FxHashSet<_>>();
1104        assert_eq!(all_slots.len(), 16);
1105
1106        assert_eq!(fd.oldest_rooted_slot(), Some(2));
1107    }
1108
1109    #[test]
1110    fn forking_treeroot_should_make_whole_tree_forked() {
1111        // test case
1112        //
1113        //          14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20
1114        //           ^
1115        //           |
1116        // 1 -> 2 -> 3 -> 4  -> 5 -> 6 -> 7 -> 8 -> 9 -> 10
1117        // |
1118        // v
1119        // 11 -> 12 -> 13
1120        //
1121        // Forking 1 should mark all slots as forks
1122
1123        let mut fd = Forks::with_max_capacity(1);
1124
1125        let mut forked_detected = FxHashSet::default();
1126        fd.add_slot_with_parent(2, 1, &mut forked_detected);
1127        fd.add_slot_with_parent(3, 2, &mut forked_detected);
1128        fd.add_slot_with_parent(4, 3, &mut forked_detected);
1129        fd.add_slot_with_parent(5, 4, &mut forked_detected);
1130        fd.add_slot_with_parent(6, 5, &mut forked_detected);
1131        fd.add_slot_with_parent(7, 6, &mut forked_detected);
1132        fd.add_slot_with_parent(8, 7, &mut forked_detected);
1133        fd.add_slot_with_parent(9, 8, &mut forked_detected);
1134        fd.add_slot_with_parent(10, 9, &mut forked_detected);
1135
1136        // 1 -> 11 -> 12 -> 13
1137        fd.add_slot_with_parent(11, 1, &mut forked_detected);
1138        fd.add_slot_with_parent(12, 11, &mut forked_detected);
1139        fd.add_slot_with_parent(13, 12, &mut forked_detected);
1140
1141        // 3 -> 14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20
1142        fd.add_slot_with_parent(14, 3, &mut forked_detected);
1143        fd.add_slot_with_parent(15, 14, &mut forked_detected);
1144        fd.add_slot_with_parent(16, 15, &mut forked_detected);
1145        fd.add_slot_with_parent(17, 16, &mut forked_detected);
1146        fd.add_slot_with_parent(18, 17, &mut forked_detected);
1147        fd.add_slot_with_parent(19, 18, &mut forked_detected);
1148        fd.add_slot_with_parent(20, 19, &mut forked_detected);
1149
1150        assert!(forked_detected.is_empty());
1151
1152        // Forking the tree root
1153        fd.mark_slot_as_forked(1, &mut forked_detected);
1154        let expected_forks_detected = FxHashSet::from_iter([
1155            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
1156        ]);
1157        assert_eq!(forked_detected, expected_forks_detected);
1158
1159        // It should be idempotent
1160        forked_detected.clear();
1161        fd.mark_slot_as_forked(1, &mut forked_detected);
1162        assert!(forked_detected.is_empty());
1163    }
1164
1165    #[test]
1166    fn test_forking_subree() {
1167        // test case
1168        //
1169        //          14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20
1170        //           ^
1171        //           |
1172        // 1 -> 2 -> 3 (mark as fork) -> 4  -> 5 -> 6 -> 7 -> 8 -> 9 -> 10
1173        // |
1174        // v
1175        // 11 -> 12 -> 13
1176        //
1177        // Forking 3 should mark all its descendants as forks
1178        //
1179
1180        let mut fd = Forks::with_max_capacity(1);
1181
1182        let mut forked_detected = FxHashSet::default();
1183        fd.add_slot_with_parent(2, 1, &mut forked_detected);
1184        fd.add_slot_with_parent(3, 2, &mut forked_detected);
1185        fd.add_slot_with_parent(4, 3, &mut forked_detected);
1186        fd.add_slot_with_parent(5, 4, &mut forked_detected);
1187        fd.add_slot_with_parent(6, 5, &mut forked_detected);
1188        fd.add_slot_with_parent(7, 6, &mut forked_detected);
1189        fd.add_slot_with_parent(8, 7, &mut forked_detected);
1190        fd.add_slot_with_parent(9, 8, &mut forked_detected);
1191        fd.add_slot_with_parent(10, 9, &mut forked_detected);
1192
1193        // 1 -> 11 -> 12 -> 13
1194        fd.add_slot_with_parent(11, 1, &mut forked_detected);
1195        fd.add_slot_with_parent(12, 11, &mut forked_detected);
1196        fd.add_slot_with_parent(13, 12, &mut forked_detected);
1197
1198        // 3 -> 14 -> 15 -> 16 -> 17 -> 18 -> 19 -> 20
1199        fd.add_slot_with_parent(14, 3, &mut forked_detected);
1200        fd.add_slot_with_parent(15, 14, &mut forked_detected);
1201        fd.add_slot_with_parent(16, 15, &mut forked_detected);
1202        fd.add_slot_with_parent(17, 16, &mut forked_detected);
1203        fd.add_slot_with_parent(18, 17, &mut forked_detected);
1204        fd.add_slot_with_parent(19, 18, &mut forked_detected);
1205        fd.add_slot_with_parent(20, 19, &mut forked_detected);
1206
1207        assert!(forked_detected.is_empty());
1208
1209        // Forking the tree root
1210        fd.mark_slot_as_forked(3, &mut forked_detected);
1211        let expected_forks_detected =
1212            FxHashSet::from_iter([3, 4, 5, 6, 7, 8, 9, 10, 14, 15, 16, 17, 18, 19, 20]);
1213        assert_eq!(forked_detected, expected_forks_detected);
1214
1215        // It should be idempotent
1216        forked_detected.clear();
1217        fd.mark_slot_as_forked(3, &mut forked_detected);
1218        assert!(forked_detected.is_empty());
1219    }
1220}