1 files changed, 114 insertions, 11 deletions

diff --git a/chain/src/item/default.rs b/chain/src/item/default.rs
index b71f940..f9d26ec 100644
--- a/chain/src/item/default.rs
+++ b/chain/src/item/default.rs
@@ -2,6 +2,8 @@
 //! forest.
 
 use super::*;
+use crate::atom::default::DefaultAtom;
+use grammar::{GrammarLabel, GrammarLabelType};
 use graph::{
     builder::BuilderMut, labelled::binary::PLGBuilderMut, Graph, LabelGraph, PLGraph, RedirectGraph,
 };
@@ -125,6 +127,11 @@ impl<T: GraphLabel> Graph for DefaultForest<T> {
     fn replace_by_builder(&mut self, _builder: impl graph::Builder<Result = Self>) {
         unimplemented!()
     }
+
+    #[inline]
+    fn print_viz(&self, filename: &str) -> Result<(), std::io::Error> {
+        self.graph.print_viz(filename)
+    }
 }
 
 impl<T: GraphLabel> ParentsGraph for DefaultForest<T> {
@@ -360,9 +367,18 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<ForestLabel<T>> {
 
         stack.push(root);
 
+        let mut seen_nodes = std::collections::HashSet::<usize>::new();
+
         while let Some(top) = stack.pop() {
+            seen_nodes.insert(top);
+
             for child in fragment.children_of(top)? {
                 builder.add_edge(conversion!(top), conversion!(child), root_label)?;
+
+                if !seen_nodes.contains(&child) {
+                    seen_nodes.insert(child);
+                    stack.push(child);
+                }
             }
         }
 
@@ -451,22 +467,82 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<ForestLabel<T>> {
     }
 }
 
+impl<T: GraphLabel> PartialEq for DefaultForest<ForestLabel<T>> {
+    fn eq(&self, other: &Self) -> bool {
+        let self_root = self.root();
+        let other_root = other.root();
+
+        if (self_root.is_some() && other_root.is_none())
+            || (self_root.is_none() && other_root.is_some())
+        {
+            return false;
+        } else if self_root.is_none() && other_root.is_none() {
+            return true;
+        }
+
+        // both roots are not empty
+
+        let self_root = self_root.unwrap();
+        let other_root = other_root.unwrap();
+
+        self.is_prefix(self_root, other).unwrap_or(false)
+            && other.is_prefix(other_root, self).unwrap_or(false)
+    }
+}
+
+impl<T: GraphLabel> Eq for DefaultForest<ForestLabel<T>> {}
+
+/// Print the labels found in the forest, so that we can easily
+/// understand what those labels mean.
+pub fn print_labels(
+    atom: impl Borrow<DefaultAtom>,
+    forest: impl Borrow<DefaultForest<ForestLabel<GrammarLabel>>>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let forest = forest.borrow();
+    let atom = atom.borrow();
+
+    for node in forest.nodes() {
+        let label = forest.vertex_label(node)?;
+
+        if let Some(label) = label {
+            let label = label.label.label();
+
+            match label {
+                GrammarLabelType::TNT(tnt) => {
+                    println!("{tnt} = {}", atom.name_of_tnt(tnt)?);
+                }
+                GrammarLabelType::Rule(pos) => {
+                    println!("pos {pos} = {}", atom.rule_pos_string(pos)?);
+                }
+            }
+        } else {
+            return Err(Error::NodeNoLabel(node).into());
+        }
+    }
+
+    Ok(())
+}
+
+#[allow(unused_macros)]
+macro_rules! leaf (
+    ($label:expr, $type:tt) =>{
+        DefaultForest::<ForestLabel<$type>>::new_leaf($label)
+    };
+    ($label:expr) => {
+        DefaultForest::<ForestLabel<usize>>::new_leaf($label)
+    }
+);
+
+#[allow(unused_imports)]
+pub(crate) use leaf;
+
 #[cfg(test)]
 mod item_test {
     use super::*;
 
-    macro_rules! leaf (
-        ($label:expr, $type:tt) =>{
-            DefaultForest::<ForestLabel<$type>>::new_leaf($label)
-        };
-        ($label:expr) => {
-            DefaultForest::<ForestLabel<usize>>::new_leaf($label)
-        }
-    );
-
     #[test]
     fn test_forest_api() -> Result<(), Box<dyn std::error::Error>> {
-        let forest: DefaultForest<usize> = Default::default();
+        let forest: DefaultForest<ForestLabel<usize>> = Default::default();
 
         // empty forest
 
@@ -537,11 +613,38 @@ mod item_test {
 
         forest.clone_node(5)?;
 
-        assert_eq!(forest.nodes_len(), 7);
+        assert_eq!(forest.nodes_len(), 8);
 
         #[cfg(feature = "test-print-viz")]
         forest.print_viz("forest.gv")?;
 
         Ok(())
     }
+
+    #[test]
+    fn test_eq() -> Result<(), Box<dyn std::error::Error>> {
+        let mut forest = leaf!(0, usize);
+
+        forest.plant(0, leaf!(1), false)?;
+        forest.plant(0, leaf!(2), false)?;
+        forest.plant(0, leaf!(3), false)?;
+        forest.plant(0, leaf!(4), false)?;
+        forest.plant(2, leaf!(5), false)?;
+
+        let mut test_forest = leaf!(0);
+        test_forest.plant(0, leaf!(1), false)?;
+        test_forest.plant(0, leaf!(2), false)?;
+        test_forest.plant(0, leaf!(3), false)?;
+        test_forest.plant(2, leaf!(5), false)?;
+
+        assert_ne!(forest, test_forest);
+        assert_ne!(test_forest, forest);
+
+        test_forest.plant(0, leaf!(4), false)?;
+
+        assert_eq!(forest, test_forest);
+        assert_eq!(test_forest, forest);
+
+        Ok(())
+    }
 }