diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 308cdb76e..9c933af43 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -1,4 +1,5 @@ import functools +import itertools import typing from collections.abc import Callable, Mapping, Sequence from types import EllipsisType @@ -10,10 +11,12 @@ except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") -from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, fvsof, typeof from effectful.ops.syntax import ( + ConstructorOperation, + PureInterpretation, Scoped, + _BaseTerm, _CustomSingleDispatchCallable, defdata, deffn, @@ -39,7 +42,71 @@ def is_eager_array(x): ) -def sizesof(value) -> Mapping[Operation[[], jax.Array], int]: +@functools.cache +def _sizesof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``sizesof``.""" + from effectful.internals.product_n import argsof, productN + + _sizes = defop(object, name="sizes") + _getitem_term = defop(object, name="getitem_args") + + def _retain(op, *args, **kwargs): + # Non-getitem subterms are opaque to this analysis. Keeping their + # arguments would retain the entire input term unnecessarily. + return _BaseTerm(op) + + def _retain_getitem(*args, **kwargs): + return defdata(jax_getitem, *args, **kwargs) + + def _merge(s1, s2): + s3 = s1.copy() + for k, v in s2.items(): + if k in s3 and s3[k] != v: + raise ValueError( + f"Named index {k} used in incompatible dimensions of size {s3[k]} and {v}" + ) + s3[k] = v + return s3 + + def _apply_sizes(op, *args, **kwargs): + analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict)) + return functools.reduce(_merge, analyses, {}) + + def _getitem(arr, index): + # Inspect this getitem's arguments in the term projection without + # forcing that projection to retain the getitem result. + term_args, _ = argsof(_getitem_term) + term_arr, term_index = term_args + + arg_sizes = (x for x in (arr, index) if isinstance(x, dict)) + if not is_eager_array(term_arr): + return functools.reduce(_merge, arg_sizes, {}) + + sizes = ( + {k.op: term_arr.shape[i]} + for i, k in enumerate(term_index) + if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0 + ) + return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) + + return ( + PureInterpretation( + productN( + { + _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, + _getitem_term: { + apply: _retain, + jax_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } + ) + ), + _sizes, + ) + + +def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: """Return the sizes of named dimensions in an array expression. Sizes are inferred from the array shape. @@ -53,30 +120,14 @@ def sizesof(value) -> Mapping[Operation[[], jax.Array], int]: >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()])) >>> assert sizes[a] == 2 and sizes[b] == 3 """ - sizes: dict[Operation[[], jax.Array], int] = {} - - def update_sizes(sizes, op, size): - old_size = sizes.get(op) - if old_size is not None and size != old_size: - raise ValueError( - f"Named index {op} used in incompatible dimensions of size {old_size} and {size}" - ) - sizes[op] = size - - def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]): - if is_eager_array(x): - for i, k in enumerate(key): - if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0: - update_sizes(sizes, k.op, x.shape[i]) - return defdata(jax_getitem, x, key) - - def _apply(op, *args, **kwargs): - return defdata(op, *args, **kwargs) - - with interpreter({jax_getitem: _getitem_sizeof, apply: _apply}): - evaluate(value) - - return sizes + from effectful.internals.product_n import _unpack + + intp, prompt = _sizesof_intp() + result = evaluate(term, intp=intp) + fvs = _unpack(result, prompt) + if not isinstance(fvs, dict): + return {} + return fvs def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index 990830d27..cc20d7498 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -1,3 +1,4 @@ +import types from typing import TYPE_CHECKING import jax.numpy @@ -7,17 +8,28 @@ _no_overload = ["array", "asarray"] for name, op in jax.numpy.__dict__.items(): - if not callable(op): + if isinstance(op, types.ModuleType): continue - jax_op = ( - _register_jax_op_no_partial_eval(op) - if name in _no_overload - else _register_jax_op(op) - ) - globals()[name] = jax_op + # copy constants + if isinstance(op, float | types.NoneType): + globals()[name] = op -pi = jax.numpy.pi + if callable(op): + if name == "__getattr__": + continue + + elif name in _no_overload: + globals()[name] = _register_jax_op_no_partial_eval(op) + + else: + globals()[name] = _register_jax_op(op) + jax_op = ( + _register_jax_op_no_partial_eval(op) + if name in _no_overload + else _register_jax_op(op) + ) + globals()[name] = jax_op # Tell mypy about our wrapped functions. if TYPE_CHECKING: diff --git a/effectful/handlers/pyro.py b/effectful/handlers/pyro.py index 2e7924210..791564afc 100644 --- a/effectful/handlers/pyro.py +++ b/effectful/handlers/pyro.py @@ -28,7 +28,7 @@ ) from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, handler, typeof -from effectful.ops.syntax import defdata, defop +from effectful.ops.syntax import ConstructorOperation, defdata, defop from effectful.ops.types import NotHandled, Operation, Term @@ -368,7 +368,9 @@ def _to_named(a): return a # Convert to a term in a context that does not evaluate distribution constructors. - with handler({apply: defdata}): + with handler( + {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + ): d = typing.cast(TorchDistribution, evaluate(value)) if not (isinstance(d, Term) and typeof(d) is TorchDistribution): @@ -403,7 +405,9 @@ def _to_positional(a, indices): else: return a - with handler({apply: defdata}): + with handler( + {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + ): d = typing.cast(TorchDistribution, evaluate(value)) if not (isinstance(d, Term) and typeof(d) is TorchDistribution): diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index c9b8f640b..62fa8742e 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -1,4 +1,5 @@ import functools +import itertools import typing from collections.abc import Callable, Mapping, Sequence from types import EllipsisType @@ -11,10 +12,17 @@ import torch.utils._pytree as pytree -from effectful.internals.runtime import interpreter from effectful.internals.tensor_utils import _desugar_tensor_index from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof -from effectful.ops.syntax import Scoped, defdata, defop, syntactic_eq +from effectful.ops.syntax import ( + ConstructorOperation, + PureInterpretation, + Scoped, + _BaseTerm, + defdata, + defop, + syntactic_eq, +) from effectful.ops.types import Expr, NotHandled, Operation, Term # + An element of a tensor index expression. @@ -34,6 +42,74 @@ def _getitem_ellipsis_and_none( return torch.reshape(x, new_shape), new_key +@functools.cache +def _sizesof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``sizesof``.""" + from effectful.internals.product_n import argsof, productN + + sizes = defop(object, name="sizes") + getitem_term = defop(object, name="getitem_args") + + def _retain(op, *args, **kwargs): + # Non-getitem subterms are opaque to this analysis. Keeping their + # arguments would retain the entire input term unnecessarily. + return _BaseTerm(op) + + def _retain_getitem(*args, **kwargs): + return defdata(torch_getitem, *args, **kwargs) + + def _merge(s1, s2): + result = s1.copy() + for k, v in s2.items(): + if k in result and result[k] != v: + raise ValueError( + f"Named index {k} used in incompatible dimensions of size {result[k]} and {v}" + ) + result[k] = v + return result + + def _apply_sizes(op, *args, **kwargs): + analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict)) + return functools.reduce(_merge, analyses, {}) + + def _getitem(x, key): + # Inspect this getitem's arguments in the term projection without + # forcing that projection to retain the getitem result. + term_args, _ = argsof(getitem_term) + term_x, term_key = term_args + + arg_sizes = (value for value in (x, key) if isinstance(value, dict)) + if not isinstance(term_x, torch.Tensor): + return functools.reduce(_merge, arg_sizes, {}) + + shape, desugared_key = _desugar_tensor_index(term_x.shape, term_key) + index_sizes = ( + {k.op: shape[i]} + for i, k in enumerate(desugared_key) + if isinstance(k, Term) + and not k.args + and not k.kwargs + and issubclass(typeof(k), torch.Tensor) + ) + return functools.reduce(_merge, itertools.chain(arg_sizes, index_sizes), {}) + + return ( + PureInterpretation( + productN( + { + sizes: {apply: _apply_sizes, torch_getitem: _getitem}, + getitem_term: { + apply: _retain, + torch_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } + ) + ), + sizes, + ) + + def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]: """Return the sizes of named dimensions in a tensor expression. @@ -48,35 +124,13 @@ def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]: >>> sizes = sizesof(torch.ones(2, 3)[a(), b()]) >>> assert sizes[a] == 2 and sizes[b] == 3 """ - sizes: dict[Operation[[], torch.Tensor], int] = {} - - def _torch_getitem_sizeof( - x: Expr[torch.Tensor], key: tuple[Expr[IndexElement], ...] - ) -> Expr[torch.Tensor]: - if isinstance(x, torch.Tensor): - shape, key_ = _desugar_tensor_index(x.shape, key) - - for i, k in enumerate(key_): - if ( - isinstance(k, Term) - and len(k.args) == 0 - and len(k.kwargs) == 0 - and issubclass(typeof(k), torch.Tensor) - ): - if k.op in sizes and sizes[k.op] != shape[i]: - raise ValueError( - f"Named index {k.op} used in incompatible dimensions of size {sizes[k.op]} and {shape[i]}" - ) - sizes[k.op] = shape[i] - - return defdata(torch_getitem, x, key) - - def _apply(op, *args, **kwargs): - return defdata(op, *args, **kwargs) - - with interpreter({torch_getitem: _torch_getitem_sizeof, apply: _apply}): - evaluate(value) + from effectful.internals.product_n import _unpack + intp, prompt = _sizesof_intp() + result = evaluate(value, intp=intp) + sizes = _unpack(result, prompt) + if not isinstance(sizes, dict): + return {} return sizes diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 43a3135ac..e721c87a1 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -10,9 +10,12 @@ from typing import Any from effectful.ops.syntax import ( + ConstructorOperation, + DataclassConstructorOperation, + ObjectInterpretation, PureInterpretation, + _BaseTerm, _CustomSingleDispatchCallable, - defdata, defop, implements, ) @@ -122,6 +125,11 @@ def handler(intp: Interpretation): yield intp +@ConstructorOperation.define +def as_tuple(*args) -> tuple: + return tuple(args) + + @_CustomSingleDispatchCallable def evaluate[T]( __dispatch: Callable[[type], Callable[..., Expr[T]]], @@ -160,19 +168,21 @@ def evaluate[T]( @evaluate.register(bytes) def _evaluate_object[T](expr: T, **kwargs) -> T: if dataclasses.is_dataclass(expr) and not isinstance(expr, type): - return typing.cast( - T, - dataclasses.replace( - expr, - **{ - field.name: evaluate(getattr(expr, field.name)) - for field in dataclasses.fields(expr) - }, - ), - ) + return _evaluate_dataclass(expr, **kwargs) return expr +def _evaluate_dataclass[T](expr: T, **kwargs) -> T: + subst = { + field.name: evaluate(getattr(expr, field.name)) + for field in dataclasses.fields(expr) # type: ignore[arg-type] + } + return typing.cast( + T, + DataclassConstructorOperation.define(type(expr))(**subst), # type: ignore[arg-type] + ) + + _EVALUATION_CACHE_ATTR = "__effectful_evaluation_cache__" @@ -223,17 +233,24 @@ def _evaluate_operation(expr: Operation, **kwargs) -> Operation: @evaluate.register(collections.defaultdict) def _evaluate_defaultdict(expr, **kwargs): - return type(expr)(expr.default_factory, evaluate(tuple(expr.items()))) + return ConstructorOperation.define(type(expr))( + expr.default_factory, + as_tuple(*(evaluate(item) for item in expr.items())), + ) @evaluate.register(types.MappingProxyType) def _evaluate_mappingproxytype(expr, **kwargs): - return type(expr)(dict(evaluate(tuple(expr.items())))) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr.items())) + ) @evaluate.register(collections.abc.Mapping) def _evaluate_mapping(expr, **kwargs): - return type(expr)(evaluate(tuple(expr.items()))) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr.items())) + ) @evaluate.register(tuple) @@ -243,27 +260,35 @@ def _evaluate_tuple(expr, **kwargs): and hasattr(expr, "_fields") and all(hasattr(expr, field) for field in getattr(expr, "_fields")) ): # namedtuple - return type(expr)( + return ConstructorOperation.define(type(expr))( **{field: evaluate(getattr(expr, field)) for field in expr._fields} ) else: - return type(expr)(evaluate(item) for item in expr) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.Sequence) def _evaluate_sequence(expr, **kwargs): - return type(expr)(evaluate(item) for item in expr) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.ItemsView) @evaluate.register(collections.abc.KeysView) def _evaluate_set_view(expr, **kwargs): - return {evaluate(item) for item in expr} + return ConstructorOperation.define(set)( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.ValuesView) def _evaluate_list_view(expr, **kwargs): - return [evaluate(item) for item in expr] + return ConstructorOperation.define(list)( + as_tuple(*(evaluate(item) for item in expr)) + ) def _simple_type(tp: type) -> type: @@ -288,15 +313,19 @@ def _simple_type(tp: type) -> type: return typing.get_origin(tp) or tp -class _TypeofIntp(PureInterpretation): +class _TypeofIntp(ObjectInterpretation): @implements(apply) def _(self, op, *args, **kwargs): from effectful.internals.unification import Box return Box(op.__type_rule__(*args, **kwargs)) + @implements(ConstructorOperation.__apply__) + def _constructor(self, op, *args, **kwargs): + return op.__default_rule__(*args, **kwargs) + -_TYPEOF_INTP = _TypeofIntp() +_TYPEOF_INTP = PureInterpretation(_TypeofIntp()) def _typeof(term: Expr): @@ -334,32 +363,111 @@ def typeof[T](term: Expr[T]) -> type[T]: return typing.cast(type[T], type(type_or_value)) +@functools.cache +def _fvsof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``fvsof``.""" + from effectful.internals.product_n import argsof, productN + + def _apply_collection_binders(op, *args, **kwargs): + return frozenset().union( + *( + {x} + if isinstance(x, Operation) + else x + if isinstance(x, frozenset) + else set() + for x in (*args, *kwargs.values()) + ) + ) + + def _apply_binders(op, *args, **kwargs): + # Parent operations only need to know that this child is a term. Its + # arguments are available through argsof while this node is analyzed. + return _BaseTerm(op) + + def _apply_passthrough_fvs(op, *args, **kwargs): + return frozenset().union( + *(x for x in (*args, *kwargs.values()) if isinstance(x, frozenset)) + ) + + def _apply_fvs(op, *args, **kwargs): + binder_args, binder_kwargs = argsof(_fvsof_binders) + # This rule handles Operation.__apply__ directly, so its first argument + # is the operation being applied rather than an argument to that + # operation. + binder_args = tuple( + frozenset() if isinstance(x, Term) else x for x in binder_args[1:] + ) + binder_kwargs = { + k: frozenset() if isinstance(v, Term) else v + for k, v in binder_kwargs.items() + } + bindings = op.__fvs_rule__(*binder_args, **binder_kwargs) + binders = frozenset().union(*(*bindings.args, *bindings.kwargs.values())) + + fvs = frozenset().union( + {op}, + *( + x if isinstance(x, frozenset) else frozenset() + for x in (*args, *kwargs.values()) + ), + ) + fvs -= binders + return fvs + + _fvsof_fvs = defop(object, name="fvsof_fvs") + _fvsof_binders = defop(object, name="fvsof_binders") + + return ( + PureInterpretation( + productN( + { + _fvsof_fvs: { + apply: _apply_fvs, + ConstructorOperation.__apply__: _apply_passthrough_fvs, + }, + _fvsof_binders: { + apply: _apply_binders, + ConstructorOperation.__apply__: _apply_collection_binders, + }, + } + ) + ), + _fvsof_fvs, + ) + + def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: - """Return the free variables of an expression. + """Return the free operations in a term. + + An operation belongs to `fvsof(t)` when it appears free in the term `t`. + This excludes operations like `apply` or collection constructors that are + raised during `evaluate` but do not appear in `t`. It also excludes + operations that are bound by a `Scoped` operation. However, it is not + restricted to the nullary operations in `t`. **Example usage**: + `fvsof` includes all unbound operations in a term: + + >>> a = defop(int) >>> @defop ... def f(x: int, y: int) -> int: ... raise NotHandled - >>> fvs = fvsof(f(1, 2)) - >>> assert f in fvs - >>> assert len(fvs) == 1 - """ - from effectful.internals.runtime import interpreter + >>> fvs = fvsof(f(a(), 1)) + >>> assert fvs >= {f, a} - _fvs: set[Operation] = set() + `fvsof` accepts the same values as `evaluate`, including collections: - def _update_fvs(op, *args, **kwargs): - _fvs.add(op) - bindings = op.__fvs_rule__(*args, **kwargs) - for bound_var in set().union(*(*bindings.args, *bindings.kwargs.values())): - assert isinstance(bound_var, Operation) - if bound_var in _fvs: - _fvs.remove(bound_var) - return defdata(op, *args, **kwargs) + >>> fvs = fvsof([a(), {'k': f(0, 1)}]) + >>> assert fvs >= {f, a} - with interpreter({apply: _update_fvs}): - evaluate(term) - - return _fvs + """ + from effectful.internals.product_n import _unpack + + intp, prompt = _fvsof_intp() + result = evaluate(term, intp=intp) + fvs = _unpack(result, prompt) + if not isinstance(fvs, frozenset): + return frozenset() + return fvs diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 212ccc34b..65b70be54 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -11,6 +11,7 @@ from effectful.ops.types import ( Annotation, Expr, + Interpretation, NotHandled, Operation, Term, @@ -508,7 +509,10 @@ 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. - with interpreter({apply: defdata} | renaming_ctx): + with interpreter( + {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + | renaming_ctx + ): return evaluate(expr) renamed_args = op.__signature__.bind(*args, **kwargs) @@ -949,14 +953,23 @@ 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())) +class PureInterpretation[T, V](Mapping[Operation[..., T], Callable[..., V]]): + """Mark an interpretation as pure so its evaluation results can be cached.""" + + def __init__(self, intp: Interpretation[T, V]): + self.intp = intp + + def __iter__(self): + return iter(self.intp) + + def __len__(self): + return len(self.intp) + + def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]: + return self.intp[item] - def __eq__(self, other): - return isinstance(other, PureInterpretation) and frozenset( - self.implementations.items() - ) == frozenset(other.implementations.items()) + __hash__ = object.__hash__ + __eq__ = object.__eq__ class _ImplementedOperation[**P, **Q, T, V]: @@ -1357,3 +1370,23 @@ class _IntegralTerm[T: numbers.Integral](_RationalTerm[T]): @defdata.register(bool) class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass + + +class ConstructorOperation[**Q, V](Operation[Q, V]): + @classmethod + @functools.cache + def define[T]( + cls, constructor: type[T] | Callable[..., T] + ) -> "ConstructorOperation[Any, T]": + if not isinstance(constructor, type): + return typing.cast( + ConstructorOperation[Any, T], super().define(constructor) + ) + + def _as_typ(*args, **kwargs) -> constructor: # type: ignore[valid-type] + return constructor(*args, **kwargs) + + return typing.cast(ConstructorOperation[Any, T], super().define(_as_typ)) + + +class DataclassConstructorOperation[**Q, V](ConstructorOperation): ... diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 2875387d1..bd6c5d0de 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -481,7 +481,7 @@ def node(x: object) -> object: term = node(node(1)) - class Intp(PureInterpretation): + class Intp(ObjectInterpretation): def __init__(self): self.calls = 0 @@ -490,32 +490,34 @@ def _(self, op, *args, **kwargs): self.calls += 1 return (op.__name__, args, kwargs) - intp = Intp() + intp_impl = Intp() + intp = PureInterpretation(intp_impl) expected = ("node", (("node", (1,), {}),), {}) assert interpreter(intp)(evaluate)(term) == expected - assert intp.calls == 2 + assert intp_impl.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 + assert intp_impl.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 + assert intp_impl.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 + assert intp_impl.calls == 4 - # An identical memoized interpretation is in the same cache namespace. - other_intp = Intp() + # A separate pure interpretation has its own cache namespace. + other_intp_impl = Intp() + other_intp = PureInterpretation(other_intp_impl) assert interpreter(other_intp)(evaluate)(term) == expected - assert intp.calls == 4 + assert other_intp_impl.calls == 2 def test_memoized_interpretation_does_not_cache_failures(): @@ -527,7 +529,7 @@ def node() -> object: term = node() - class Intp(PureInterpretation): + class Intp(ObjectInterpretation): def __init__(self): self.calls = 0 @@ -538,28 +540,38 @@ def _(self, op, *args, **kwargs): raise ValueError("failed analysis") return "success" - intp = Intp() + intp_impl = Intp() + intp = PureInterpretation(intp_impl) with pytest.raises(ValueError, match="failed analysis"): interpreter(intp)(evaluate)(term) assert interpreter(intp)(evaluate)(term) == "success" - assert intp.calls == 2 + assert intp_impl.calls == 2 assert interpreter(intp)(evaluate)(term) == "success" - assert intp.calls == 2 + assert intp_impl.calls == 2 -def test_ctxof(): - x = defop(object) - y = defop(object) +@pytest.mark.parametrize( + "build_args", + [ + lambda x, y: (x(), y()), + lambda x, y: ([x()], y()), + lambda x, y: ([x()], [y()]), + lambda x, y: (([x()], [y()]),), + ], +) +def test_ctxof(build_args): + x = defop(object, name="x") + y = defop(object, name="y") @defop def Nested(*args, **kwargs): raise NotHandled - assert fvsof(Nested(x(), y())) >= {x, y} - assert fvsof(Nested([x()], y())) >= {x, y} - assert fvsof(Nested([x()], [y()])) >= {x, y} - assert fvsof(Nested((x(), y()))) >= {x, y} + term = Nested(*build_args(x, y)) + actual = fvsof(term) + expected = {x, y, Nested} + assert actual >= expected def test_handler_typing() -> None: @@ -813,8 +825,15 @@ def Lam2[A, B]( raise NotHandled term = Lam2(add(x(), add(y(), z())), x, y) - assert not {x, y} <= fvsof(term) - assert fvsof(term) == {z, Lam2, add} + actual = fvsof(term) + assert not ({x, y} & actual) + assert actual >= {z, Lam2, add} + + +def test_fvsof_collection_does_not_include_apply(): + x = defop(int, name="x") + + assert fvsof((x(),)) == {x} def test_interpretation_typing(): @@ -875,7 +894,29 @@ def __init__(self, x: int): self.x = x v = Operation.define(int) - assert fvsof(A(v())) == {v} + actual = fvsof(A(v())) + assert actual == {v} + + +def test_defdata_dataclass_init_effects() -> None: + @Operation.define + def f(x: int): + raise NotHandled + + @dataclasses.dataclass + class A: + x: int + + def __init__(self, x: int): + self.x = f(x) + + @Operation.define + def g(a: A): + raise NotHandled + + v = Operation.define(int) + t = g(A(v())) + assert isinstance(t.args[0].x, Term) def test_instanceop_super() -> None: @@ -909,6 +950,31 @@ def f(self): assert b.f() == "*B*" +def test_instanceop_dataclass() -> None: + """Dataclasses with no free variables get instance operations.""" + + @dataclasses.dataclass + class A: + @Operation.define + def f(self): + raise NotHandled + + assert isinstance(A.f, Operation) + assert isinstance(A().f, Operation) + + @dataclasses.dataclass + class B: + x: int + + @Operation.define + def g(self): + raise NotHandled + + assert isinstance(B.g, Operation) + fv = Operation.define(int)() + assert not isinstance(B(fv).g, Operation) + + def test_coproduct_fwd_chain(benchmark): """Benchmark coproduct + fwd over a deep chain of forwarding handlers.