substitute(self) -> substitute(&self) calling subst_mut(&mut self)

CQCL · Dec 4, 2024 · 3e755d1 · 3e755d1
1 parent f7a6dca
commit 3e755d1
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 130 deletions.
diff --git a/hugr-core/src/ops.rs b/hugr-core/src/ops.rs
@@ -350,7 +350,7 @@ pub trait StaticTag {
 
 #[enum_dispatch]
 /// Trait implemented by all OpType variants.
-pub trait OpTrait: Sized {
+pub trait OpTrait: Sized + Clone {
     /// A human-readable description of the operation.
     fn description(&self) -> &str;
 
@@ -415,11 +415,19 @@ pub trait OpTrait: Sized {
         .is_some() as usize
     }
 
+    /// Like [Self::subst_mut] but returns a substituted clone of `self`.
+    /// The default impl is correct so long as `subst_mut` is, so trait
+    /// `impl`s should/must override that if the op may contain type variables.
+    fn substitute(&self, subst: &Substitution) -> Self {
+        let mut s = self.clone();
+        s.subst_mut(subst);
+        s
+    }
+
     /// Apply a type-level substitution to this OpType, i.e. replace
     /// [type variables](crate::types::TypeArg::new_var_use) with new types.
-    fn substitute(self, _subst: &Substitution) -> Self {
-        self
-    }
+    /// The default is appropriate when there are no type (variables) in the op.
+    fn subst_mut(&mut self, _subst: &Substitution) {}
 }
 
 /// Properties of child graphs of ops, if the op has children.

diff --git a/hugr-core/src/ops/controlflow.rs b/hugr-core/src/ops/controlflow.rs
@@ -1,7 +1,7 @@
 //! Control flow operations.
 
 use crate::extension::ExtensionSet;
-use crate::types::{EdgeKind, Signature, Type, TypeRow};
+use crate::types::{EdgeKind, Signature, Substitution, Type, TypeRow};
 use crate::Direction;
 
 use super::dataflow::{DataflowOpTrait, DataflowParent};
@@ -37,13 +37,11 @@ impl DataflowOpTrait for TailLoop {
         Signature::new(inputs, outputs).with_extension_delta(self.extension_delta.clone())
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            just_inputs: self.just_inputs.substitute(subst),
-            just_outputs: self.just_outputs.substitute(subst),
-            rest: self.rest.substitute(subst),
-            extension_delta: self.extension_delta.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.just_inputs = self.just_inputs.substitute(subst);
+        self.just_outputs = self.just_outputs.substitute(subst);
+        self.rest = self.rest.substitute(subst);
+        self.extension_delta = self.extension_delta.substitute(subst);
     }
 }
 
@@ -110,17 +108,13 @@ impl DataflowOpTrait for Conditional {
             .with_extension_delta(self.extension_delta.clone())
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            sum_rows: self
-                .sum_rows
-                .into_iter()
-                .map(|r| r.substitute(subst))
-                .collect(),
-            other_inputs: self.other_inputs.substitute(subst),
-            outputs: self.outputs.substitute(subst),
-            extension_delta: self.extension_delta.substitute(subst),
+    fn subst_mut(&mut self, subst: &Substitution) {
+        for row in self.sum_rows.iter_mut() {
+            *row = row.substitute(subst)
         }
+        self.other_inputs = self.other_inputs.substitute(subst);
+        self.outputs = self.outputs.substitute(subst);
+        self.extension_delta = self.extension_delta.substitute(subst);
     }
 }
 
@@ -152,10 +146,8 @@ impl DataflowOpTrait for CFG {
         self.signature.clone()
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            signature: self.signature.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.signature = self.signature.substitute(subst);
     }
 }
 
@@ -238,17 +230,13 @@ impl OpTrait for DataflowBlock {
         }
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            inputs: self.inputs.substitute(subst),
-            other_outputs: self.other_outputs.substitute(subst),
-            sum_rows: self
-                .sum_rows
-                .into_iter()
-                .map(|r| r.substitute(subst))
-                .collect(),
-            extension_delta: self.extension_delta.substitute(subst),
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.inputs = self.inputs.substitute(subst);
+        self.other_outputs = self.other_outputs.substitute(subst);
+        for r in self.sum_rows.iter_mut() {
+            *r = r.substitute(subst);
         }
