Memoize proxy locations in produce() with epoch invalidation

benchmarks/benchmark.py

@@ -526,6 +526,71 @@ def run():
     benchmark(run)
 
 
+# --- Deep-Path Benchmarks ---
+
+
+DEEP_DEPTH = 12
+
+
+def _make_deep_dict(depth: int) -> dict:
+    """A chain of nested dicts with a small payload at the leaf."""
+    node: dict = {"value": 0, "tag": "leaf"}
+    for level in range(depth):
+        node = {"child": node, "level": level}
+    return node
+
+
+DEEP_BASE = _make_deep_dict(DEEP_DEPTH)
+
+
+def deep_leaf_writes_recipe(draft):
+    """Descend once, then write the leaf repeatedly: every write records
+    a patch whose path is DEEP_DEPTH+1 tokens long."""
+    node = draft
+    for _ in range(DEEP_DEPTH):
+        node = node["child"]
+    for i in range(100):
+        node["value"] = i + 1
+
+
+@pytest.mark.parametrize("in_place", [False, True], ids=["copy", "in_place"])
+@pytest.mark.benchmark(group="produce-deep")
+def test_produce_deep_leaf_writes(benchmark, in_place):
+    """Benchmark: produce() writing repeatedly at the bottom of a deep tree."""
+
+    def run():
+        data = copy.deepcopy(DEEP_BASE)
+        return produce(data, deep_leaf_writes_recipe, in_place=in_place)
+
+    benchmark(run)
+
+
+ITEM_LIST_BASE = {
+    "items": [{"id": i, "done": False, "meta": {"prio": i % 3}} for i in range(100)]
+}
+
+
+def item_field_updates_recipe(draft):
+    """Touch a field on every item of a list of dicts — the classic
+    'mark everything done' shape."""
+    for item in draft["items"]:
+        item["done"] = True
+        item["meta"]["prio"] = 0
+
+
+@pytest.mark.parametrize("in_place", [False, True], ids=["copy", "in_place"])
+@pytest.mark.benchmark(group="produce-deep")
+def test_produce_item_field_updates(benchmark, in_place):
+    """Benchmark: produce() updating a field on every element of a list
+    of nested dicts."""
+
+    def run():
+        data = copy.deepcopy(ITEM_LIST_BASE)
+        return produce(data, item_field_updates_recipe, in_place=in_place)
+
+    benchmark(run)
+
+
 # --- Set Benchmarks ---
 
 SET_BASE = set(range(500))

docs/internals/architecture.md

@@ -42,6 +42,8 @@ The reverse operations are built from the same hunks with input/output roles swa
 
 A proxy stores only its parent proxy and its key within that parent — never an absolute path. Its location is computed on demand by `_location()`, walking up to the root and collecting keys. This is what keeps recorded paths correct when the tree changes under a handed-out reference: when `insert(0, …)` shifts list elements, the list proxy renumbers the keys of its children (`_shift_cache`), and a child that was at index 0 now reports index 1 — no stored path to invalidate.
 
