Add cached interpretations and use for typeof

effectful/internals/runtime.py

@@ -27,7 +27,7 @@ def interpreter(intp: "Interpretation"):
     r = get_runtime()
     old_intp = r.interpretation
     try:
-        old_intp, r.interpretation = r.interpretation, dict(intp)
+        old_intp, r.interpretation = r.interpretation, intp
         yield intp
     finally:
         r.interpretation = old_intp

effectful/ops/semantics.py

@@ -5,10 +5,17 @@
 import operator
 import types
 import typing
+import weakref
 from collections.abc import Callable
 from typing import Any
 
-from effectful.ops.syntax import _CustomSingleDispatchCallable, defdata, defop
+from effectful.ops.syntax import (
+    PureInterpretation,
+    _CustomSingleDispatchCallable,
+    defdata,
+    defop,
+    implements,
+)
 from effectful.ops.types import (
     Expr,
     Interpretation,
@@ -166,11 +173,44 @@ def _evaluate_object[T](expr: T, **kwargs) -> T:
     return expr
 
 
+_EVALUATION_CACHE_ATTR = "__effectful_evaluation_cache__"
+
+
+def _term_cache(
+    expr: Term,
+) -> weakref.WeakKeyDictionary[PureInterpretation, Any] | None:
+    """Return the cache owned by ``expr``, or ``None`` if it cannot store one."""
+    try:
+        return getattr(expr, _EVALUATION_CACHE_ATTR)
+    except AttributeError:
+        cache: weakref.WeakKeyDictionary[PureInterpretation, Any] = (
+            weakref.WeakKeyDictionary()
+        )
+        try:
+            setattr(expr, _EVALUATION_CACHE_ATTR, cache)
+        except (AttributeError, TypeError):
+            return None
+        return cache
+
+
 @evaluate.register(Term)
 def _evaluate_term(expr: Term, **kwargs):
+    from effectful.internals.runtime import get_interpretation
+
+    current_intp = get_interpretation()
+    if isinstance(current_intp, PureInterpretation):
+        cache = _term_cache(expr)
+        if cache is not None and current_intp in cache:
+            return cache[current_intp]
+    else:
+        cache = None
+
     args = tuple(evaluate(arg) for arg in expr.args)
     kwargs = {k: evaluate(v) for k, v in expr.kwargs.items()}
-    return expr.op(*args, **kwargs)
+    result = expr.op(*args, **kwargs)
+    if cache is not None:
+        cache[current_intp] = result
+    return result
 
 
 @evaluate.register(Operation)
@@ -248,6 +288,22 @@ def _simple_type(tp: type) -> type:
     return typing.get_origin(tp) or tp
 
 
+class _TypeofIntp(PureInterpretation):
+    @implements(apply)
+    def _(self, op, *args, **kwargs):
+        from effectful.internals.unification import Box
+
+        return Box(op.__type_rule__(*args, **kwargs))
+
+
+_TYPEOF_INTP = _TypeofIntp()
+
+
+def _typeof(term: Expr):
+    """Evaluate the cached type analysis without unwrapping its result."""
+    return evaluate(term, intp=_TYPEOF_INTP)
+
+
 def typeof[T](term: Expr[T]) -> type[T]:
     """Return the type of an expression.
 
@@ -270,17 +326,12 @@ def typeof[T](term: Expr[T]) -> type[T]:
     <class 'int'>
 
     """
-    from effectful.internals.runtime import interpreter
     from effectful.internals.unification import Box
 
-    def _apply(op, *args, **kwargs):
-        return Box(op.__type_rule__(*args, **kwargs))
-
-    with interpreter({apply: _apply}):
-        type_or_value = evaluate(term)
-        if isinstance(type_or_value, Box):
-            return _simple_type(type_or_value.value)
-        return typing.cast(type[T], type(type_or_value))
+    type_or_value = _typeof(term)
+    if isinstance(type_or_value, Box):
+        return _simple_type(type_or_value.value)
+    return typing.cast(type[T], type(type_or_value))
 
 
 def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]:

effectful/ops/syntax.py

