Forwarding proceeds through apply to reach default rule

effectful/ops/types.py

@@ -510,10 +510,15 @@ def _instance_op(instance, *args, **kwargs):
             return self
 
     @functools.cached_property
-    def _default_rule_with_args(self):
+    def _next_rule_with_args(self):
         from effectful.internals.runtime import _restore_args
 
-        return _restore_args(self.__default_rule__)
+        rule = (
+            self.__default_rule__
+            if isinstance(self, ApplyOperation)
+            else functools.partial(self.__apply__, self)
+        )
+        return _restore_args(rule)
 
     def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V:
         from effectful.internals.runtime import get_interpretation
@@ -523,22 +528,23 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V:
 
         self_handler = intp.get(self)
         if self_handler is not None:
-            # ensure that fwd is bound to the default rule. if this handler has
-            # a bound fwd, it will override this binding
-            fwd_intp = typing.cast(Interpretation, {fwd: self._default_rule_with_args})
+            # Operation handlers forward through apply before reaching the default.
+            # Apply handlers forward to their own generated/default implementation.
+            fwd_intp = typing.cast(Interpretation, {fwd: self._next_rule_with_args})
             with handler(fwd_intp):
                 return self_handler(*args, **kwargs)
-        elif args and isinstance(args[0], Operation) and self is args[0].__apply__:
-            # Prevent infinite recursion when calling self.apply directly
+        elif isinstance(self, ApplyOperation):
             return self.__default__(*args, **kwargs)
         else:
             return self.__apply__(self, *args, **kwargs)
 
-    def __init_subclass__(cls, **kwargs) -> None:
-        assert "__apply__" not in cls.__dict__ or cls is Operation, (
-            "Cannot manually override apply"
-        )
-        assert isinstance(cls.__apply__, Operation)
+    def __init_subclass__(cls, *, _generate_apply: bool = True, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        if not _generate_apply:
+            return
+
+        assert "__apply__" not in cls.__dict__, "Cannot manually override apply"
+        assert isinstance(cls.__apply__, ApplyOperation)
 
         cls.__apply__ = cls.__apply__.define(
             staticmethod(
@@ -552,6 +558,10 @@ def __init_subclass__(cls, **kwargs) -> None:
         )
 
 
+class ApplyOperation[**Q, V](Operation[Q, V], _generate_apply=False):
+    """An operation that implements application for an Operation subclass."""
+
+
 def __apply__[**A, B](op: Operation[A, B], *args: A.args, **kwargs: A.kwargs) -> B:
     """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
 
@@ -584,7 +594,7 @@ def __apply__[**A, B](op: Operation[A, B], *args: A.args, **kwargs: A.kwargs) ->
     return op.__default_rule__(*args, **kwargs)  # type: ignore[return-value]
 
 
-Operation.__apply__ = Operation.define(staticmethod(__apply__))
+Operation.__apply__ = ApplyOperation.define(staticmethod(__apply__))
 del __apply__
 
 

tests/test_ops_semantics.py

@@ -23,7 +23,13 @@
     defop,
     implements,
 )
-from effectful.ops.types import Interpretation, NotHandled, Operation, Term
+from effectful.ops.types import (
+    ApplyOperation,
+    Interpretation,
+    NotHandled,
+    Operation,
+    Term,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -319,6 +325,131 @@ def plus_1_fwd(x):
         assert plus_1(1) == 2
 
 
+def test_fwd_from_operation_handler_to_apply_handler():
+    @Operation.define
+    def f(x: int) -> int:
+        return x + 1
+
+    calls = []
+
+    def f_handler(x):
+        calls.append(("f", x))
+        return fwd()
+
+    def apply_handler(op, *args, **kwargs):
+        calls.append(("apply", op, args, kwargs))
+        return fwd()
+
+    with handler({apply: apply_handler}), handler({f: f_handler}):
+        assert f(1) == 2
+
+    assert calls == [("f", 1), ("apply", f, (1,), {})]
+
+
+def test_fwd_from_operation_handler_to_apply_handler_with_replacement_args():
+    @Operation.define
+    def f(x: int) -> int:
+        return x
+
+    apply_args = []
+
+    def apply_handler(op, *args, **kwargs):
+        apply_args.append((op, args, kwargs))
+        return fwd(op, *args, **kwargs)
+
+    with handler({apply: apply_handler}), handler({f: lambda x: fwd(x + 1)}):
+        assert f(1) == 2
+
+    assert apply_args == [(f, (2,), {})]
+
+
+def test_fwd_through_apply_handlers_is_associative():
+    calls = []
+
+    @Operation.define
+    def f() -> int:
+        calls.append("default")
+        return 1
+
+    def forwarding(name):
+        def impl(*args, **kwargs):
+            calls.append(name)
+            return fwd()
+
+        return impl
+
+    h0 = {apply: forwarding("left apply")}
+    h1 = {f: forwarding("exact")}
+    h2 = {apply: forwarding("right apply")}
+    expected = ["exact", "right apply", "left apply", "default"]
+
+    for intp in (
+        coproduct(coproduct(h0, h1), h2),
+        coproduct(h0, coproduct(h1, h2)),
+    ):
+        calls.clear()
+        with handler(intp):
+            assert f() == 1
+        assert calls == expected
+
+    calls.clear()
+    with handler(h0), handler(h1), handler(h2):
+        assert f() == 1
+    assert calls == expected
+
+
+def test_fwd_through_apply_operation_subtypes():
+    calls = []
+
+    class BaseOperation(Operation):
+        pass
+
+    class DerivedOperation(BaseOperation):
+        pass
+
+    @DerivedOperation.define
+    def f(x: int) -> int:
+        calls.append("default")
+        return x + 1
+
+    def forwarding(name):
+        def impl(*args, **kwargs):
+            calls.append(name)
+            return fwd()
+
+        return impl
+
+    assert isinstance(Operation.__apply__, ApplyOperation)
+    assert isinstance(BaseOperation.__apply__, ApplyOperation)
+    assert isinstance(DerivedOperation.__apply__, ApplyOperation)
+    assert not isinstance(f, ApplyOperation)
+
+    with handler(
+        {
+            Operation.__apply__: forwarding("apply"),
+            BaseOperation.__apply__: forwarding("base apply"),
+            DerivedOperation.__apply__: forwarding("derived apply"),
+            f: forwarding("exact"),
+        }
+    ):
+        assert f(1) == 2
+
+    assert calls == ["exact", "derived apply", "base apply", "apply", "default"]
+
+    # Unhandled intermediate apply operations proceed directly to their defaults,
+    # so the base apply handler sees the original operation exactly once.
+    calls.clear()
+    with handler(
+        {
+            Operation.__apply__: forwarding("apply"),
+            f: forwarding("exact"),
+        }
+    ):
+        assert f(1) == 2
+
+    assert calls == ["exact", "apply", "default"]
+
+
 @pytest.mark.parametrize("op,args", OPERATION_CASES)
 @pytest.mark.parametrize("n1", N_CASES)
 @pytest.mark.parametrize("n2", N_CASES)