diff --git a/effectful/internals/runtime.py b/effectful/internals/runtime.py index 4c9ebd7b1..26deb1f4d 100644 --- a/effectful/internals/runtime.py +++ b/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 diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index de041b61f..43a3135ac 100644 --- a/effectful/ops/semantics.py +++ b/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]: """ - 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]: diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 764016752..fefb01768 100644 --- a/effectful/ops/syntax.py +++ b/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] diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 75d040162..2875387d1 100644 --- a/tests/test_ops_semantics.py +++ b/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)