diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 8bf231573..d73546028 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -433,6 +433,26 @@ def deffn[T, A, B]( raise NotHandled +def _build_term[T]( + __dispatch: Callable[[type], Callable[..., Expr[T]]], + op: Operation[..., T], + *args, + **kwargs, +) -> Expr[T]: + """Build a single node from arguments whose bound variables are already renamed. + + This is :func:`defdata` without the renaming step: it computes the node's + type from the types of its arguments and dispatches on that type to pick a + constructor. + """ + from effectful.ops.semantics import _simple_type, _typeof + + typed_args = tuple(_typeof(arg) for arg in args) + typed_kwargs = {k: _typeof(v) for k, v in kwargs.items()} + dispatch_type = _simple_type(op.__type_rule__(*typed_args, **typed_kwargs)) + return __dispatch(dispatch_type)(dispatch_type, op, *args, **kwargs) + + @_CustomSingleDispatchCallable def defdata[T]( __dispatch: Callable[[type], Callable[..., Expr[T]]], @@ -492,7 +512,7 @@ def __call__(self: collections.abc.Callable[P, T], *args: P.args, **kwargs: P.kw it is reconstructed as a :class:`_CallableTerm`, which implements the :func:`__call__` method. """ from effectful.internals.runtime import interpreter - from effectful.ops.semantics import _simple_type, _typeof, apply, evaluate + from effectful.ops.semantics import apply, evaluate # If this operation binds variables, we need to rename them in the # appropriate parts of the child term. @@ -515,8 +535,14 @@ def evaluate_with_renaming(expr, ctx): # 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. + # + # Rebuild with ``_build_term`` rather than ``defdata``: the subterm was + # already renamed when it was first built, so re-entering ``defdata`` + # here would recompute binders and rename each child again at every + # level, re-traversing the subtree once per level of nesting. + rebuild = functools.partial(_build_term, __dispatch) with interpreter( - {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + {apply: rebuild, ConstructorOperation.__apply__: apply.__default_rule__} | renaming_ctx ): return evaluate(expr) @@ -533,12 +559,7 @@ def evaluate_with_renaming(expr, ctx): for (k, v) in renamed_args.kwargs.items() } - # 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_) + return _build_term(__dispatch, op, *args_, **kwargs_) def _construct_dataclass_term[T]( diff --git a/tests/test_ops_syntax.py b/tests/test_ops_syntax.py index 413f701e3..a5fdb749c 100644 --- a/tests/test_ops_syntax.py +++ b/tests/test_ops_syntax.py @@ -1457,3 +1457,45 @@ def _make_benchmark_term(size: int) -> Term[int]: result = benchmark(_make_benchmark_term, 25) assert isinstance(result, Term) + + +def test_bench_nested_binder_construction(benchmark): + """Benchmark term construction under *nested* binders. + + Constructing an operation that binds a variable renames that variable + throughout the body, which rebuilds the body. Nesting binders means each + level rebuilds everything beneath it, so a rebuild that is not single-pass + compounds multiplicatively with depth rather than adding to it. + + ``test_bench_term_construction`` builds a binder-free term, so it never + reaches this path -- it is fast even when nested construction is + exponential in depth. + """ + + @defop + def _benchmark_let[S, T, A]( + var: Annotated[Operation[[], S], Scoped[A]], + val: S, + body: Annotated[T, Scoped[A]], + ) -> T: + raise NotHandled + + @defop + def _benchmark_add(x: int, y: int) -> int: + raise NotHandled + + def _make_nested_term(depth: int) -> Term[int]: + """A term of ``depth`` nested binders, each used in the body below it.""" + if depth < 1: + raise ValueError("depth must be positive") + + body: Expr[int] = 0 + for index in range(depth): + var = defop(int, name=f"_benchmark_var_{index}") + body = _benchmark_let(var, 1, _benchmark_add(var(), body)) + + assert isinstance(body, Term) + return body + + result = benchmark(_make_nested_term, 10) + assert isinstance(result, Term)