feat: Result type for ComposablePasses

hugr-py/src/hugr/hugr/base.py

@@ -934,6 +934,14 @@ def insert_hugr(self, hugr: Hugr, parent: ToNode | None = None) -> dict[Node, No
             )
         return mapping
 
+    def _overwrite_hugr(self, new_hugr: Hugr) -> None:
+        """Modify a Hugr in place by replacing contents with those from a new Hugr."""
+        self.module_root = new_hugr.module_root
+        self.entrypoint = new_hugr.entrypoint
+        self._nodes = new_hugr._nodes
+        self._links = new_hugr._links
+        self._free_nodes = new_hugr._free_nodes
+
     def _to_serial(self) -> SerialHugr:
         """Serialize the HUGR."""
 

hugr-py/src/hugr/passes/_composable_pass.py

@@ -5,41 +5,84 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from hugr.hugr.base import Hugr
 
 
+# Type alias for a pass name
+PassName = str
+
+
 @runtime_checkable
 class ComposablePass(Protocol):
     """A Protocol which represents a composable Hugr transformation."""
 
-    def __call__(self, hugr: Hugr) -> None:
-        """Call the pass to transform a HUGR."""
-        ...
+    def __call__(self, hugr: Hugr, *, inplace: bool = True) -> Hugr:
+        """Call the pass to transform a HUGR, returning a Hugr."""
+        return self.run(hugr, inplace=inplace).hugr
+
+    def run(self, hugr: Hugr, *, inplace: bool = True) -> PassResult:
+        """Run the pass to transform a HUGR, returning a PassResult.
+
+        See :func:`_impl_pass_run` for a helper function to implement this method.
+        """
 
     @property
-    def name(self) -> str:
+    def name(self) -> PassName:
         """Returns the name of the pass."""
         return self.__class__.__name__
 
     def then(self, other: ComposablePass) -> ComposablePass:
         """Perform another composable pass after this pass."""
-        # Provide a default implementation for composing passes.
-        pass_list = []
-        if isinstance(self, ComposedPass):
-            pass_list.extend(self.passes)
-        else:
-            pass_list.append(self)
-
-        if isinstance(other, ComposedPass):
-            pass_list.extend(other.passes)
-        else:
-            pass_list.append(other)
-
-        return ComposedPass(pass_list)
+        return ComposedPass(self, other)
+
+
+def impl_pass_run(
+    pass_: ComposablePass,
+    *,
+    hugr: Hugr,
+    inplace: bool,
+    inplace_call: Callable[[Hugr], PassResult] | None = None,
+    copy_call: Callable[[Hugr], PassResult] | None = None,
+) -> PassResult:
+    """Helper function to implement a ComposablePass.run method, given an
+    inplace or copy-returning pass methods.
+
+    At least one of the `inplace_call` or `copy_call` arguments must be provided.
+
+    :param pass_: The pass being run. Used for error messages.
+    :param hugr: The Hugr to apply the pass to.
+    :param inplace: Whether to apply the pass inplace.
+    :param inplace_call: The method to apply the pass inplace.
+    :param copy_call: The method to apply the pass by copying the Hugr.
+    :return: The result of the pass application.
+    :raises ValueError: If neither `inplace_call` nor `copy_call` is provided.
+    """
+    if inplace and inplace_call is not None:
+        return inplace_call(hugr)
+    elif inplace and copy_call is not None:
+        pass_result = copy_call(hugr)
+        pass_result.hugr = hugr
+        if pass_result.modified:
+            hugr._overwrite_hugr(pass_result.hugr)
+            pass_result.original_dirty = True
+        return pass_result
+    elif not inplace and copy_call is not None:
+        return copy_call(hugr)
+    elif not inplace and inplace_call is not None:
+        new_hugr = deepcopy(hugr)
+        pass_result = inplace_call(new_hugr)
+        pass_result.original_dirty = False
+        return pass_result
+    else:
+        msg = f"{pass_.name} needs to implement at least an inplace or copy run method"
+        raise ValueError(msg)
 
 
 @dataclass
