Acl/order effects

guppylang-internals/src/guppylang_internals/compiler/builder.py → guppylang-internals/src/guppylang_internals/compiler/builder/__init__.py

@@ -1,24 +1,33 @@
 from abc import ABC, abstractmethod, abstractproperty
-from collections.abc import Iterator, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from types import TracebackType
-from typing import Generic, TypeVar
+from typing import Generic, NamedTuple, TypeAlias, TypeVar
 
-from hugr import Node, Wire, ops, val
+from hugr import Node, Wire, val
 from hugr import tys as ht
 from hugr.build import Block, Case, Cfg, Conditional, TailLoop
 from hugr.build import function as hf
 from hugr.hugr.node_port import ToNode
 from hugr.metadata import HugrDebugInfo
+from hugr.ops import DataflowOp, Output
 from typing_extensions import Self, override
 
 from guppylang_internals.ast_util import AstNode
-from guppylang_internals.compiler.core import may_have_side_effect
 from guppylang_internals.metadata.debug_info_util import (
     debug_conditions_fulfilled,
     make_location_record,
 )
+from guppylang_internals.tys import Effect
+
+
+@dataclass(frozen=True)
+class Pure:
+    op: DataflowOp
+
+
+OpWithEffects: TypeAlias = Pure | tuple[DataflowOp, Sequence[Effect]]
 
 
 @dataclass
