Use new internals.weak module to unify caching mechanisms in evaluate

effectful/handlers/jax/_handlers.py

@@ -14,7 +14,6 @@
 from effectful.ops.semantics import apply, evaluate, fvsof, typeof
 from effectful.ops.syntax import (
     ConstructorOperation,
-    PureInterpretation,
     Scoped,
     _BaseTerm,
     _CustomSingleDispatchCallable,
@@ -23,7 +22,7 @@
     defop,
     syntactic_eq,
 )
-from effectful.ops.types import Expr, NotHandled, Operation, Term
+from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term
 
 # + An element of an array index expression.
 IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array
@@ -43,7 +42,7 @@ def is_eager_array(x):
 
 
 @functools.cache
-def _sizesof_intp() -> tuple[PureInterpretation, Operation]:
+def _sizesof_intp() -> tuple[Interpretation, Operation]:
     """Construct the singleton interpretation used by ``sizesof``."""
     from effectful.internals.product_n import argsof, productN
 
@@ -90,17 +89,15 @@ def _getitem(arr, index):
         return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {})
 
     return (
-        PureInterpretation(
-            productN(
-                {
-                    _sizes: {apply: _apply_sizes, jax_getitem: _getitem},
-                    _getitem_term: {
-                        apply: _retain,
-                        jax_getitem: _retain_getitem,
-                        ConstructorOperation.__apply__: apply.__default_rule__,
-                    },
-                }
-            )
+        productN(
+            {
+                _sizes: {apply: _apply_sizes, jax_getitem: _getitem},
+                _getitem_term: {
+                    apply: _retain,
+                    jax_getitem: _retain_getitem,
+                    ConstructorOperation.__apply__: apply.__default_rule__,
+                },
+            }
         ),
         _sizes,
     )

effectful/handlers/torch.py

@@ -16,14 +16,13 @@
 from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof
 from effectful.ops.syntax import (
     ConstructorOperation,
-    PureInterpretation,
     Scoped,
     _BaseTerm,
     defdata,
     defop,
     syntactic_eq,
 )
-from effectful.ops.types import Expr, NotHandled, Operation, Term
+from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term
 
 # + An element of a tensor index expression.
 IndexElement = None | int | slice | Sequence[int] | EllipsisType | torch.Tensor
@@ -43,7 +42,7 @@ def _getitem_ellipsis_and_none(
 
 
 @functools.cache
-def _sizesof_intp() -> tuple[PureInterpretation, Operation]:
+def _sizesof_intp() -> tuple[Interpretation, Operation]:
     """Construct the singleton interpretation used by ``sizesof``."""
     from effectful.internals.product_n import argsof, productN
 
@@ -94,17 +93,15 @@ def _getitem(x, key):
         return functools.reduce(_merge, itertools.chain(arg_sizes, index_sizes), {})
 
     return (
-        PureInterpretation(
-            productN(
-                {
-                    sizes: {apply: _apply_sizes, torch_getitem: _getitem},
-                    getitem_term: {
-                        apply: _retain,
-                        torch_getitem: _retain_getitem,
-                        ConstructorOperation.__apply__: apply.__default_rule__,
-                    },
-                }
-            )
+        productN(
+            {
+                sizes: {apply: _apply_sizes, torch_getitem: _getitem},
+                getitem_term: {
+                    apply: _retain,
+                    torch_getitem: _retain_getitem,
+                    ConstructorOperation.__apply__: apply.__default_rule__,
+                },
+            }
         ),
         sizes,
     )

effectful/internals/runtime.py

@@ -1,49 +1,111 @@
 import contextlib
-import dataclasses
+import contextvars
 import functools
 import inspect
 import typing
-from collections.abc import Callable, Mapping, MutableMapping
-from threading import local
+from collections.abc import Callable, Mapping
 
+from effectful.internals.weak import (
+    AutoIdKeyDictionary,
+    WeakIdKeyDictionary,
+    weak_memoize,
+)
 from effectful.ops.types import Interpretation, Operation
 
+type CacheEntry = AutoIdKeyDictionary[Interpretation, typing.Any]
+type EvalCache = AutoIdKeyDictionary[typing.Any, CacheEntry]
 
-@dataclasses.dataclass
-class Runtime[S, T](local):
-    interpretation: "Interpretation[S, T]"
-    cache: MutableMapping[int, typing.Any] | None
+EVAL_CACHE: contextvars.ContextVar[EvalCache | None] = contextvars.ContextVar(
+    "EVAL_CACHE", default=None
+)
 
 
-@functools.lru_cache(maxsize=1)
-def get_runtime() -> Runtime:
-    return Runtime(interpretation={}, cache=None)
+INTERPRETATION: contextvars.ContextVar[Interpretation] = contextvars.ContextVar(
+    "INTERPRETATION", default=typing.cast(Interpretation, {})
+)
 
 
-def get_interpretation():
-    return get_runtime().interpretation
+get_interpretation = INTERPRETATION.get
 
 
 @contextlib.contextmanager
 def interpreter(intp: "Interpretation"):
-    r = get_runtime()
-    old_intp, old_cache = r.interpretation, r.cache
+    token = INTERPRETATION.set(intp)
     try:
-        old_intp, r.interpretation = r.interpretation, intp
-        old_cache, r.cache = (
-            r.cache,
-            old_cache if old_intp is intp and old_cache is not None else {},
-        )
         yield intp
     finally:
-        r.interpretation, r.cache = old_intp, old_cache
+        INTERPRETATION.reset(token)
+
+
+@contextlib.contextmanager
+def cache(store: EvalCache | None = None):
+    """Memoize evaluation under any interpretation for the duration of this block.
+
+    Installs ``store``, or a fresh cache if none is given, and yields it so that a
+    later block can reuse it::
+
+        with cache() as store:
+            ...
+        with cache(store):
+            ...
+
+    :func:`effectful.ops.semantics.evaluate` installs one for the duration of a
+    call when none is active, so a lone call is memoized internally whether or not
+    a scope is open. Holding a scope is what shares that work *between* calls.
+    """
+    store = AutoIdKeyDictionary() if store is None else store
+    token = EVAL_CACHE.set(store)
+    try:
+        yield store
+    finally:
+        EVAL_CACHE.reset(token)
+
+
+def cache_get(
+    store: EvalCache,
+    expr: typing.Any,
+    intp: "Interpretation",
+    default: typing.Any = None,
+) -> typing.Any:
+    """Look ``expr`` up under ``intp``, returning ``default`` if it is not cached."""
+    inner = store.get(expr)
+    return default if inner is None else inner.get(intp, default)
+
+
+def cache_put(
+    store: EvalCache, expr: typing.Any, intp: "Interpretation", value: typing.Any
+) -> None:
+    """Record that ``expr`` evaluates to ``value`` under ``intp``."""
+    inner = store.get(expr)
+    if inner is None:
+        # Same flavour as the outer store, so the inner map accepts the plain
+        # ``dict`` interpretations that ``coproduct`` builds.
+        inner = type(store)()
+        store[expr] = inner
+    inner[intp] = value
+
+
+def copy_cache_entries(src, dst) -> None:
+    """Copy everything cached for ``src`` onto ``dst``.
+
+    :func:`effectful.ops.syntax._build_term` computes a node's type analysis on a
+    throwaway term and then needs it attributed to the term it actually returns.
+    """
+    store = EVAL_CACHE.get()
+    if store is None:
+        return
+    inner = store.get(src)
+    if inner is not None:
+        for intp, value in inner.items():
+            cache_put(store, dst, intp, value)
 
 
 @Operation.define
 def _get_args() -> tuple[tuple, Mapping]:
     return ((), {})
 
 
+@weak_memoize(cache=WeakIdKeyDictionary())
 def _restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
     sig = inspect.signature(fn)
     if not sig.parameters:
@@ -57,6 +119,7 @@ def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
     return _cont_wrapper
 
 
+@weak_memoize(cache=WeakIdKeyDictionary())
 def _save_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
     from effectful.ops.semantics import handler
 
@@ -72,14 +135,32 @@ def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
     return _cont_wrapper
 
 
+@weak_memoize(cache=WeakIdKeyDictionary())
+def _save_then_restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
+    # should be equivalent to _restore_args(_save_args(fn)), just fused
+    from effectful.ops.semantics import handler
+
+    sig = inspect.signature(fn)
+    if not sig.parameters:
+        return fn
+
+    @functools.wraps(fn)
+    def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
+        a, k = (a, k) if a or k else _get_args()
+        with handler({_get_args: lambda: (a, k)}):
+            return fn(*a, **k)
+
+    return _cont_wrapper
+
+
 def _set_prompt[**P, T](
     prompt: Operation[P, T], cont: Callable[P, T], body: Callable[P, T]
 ) -> Callable[P, T]:
     from effectful.ops.semantics import handler
 
     @functools.wraps(body)
     def bound_body(*a: P.args, **k: P.kwargs) -> T:
-        next_cont = get_interpretation().get(prompt, prompt.__default_rule__)
+        next_cont = INTERPRETATION.get().get(prompt, prompt.__default_rule__)
         with handler({prompt: handler({prompt: next_cont})(cont)}):
             return body(*a, **k)