From 5244d1657ad71f1dc84c6d263f0d7bdd5bf27b28 Mon Sep 17 00:00:00 2001
From: Tristan Simas <tristan.simas@mail.mcgill.ca>
Date: Thu, 12 Nov 2026 17:51:50 -0500
Subject: [PATCH] Collect every declared nested configuration in one traversal

---
 src/objectstate/context_manager.py | 79 +++++++++++-------------------
 tests/test_context_manager.py      | 37 ++++++++++++++
 2 files changed, 66 insertions(+), 50 deletions(-)

diff --git a/src/objectstate/context_manager.py b/src/objectstate/context_manager.py
index 5218869..03f12e2 100644
--- a/src/objectstate/context_manager.py
+++ b/src/objectstate/context_manager.py
@@ -20,6 +20,7 @@
 import functools
 import logging
 import threading
+from collections import deque
 from collections.abc import Sequence
 from contextlib import ExitStack, contextmanager
 from typing import Any, Dict, get_type_hints
@@ -755,59 +756,37 @@ def _declared_config_fields(dataclass_type: type) -> tuple[tuple[str, type], ...
 
 
 def extract_all_configs(context_obj) -> Dict[type, Any]:
-    """
-    Extract all config instances from a context object using type-driven approach.
-
-    A context layer keeps its own result (``ContextLayer.configs``); this function
-    itself keeps nothing.
-
-    This function leverages dataclass field type annotations to efficiently extract
-    config instances, avoiding string matching and runtime attribute scanning.
-
-    Args:
-        context_obj: Object to extract configs from (orchestrator, merged config, etc.)
+    """Read every declared nested configuration, retaining its nominal identity.
 
-    Returns:
-        Dict mapping nominal config types to config instances.
+    Each stored object is visited once. Repeated paths to the same object are
+    shared; different objects claiming the same config type still refuse.
+    Stored values are read without invoking lazy resolution in the layer being built.
     """
-    if context_obj is None:
-        return {}
-
     configs: Dict[type, Any] = {}
-
-    # Include the context object itself if it's a dataclass
-    if is_dataclass(context_obj):
-        _record_config_instance(configs, context_obj, source="context object")
-
-    # Type-driven extraction: Use dataclass field annotations to find config fields
-    if is_dataclass(type(context_obj)):
-        for field_name, declared_type in _declared_config_fields(type(context_obj)):
-            # Read the stored value: a None field holds no config in this object,
-            # and resolving it here would ask the context that is being read.
-            field_value = object.__getattribute__(context_obj, field_name)
-            if field_value is None:
-                continue
-            if not _is_compatible_config_type(field_value, declared_type):
-                raise TypeError(
-                    f"{type(context_obj).__name__}.{field_name} declares "
-                    f"{declared_type.__name__}, got {type(field_value).__name__}."
-                )
-            _record_config_instance(
-                configs,
-                field_value,
-                source=f"{type(context_obj).__name__}.{field_name}",
-            )
-
-    # For non-dataclass objects (orchestrators, etc.), extract dataclass attributes
-    else:
-        _extract_from_object_attributes_typed(context_obj, configs)
-
-    if logger.isEnabledFor(logging.DEBUG):
-        logger.debug(
-            "Extracted %d configs: %s",
-            len(configs),
-            [config_type.__name__ for config_type in configs],
-        )
+    pending = deque([(context_obj, "context object")])
+    visited: set[int] = set()
+    while pending:
+        current, source = pending.popleft()
+        if current is None or id(current) in visited:
+            continue
+        visited.add(id(current))
+        if is_dataclass(current):
+            _record_config_instance(configs, current, source=source)
+            for field_name, declared_type in _declared_config_fields(type(current)):
+                value = object.__getattribute__(current, field_name)
+                if value is None:
+                    continue
+                if not _is_compatible_config_type(value, declared_type):
+                    raise TypeError(
+                        f"{type(current).__name__}.{field_name} declares "
+                        f"{declared_type.__name__}, got {type(value).__name__}."
+                    )
+                pending.append((value, f"{source}.{field_name}"))
+        else:
+            children: Dict[type, Any] = {}
+            _extract_from_object_attributes_typed(current, children)
+            pending.extend((child, f"{source}.{kind.__name__}")
+                           for kind, child in children.items())
     return configs
 
 
diff --git a/tests/test_context_manager.py b/tests/test_context_manager.py
index 4c8f518..cf9bc4d 100644
--- a/tests/test_context_manager.py
+++ b/tests/test_context_manager.py
@@ -194,3 +194,40 @@ class ParentConfig:
         ParentConfig: parent,
         ChildConfig: concrete,
     }
+
+
+def test_extract_all_configs_visits_nested_and_shared_records() -> None:
+    @dataclass
+    class Leaf:
+        value: str = "original"
+
+    @dataclass
+    class Branch:
+        first: Leaf
+        shared: Leaf
+
+    @dataclass
+    class Root:
+        branch: Branch
+
+    leaf = Leaf()
+    root = Root(Branch(leaf, leaf))
+    assert extract_all_configs(root) == {Root: root, Branch: root.branch, Leaf: leaf}
+
+
+def test_extract_all_configs_refuses_collision_below_a_direct_child() -> None:
+    @dataclass
+    class Leaf:
+        value: str
+
+    @dataclass
+    class Branch:
+        first: Leaf
+        second: Leaf
+
+    @dataclass
+    class Root:
+        branch: Branch
+
+    with pytest.raises(ValueError, match="second Leaf instance"):
+        extract_all_configs(Root(Branch(Leaf("first"), Leaf("second"))))