script/research/dark-forest: appending functionality along with optional capacity added

script/research/dark-forest/src/error.rs

@@ -30,4 +30,7 @@ pub enum DarkTreeError {
 
     #[error("Invalid DarkLeaf children index found for leaf: {0}")]
     InvalidLeafChildrenIndexes(usize),
+
+    #[error("DarkTree capacity have been exceeded")]
+    CapacityExceeded,
 }

script/research/dark-forest/src/lib.rs

@@ -74,25 +74,75 @@ impl<T> DarkLeaf<T> {
 /// where when we iterate through the tree, we first process tree
 /// node's children, and then the node itself, recursively.
 /// Based on this, initial tree node (leaf), known as the root,
-/// will always show up at the end of iteration.
+/// will always show up at the end of iteration. It is advised
+/// to always execute .build() after finishing setting up the
+/// Tree, to properly index it and check its integrity.
 #[derive(Debug, PartialEq)]
 struct DarkTree<T> {
     /// This tree's leaf information, along with its data
     leaf: DarkLeaf<T>,
     /// Vector containing all tree's branches(children tree)
     children: Vec<DarkTree<T>>,
+    /// Optional max capacity of the tree, including all children
+    /// nodes recursively from the root. None indicates no
+    /// capacity restrictions. This is enforced by the root,
+    /// so children nodes don't have to set it up. If children
+    /// nodes children(recursively) make us exceed that capacity,
+    /// we will be able to catch it using .check_capacity() or
+    /// .integrity_check().
+    capacity: Option<usize>,
 }
 
 impl<T> DarkTree<T> {
     /// Initialize a [`DarkTree`], using provided data to
     /// generate its root.
-    fn new(data: T, children: Vec<DarkTree<T>>) -> DarkTree<T> {
+    fn new(data: T, children: Vec<DarkTree<T>>, capacity: Option<usize>) -> DarkTree<T> {
         let leaf = DarkLeaf::new(data);
-        Self { leaf, children }
+        Self { leaf, children, capacity }
+    }
+
+    /// Build the [`DarkTree`] indexes and perform an
+    /// integrity check on them. This should be used
+    /// after we have appended all child nodes, so we
+    /// don't have to call .index() and .integrity_check()
+    /// manually.
+    fn build(&mut self) -> DarkTreeResult<()> {
+        self.index();
+        self.integrity_check()
+    }
+
+    /// Return the count of all [`DarkTree`] leafs.
+    fn len(&self) -> usize {
+        self.iter().count()
+    }
+
+    /// Check if configured capacity have been exceeded.
+    fn check_capacity(&self) -> DarkTreeResult<()> {
+        if let Some(capacity) = self.capacity {
+            if self.len() >= capacity {
+                return Err(DarkTreeError::CapacityExceeded)
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Append a new child node to the [`DarkTree`],
+    /// if capacity has not been exceeded. This call
+    /// doesn't update the indexes, so either .index()
+    /// or .build() must be called after it.
+    fn append(&mut self, child: DarkTree<T>) -> DarkTreeResult<()> {
+        // Check current capacity
+        self.check_capacity()?;
+
+        // Append the new child
+        self.children.push(child);
+
+        Ok(())
     }
 
     /// Set [`DarkTree`]'s leaf parent and children indexes,
-    /// and trigger the setup of its children indexes
+    /// and trigger the setup of its children indexes.
     fn set_parent_children_indexes(&mut self, parent_index: Option<usize>) {
         // Set our leafs parent index
         self.leaf.set_parent_index(parent_index);
@@ -123,7 +173,7 @@ impl<T> DarkTree<T> {
     }
 
     /// Verify [`DarkTree`]'s leaf parent and children indexes validity,
-    /// and trigger the check of its children indexes
+    /// and trigger the check of its children indexes.
     fn check_parent_children_indexes(&self, parent_index: Option<usize>) -> DarkTreeResult<()> {
         // Check our leafs parent index
         if self.leaf.parent_index != parent_index {
@@ -147,18 +197,22 @@ impl<T> DarkTree<T> {
     }
 
     /// Verify current [`DarkTree`]'s leafs indexes validity,
-    /// based on DFS post-order traversal order. This call
+    /// based on DFS post-order traversal order. Additionally,
+    /// check that capacity has not been exceeded. This call
     /// assumes it was triggered for the root of the tree,
     /// which has no parent index.
     fn integrity_check(&self) -> DarkTreeResult<()> {
-        // First we check each leaf index
+        // Check current capacity
+        self.check_capacity()?;
+
+        // Check each leaf index
         for (index, leaf) in self.iter().enumerate() {
             if index != leaf.index {
                 return Err(DarkTreeError::InvalidLeafIndex(leaf.index, index))
             }
         }
 
-        // Now we trigger recursion to check each nodes rest indexes
+        // Trigger recursion to check each nodes rest indexes
         self.check_parent_children_indexes(None)
     }
 

script/research/dark-forest/src/tests.rs

@@ -38,29 +38,47 @@ fn generate_tree() -> DarkTreeResult<(DarkTree<i32>, Vec<i32>)> {
             DarkTree::new(
                 10,
                 vec![
-                    DarkTree::new(2, vec![DarkTree::new(0, vec![]), DarkTree::new(1, vec![])]),
-                    DarkTree::new(4, vec![DarkTree::new(3, vec![])]),
-                    DarkTree::new(6, vec![DarkTree::new(5, vec![])]),
-                    DarkTree::new(9, vec![DarkTree::new(7, vec![]), DarkTree::new(8, vec![])]),
+                    DarkTree::new(
+                        2,
+                        vec![DarkTree::new(0, vec![], None), DarkTree::new(1, vec![], None)],
+                        None,
+                    ),
+                    DarkTree::new(4, vec![DarkTree::new(3, vec![], None)], None),
+                    DarkTree::new(6, vec![DarkTree::new(5, vec![], None)], None),
+                    DarkTree::new(
+                        9,
+                        vec![DarkTree::new(7, vec![], None), DarkTree::new(8, vec![], None)],
+                        None,
+                    ),
                 ],
+                None,
             ),
             DarkTree::new(
                 14,
-                vec![DarkTree::new(12, vec![DarkTree::new(11, vec![])]), DarkTree::new(13, vec![])],
+                vec![
+                    DarkTree::new(12, vec![DarkTree::new(11, vec![], None)], None),
+                    DarkTree::new(13, vec![], None),
+                ],
+                None,
             ),
             DarkTree::new(
                 21,
                 vec![
-                    DarkTree::new(17, vec![DarkTree::new(15, vec![]), DarkTree::new(16, vec![])]),
-                    DarkTree::new(18, vec![]),
-                    DarkTree::new(20, vec![DarkTree::new(19, vec![])]),
+                    DarkTree::new(
+                        17,
+                        vec![DarkTree::new(15, vec![], None), DarkTree::new(16, vec![], None)],
+                        None,
+                    ),
+                    DarkTree::new(18, vec![], None),
+                    DarkTree::new(20, vec![DarkTree::new(19, vec![], None)], None),
                 ],
+                None,
             ),
         ],
+        None,
     );
 
-    tree.index();
-    tree.integrity_check()?;
+    tree.build()?;
 
     let traversal_order = (0..23).collect();
 
@@ -195,6 +213,7 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
                                 DarkTree {
                                     leaf: DarkLeaf {
@@ -204,8 +223,10 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
-                            ]
+                            ],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -222,7 +243,9 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                     children_indexes: vec![]
                                 },
                                 children: vec![],
+                                capacity: None,
                             },],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -239,7 +262,9 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                     children_indexes: vec![]
                                 },
                                 children: vec![],
+                                capacity: None,
                             },],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -257,6 +282,7 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
                                 DarkTree {
                                     leaf: DarkLeaf {
@@ -266,10 +292,13 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
                             ],
+                            capacity: None,
                         },
                     ],
+                    capacity: None,
                 },
                 DarkTree {
                     leaf: DarkLeaf {
@@ -294,7 +323,9 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                     children_indexes: vec![]
                                 },
                                 children: vec![],
+                                capacity: None,
                             },],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -304,8 +335,10 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                 children_indexes: vec![]
                             },
                             children: vec![],