@@ -48,11 +91,90 @@ class ComposedPass(ComposablePass):
 
     passes: list[ComposablePass]
 
-    def __call__(self, hugr: Hugr):
-        """Call all of the passes in sequence."""
-        for comp_pass in self.passes:
-            comp_pass(hugr)
+    def __init__(self, *passes: ComposablePass) -> None:
+        self.passes = []
+        for pass_ in passes:
+            if isinstance(pass_, ComposedPass):
+                self.passes.extend(pass_.passes)
+            else:
+                self.passes.append(pass_)
+
+    def run(self, hugr: Hugr, *, inplace: bool = True) -> PassResult:
+        def apply(hugr: Hugr) -> PassResult:
+            pass_result = PassResult(hugr=hugr)
+            for comp_pass in self.passes:
+                new_result = comp_pass.run(pass_result.hugr, inplace=inplace)
+                pass_result = pass_result.then(new_result)
+            return pass_result
+
+        return impl_pass_run(
+            self,
+            hugr=hugr,
+            inplace=inplace,
+            inplace_call=apply,
+            copy_call=apply,
+        )
 
     @property
-    def name(self) -> str:
+    def name(self) -> PassName:
         return f"Composed({ ', '.join(pass_.name for pass_ in self.passes) })"
+
+
+@dataclass
+class PassResult:
+    """The result of a series of composed passes applied to a HUGR.
+
+    Includes a flag indicating whether the passes modified the HUGR, and an
+    arbitrary result object for each pass.
+
+    In some cases, `modified` may be set to `True` even if the pass did not
+    modify the program.
+
+    :attr hugr: The transformed Hugr.
+    :attr original_dirty: Whether the original HUGR was modified by the pass.
+    :attr modified: Whether the pass made changes to the HUGR.
+    :attr results: The result of each applied pass, as a tuple of the pass and
+        the result.
+    """
+
+    hugr: Hugr
+    original_dirty: bool = False
+    modified: bool = False
+    results: list[tuple[PassName, Any]] = field(default_factory=list)
+
+    @classmethod
+    def for_pass(
+        cls,
+        pass_: ComposablePass,
+        hugr: Hugr,
+        *,
+        result: Any,
+        inline: bool,
+        modified: bool = True,
+    ) -> PassResult:
+        """Create a new PassResult after a pass application.
+
+        :param hugr: The Hugr that was transformed.
+        :param pass_: The pass that was applied.
+        :param result: The result of the pass application.
+        :param inline: Whether the pass was applied inplace.
+        :param modified: Whether the pass modified the HUGR.
+        """
+        return cls(
+            hugr=hugr,
+            original_dirty=inline and modified,
+            modified=modified,
+            results=[(pass_.name, result)],
+        )
+
+    def then(self, other: PassResult) -> PassResult:
+        """Extend the PassResult with the results of another PassResult.
+
+        Keeps the hugr returned by the last pass.
+        """
+        return PassResult(
+            hugr=other.hugr,
+            original_dirty=self.original_dirty or other.original_dirty,
+            modified=self.modified or other.modified,
+            results=self.results + other.results,
+        )

hugr-py/tests/test_passes.py

@@ -1,26 +1,105 @@
+from copy import deepcopy
+
+import pytest
+
 from hugr.hugr.base import Hugr
-from hugr.passes._composable_pass import ComposablePass, ComposedPass
+from hugr.passes._composable_pass import (
+    ComposablePass,
+    ComposedPass,
+    PassResult,
+    impl_pass_run,
+)
 
 
 def test_composable_pass() -> None:
-    class MyDummyPass(ComposablePass):
-        def __call__(self, hugr: Hugr) -> None:
-            return self(hugr)
+    class MyDummyInlinePass(ComposablePass):
+        def run(self, hugr: Hugr, inplace: bool = True) -> PassResult:
+            return impl_pass_run(
+                self,
+                hugr=hugr,
+                inplace=inplace,
+                inplace_call=lambda hugr: PassResult.for_pass(
+                    self,
+                    hugr,
+                    result=None,
+                    inline=True,
+                    # Say that we modified the HUGR even though we didn't
+                    modified=True,
+                ),
+            )
 
-    dummy = MyDummyPass()
+    class MyDummyCopyPass(ComposablePass):
+        def run(self, hugr: Hugr, inplace: bool = True) -> PassResult:
+            return impl_pass_run(
+                self,
+                hugr=hugr,
+                inplace=inplace,
+                copy_call=lambda hugr: PassResult.for_pass(
+                    self,
+                    deepcopy(hugr),
+                    result=None,
+                    inline=False,
+                    # Say that we modified the HUGR even though we didn't
+                    modified=True,
+                ),
+            )
 
-    composed_dummies = dummy.then(dummy)
+    dummy_inline = MyDummyInlinePass()
+    dummy_copy = MyDummyCopyPass()
 
-    my_composed_pass = ComposedPass([dummy, dummy])
-    assert my_composed_pass.passes == [dummy, dummy]
+    composed_dummies = dummy_inline.then(dummy_copy)
+    assert isinstance(composed_dummies, ComposedPass)
 
-    assert isinstance(composed_dummies, ComposablePass)
-    assert composed_dummies == my_composed_pass
+    assert dummy_inline.name == "MyDummyInlinePass"
+    assert dummy_copy.name == "MyDummyCopyPass"
+    assert composed_dummies.name == "Composed(MyDummyInlinePass, MyDummyCopyPass)"
+    assert composed_dummies.then(dummy_inline).then(composed_dummies).name == (
+        "Composed("
+        + "MyDummyInlinePass, MyDummyCopyPass, "
+        + "MyDummyInlinePass, "
+        + "MyDummyInlinePass, MyDummyCopyPass)"
+    )
 
-    assert dummy.name == "MyDummyPass"
-    assert composed_dummies.name == "Composed(MyDummyPass, MyDummyPass)"
+    # Apply the passes
+    hugr: Hugr = Hugr()
+    new_hugr = composed_dummies(hugr, inplace=False)
+    assert hugr == new_hugr
+    assert new_hugr is not hugr
 
-    assert (
-        composed_dummies.then(my_composed_pass).name
-        == "Composed(MyDummyPass, MyDummyPass, MyDummyPass, MyDummyPass)"
-    )
+    # Verify the pass results
+    hugr = Hugr()
+    inplace_result = composed_dummies.run(hugr, inplace=True)
+    assert inplace_result.modified
+    assert inplace_result.original_dirty
+    assert inplace_result.results == [
+        ("MyDummyInlinePass", None),
+        ("MyDummyCopyPass", None),
+    ]
+    assert inplace_result.hugr is hugr
+
+    hugr = Hugr()
+    copy_result = composed_dummies.run(hugr, inplace=False)
+    assert copy_result.modified
+    assert not copy_result.original_dirty
+    assert copy_result.results == [
+        ("MyDummyInlinePass", None),
+        ("MyDummyCopyPass", None),
+    ]
+    assert copy_result.hugr is not hugr
+
+
+def test_invalid_composable_pass() -> None:
+    class MyDummyInvalidPass(ComposablePass):
+        def run(self, hugr: Hugr, inplace: bool = True) -> PassResult:
+            return impl_pass_run(
+                self,
+                hugr=hugr,
+                inplace=inplace,
+            )
+
+    dummy_invalid = MyDummyInvalidPass()
+    with pytest.raises(
+        ValueError,
+        match="MyDummyInvalidPass needs to implement at least an inplace or copy run method",  # noqa: E501
+    ):
+        dummy_invalid.run(Hugr())