Cache Fragment._component_parent() lookups

The parent-chain walk ran on every reactive attribute update (via
_set_attr/_rem_attr calling component.updated()) and on every ref
(un)registration, making update cost scale with fragment depth.

The result is now cached per fragment. The cache is invalidated in the
parent property setter, on unmount(destroy=True), and - via a new
recursive helper - after DynamicFragment._create_fragment_for_tag()
transfers live children to a new active fragment on a <component :is>
tag switch, which reassigns _parent directly and changes the owning
component of every transferred descendant.

Two regression tests cover the tag-switch reparenting: updated()
attribution and string-ref registration must follow the new owner.
Both fail when the invalidation hook is removed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: berendkleinhaneveld

collagraph/fragment.py

@@ -11,6 +11,9 @@
 from .renderers import Renderer
 from .weak import weak
 
+# Sentinel for an unpopulated Fragment._component_parent_cache
+_UNSET = object()
+
 
 class Fragment:
     """
@@ -45,6 +48,10 @@ def __init__(
 
         # Weak ref to parent fragment
         self._parent: ref[Fragment] | None = ref(parent) if parent else None
+        # Cache for _component_parent(), invalidated whenever this fragment
+        # is reparented to a different owner (see the `parent` setter and
+        # DynamicFragment._create_fragment_for_tag)
+        self._component_parent_cache: Component | None = _UNSET
         # Static attributes for the DOM element
         self._attributes: dict[str, str] = {}
         # Events for the DOM element
@@ -88,22 +95,31 @@ def _component_parent(self) -> Component | None:
         Returns the component of the first parent ComponentFragment that
         has a component property. Or None, if not there.
 
+        Cached, since parenthood is stable once a fragment is mounted. The
+        cache is invalidated whenever a fragment is reparented to a
+        different owner (see the `parent` setter and
+        DynamicFragment._create_fragment_for_tag, the only place a live,
+        already-mounted fragment's parent is reassigned directly).
         """
-        # TODO: would be nice if we could cache the _component_parent in a clever way
+        if self._component_parent_cache is not _UNSET:
+            return self._component_parent_cache
+
         parent = self.parent
         while parent and (
             not isinstance(parent, ComponentFragment) or not parent.component
         ):
             parent = parent.parent
 
-        if parent:
-            return parent.component
+        result = parent.component if parent else None
+        self._component_parent_cache = result
+        return result
 
     @parent.setter
     def parent(self, parent: Fragment | None):
         # TODO: should this also check that this item is
         # now in the list of the parent's children?
         self._parent = ref(parent) if parent else None
+        self._component_parent_cache = _UNSET
 
     def register_child(self, child: Fragment) -> None:
         self.children.append(child)
@@ -448,6 +464,7 @@ def unmount(self, destroy=True):
             self.tag = None
             self._ref_name = None
             self._ref_is_dynamic = False
+            self._component_parent_cache = _UNSET
         else:
             self.element = None
             # Disable the fn and callback of the watcher to disable
@@ -1038,6 +1055,13 @@ def _create_fragment_for_tag(self, tag):
         # Create the fragment
         self._active_fragment.create()
 
+        # `existing_children` (and their descendants) may have been mounted
+        # under a different owning component before this tag switch (e.g.
+        # transferred out of a previous ComponentFragment's slot_contents).
+        # Their cached _component_parent() result would now be stale, since
+        # their effective parent chain just changed.
+        _invalidate_component_parent_cache(self._active_fragment)
+
     def mount(self, target: Any, anchor: Any | None = None):
         if self._mounted:
             return
@@ -1103,6 +1127,36 @@ def _remove(self):
             self._active_fragment._remove()
 
 
+def _invalidate_component_parent_cache(fragment: Fragment | None) -> None:
+    """
+    Recursively reset the cached _component_parent() result for `fragment`
+    and its descendants.
+
+    Needed after a live, already-mounted fragment is reparented to a
+    different owner without going through the `parent` property setter
+    (currently only DynamicFragment._create_fragment_for_tag does this,
+    when transferring existing children to a newly created active
+    fragment on a `<component :is=...>` tag switch).
+    """
+    if fragment is None:
+        return
+
+    fragment._component_parent_cache = _UNSET
+
+    for child in fragment.children:
+        _invalidate_component_parent_cache(child)
+
+    if isinstance(fragment, ComponentFragment):
+        # Note: fragment.fragment (the component's rendered template root)
+        # is already reachable via fragment.children, since create() always
+        # appends it there.
+        for slot_content in fragment.slot_contents:
+            _invalidate_component_parent_cache(slot_content)
+
+    if isinstance(fragment, DynamicFragment):
+        _invalidate_component_parent_cache(fragment._active_fragment)
+
+
 def move_fragment_dom(
     fragment: Fragment, renderer: Renderer, target: Any, anchor: Any | None
 ) -> None:

tests/test_tag_dynamic.py

@@ -984,6 +984,145 @@ class App(Component):
     assert root["children"][0]["type"] == "content"
 
 
+def test_dynamic_component_tag_switch_updates_correct_owner(parse_source):
+    """
+    Regression test for Fragment._component_parent() caching: when a
+    <component :is="tag"> switches between a plain element and a
+    Component class, the transferred slot content's `updated()`
+    notifications must go to the *new* owning component, not a stale
+    cached owner from before the switch.
+    """
+    Wrapper, _ = parse_source(
+        """
+        <wrapper>
+          <slot />
+        </wrapper>
+        <script>
+        from collagraph import Component
+        class Wrapper(Component):
+            updates = 0
+
+            def updated(self):
+                Wrapper.updates += 1
+        </script>
+        """
+    )
+
+    App, _ = parse_source(
+        """
+        <component :is="tag">
+          <content :text="text" />
+        </component>
+        <script>
+        from collagraph import Component
+        class App(Component):
+            updates = 0
+
+            def updated(self):
+                App.updates += 1
+        </script>
+        """,
+        namespace={"Wrapper": Wrapper},
+    )
+
+    state = reactive({"tag": "div", "text": "Hello"})
+    container = {"type": "root"}
+    gui = Collagraph(
+        renderer=DictRenderer(),
+        event_loop_type=EventLoopType.SYNC,
+    )
+    gui.render(App, container, state=state)
+
+    # Plain element: updates to "content" should be attributed to App
+    App.updates = 0
+    Wrapper.updates = 0
+    state["text"] = "World"
+    assert App.updates == 1
+    assert Wrapper.updates == 0
+
+    # Switch to a Component tag: updates to "content" should now be
+    # attributed to Wrapper, not the stale cached App
+    state["tag"] = Wrapper
+    App.updates = 0
+    Wrapper.updates = 0
+    state["text"] = "Wrapped"
+    assert Wrapper.updates == 1
+    assert App.updates == 0
+
+    # Switch back to a plain element: updates should flow to App again
+    state["tag"] = "section"
+    App.updates = 0
+    Wrapper.updates = 0
+    state["text"] = "Plain again"
+    assert App.updates == 1
+    assert Wrapper.updates == 0
+
+
+def test_dynamic_component_tag_switch_moves_string_ref(parse_source):
+    """
+    Regression test for Fragment._component_parent() caching: a string
+    ref on slot content of <component :is="tag"> must register with the
+    component that structurally owns the content after each tag switch,
+    not with a stale cached owner from before the switch.
+    """
+    Wrapper, _ = parse_source(
+        """
+        <wrapper>
+          <slot />
+        </wrapper>
+        <script>
+        from collagraph import Component
+        class Wrapper(Component):
+            instance = None
+
+            def mounted(self):
+                Wrapper.instance = self
+        </script>
+        """
+    )
+
+    App, _ = parse_source(
+        """
+        <component :is="tag">
+          <content ref="content_el" />
+        </component>
+        <script>
+        from collagraph import Component
+        class App(Component):
+            instance = None
+
+            def mounted(self):
+                App.instance = self
+        </script>
+        """,
+        namespace={"Wrapper": Wrapper},
+    )
+
+    state = reactive({"tag": "div"})
+    container = {"type": "root"}
+    gui = Collagraph(
+        renderer=DictRenderer(),
+        event_loop_type=EventLoopType.SYNC,
+    )
+    gui.render(App, container, state=state)
+
+    # Plain element tag: the content is owned by App
+    assert "content_el" in App.instance.refs
+    assert App.instance.refs["content_el"]["type"] == "content"
+
+    # Switch to a Component tag: the content is transferred into
+    # Wrapper's slot and remounted, so the ref must register with the
+    # new owner
+    state["tag"] = Wrapper
+    assert "content_el" in Wrapper.instance.refs
+    assert Wrapper.instance.refs["content_el"]["type"] == "content"
+
+    # Switch back to a plain element: the ref must follow back to App
+    state["tag"] = "section"
+    assert "content_el" in App.instance.refs
+    assert App.instance.refs["content_el"]["type"] == "content"
+
+
 def test_dynamic_component_with_named_slots(parse_source):
     """Test dynamic component with Component class that has named slots"""
     Layout, _ = parse_source(