improve memoize input equality check logic

ipyflow · Nov 29, 2023 · 3beb795 · 3beb795
1 parent 0914a30
commit 3beb795
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 2 deletions.
diff --git a/core/ipyflow/data_model/cell.py b/core/ipyflow/data_model/cell.py
@@ -289,10 +289,12 @@ def _maybe_memoize_params(self) -> None:
                         inputs[sym] = MemoizedInput(
                             sym,
                             sym.timestamp,
+                            sym.obj_id,
                             sym.make_memoize_comparable()[0],
                         )
         outputs: Dict["Symbol", MemoizedOutput] = {}
         for sym in flow().updated_symbols:
+            sym.last_updated_timestamp_by_obj_id[sym.obj_id] = sym.timestamp
             if not sym.is_user_accessible or not sym.containing_scope.is_global:
                 continue
             outputs[sym] = MemoizedOutput(

diff --git a/core/ipyflow/data_model/symbol.py b/core/ipyflow/data_model/symbol.py
@@ -153,6 +153,8 @@ def __init__(
         self.timestamp_by_liveness_time: Dict[Timestamp, Timestamp] = {}
         # All timestamps associated with updates to this symbol
         self.updated_timestamps: Set[Timestamp] = set()
+        # The most recent timestamp associated with a particular object id
+        self.last_updated_timestamp_by_obj_id: Dict[int, Timestamp] = {}
 
         self.fresher_ancestors: Set["Symbol"] = set()
         self.fresher_ancestor_timestamps: Set[Timestamp] = set()

diff --git a/core/ipyflow/memoization/__init__.py b/core/ipyflow/memoization/__init__.py
@@ -9,6 +9,7 @@
 class MemoizedInput(NamedTuple):
     symbol: "Symbol"
     ts_at_execution: "Timestamp"
+    obj_id_at_execution: int
     comparable: Any
 
 

diff --git a/core/ipyflow/shell/interactiveshell.py b/core/ipyflow/shell/interactiveshell.py
@@ -490,9 +490,14 @@ def before_run_cell(
             for content, inputs, outputs, ctr in prev_cell.memoized_executions:
                 if content != cell.executed_content:
                     continue
-                for (sym, in_ts, comparable) in inputs:
+                for sym, in_ts, obj_id, comparable in inputs:
                     if sym.timestamp.cell_num == in_ts.cell_num:
                         continue
+                    elif (
+                        sym.obj_id == obj_id
+                        and sym.last_updated_timestamp_by_obj_id.get(obj_id) == in_ts
+                    ):
+                        continue
                     elif comparable is Symbol.NULL:
                         break
                     current_comp, eq = sym.make_memoize_comparable()
@@ -508,7 +513,7 @@ def before_run_cell(
             print_purple(
                 "Detected identical symbol usages to previous run; reusing memoized result..."
             )
-            for (sym, out_ts, value) in memoized_outputs:
+            for sym, out_ts, value in memoized_outputs:
                 if sym.obj is not value:
                     self.user_ns[sym.name] = value
                     sym.update_obj_ref(value)