+        self.extension_delta = self.extension_delta.substitute(subst);
     }
 }
 
@@ -276,10 +264,8 @@ impl OpTrait for ExitBlock {
         }
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            cfg_outputs: self.cfg_outputs.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.cfg_outputs = self.cfg_outputs.substitute(subst);
     }
 }
 
@@ -345,10 +331,8 @@ impl OpTrait for Case {
         <Self as StaticTag>::TAG
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            signature: self.signature.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.signature = self.signature.substitute(subst);
     }
 }
 

diff --git a/hugr-core/src/ops/custom.rs b/hugr-core/src/ops/custom.rs
@@ -11,8 +11,11 @@ use {
     ::proptest_derive::Arbitrary,
 };
 
-use crate::extension::{ConstFoldResult, ExtensionId, ExtensionRegistry, OpDef, SignatureError};
 use crate::types::{type_param::TypeArg, Signature};
+use crate::{
+    extension::{ConstFoldResult, ExtensionId, ExtensionRegistry, OpDef, SignatureError},
+    types::Substitution,
+};
 use crate::{ops, IncomingPort, Node};
 
 use super::dataflow::DataflowOpTrait;
@@ -161,24 +164,17 @@ impl DataflowOpTrait for ExtensionOp {
         self.signature.clone()
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        let args = self
-            .args
-            .into_iter()
-            .map(|ta| ta.substitute(subst))
-            .collect::<Vec<_>>();
-        let signature = self.signature.substitute(subst);
+    fn subst_mut(&mut self, subst: &Substitution) {
+        for ta in self.args.iter_mut() {
+            *ta = ta.substitute(subst);
+        }
+        self.signature = self.signature.substitute(subst);
         debug_assert_eq!(
             self.def
-                .compute_signature(&args, subst.extension_registry())
+                .compute_signature(&self.args, subst.extension_registry())
                 .as_ref(),
-            Ok(&signature)
+            Ok(&self.signature)
         );
-        Self {
-            def: self.def,
-            args,
-            signature,
-        }
     }
 }
 
