feat: Always require a signature in OpaqueOp (#732)

Closes #689 

Adds a previously-failing test for the serialiser, where the signature
of an opaque op was discarded and couldn't be recovered. Depends on
#730.

src/builder/circuit.rs

@@ -179,7 +179,7 @@ mod test {
                 "MyOp",
                 "unknown op".to_string(),
                 vec![],
-                Some(FunctionType::new(vec![QB, NAT], vec![QB])),
+                FunctionType::new(vec![QB, NAT], vec![QB]),
             ))
             .into(),
         );

src/extension/infer/test.rs

@@ -450,8 +450,7 @@ fn extension_adding_sequence() -> Result<(), Box<dyn Error>> {
 }
 
 fn make_opaque(extension: impl Into<ExtensionId>, signature: FunctionType) -> ops::LeafOp {
-    let opaque =
-        ops::custom::OpaqueOp::new(extension.into(), "", "".into(), vec![], Some(signature));
+    let opaque = ops::custom::OpaqueOp::new(extension.into(), "", "".into(), vec![], signature);
     ops::custom::ExternalOp::from(opaque).into()
 }
 
@@ -874,7 +873,7 @@ fn plus_on_self() -> Result<(), Box<dyn std::error::Error>> {
         "2qb_op",
         String::new(),
         vec![],
-        Some(ft),
+        ft,
     ))
     .into();
     let unary_sig = FunctionType::new_endo(type_row![QB_T])
@@ -884,7 +883,7 @@ fn plus_on_self() -> Result<(), Box<dyn std::error::Error>> {
         "1qb_op",
         String::new(),
         vec![],
-        Some(unary_sig),
+        unary_sig,
     ))
     .into();
     // Constrain q1,q2 as PLUS(ext1, inputs):

src/extension/op_def.rs

@@ -349,13 +349,6 @@ impl OpDef {
         self.signature_func.compute_signature(self, args, exts)
     }
 
