2 files changed, 79 insertions, 27 deletions

diff --git a/chain/src/item/default/mod.rs b/chain/src/item/default/mod.rs
index 7369c6f..5971985 100644
--- a/chain/src/item/default/mod.rs
+++ b/chain/src/item/default/mod.rs
@@ -3,7 +3,7 @@
 //! forest.
 
 use super::*;
-use crate::atom::default::DefaultAtom;
+use crate::atom::{default::DefaultAtom, Atom};
 use grammar::{GrammarLabel, GrammarLabelType, TNT};
 use graph::{
     builder::BuilderMut, labelled::binary::PLGBuilderMut, Graph, LabelGraph, PLGraph, RedirectGraph,
@@ -1174,11 +1174,82 @@ impl DefaultForest<ForestLabel<GrammarLabel>> {
     /// Set the position of every node.
     ///
     /// If ALLP is non-nil or if the node is a terminal node, also set
-    /// the ending position.
-    pub(crate) fn set_pos(&mut self, pos: usize, allp: bool) -> Result<(), Error> {
+    /// the ending position whereever reasonable.
+    pub(crate) fn set_pos(
+        &mut self,
+        atom: &DefaultAtom,
+        pos: usize,
+        allp: bool,
+    ) -> Result<(), Error> {
         let mut builder = PLGBuilderMut::from_graph_mut(&mut self.graph);
 
-        for node in 0..builder.nodes_len() {
+        let nodes_len = builder.nodes_len();
+
+        if !allp {
+            for node in 0..nodes_len {
+                let label = builder
+                    .vertex_label(node)?
+                    .ok_or(Error::NodeNoLabel(node))?;
+
+                let mut label_label = label.label();
+
+                label_label.set_start(pos);
+
+                let new_label = ForestLabel::new(label_label, label.status);
+
+                builder.set_label(node, new_label)?;
+            }
+
+            return Ok(());
+        }
+
+        // We assign every node to be open first, and then start from
+        // any terminal node and assign them to be closed, then their
+        // parents are rules.  If those rules are accepting, we assign
+        // the rules and the parents of rules to be closed, and so on
+        // and so forth.
+        //
+        // We need to test this behaviour, though.
+
+        let mut closed_nodes: Vec<bool> = std::iter::repeat(false).take(nodes_len).collect();
+
+        let mut closed_stack: Vec<usize> = Vec::with_capacity(nodes_len);
+
+        for (node, closed_nodes_ptr) in (0..nodes_len).zip(closed_nodes.iter_mut()) {
+            let label = builder
+                .vertex_label(node)?
+                .ok_or(Error::NodeNoLabel(node))?;
+
+            if matches!(label.label().label().tnt(), Some(TNT::Ter(_))) {
+                *closed_nodes_ptr = true;
+                closed_stack.push(node);
+            }
+        }
+
+        while let Some(node) = closed_stack.pop() {
+            for parent in builder.parents_of(node)?.map(|p| p.node()) {
+                *closed_nodes
+                    .get_mut(parent)
+                    .ok_or(Error::IndexOutOfBounds(parent, nodes_len))? = true;
+
+                let label = builder
+                    .vertex_label(parent)?
+                    .ok_or(Error::NodeNoLabel(parent))?;
+
+                if let Some(rule) = label.label().label().rule() {
+                    if atom
+                        .is_accepting(2 * rule + 1)
+                        .map_err(|_| Error::IndexOutOfBounds(2 * rule + 1, atom.accepting_len()))?
+                    {
+                        closed_stack.push(parent);
+                    }
+                } else {
+                    closed_stack.push(parent);
+                }
+            }
+        }
+
+        for (node, closed_p) in (0..nodes_len).zip(closed_nodes.iter().copied()) {
             let label = builder
                 .vertex_label(node)?
                 .ok_or(Error::NodeNoLabel(node))?;
@@ -1187,22 +1258,8 @@ impl DefaultForest<ForestLabel<GrammarLabel>> {
 
             label_label.set_start(pos);
 
-            if allp || matches!(label_label.label().tnt(), Some(TNT::Ter(_))) {
+            if closed_p {
                 label_label.set_end(pos + 1);
-            } else if builder.degree(node)? == 1 && label_label.label().rule().is_some() {
-                let child = builder.children_of(node)?.next().unwrap();
-
-                if matches!(
-                    builder
-                        .vertex_label(child)?
-                        .ok_or(Error::NodeNoLabel(child))?
-                        .label()
-                        .label()
-                        .tnt(),
-                    Some(TNT::Ter(_))
-                ) {
-                    label_label.set_end(pos + 1);
-                }
             }
 
             let new_label = ForestLabel::new(label_label, label.status);
diff --git a/chain/src/item/genins.rs b/chain/src/item/genins.rs
index 07dbdfb..1996545 100644
--- a/chain/src/item/genins.rs
+++ b/chain/src/item/genins.rs
@@ -335,7 +335,7 @@ impl DefaultForest<ForestLabel<GrammarLabel>> {
                             if let Some(frag) = virtual_frag {
                                 let mut frag = (*frag.get(0).unwrap()).clone();
 
-                                frag.set_pos(node_label.label().start(), false)?;
+                                frag.set_pos(atom, node_label.label().start(), false)?;
 
                                 let frag_nodes_len = frag.nodes_len();
 
@@ -720,15 +720,10 @@ impl DefaultForest<ForestLabel<GrammarLabel>> {
                 //     return Err(Error::CannotClose(nt, t, node, node_label_start));
                 // }
 
-                // FIXME: It is not correct to set the end of all
-                // nodes unconditionally here.
-                //
-                // We need to figure out the correct way to set the
-                // ending positions.
-
                 for frag in reduction_fragment.into_iter().flatten() {
                     let mut frag = frag.clone();
-                    frag.set_pos(node_label_start, true)?;
+
+                    frag.set_pos(atom, node_label_start, true)?;
 
                     node = self.plant_at_start(node, frag)?;
                 }