@@ -33,7 +42,7 @@ class DFBuilder(ABC, ToNode):
     """
 
     current_ast_node: AstNode | None = field(default=None, kw_only=True)
-    _last_side_effect: Node | None = field(default=None, init=False)
+    _last_side_effect: dict[Effect, Node] = field(default_factory=dict, init=False)
 
     @abstractproperty
     def _raw(self) -> hf.Function | Case | TailLoop | Block:
@@ -76,23 +85,24 @@ def inputs(self) -> Sequence[Wire]:
 
     def set_outputs(self, *outputs: Wire) -> hf.Function | Case | TailLoop | Block:
         self._raw.set_outputs(*outputs)
-        if self._last_side_effect is not None:
-            self._handle_side_effects(self._raw.output_node)
+        self._handle_side_effects(
+            self._raw.output_node, list(self._last_side_effect.keys())
+        )
         return self._raw
 
     def add_op(
         self,
-        op: ops.DataflowOp,
+        op: OpWithEffects,
         /,
         *args: Wire,
         set_debug_info: bool = True,
     ) -> Node:
         """Adds an op to the dataflow graph builder. Set `set_debug_info=False` to
         avoid automatic debug information attachment.
         """
+        op, effects = (op.op, []) if isinstance(op, Pure) else op
         op_node = self._raw.add_op(op, *args)
-        if may_have_side_effect(op):
-            self._handle_side_effects(op_node)
+        self._handle_side_effects(op_node, effects)
 
         if set_debug_info and debug_conditions_fulfilled(self.current_ast_node):
             assert self.current_ast_node is not None  # for type-checker
@@ -101,26 +111,44 @@ def add_op(
             )
         return op_node
 
-    def _handle_side_effects(self, op_node: ToNode) -> None:
-        if self._last_side_effect is None:
-            self._propagate_side_effects()
-            self._last_side_effect = self.input_node
-        else:
-            assert not isinstance(self._raw.hugr[self._last_side_effect].op, ops.Output)
+    def _handle_side_effects(self, op_node: ToNode, effects: Iterable[Effect]) -> None:
+        """Updates Hugr to reflect `op_node` having effects `effects`.
+        Does nothing if effects is empty (or the node already has those effects)."""
         node = op_node.to_node()
-        if self._last_side_effect != node:  # avoid self-loops when propagating
-            self._raw.add_state_order(self._last_side_effect, node)
-            self._last_side_effect = node
+        to_propagate = set()  # Effects newly added to our container
+
+        def get_last_node(e: Effect) -> Node:
+            last = self._last_side_effect.get(e)
+            if last is None:
+                to_propagate.add(e)
+                last = self.input_node
+            else:
+                assert not isinstance(self._raw.hugr[last].op, Output)
+            self._last_side_effect[e] = node
+            return last
+
+        prev_nodes = {get_last_node(e) for e in effects}
+        # Avoid cycles and duplicate edges:
+        prev_nodes.discard(node)
+        for prev in self._raw.hugr.incoming_order_links(node):
+            prev_nodes.discard(prev)
+
+        for prev in prev_nodes:
+            self._raw.add_state_order(prev, node)
+
+        if to_propagate:
+            self._propagate_side_effects(to_propagate)
 
     @abstractmethod
-    def _propagate_side_effects(self) -> None:
+    def _propagate_side_effects(self, effects: Iterable[Effect]) -> None:
         """Subclasses must implement to mark the container node
         as side-effecting within any parent/ancestor builder"""
 
     def call(
         self,
         func: ToNode,
         *args: Wire,
+        effects: Sequence[Effect],
         instantiation: ht.FunctionType | None = None,
         type_args: Sequence[ht.TypeArg] | None = None,
         set_debug_info: bool = True,
@@ -131,7 +159,7 @@ def call(
         call = self._raw.call(
             func, *args, instantiation=instantiation, type_args=type_args
         )
-        self._handle_side_effects(call)
+        self._handle_side_effects(call, effects)
         if set_debug_info and debug_conditions_fulfilled(self.current_ast_node):
             assert self.current_ast_node is not None  # for type-checker
             call.metadata[HugrDebugInfo] = make_location_record(self.current_ast_node)
@@ -190,16 +218,17 @@ class TailLoopBuilder(_DFBuilderRaw[TailLoop]):
 
     def set_loop_outputs(self, predicate: Wire, *outputs: Wire) -> None:
         self._raw.set_loop_outputs(predicate, *outputs)
-        if self._last_side_effect is not None:
-            self._handle_side_effects(self._raw.output_node)
+        self._handle_side_effects(
+            self._raw.output_node, list(self._last_side_effect.keys())
+        )
 
-    def _propagate_side_effects(self) -> None:
-        self.parent._handle_side_effects(self._raw)
+    def _propagate_side_effects(self, effects: Iterable[Effect]) -> None:
+        self.parent._handle_side_effects(self._raw, effects)
 
 
 @dataclass
 class FunctionBuilder(_DFBuilderRaw[hf.Function]):
-    def _propagate_side_effects(self) -> None:
+    def _propagate_side_effects(self, effects: Iterable[Effect]) -> None:
         pass  # No parent
 
     @override
@@ -213,10 +242,10 @@ class CaseBuilder(_DFBuilderRaw[Case]):
     parent: Conditional
     grandparent: DFBuilder
 
-    def _propagate_side_effects(self) -> None:
+    def _propagate_side_effects(self, effects: Iterable[Effect]) -> None:
         # No need to do anything in the Conditional,
         # but the Conditional itself needs to be ordered inside its parent
-        self.grandparent._handle_side_effects(self.parent)
+        self.grandparent._handle_side_effects(self.parent, effects)
 
 
 @dataclass
@@ -251,12 +280,13 @@ class BlockBuilder(_DFBuilderRaw[Block]):
     parent: Cfg
     grandparent: DFBuilder
 
-    def _propagate_side_effects(self) -> None:
+    def _propagate_side_effects(self, effects: Iterable[Effect]) -> None:
         # No need to do anything in the CFG, but the CFG itself
         # needs to be ordered inside its parent,
-        self.grandparent._handle_side_effects(self.parent)
+        self.grandparent._handle_side_effects(self.parent, effects)
 
     def set_block_outputs(self, branching: Wire, *other_outputs: Wire) -> None:
         self._raw.set_outputs(branching, *other_outputs)
-        if self._last_side_effect is not None:
-            self._handle_side_effects(self._raw.output_node)
+        self._handle_side_effects(
+            self._raw.output_node, list(self._last_side_effect.keys())
+        )

guppylang-internals/src/guppylang_internals/compiler/builder/ops.py

@@ -0,0 +1,20 @@
+from hugr import ops
+from hugr.tys import Sum, Type, TypeRow
+
+from guppylang_internals.compiler.builder import OpWithEffects, Pure
+
+
+def make_tuple(tys: TypeRow | None = None) -> OpWithEffects:
+    return Pure(ops.MakeTuple(tys))
+
+
+def unpack_tuple(tys: TypeRow | None = None) -> OpWithEffects:
+    return Pure(ops.UnpackTuple(tys))
+
+
+def tag(tag: int, rows: Sum) -> OpWithEffects:
+    return Pure(ops.Tag(tag, rows))
+
+
+def some(ty: Type) -> OpWithEffects:
+    return Pure(ops.Some(ty))

guppylang-internals/src/guppylang_internals/compiler/cfg_compiler.py

@@ -1,7 +1,7 @@
 import functools
 from collections.abc import Sequence
 
-from hugr import Wire, ops
+from hugr import Wire
 from hugr import tys as ht
 from hugr.build import cfg as hc
 from hugr.hugr.node_port import ToNode
@@ -13,7 +13,7 @@
     Signature,
 )
 from guppylang_internals.checker.core import Place, Variable
-from guppylang_internals.compiler.builder import BlockBuilder, DFBuilder
+from guppylang_internals.compiler.builder import BlockBuilder, DFBuilder, ops
 from guppylang_internals.compiler.core import (
     CompilerContext,
     DFContainer,
@@ -109,7 +109,7 @@ def compile_bb(
     else:
         # Even if we don't branch, we still have to add a `Sum(())` predicates
         branch_port = dfg.builder.add_op(
-            ops.Tag(0, ht.UnitSum(1)), set_debug_info=False
+            ops.tag(0, ht.UnitSum(1)), set_debug_info=False
         )
 
     # Finally, we have to add the block output.
@@ -206,7 +206,7 @@ def choose_vars_for_tuple_sum(
         for i, var_row in enumerate(output_vars):
             case = conditional.add_case(i)
             outputs = [case.inputs()[all_vars_idxs[v.id]] for v in var_row]
-            tag = case.add_op(ops.Tag(i, sum_type), *outputs)
+            tag = case.add_op(ops.tag(i, sum_type), *outputs)
             case.set_outputs(tag)
         return conditional
 

guppylang-internals/src/guppylang_internals/compiler/core.py

@@ -3,11 +3,12 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-from hugr import Hugr, Wire, ops
+from hugr import Hugr, Wire
 from hugr import tys as ht
 from hugr.build import function as hf
 from hugr.build.dfg import DefinitionBuilder
 from hugr.hugr.base import OpVarCov
+from hugr.ops import Module
 from hugr.std import PRELUDE
 from hugr.std.collections.array import EXTENSION as ARRAY_EXTENSION
 from hugr.std.collections.borrow_array import EXTENSION as BORROW_ARRAY_EXTENSION
@@ -19,6 +20,7 @@
     TupleAccess,
     Variable,
 )
+from guppylang_internals.compiler.builder import ops
 from guppylang_internals.definition.common import (
     CompilableDef,
     CompiledDef,
@@ -85,7 +87,7 @@ class CompilerContext(ToHugrContext):
     themselves (i.e. `compile_inner` has not yet been called).
     """
 