+The walk is **memoized**: each proxy caches its location (reusing the parent's cached location recursively), and every structural change — a detach, a re-attach, a list index shift, `sort`/`reverse` re-keying, finalize — bumps an epoch counter on the recorder that invalidates all cached locations at once. Repeated writes into a stable tree therefore pay O(1) per write instead of an O(depth) walk, while any change that could move a proxy falls back to recomputation on the next write.
+
 ### Detachment stops recording
 
 Each proxy keeps a registry (`_proxies`) of the child proxies it has handed out. Removing or replacing an element marks the corresponding child proxy *detached*: it still mutates its underlying data (so user code holding it keeps working), but `_location()` returns `None` and nothing gets recorded — its data will be captured by the snapshot of whatever write re-inserts it. Re-inserting a *detached* proxy directly is a move: `_adopt` re-attaches it at the new location, and later mutations through the held reference record at the new path. Adopting a still-attached proxy (or one wrapping duck-typed data) snapshots instead, since a value can only live in one place.

patchdiff/produce.py

@@ -120,6 +120,10 @@ class PatchRecorder:
     def __init__(self):
         self.patches: list[Operation] = []
         self.reverse_patches: list[Operation] = []
+        # Bumped on every structural change to the proxy tree (detach,
+        # re-attach, list index shifts); proxies use it to invalidate
+        # their memoized _location() results.
+        self.epoch: int = 0
 
     def finalize(self, root: "_Proxy") -> None:
         """Put the reverse patches in reverse application order and
@@ -132,6 +136,7 @@ def finalize(self, root: "_Proxy") -> None:
         """
         self.reverse_patches.reverse()
         root._detached = True
+        self.epoch += 1
 
     def record_add(
         self, path: Pointer, value: Any, reverse_path: Pointer | None = None
@@ -210,6 +215,8 @@ class _Proxy:
         "_data",
         "_detached",
         "_key",
+        "_loc_epoch",
+        "_loc_tokens",
         "_parent",
         "_proxies",
         "_recorder",
@@ -219,6 +226,8 @@ class _Proxy:
     _data: Any
     _detached: bool
     _key: Any
+    _loc_epoch: int
+    _loc_tokens: tuple | None
     _parent: Any  # weakref.ref[_Proxy] | None
     _proxies: dict[Any, _Proxy] | None
     _recorder: PatchRecorder
@@ -235,6 +244,10 @@ def __init__(
         self._parent = ref(parent) if parent is not None else None
         self._key = key
         self._detached = False
+        # Memoized _location() result, valid while _loc_epoch matches the
+        # recorder's epoch (-1 = never computed)
+        self._loc_epoch = -1
+        self._loc_tokens = None
         # Child-proxy registry, allocated lazily on first wrap: most
         # proxies are leaves that never hand out children
         self._proxies = None
@@ -248,29 +261,38 @@ def _location(self) -> tuple | None:
         or any of its ancestors has been detached from the draft (a dead
         parent reference counts as detached).
 
+        The result is memoized per proxy and recursively reuses the
+        parent's memoized location, so repeated writes into a stable tree
+        pay O(1) per write instead of walking to the root every time.
+        Every structural change (detach, re-attach, list index shifts)
+        bumps the recorder's epoch, which invalidates all memoized
+        locations at once.
+
         Returns a tuple so callers can build a child Pointer in a single
         allocation: Pointer((*tokens, key))."""
+        recorder = self._recorder
+        if self._loc_epoch == recorder.epoch:
+            return self._loc_tokens
+        tokens: tuple | None
         if self._detached:
-            return None
-        parent_ref = self._parent
-        if parent_ref is None:
-            return ()
-        node = parent_ref()
-        if node is None:
-            return None
-        tokens = [self._key]
-        while True:
-            if node._detached:
-                return None
-            parent_ref = node._parent
+            tokens = None
+        else:
+            parent_ref = self._parent
             if parent_ref is None:
-                break
-            tokens.append(node._key)
-            node = parent_ref()
-            if node is None:
-                return None
-        tokens.reverse()
-        return tuple(tokens)
+                tokens = ()
+            else:
+                node = parent_ref()
+                if node is None:
+                    tokens = None
+                else:
+                    parent_tokens = node._location()
+                    if parent_tokens is None:
+                        tokens = None
+                    else:
+                        tokens = (*parent_tokens, self._key)
+        self._loc_tokens = tokens
+        self._loc_epoch = recorder.epoch
+        return tokens
 
     def _wrap(self, key: Hashable, value: Any) -> Any:
         """Wrap nested structures in proxies using duck typing."""
@@ -303,13 +325,15 @@ def _detach(self, key: Hashable) -> None:
         proxy = proxies.pop(key, None)
         if proxy is not None:
             proxy._detached = True
+            self._recorder.epoch += 1
 
     def _detach_all(self) -> None:
         if not self._proxies:
             return
         for proxy in self._proxies.values():
             proxy._detached = True
         self._proxies.clear()
+        self._recorder.epoch += 1
 
     def _in_parent_chain(self, proxy: "_Proxy") -> bool:
         """True when proxy is self or an ancestor of self."""
@@ -343,6 +367,7 @@ def _adopt(self, key: Hashable, value: Any) -> Any:
                 value._detached = False
                 value._parent = ref(self)
                 value._key = key
+                self._recorder.epoch += 1
                 if self._proxies is None:
                     self._proxies = {}
                 self._proxies[key] = value
@@ -497,6 +522,7 @@ def _shift_cache(self, start: int, delta: int) -> None:
         after elements shifted in the underlying list."""
         if not self._proxies:
             return
+        self._recorder.epoch += 1
         shifted = {}
         for index, proxy in self._proxies.items():
             if index >= start:
@@ -739,6 +765,7 @@ def reverse(self) -> None:
         self._data.reverse()
         # Reindex handed-out child proxies to their new positions
         if self._proxies:
+            self._recorder.epoch += 1
             remapped = {}
             for index, proxy in self._proxies.items():
                 new_index = n - 1 - index
@@ -766,6 +793,7 @@ def sort(self, *args, **kwargs) -> None:
         # Reindex handed-out child proxies by following their element's
         # identity to its new position
         if self._proxies:
+            self._recorder.epoch += 1
             positions = {}
             for i, item in enumerate(self._data):
                 positions.setdefault(id(item), []).append(i)