+                            capacity: None,
                         },
                     ],
+                    capacity: None,
                 },
                 DarkTree {
                     leaf: DarkLeaf {
@@ -331,6 +364,7 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
                                 DarkTree {
                                     leaf: DarkLeaf {
@@ -340,8 +374,10 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                         children_indexes: vec![]
                                     },
                                     children: vec![],
+                                    capacity: None,
                                 },
                             ],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -351,6 +387,7 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                 children_indexes: vec![]
                             },
                             children: vec![],
+                            capacity: None,
                         },
                         DarkTree {
                             leaf: DarkLeaf {
@@ -367,11 +404,15 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
                                     children_indexes: vec![]
                                 },
                                 children: vec![],
+                                capacity: None,
                             },],
+                            capacity: None,
                         },
                     ],
+                    capacity: None,
                 },
             ],
+            capacity: None,
         }
     );
 
@@ -388,3 +429,54 @@ fn test_darktree_mut_iterator() -> DarkTreeResult<()> {
     // Thanks for reading
     Ok(())
 }
+
+#[test]
+pub fn test_darktree_max_capacity() -> DarkTreeResult<()> {
+    // Generate a new [`DarkTree`] with capacity 2
+    let mut tree = DarkTree::new(0, vec![], Some(2));
+
+    // Append a new node
+    tree.append(DarkTree::new(1, vec![], None))?;
+
+    // Try to append a new node
+    assert!(tree.append(DarkTree::new(2, vec![], None)).is_err());
+
+    // Generate a new [`DarkTree`] with capacity 2
+    let mut new_tree = DarkTree::new(3, vec![], Some(2));
+
+    // Append the previous tree as a new node
+    new_tree.append(tree)?;
+
+    // Check that capacity has been exceeded
+    assert!(new_tree.check_capacity().is_err());
+
+    // Generate a new [`DarkTree`] manually with
+    // capacity 1
+    let mut tree = DarkTree {
+        leaf: DarkLeaf { data: 0, index: 0, parent_index: None, children_indexes: vec![] },
+        children: vec![
+            DarkTree {
+                leaf: DarkLeaf { data: 0, index: 0, parent_index: None, children_indexes: vec![] },
+                children: vec![],
+                capacity: None,
+            },
+            DarkTree {
+                leaf: DarkLeaf { data: 0, index: 0, parent_index: None, children_indexes: vec![] },
+                children: vec![],
+                capacity: None,
+            },
+            DarkTree {
+                leaf: DarkLeaf { data: 0, index: 0, parent_index: None, children_indexes: vec![0] },
+                children: vec![],
+                capacity: None,
+            },
+        ],
+        capacity: Some(1),
+    };
+
+    // Verify that building it will fail
+    assert!(tree.build().is_err());
+
+    // Thanks for reading
+    Ok(())
+}