fix: assembly_substituent_tree() now accepts precomputed trace_segments, so callers do not recompute assembly_trace_segments(parts) when both trace and tree are requested, component_namer.name_component() and recursive name_subgraph() now compute trace segments once and reuse them for both return value and tree. Grouped same-name substituents no longer silently overwrite substituent_tree. If multiple instances have different trees, they are preserved

Author: r-fedorov

src/bluenamer/component_namer.py

@@ -593,17 +593,21 @@ def name_component_again(next_mol: Molecule, next_atoms: set[int], is_substituen
             "name_rewrite_history": parts.name_rewrite_history,
         },
     )
-    tree = assembly_substituent_tree(
-        parts,
-        name=name,
-        atom_ids=state.component_atoms,
-        bond_ids=bond_ids_within(mol, state.component_atoms),
-    )
-    tree["kind"] = "component"
+    trace_segments = assembly_trace_segments(parts) if return_trace or return_tree else []
+    tree = None
+    if return_tree:
+        tree = assembly_substituent_tree(
+            parts,
+            name=name,
+            atom_ids=state.component_atoms,
+            bond_ids=bond_ids_within(mol, state.component_atoms),
+            trace_segments=trace_segments,
+        )
+        tree["kind"] = "component"
     if return_trace and return_tree:
-        return name, assembly_trace_segments(parts), tree
+        return name, trace_segments, tree
     if return_trace:
-        return name, assembly_trace_segments(parts)
+        return name, trace_segments
     if return_tree:
         return name, tree
     return name

src/bluenamer/namer.py

@@ -1236,6 +1236,7 @@ def name_subgraph(
     )
 
     name = _assemble_parent_name(mol, parts, numbered_path, get_loc, finalize_subgraph=True)
+    trace_segments = _assembly_trace_segments(parts) if return_trace or return_tree else []
     trace_decision(
         decision_trace,
         TracePhase.ASSEMBLY,
@@ -1245,11 +1246,10 @@ def name_subgraph(
         bonds=_bond_ids_within(mol, component),
         data={
             "name": name,
-            "trace_segment_count": len(_assembly_trace_segments(parts)) if return_trace else 0,
+            "trace_segment_count": len(trace_segments),
         },
     )
     if return_trace:
-        trace_segments = _assembly_trace_segments(parts)
         if return_tree:
             return (
                 name,
@@ -1260,6 +1260,7 @@ def name_subgraph(
                     atom_ids=component,
                     bond_ids=_bond_ids_within(mol, component),
                     decisions=decision_trace_data(decision_trace),
+                    trace_segments=trace_segments,
                 ),
             )
         return name, trace_segments
@@ -1272,6 +1273,7 @@ def name_subgraph(
                 atom_ids=component,
                 bond_ids=_bond_ids_within(mol, component),
                 decisions=decision_trace_data(decision_trace),
+                trace_segments=trace_segments,
             ),
         )
     return name

src/bluenamer/tests/test_analysis.py

@@ -290,6 +290,29 @@ def test_add_substituent_trace_preserves_nested_decisions_into_segments():
     assert assembly_trace_segments(parts)[0]["nested_decisions"] == nested_decisions
 
 
+def test_add_substituent_trace_preserves_grouped_tree_instances():
+    parts = AssemblyParts(parent_length=1, parent_atom_ids={0})
+
+    add_substituent_trace(
+        parts,
+        "ethyl",
+        "1",
+        atom_ids={1, 2},
+        substituent_tree={"kind": "substituent", "name": "ethyl", "atoms": [1, 2]},
+    )
+    add_substituent_trace(
+        parts,
+        "ethyl",
+        "2",
+        atom_ids={3, 4},
+        substituent_tree={"kind": "substituent", "name": "ethyl", "atoms": [3, 4]},
+    )
+
+    tree = parts.substituents[0].substituent_tree
+    assert tree["kind"] == "grouped_substituent_instances"
+    assert [instance["atoms"] for instance in tree["instances"]] == [[1, 2], [3, 4]]
+
+
 def test_principal_suffix_tokens_are_emitted_from_functional_group_renderer():
     parts = AssemblyParts(
         parent_length=3,

src/bluenamer/trace_helpers.py

@@ -113,7 +113,11 @@ def add_substituent_trace(
         existing.trace_segments.extend(trace_segments)
         existing.nested_decisions.extend(nested_decisions)
         if substituent_tree:
-            existing.substituent_tree = substituent_tree
+            existing.substituent_tree = _merge_substituent_tree_instances(
+                existing.substituent_tree,
+                substituent_tree,
+                existing.name,
+            )
         existing.emitted_tokens = existing.emitted_tokens + tuple(emitted_tokens)
     else:
         parts.substituents.append(
@@ -186,7 +190,11 @@ def assembly_trace_segments(parts: AssemblyParts) -> list[dict]:
             target.trace_segments.extend(item.trace_segments)
             target.nested_decisions.extend(item.nested_decisions)
             if item.substituent_tree:
-                target.substituent_tree = item.substituent_tree
+                target.substituent_tree = _merge_substituent_tree_instances(
+                    target.substituent_tree,
+                    item.substituent_tree,
+                    target.name,
+                )
         for item in grouped.values():
             if item.trace_segments and strip_outer_parentheses(item.name) != "methyl":
                 for segment in item.trace_segments:
@@ -305,9 +313,12 @@ def assembly_substituent_tree(
     atom_ids=None,
     bond_ids=None,
     decisions=None,
+    trace_segments=None,
 ) -> dict:
     """Return a nested substituent tree from the graph-bound assembly parts."""
 
+    if trace_segments is None:
+        trace_segments = assembly_trace_segments(parts)
     component_atoms = set(atom_ids or parts.parent_atom_ids)
     component_bonds = set(bond_ids or parts.parent_bond_ids)
     parent_node = {
@@ -352,12 +363,32 @@ def assembly_substituent_tree(
             }
             for item in parts.parent_charges
         ],
-        "trace_segments": assembly_trace_segments(parts),
+        "trace_segments": list(trace_segments),
         "nested_decisions": list(decisions or ()),
     }
     return parent_node
 
 
+def _merge_substituent_tree_instances(existing: dict | None, new: dict, name: str) -> dict:
+    """Preserve all tree instances when same-name substituents are grouped."""
+
+    if existing is None:
+        return new
+    if existing == new:
+        merged = dict(existing)
+        merged["instance_count"] = int(merged.get("instance_count", 1)) + 1
+        return merged
+    if existing.get("kind") == "grouped_substituent_instances":
+        merged = dict(existing)
+        merged["instances"] = [*existing.get("instances", ()), new]
+        return merged
+    return {
+        "kind": "grouped_substituent_instances",
+        "name": name,
+        "instances": [existing, new],
+    }
+
+
 def _parent_tree_node(parts: AssemblyParts) -> dict:
     return {
         "kind": "parent",