-    pub(crate) fn should_serialize_signature(&self) -> bool {
-        match self.signature_func {
-            SignatureFunc::TypeScheme { .. } => false,
-            SignatureFunc::CustomFunc { .. } => true,
-        }
-    }
-
     /// Fallibly returns a Hugr that may replace an instance of this OpDef
     /// given a set of available extensions that may be used in the Hugr.
     pub fn try_lower(&self, args: &[TypeArg], available_extensions: &ExtensionSet) -> Option<Hugr> {

src/hugr/rewrite/replace.rs

@@ -652,7 +652,7 @@ mod test {
                 s,
                 String::new(),
                 vec![],
-                Some(utou.clone()),
+                utou.clone(),
             )))
         };
         let mut h = DFGBuilder::new(FunctionType::new(

src/hugr/serialize.rs

@@ -257,10 +257,13 @@ pub mod test {
         DataflowSubContainer, HugrBuilder, ModuleBuilder,
     };
     use crate::extension::prelude::BOOL_T;
+    use crate::extension::simple_op::MakeRegisteredOp;
     use crate::extension::{EMPTY_REG, PRELUDE_REGISTRY};
     use crate::hugr::hugrmut::sealed::HugrMutInternals;
     use crate::hugr::NodeType;
+    use crate::ops::custom::{ExtensionOp, OpaqueOp};
     use crate::ops::{dataflow::IOTrait, Input, LeafOp, Module, Output, DFG};
+    use crate::std_extensions::logic::NotOp;
     use crate::types::{FunctionType, Type};
     use crate::OutgoingPort;
     use itertools::Itertools;
@@ -298,9 +301,23 @@ pub mod test {
 
         assert_eq!(new_hugr.root, h_canon.root);
         assert_eq!(new_hugr.hierarchy, h_canon.hierarchy);
-        assert_eq!(new_hugr.op_types, h_canon.op_types);
         assert_eq!(new_hugr.metadata, h_canon.metadata);
 
+        // Extension operations may have been downgraded to opaque operations.
+        for node in new_hugr.nodes() {
+            let new_op = new_hugr.get_optype(node);
+            let old_op = h_canon.get_optype(node);
+            if let OpType::LeafOp(LeafOp::CustomOp(new_ext)) = new_op {
+                if let OpType::LeafOp(LeafOp::CustomOp(old_ext)) = old_op {
+                    assert_eq!(new_ext.clone().as_opaque(), old_ext.clone().as_opaque());
+                } else {
+                    panic!("Expected old_op to be a custom op");
+                }
+            } else {
+                assert_eq!(new_op, old_op);
+            }
+        }
+
         // Check that the graphs are equivalent up to port renumbering.
         let new_graph = &new_hugr.graph;
         let old_graph = &h_canon.graph;
@@ -423,7 +440,30 @@ pub mod test {
     }
 
     #[test]
-    fn canonicalisation() {
+    fn opaque_ops() -> Result<(), Box<dyn std::error::Error>> {
+        let tp: Vec<Type> = vec![BOOL_T; 1];
+        let mut dfg = DFGBuilder::new(FunctionType::new_endo(tp))?;
+        let [wire] = dfg.input_wires_arr();
+
+        // Add an extension operation
+        let extension_op: ExtensionOp = NotOp.to_extension_op().unwrap();
+        let wire = dfg
+            .add_dataflow_op(extension_op.clone(), [wire])
+            .unwrap()
+            .out_wire(0);
+
+        // Add an unresolved opaque operation
+        let opaque_op: OpaqueOp = extension_op.into();
+        let wire = dfg.add_dataflow_op(opaque_op, [wire]).unwrap().out_wire(0);
+
+        let hugr = dfg.finish_hugr_with_outputs([wire], &PRELUDE_REGISTRY)?;
+
+        check_hugr_roundtrip(&hugr);
+        Ok(())
+    }
+
+    #[test]
+    fn hierarchy_order() {
         let mut hugr = closed_dfg_root_hugr(FunctionType::new(vec![QB], vec![QB]));
         let [old_in, out] = hugr.get_io(hugr.root()).unwrap();
         hugr.connect(old_in, 0, out, 0).unwrap();

src/ops/custom.rs

@@ -36,6 +36,23 @@ impl ExternalOp {
             Self::Extension(op) => op.args(),
         }
     }
+
+    /// Name of the ExternalOp
+    pub fn name(&self) -> SmolStr {
+        let (res_id, op_name) = match self {
+            Self::Opaque(op) => (&op.extension, &op.op_name),
+            Self::Extension(ExtensionOp { def, .. }) => (def.extension(), def.name()),
+        };
+        qualify_name(res_id, op_name)
+    }
+
+    /// Downgrades this ExternalOp into an OpaqueOp
+    pub fn as_opaque(self) -> OpaqueOp {
+        match self {
+            Self::Opaque(op) => op,
+            Self::Extension(op) => op.into(),
+        }
+    }
 }
 
 impl From<ExternalOp> for OpaqueOp {
@@ -66,34 +83,19 @@ impl From<ExternalOp> for OpType {
     }
 }
 
-impl ExternalOp {
-    /// Name of the ExternalOp
-    pub fn name(&self) -> SmolStr {
-        let (res_id, op_name) = match self {
-            Self::Opaque(op) => (&op.extension, &op.op_name),
-            Self::Extension(ExtensionOp { def, .. }) => (def.extension(), def.name()),
-        };
-        qualify_name(res_id, op_name)
-    }
-}
+impl DataflowOpTrait for ExternalOp {
+    const TAG: OpTag = OpTag::Leaf;
 
-impl ExternalOp {
-    /// A description of the external op.
-    pub fn description(&self) -> &str {
+    fn description(&self) -> &str {
         match self {
-            Self::Opaque(op) => op.description.as_str(),
+            Self::Opaque(op) => DataflowOpTrait::description(op),
             Self::Extension(ext_op) => DataflowOpTrait::description(ext_op),
         }
     }
 
-    /// Note the case of an OpaqueOp without a signature should already
-    /// have been detected in [resolve_extension_ops]
-    pub fn dataflow_signature(&self) -> FunctionType {
+    fn signature(&self) -> FunctionType {
         match self {
-            Self::Opaque(op) => op
-                .signature
-                .clone()
-                .expect("Op should have been serialized with signature."),
+            Self::Opaque(op) => op.signature.clone(),
             Self::Extension(ext_op) => ext_op.signature(),
         }
     }
@@ -144,17 +146,12 @@ impl From<ExtensionOp> for OpaqueOp {
             args,
             signature,
         } = op;
-        let opt_sig = if def.should_serialize_signature() {
-            Some(signature)
-        } else {
-            None
-        };
         OpaqueOp {
             extension: def.extension().clone(),
             op_name: def.name().clone(),
             description: def.description().into(),
             args,
-            signature: opt_sig,
+            signature,
         }
     }
 }
@@ -199,7 +196,7 @@ pub struct OpaqueOp {
     op_name: SmolStr,
     description: String, // cache in advance so description() can return &str
     args: Vec<TypeArg>,
-    signature: Option<FunctionType>,
+    signature: FunctionType,
 }
 
 fn qualify_name(res_id: &ExtensionId, op_name: &SmolStr) -> SmolStr {
@@ -213,7 +210,7 @@ impl OpaqueOp {
         op_name: impl Into<SmolStr>,
         description: String,
         args: impl Into<Vec<TypeArg>>,
-        signature: Option<FunctionType>,
+        signature: FunctionType,
     ) -> Self {
         Self {
             extension,
@@ -255,6 +252,18 @@ impl From<OpaqueOp> for OpType {
     }
 }
 
+impl DataflowOpTrait for OpaqueOp {
+    const TAG: OpTag = OpTag::Leaf;
+
+    fn description(&self) -> &str {
+        &self.description
+    }
+
+    fn signature(&self) -> FunctionType {
+        self.signature.clone()
+    }
+}
+
 /// Resolve serialized names of operations into concrete implementation (OpDefs) where possible
 #[allow(dead_code)]
 pub fn resolve_extension_ops(
@@ -291,7 +300,7 @@ pub fn resolve_extension_ops(
 /// # Errors
 /// If the serialized opaque resolves to a definition that conflicts with what was serialized
 pub fn resolve_opaque_op(
-    n: Node,
+    _n: Node,
     opaque: &OpaqueOp,
     extension_registry: &ExtensionRegistry,
 ) -> Result<Option<ExtensionOp>, CustomOpError> {
@@ -305,22 +314,15 @@ pub fn resolve_opaque_op(
         };
         let ext_op =
             ExtensionOp::new(def.clone(), opaque.args.clone(), extension_registry).unwrap();
-        if let Some(stored_sig) = &opaque.signature {
-            if stored_sig != &ext_op.signature {
-                return Err(CustomOpError::SignatureMismatch {
-                    extension: opaque.extension.clone(),
-                    op: def.name().clone(),
-                    computed: ext_op.signature.clone(),
-                    stored: stored_sig.clone(),
-                });
-            };
-        }
+        if opaque.signature != ext_op.signature {
+            return Err(CustomOpError::SignatureMismatch {
+                extension: opaque.extension.clone(),
+                op: def.name().clone(),
+                computed: ext_op.signature.clone(),
+                stored: opaque.signature.clone(),
+            });
+        };
         Ok(Some(ext_op))
-    } else if opaque.signature.is_none() {
-        Err(CustomOpError::NoStoredSignature(
-            ExternalOp::Opaque(opaque.clone()).name(),
-            n,
-        ))
     } else {
         Ok(None)
     }
@@ -330,9 +332,6 @@ pub fn resolve_opaque_op(
 /// when trying to resolve the serialized names against a registry of known Extensions.
 #[derive(Clone, Debug, Error, PartialEq)]
 pub enum CustomOpError {
-    /// Extension not found, and no signature
-    #[error("Unable to resolve operation {0} for node {1:?} with no saved signature")]
-    NoStoredSignature(SmolStr, Node),
     /// The Extension was found but did not contain the expected OpDef
     #[error("Operation {0} not found in Extension {1}")]
     OpNotFoundInExtension(SmolStr, ExtensionId),
@@ -350,22 +349,24 @@ pub enum CustomOpError {
 #[cfg(test)]
 mod test {
 
-    use crate::extension::prelude::USIZE_T;
+    use crate::extension::prelude::{QB_T, USIZE_T};
 
     use super::*;
 
     #[test]
     fn new_opaque_op() {
+        let sig = FunctionType::new_endo(vec![QB_T]);
         let op = OpaqueOp::new(
             "res".try_into().unwrap(),
             "op",
             "desc".into(),
             vec![TypeArg::Type { ty: USIZE_T }],
-            None,
+            sig.clone(),
         );
         let op: ExternalOp = op.into();
         assert_eq!(op.name(), "res.op");
-        assert_eq!(op.description(), "desc");
+        assert_eq!(DataflowOpTrait::description(&op), "desc");
         assert_eq!(op.args(), &[TypeArg::Type { ty: USIZE_T }]);
+        assert_eq!(op.signature(), sig);
     }
 }

src/ops/leaf.rs

@@ -170,7 +170,7 @@ impl DataflowOpTrait for LeafOp {
 
         match self {
             LeafOp::Noop { ty: typ } => FunctionType::new(vec![typ.clone()], vec![typ.clone()]),
-            LeafOp::CustomOp(ext) => ext.dataflow_signature(),
+            LeafOp::CustomOp(ext) => ext.signature(),
             LeafOp::MakeTuple { tys: types } => {
                 FunctionType::new(types.clone(), vec![Type::new_tuple(types.clone())])
             }