@@ -270,16 +266,11 @@ impl DataflowOpTrait for OpaqueOp {
             .with_extension_delta(self.extension().clone())
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        Self {
-            args: self
-                .args
-                .into_iter()
-                .map(|ta| ta.substitute(subst))
-                .collect(),
-            signature: self.signature.substitute(subst),
-            ..self
+    fn subst_mut(&mut self, subst: &Substitution) {
+        for ta in self.args.iter_mut() {
+            *ta = ta.substitute(subst);
         }
+        self.signature = self.signature.substitute(subst);
     }
 }
 

diff --git a/hugr-core/src/ops/dataflow.rs b/hugr-core/src/ops/dataflow.rs
@@ -11,7 +11,7 @@ use crate::IncomingPort;
 use ::proptest_derive::Arbitrary;
 
 /// Trait implemented by all dataflow operations.
-pub trait DataflowOpTrait: Sized {
+pub trait DataflowOpTrait: Sized + Clone {
     /// Tag identifying the operation.
     const TAG: OpTag;
 
@@ -52,7 +52,7 @@ pub trait DataflowOpTrait: Sized {
 
     /// Apply a type-level substitution to this OpType, i.e. replace
     /// [type variables](TypeArg::new_var_use) with new types.
-    fn substitute(self, _subst: &Substitution) -> Self;
+    fn subst_mut(&mut self, _subst: &Substitution) {}
 }
 
 /// Helpers to construct input and output nodes
@@ -113,10 +113,8 @@ impl DataflowOpTrait for Input {
         Signature::new(TypeRow::new(), self.types.clone())
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        Self {
-            types: self.types.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.types = self.types.substitute(subst);
     }
 }
 impl DataflowOpTrait for Output {
@@ -136,10 +134,8 @@ impl DataflowOpTrait for Output {
         None
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        Self {
-            types: self.types.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.types = self.types.substitute(subst);
     }
 }
 
@@ -168,8 +164,8 @@ impl<T: DataflowOpTrait> OpTrait for T {
         DataflowOpTrait::static_input(self)
     }
 
-    fn substitute(self, subst: &crate::types::Substitution) -> Self {
-        DataflowOpTrait::substitute(self, subst)
+    fn subst_mut(&mut self, subst: &Substitution) {
+        DataflowOpTrait::subst_mut(self, subst)
     }
 }
 impl<T: DataflowOpTrait> StaticTag for T {
@@ -208,24 +204,17 @@ impl DataflowOpTrait for Call {
         Some(EdgeKind::Function(self.called_function_type().clone()))
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        let type_args = self
-            .type_args
-            .into_iter()
-            .map(|ta| ta.substitute(subst))
-            .collect::<Vec<_>>();
-        let instantiation = self.instantiation.substitute(subst);
+    fn subst_mut(&mut self, subst: &Substitution) {
+        for ta in self.type_args.iter_mut() {
+            *ta = ta.substitute(subst);
+        }
+        self.instantiation = self.instantiation.substitute(subst);
         debug_assert_eq!(
             self.func_sig
-                .instantiate(&type_args, subst.extension_registry())
+                .instantiate(&self.type_args, subst.extension_registry())
                 .as_ref(),
-            Ok(&instantiation)
+            Ok(&self.instantiation)
         );
-        Self {
-            type_args,
-            instantiation,
-            func_sig: self.func_sig,
-        }
     }
 }
 impl Call {
@@ -320,10 +309,8 @@ impl DataflowOpTrait for CallIndirect {
         s
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        Self {
-            signature: self.signature.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.signature = self.signature.substitute(subst);
     }
 }
 
@@ -335,6 +322,7 @@ pub struct LoadConstant {
     pub datatype: Type,
 }
 impl_op_name!(LoadConstant);
+// Constants cannot contain type variables, so no substitution required
 impl DataflowOpTrait for LoadConstant {
     const TAG: OpTag = OpTag::LoadConst;
 
@@ -349,11 +337,6 @@ impl DataflowOpTrait for LoadConstant {
     fn static_input(&self) -> Option<EdgeKind> {
         Some(EdgeKind::Const(self.constant_type().clone()))
     }
-
-    fn substitute(self, _subst: &Substitution) -> Self {
-        // Constants cannot refer to TypeArgs, so neither can loading them
-        self
-    }
 }
 
 impl LoadConstant {
@@ -411,24 +394,17 @@ impl DataflowOpTrait for LoadFunction {
         Some(EdgeKind::Function(self.func_sig.clone()))
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        let type_args = self
-            .type_args
-            .into_iter()
-            .map(|ta| ta.substitute(subst))
-            .collect::<Vec<_>>();
-        let signature = self.signature.substitute(subst);
+    fn subst_mut(&mut self, subst: &Substitution) {
+        for ta in self.type_args.iter_mut() {
+            *ta = ta.substitute(subst);
+        }
+        self.signature = self.signature.substitute(subst);
         debug_assert_eq!(
             self.func_sig
-                .instantiate(&type_args, subst.extension_registry())
+                .instantiate(&self.type_args, subst.extension_registry())
                 .as_ref(),
-            Ok(&signature)
+            Ok(&self.signature)
         );
-        Self {
-            func_sig: self.func_sig,
-            type_args,
-            signature,
-        }
     }
 }
 impl LoadFunction {
@@ -520,9 +496,7 @@ impl DataflowOpTrait for DFG {
         self.inner_signature()
     }
 
-    fn substitute(self, subst: &Substitution) -> Self {
-        Self {
-            signature: self.signature.substitute(subst),
-        }
+    fn subst_mut(&mut self, subst: &Substitution) {
+        self.signature = self.signature.substitute(subst)
     }
 }