-    module: DefinitionBuilder[ops.Module]
+    module: DefinitionBuilder[Module]
 
     #: The definitions compiled so far. For generic definitions, their id can occur
     #: multiple times here with respectively different monomorphizations. See
@@ -107,7 +109,7 @@ class CompilerContext(ToHugrContext):
 
     def __init__(
         self,
-        module: DefinitionBuilder[ops.Module],
+        module: DefinitionBuilder[Module],
         exported_defs: set[DefId],
         file_table: StringTable | None = None,
     ) -> None:
@@ -227,7 +229,7 @@ def __getitem__(self, place: Place) -> Wire:
             raise InternalGuppyError(f"Couldn't obtain a port for `{place}`")
         child_types = [child.ty.to_hugr(self.ctx) for child in children]
         child_wires = [self[child] for child in children]
-        wire = self.builder.add_op(ops.MakeTuple(child_types), *child_wires)[0]
+        wire = self.builder.add_op(ops.make_tuple(child_types), *child_wires)[0]
         for child in children:
             if child.ty.linear:
                 self.locals.pop(child.id)
@@ -240,7 +242,7 @@ def __setitem__(self, place: Place, port: Wire) -> None:
         is_return = isinstance(place, Variable) and is_return_var(place.name)
         if isinstance(place.ty, StructType) and not is_return:
             hugr_fields_ty = [t.ty.to_hugr(self.ctx) for t in place.ty.fields]