@@ -484,9 +484,8 @@ def __call__(self: collections.abc.Callable[P, T], *args: P.args, **kwargs: P.kw
     When an Operation whose return type is `Callable` is passed to :func:`defdata`,
     it is reconstructed as a :class:`_CallableTerm`, which implements the :func:`__call__` method.
     """
-    from effectful.internals.product_n import _unpack, productN
     from effectful.internals.runtime import interpreter
-    from effectful.ops.semantics import _simple_type, apply, evaluate
+    from effectful.ops.semantics import _simple_type, _typeof, apply, evaluate
 
     # If this operation binds variables, we need to rename them in the
     # appropriate parts of the child term.
@@ -497,38 +496,20 @@ def __call__(self: collections.abc.Callable[P, T], *args: P.args, **kwargs: P.kw
         for var in bound_vars
     }
 
-    # Analysis for type computation and term reconstruction
-    typ = defop(object, name="typ")
-    cast = defop(object, name="cast")
-
-    def apply_type(op, *args, **kwargs):
-        from effectful.internals.unification import Box
-
-        assert isinstance(op, Operation)
-        tp = op.__type_rule__(*args, **kwargs)
-        return Box(tp)
-
-    def apply_cast(op, *args, **kwargs):
-        assert isinstance(op, Operation)
-        full_type = typ()
-        dispatch_type = _simple_type(full_type.value)
-        return __dispatch(dispatch_type)(dispatch_type, op, *args, **kwargs)
-
-    analysis = productN({typ: {apply: apply_type}, cast: {apply: apply_cast}})
-
     def evaluate_with_renaming(expr, ctx):
         """Evaluate an expression with renaming applied."""
         renaming_ctx = {
             old_var: new_var for old_var, new_var in renaming.items() if old_var in ctx
         }
 
+        if not renaming_ctx:
+            return expr
+
         # Note: coproduct cannot be used to compose these interpretations
         # because evaluate will only do operation replacement when the handler
         # is operation typed, which coproduct does not satisfy.
-        with interpreter(analysis | renaming_ctx):
-            result = evaluate(expr)
-
-        return result
+        with interpreter({apply: defdata} | renaming_ctx):
+            return evaluate(expr)
 
     renamed_args = op.__signature__.bind(*args, **kwargs)
     renamed_args.apply_defaults()
@@ -542,11 +523,12 @@ def evaluate_with_renaming(expr, ctx):
         for (k, v) in renamed_args.kwargs.items()
     }
 
-    # Build the final term with type analysis
-    with interpreter(analysis):
-        result = op(*args_, **kwargs_)
-
-    return _unpack(result, cast)
+    # Build the final term using the cached type analysis of its children.
+    typed_args = tuple(_typeof(arg) for arg in args_)
+    typed_kwargs = {k: _typeof(v) for k, v in kwargs_.items()}
+    full_type = op.__type_rule__(*typed_args, **typed_kwargs)
+    dispatch_type = _simple_type(full_type)
+    return __dispatch(dispatch_type)(dispatch_type, op, *args_, **kwargs_)
 
 
 def _construct_dataclass_term[T](
@@ -929,6 +911,16 @@ def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]:
         return self.implementations[item].__get__(self, type(self))
 
 
+class PureInterpretation[T, V](ObjectInterpretation[T, V]):
+    def __hash__(self):
+        return hash(frozenset(self.implementations.items()))
+
+    def __eq__(self, other):
+        return isinstance(other, PureInterpretation) and frozenset(
+            self.implementations.items()
+        ) == frozenset(other.implementations.items())
+
+
 class _ImplementedOperation[**P, **Q, T, V]:
     impl: Callable[Q, V] | None
     op: Operation[P, T]

tests/test_ops_semantics.py

@@ -8,8 +8,23 @@
 
 import pytest
 
-from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof
-from effectful.ops.syntax import ObjectInterpretation, Scoped, deffn, defop, implements
+from effectful.ops.semantics import (
+    apply,
+    coproduct,
+    evaluate,
+    fvsof,
+    fwd,
+    handler,
+    typeof,
+)
+from effectful.ops.syntax import (
+    ObjectInterpretation,
+    PureInterpretation,
+    Scoped,
+    deffn,
+    defop,
+    implements,
+)
 from effectful.ops.types import Interpretation, NotHandled, Operation, Term
 
 logger = logging.getLogger(__name__)
@@ -457,6 +472,82 @@ def Nested(*args, **kwargs):
         assert evaluate(t) == Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1})
 
 
+def test_memoized_interpretation():
+    from effectful.internals.runtime import interpreter
+
+    @defop
+    def node(x: object) -> object:
+        raise NotHandled
+
+    term = node(node(1))
+
+    class Intp(PureInterpretation):
+        def __init__(self):
+            self.calls = 0
+
+        @implements(apply)
+        def _(self, op, *args, **kwargs):
+            self.calls += 1
+            return (op.__name__, args, kwargs)
+
+    intp = Intp()
+    expected = ("node", (("node", (1,), {}),), {})
+
+    assert interpreter(intp)(evaluate)(term) == expected
+    assert intp.calls == 2
+
+    # The root cache is checked before its children are traversed, including
+    # when evaluation is expressed directly through a handler.
+    with interpreter(intp):
+        assert evaluate(term) == expected
+    assert intp.calls == 2
+
+    # Child results are cached independently and can be reused directly.
+    assert interpreter(intp)(evaluate)(term.args[0]) == expected[1][0]
+    assert intp.calls == 2
+
+    # A composition has a distinct identity even when its added handler is not
+    # used while evaluating this term.
+    combined_intp = coproduct(intp, {plus_1: lambda x: x})
+    assert interpreter(combined_intp)(evaluate)(term) == expected
+    assert intp.calls == 4
+
+    # An identical memoized interpretation is in the same cache namespace.
+    other_intp = Intp()
+    assert interpreter(other_intp)(evaluate)(term) == expected
+    assert intp.calls == 4
+
+
+def test_memoized_interpretation_does_not_cache_failures():
+    from effectful.internals.runtime import interpreter
+
+    @defop
+    def node() -> object:
+        raise NotHandled
+
+    term = node()
+
+    class Intp(PureInterpretation):
+        def __init__(self):
+            self.calls = 0
+
+        @implements(apply)
+        def _(self, op, *args, **kwargs):
+            self.calls += 1
+            if self.calls == 1:
+                raise ValueError("failed analysis")
+            return "success"
+
+    intp = Intp()
+    with pytest.raises(ValueError, match="failed analysis"):
+        interpreter(intp)(evaluate)(term)
+
+    assert interpreter(intp)(evaluate)(term) == "success"
+    assert intp.calls == 2
+    assert interpreter(intp)(evaluate)(term) == "success"
+    assert intp.calls == 2
+
+
 def test_ctxof():
     x = defop(object)
     y = defop(object)