diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index de041b61..eb1f1e1d 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -6,6 +6,7 @@ import types import typing from collections.abc import Callable +from dataclasses import dataclass, field from typing import Any from effectful.ops.syntax import _CustomSingleDispatchCallable, defdata, defop @@ -36,6 +37,22 @@ def fwd(*args, **kwargs) -> Any: raise RuntimeError("fwd should only be called in the context of a handler") +@dataclass +class Fwd: + next: Callable + args: tuple = () + kwargs: collections.abc.Mapping = field(default_factory=dict) + + def __init__(self, *args, **kwargs): + from effectful.internals.runtime import _get_args, get_interpretation + + self.args, self.kwargs = (args, kwargs) if args or kwargs else _get_args() + + intp = get_interpretation() + assert fwd in intp + self.next = handler({_get_args: lambda: (self.args, self.kwargs)})(intp[fwd]) + + def coproduct(intp: Interpretation, intp2: Interpretation) -> Interpretation: """The coproduct of two interpretations handles any effect that is handled by either. If both interpretations handle an effect, ``intp2`` takes diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 46419d7a..ee57e7f4 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -498,7 +498,7 @@ def _instance_op(instance, *args, **kwargs): def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: from effectful.internals.runtime import _restore_args, get_interpretation - from effectful.ops.semantics import fwd, handler + from effectful.ops.semantics import Fwd, fwd, handler intp = get_interpretation() @@ -510,7 +510,11 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: Interpretation, {fwd: _restore_args(self.__default_rule__)} ) with handler(fwd_intp): - return self_handler(*args, **kwargs) + result = self_handler(*args, **kwargs) + while isinstance(result, Fwd): # discharge fwd tail calls + result = result.next(*result.args, **result.kwargs) + return result + elif args and isinstance(args[0], Operation) and self is args[0].__apply__: # Prevent infinite recursion when calling self.apply directly return self.__default__(*args, **kwargs) diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 75d04016..352b84ab 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -8,7 +8,15 @@ import pytest -from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof +from effectful.ops.semantics import ( + Fwd, + coproduct, + evaluate, + fvsof, + fwd, + handler, + typeof, +) from effectful.ops.syntax import ObjectInterpretation, Scoped, deffn, defop, implements from effectful.ops.types import Interpretation, NotHandled, Operation, Term @@ -851,3 +859,31 @@ def f(): with pytest.raises(RuntimeError): f() + + +def test_fwd_return(): + @Operation.define + def f(x): + return x + + def f_call(x): + return fwd(x + 1) + + call_intp = functools.reduce(coproduct, [{f: f_call} for _ in range(10000)]) + with pytest.raises(RuntimeError): + with handler(call_intp): + f(0) + + def f_restore(x): + return Fwd() + + restore_intp = functools.reduce(coproduct, [{f: f_restore} for _ in range(10000)]) + with handler(restore_intp): + assert f(999) == 999 + + def f_ret(x): + return Fwd(x + 1) + + ret_intp = functools.reduce(coproduct, [{f: f_ret} for _ in range(10000)]) + with handler(ret_intp): + assert f(0) == 10000