1 files changed, 28 insertions, 3 deletions
diff --git a/chain/src/atom/default.rs b/chain/src/atom/default.rs
index 9c91296..ec53596 100644
--- a/chain/src/atom/default.rs
+++ b/chain/src/atom/default.rs
@@ -415,7 +415,7 @@ impl DefaultAtom {
                 .map(|(label, target_iter)| (*label, target_iter))
                 .collect();
 
-            type TerminalsValue = (HashSet<(LabelType<TNT>, usize)>, bool);
+            type TerminalsValue = (HashSet<(LabelType<TNT>, usize, Option<Vec<usize>>)>, bool);
 
             let mut terminals_map: HashMap<usize, TerminalsValue> = HashMap::new();
 
@@ -454,14 +454,25 @@ impl DefaultAtom {
                             // We checked this is safe above.
                             let (set, _) = terminals_map.get_mut(&t).unwrap();
 
-                            set.extend(child_children_iter.map(|target| (*child_label, target)));
+                            set.extend(child_children_iter.map(|target| {
+                                (
+                                    *child_label,
+                                    target,
+                                    grammar
+                                        .query_reduction(child, target)
+                                        .unwrap()
+                                        .map(|slice| slice.to_vec()),
+                                )
+                            }));
                         }
                     }
                 }
             }
 
             for (t, (set, accepting)) in terminals_map.into_iter() {
-                let new_index = nfa.extend(set).map_err(index_out_of_bounds_conversion)?;
+                let new_index = nfa
+                    .extend(set.iter().map(|(label, target, _)| (*label, *target)))
+                    .map_err(index_out_of_bounds_conversion)?;
 
                 if accepting_vec.get(new_index).is_none() {
                     #[cfg(debug_assertions)]
@@ -473,6 +484,20 @@ impl DefaultAtom {
                 let virtual_node = VirtualNode::new(nt, t);
 
                 virtual_nodes.insert(virtual_node, new_index);
+
+                // update the reduction information
+                for (_, target, info) in set.into_iter() {
+                    if let Some(info) = info {
+                        if !matches!(
+                            grammar.query_reduction(new_index, target)?,
+                            Some(original_reduction)
+                                if original_reduction.len()
+                                >= info.len())
+                        {
+                            grammar.set_reduction(new_index, target, info);
+                        }
+                    }
+                }
             }
         }