1 files changed, 39 insertions, 17 deletions

diff --git a/chain/src/default.rs b/chain/src/default.rs
index 97dbfa3..5f83115 100644
--- a/chain/src/default.rs
+++ b/chain/src/default.rs
@@ -207,6 +207,8 @@ impl Iterator for DerIter {
     }
 }
 
+// TODO: Implement memoization.
+
 /// A default implementation for the [`Chain`] trait.
 #[derive(Debug, Clone, Default, Graph)]
 pub struct DefaultChain {
@@ -492,18 +494,20 @@ impl Chain for DefaultChain {
             let graph_len = graph.nodes_len();
             let atom_len = atom.nodes_len();
 
-            for child in child_iter.clone() {
-                if !graph.has_node(child) {
-                    return Err(Error::IndexOutOfBounds(child, graph_len));
+            #[cfg(debug_assertions)]
+            {
+                for child in child_iter.clone() {
+                    if !graph.has_node(child) {
+                        return Err(Error::IndexOutOfBounds(child, graph_len));
+                    }
                 }
-            }
 
-            for atom_child in atom_child_iter.clone() {
-                if !atom.has_node(atom_child) {
-                    return Err(Error::IndexOutOfBounds(atom_child, atom_len));
+                for atom_child in atom_child_iter.clone() {
+                    if !atom.has_node(atom_child) {
+                        return Err(Error::IndexOutOfBounds(atom_child, atom_len));
+                    }
                 }
             }
-
             // From now on the nodes are all valid, so we can just
             // call `unwrap`.
 
@@ -511,21 +515,37 @@ impl Chain for DefaultChain {
             // repeated allocations
             let mut num = 0usize;
 
+            let child_iter_len = child_iter.len();
+
             let child_iter_total_degree = child_iter
                 .clone()
                 .map(|child| graph.degree(child).unwrap())
                 .sum::<usize>();
 
-            for atom_child in atom_child_iter.clone() {
-                let atom_child_accepting = atom.is_accepting(atom_child).unwrap();
+            num += atom_child_iter
+                .clone()
+                .map(|atom_child| {
+                    let atom_child_accepting = atom.is_accepting(atom_child).unwrap();
 
-                num += child_iter.len();
+                    if atom_child_accepting {
+                        accepting = true;
+                        child_iter_len + child_iter_total_degree
+                    } else {
+                        child_iter_len
+                    }
+                })
+                .sum::<usize>();
 
-                if atom_child_accepting {
-                    accepting = true;
-                    num += child_iter_total_degree;
-                }
-            }
+            // for atom_child in atom_child_iter.clone() {
+            //     let atom_child_accepting = atom.is_accepting(atom_child).unwrap();
+
+            //     num += child_iter.len();
+
+            //     if atom_child_accepting {
+            //         accepting = true;
+            //         num += child_iter_total_degree;
+            //     }
+            // }
 
             let num = num;
 
@@ -1124,9 +1144,11 @@ mod test_chain {
         chain.chain(0, 12, no_item)?;
 
         if !no_item {
-            let _output = chain.end_of_input(13, 0)?;
+            let output = chain.end_of_input(13, 0)?;
             chain.forest.print_viz("forest.gv")?;
 
+            output.print_viz("extracted forest.gv")?;
+
             chain.graph.print_viz("chain.gv")?;
             chain.atom.print_nfa("nfa.gv")?;
             item::default::print_labels(&chain.atom, &chain.forest)?;