-            unpack = self.builder.add_op(ops.UnpackTuple(hugr_fields_ty), port)
+            unpack = self.builder.add_op(ops.unpack_tuple(hugr_fields_ty), port)
             for field, field_port in zip(place.ty.fields, unpack, strict=True):
                 self[FieldAccess(place, field, None)] = field_port
             # If we had a previous wire assigned to this place, we need forget about it.
@@ -249,7 +251,7 @@ def __setitem__(self, place: Place, port: Wire) -> None:
         # Same for tuples.
         elif isinstance(place.ty, TupleType) and not is_return:
             hugr_elem_tys = [ty.to_hugr(self.ctx) for ty in place.ty.element_types]
-            unpack = self.builder.add_op(ops.UnpackTuple(hugr_elem_tys), port)
+            unpack = self.builder.add_op(ops.unpack_tuple(hugr_elem_tys), port)
             for idx, (elem, elem_port) in enumerate(
                 zip(place.ty.element_types, unpack, strict=True)
             ):
@@ -319,30 +321,6 @@ def get_parent_type(defn: Definition) -> "RawDef | None":
 ]
 
 
-def may_have_side_effect(op: ops.Op) -> bool:
-    """Checks whether an operation could have a side-effect.
-
-    We need to insert implicit state order edges between these kinds of nodes to ensure
-    they are executed in the correct order, even if there is no data dependency.
-    """
-    match op:
-        case ops.ExtOp() as ext_op:
-            return ext_op.op_def().qualified_name() in EXTENSION_OPS_WITH_SIDE_EFFECTS
-        case ops.Custom(op_name=op_name, extension=extension):
-            qualified_name = f"{extension}.{op_name}" if extension else op_name
-            return qualified_name in EXTENSION_OPS_WITH_SIDE_EFFECTS
-        case ops.Call() | ops.CallIndirect():
-            # Conservative choice is to assume that all calls could have side effects.
-            # In the future we could inspect the call graph to figure out a more
-            # precise answer
-            return True
-        case _:
-            # There is no need to handle TailLoop (in case of non-termination) since
-            # TailLoops are only generated for array comprehensions which must have
-            # statically-guaranteed (finite) size. TODO revisit this for lists.
-            return False
-
-
 #: List of linear extension types that correspond to affine Guppy types and thus require
 #: insertion of an explicit drop operation.
 AFFINE_EXTENSION_TYS: list[str] = [
@@ -378,13 +356,6 @@ def requires_drop(ty: ht.Type) -> bool:
             return False
 
 
-def drop_op(ty: ht.Type) -> ops.ExtOp:
-    """Returns the operation to drop affine values."""
-    return GUPPY_EXTENSION.get_op("drop").instantiate(
-        [ht.TypeTypeArg(ty)], ht.FunctionType([ty], [])
-    )
-
-
 def insert_drops(hugr: Hugr[OpVarCov]) -> None:
     """Inserts explicit drop ops for unconnected ports into the Hugr.
     TODO: This is a quick workaround until we can properly insert these drops during
@@ -403,5 +374,8 @@ def insert_drops(hugr: Hugr[OpVarCov]) -> None:
                 and isinstance(kind, ht.ValueKind)
                 and requires_drop(kind.ty)
             ):
-                drop = hugr.add_node(drop_op(kind.ty), parent=data.parent)
+                drop_op = GUPPY_EXTENSION.get_op("drop").instantiate(
+                    [ht.TypeTypeArg(kind.ty)], ht.FunctionType([kind.ty], [])
+                )
+                drop = hugr.add_node(drop_op, parent=data.parent)
                 hugr.add_link(port, drop.inp(0))