feat: LowerTypes pass allows replacing extension types and ops

hugr-core/src/hugr/internal.rs

@@ -10,7 +10,7 @@ use itertools::Itertools;
 use portgraph::{LinkMut, LinkView, MultiPortGraph, PortMut, PortOffset, PortView};
 
 use crate::ops::handle::NodeHandle;
-use crate::ops::OpTrait;
+use crate::ops::{OpTag, OpTrait};
 use crate::{Direction, Hugr, Node};
 
 use super::hugrmut::{panic_invalid_node, panic_invalid_non_root};
@@ -309,6 +309,22 @@ pub trait HugrMutInternals: RootTagged<Node = Node> {
         }
         self.hugr_mut().replace_op(node, op)
     }
+
+    /// Gets a mutable reference to the optype.
+    ///
+    /// Changing this may invalidate the ports, which may need to be resized to
+    /// match the OpType signature.
+    ///
+    /// Will panic for the root node unless [`Self::RootHandle`](RootTagged::RootHandle)
+    /// is [OpTag::Any], as mutation could invalidate the bound.
+    fn optype_mut(&mut self, node: Node) -> &mut OpType {
+        if Self::RootHandle::TAG.is_superset(OpTag::Any) {
+            panic_invalid_node(self, node);
+        } else {
+            panic_invalid_non_root(self, node);
+        }
+        self.hugr_mut().op_types.get_mut(node.pg_index())
+    }
 }
 
 /// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr.
@@ -406,8 +422,7 @@ impl<T: RootTagged<RootHandle = Node, Node = Node> + AsMut<Hugr>> HugrMutInterna
 
     fn replace_op(&mut self, node: Node, op: impl Into<OpType>) -> Result<OpType, HugrError> {
         // We know RootHandle=Node here so no need to check
-        let cur = self.hugr_mut().op_types.get_mut(node.pg_index());
-        Ok(std::mem::replace(cur, op.into()))
+        Ok(std::mem::replace(self.optype_mut(node), op.into()))
     }
 }
 

hugr-core/src/ops/custom.rs

@@ -85,6 +85,12 @@ impl ExtensionOp {
         self.def.as_ref()
     }
 
+    /// Gets an Arc to the [`OpDef`] of this instance, i.e. usable to create
+    /// new instances.
+    pub fn def_arc(&self) -> &Arc<OpDef> {
+        &self.def
+    }
+
     /// Attempt to evaluate this operation. See [`OpDef::constant_fold`].
     pub fn constant_fold(&self, consts: &[(IncomingPort, ops::Value)]) -> ConstFoldResult {
         self.def().constant_fold(self.args(), consts)

hugr-core/src/types/poly_func.rs

@@ -119,7 +119,7 @@ impl<RV: MaybeRV> PolyFuncTypeBase<RV> {
     /// # Errors
     /// If there is not exactly one [TypeArg] for each binder ([Self::params]),
     /// or an arg does not fit into its corresponding [TypeParam]
-    pub(crate) fn instantiate(&self, args: &[TypeArg]) -> Result<FuncTypeBase<RV>, SignatureError> {
+    pub fn instantiate(&self, args: &[TypeArg]) -> Result<FuncTypeBase<RV>, SignatureError> {
         // Check that args are applicable, and that we have a value for each binder,
         // i.e. each possible free variable within the body.
         check_type_args(args, &self.params)?;

hugr-passes/src/lib.rs

@@ -28,6 +28,8 @@ pub use monomorphize::remove_polyfuncs;
 #[allow(deprecated)]
 pub use monomorphize::monomorphize;
 pub use monomorphize::{MonomorphizeError, MonomorphizePass};
+pub mod lower_types;
+pub use lower_types::LowerTypes;
 pub mod nest_cfgs;
 pub mod non_local;
 pub mod validation;

hugr-passes/src/validation.rs

@@ -23,7 +23,7 @@ pub enum ValidationLevel {
     WithExtensions,
 }
 
-#[derive(Error, Debug)]
+#[derive(Error, Debug, PartialEq)]
 #[allow(missing_docs)]
 pub enum ValidatePassError {
     #[error("Failed to validate input HUGR: {err}\n{pretty_hugr}")]