Add tail call fwd interface

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

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)

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