Fix tests and stuff

hugr-cli/tests/validate.rs

@@ -126,7 +126,9 @@ fn test_mermaid_invalid(bad_hugr_string: String, mut cmd: Command) {
     cmd.arg("mermaid");
     cmd.arg("--validate");
     cmd.write_stdin(bad_hugr_string);
-    cmd.assert().failure().stderr(contains("UnconnectedPort"));
+    cmd.assert()
+        .failure()
+        .stderr(contains("has an unconnected port"));
 }
 
 #[rstest]

hugr-core/src/builder.rs

@@ -242,7 +242,6 @@ pub(crate) mod test {
     use crate::hugr::{views::HugrView, HugrMut};
     use crate::ops;
     use crate::types::{PolyFuncType, Signature};
-    use crate::utils::test_quantum_extension;
     use crate::Hugr;
 
     use super::handle::BuildHandle;

hugr-core/src/builder/build_traits.rs

@@ -52,18 +52,10 @@ pub trait Container {
         let node: OpType = node.into();
 
         // Add the extension the operation is defined in to the HUGR.
-        if let Some(op) = node.as_extension_op() {
-            // Since the operations only keep a weak reference to the extension,
-            // we have to ignore the unlikely case where the extension has been dropped.
-            if let Some(extensions) = op.def().extension().upgrade() {
-                self.use_extension(extensions);
-            };
-        }
-        // If the operation is a dataflow operation, also add the extensions
-        // used by the types.
-        if let Some(sig) = node.dataflow_signature() {
-            self.use_extensions(&sig.used_extensions());
-        }
+        let used_extensions = node
+            .used_extensions()
+            .unwrap_or_else(|e| panic!("Build-time signatures should have valid extensions. {e}"));
+        self.use_extensions(used_extensions);
 
         let parent = self.container_node();
         self.hugr_mut().add_node_with_parent(parent, node)
@@ -110,7 +102,11 @@ pub trait Container {
         });
 
         // Add the extensions used by the function types.
-        self.use_extensions(&body.used_extensions());
+        self.use_extensions(
+            body.used_extensions().unwrap_or_else(|e| {
+                panic!("Build-time signatures should have valid extensions. {e}")
+            }),
+        );
 
         let db = DFGBuilder::create_with_io(self.hugr_mut(), f_node, body)?;
         Ok(FunctionBuilder::from_dfg_builder(db))
@@ -153,7 +149,10 @@ pub trait Container {
     }
 
     /// Extend the set of extensions used by the hugr with the extensions in the registry.
-    fn use_extensions(&mut self, registry: &ExtensionRegistry) {
+    fn use_extensions<Reg>(&mut self, registry: impl IntoIterator<Item = Reg>)
+    where
+        ExtensionRegistry: Extend<Reg>,
+    {
         self.hugr_mut().extensions_mut().extend(registry);
     }
 }

hugr-core/src/builder/circuit.rs

@@ -346,7 +346,7 @@ mod test {
 
         let mut registry = test_quantum_extension::REG.clone();
         registry.register(my_ext).unwrap();
-        let build_res = module_builder.finish_hugr(&registry);
+        let build_res = module_builder.finish_hugr();
 
         assert_matches!(build_res, Ok(_));
     }

hugr-core/src/builder/dataflow.rs

@@ -324,7 +324,7 @@ pub(crate) mod test {
     use crate::std_extensions::logic::test::and_op;
     use crate::types::type_param::TypeParam;
     use crate::types::{EdgeKind, FuncValueType, RowVariable, Signature, Type, TypeBound, TypeRV};
-    use crate::utils::test_quantum_extension::{self, h_gate};
+    use crate::utils::test_quantum_extension::h_gate;
     use crate::{builder::test::n_identity, type_row, Wire};
 
     use super::super::test::simple_dfg_hugr;

hugr-core/src/builder/module.rs

@@ -109,7 +109,11 @@ impl<T: AsMut<Hugr> + AsRef<Hugr>> ModuleBuilder<T> {
         });
 
         // Add the extensions used by the function types.
-        self.use_extensions(&body.used_extensions());
+        self.use_extensions(
+            body.used_extensions().unwrap_or_else(|e| {
+                panic!("Build-time signatures should have valid extensions. {e}")
+            }),
+        );
 
         Ok(declare_n.into())
     }

hugr-core/src/extension.rs

@@ -3,14 +3,16 @@
 //! TODO: YAML declaration and parsing. This should be similar to a plugin
 //! system (outside the `types` module), which also parses nested [`OpDef`]s.
 
+use itertools::Itertools;
 pub use semver::Version;
 use serde::{Deserialize, Deserializer, Serialize};
 use std::collections::btree_map;
 use std::collections::{BTreeMap, BTreeSet};
-use std::fmt::{Debug, Display, Formatter};
+use std::fmt::Debug;
 use std::mem;
 use std::sync::{Arc, Weak};
 
+use derive_more::Display;
 use thiserror::Error;
 
 use crate::hugr::IdentList;
@@ -41,15 +43,10 @@ pub use type_def::{TypeDef, TypeDefBound};
 pub mod declarative;
 
 /// Extension Registries store extensions to be looked up e.g. during validation.
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, Display, Default, PartialEq)]
+#[display("ExtensionRegistry[{}]", _0.keys().join(", "))]
 pub struct ExtensionRegistry(BTreeMap<ExtensionId, Arc<Extension>>);
 
-impl Default for ExtensionRegistry {
-    fn default() -> Self {
-        Self(BTreeMap::new())
-    }
-}
-
 impl ExtensionRegistry {
     /// Create a new empty extension registry.
     ///
@@ -264,7 +261,7 @@ pub enum SignatureError {
     #[error("Invalid type arguments for operation")]
     InvalidTypeArgs,
     /// The Extension Registry did not contain an Extension referenced by the Signature
-    #[error("Extension '{missing}' not found. Available extensions: {}",
+    #[error("Extension '{missing}' is not part of the declared HUGR extensions [{}]",
         available.iter().map(|e| e.to_string()).collect::<Vec<_>>().join(", ")
     )]
     ExtensionNotFound {
@@ -652,7 +649,10 @@ pub enum ExtensionBuildError {
 }
 
 /// A set of extensions identified by their unique [`ExtensionId`].
-#[derive(Clone, Debug, Default, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+#[derive(
+    Clone, Debug, Display, Default, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize,
+)]
+#[display("[{}]", _0.iter().join(", "))]
 pub struct ExtensionSet(BTreeSet<ExtensionId>);
 
 /// A special ExtensionId which indicates that the delta of a non-Function
@@ -776,12 +776,6 @@ fn as_typevar(e: &ExtensionId) -> Option<usize> {
     }
 }
 
-impl Display for ExtensionSet {
-    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        f.debug_list().entries(self.0.iter()).finish()
-    }
-}
-
 impl FromIterator<ExtensionId> for ExtensionSet {
     fn from_iter<I: IntoIterator<Item = ExtensionId>>(iter: I) -> Self {
         Self(BTreeSet::from_iter(iter))

hugr-core/src/extension/prelude.rs

@@ -959,7 +959,6 @@ impl MakeRegisteredOp for Lift {
 mod test {
     use crate::builder::inout_sig;
     use crate::std_extensions::arithmetic::float_types::{float64_type, ConstF64};
-    use crate::utils::test_quantum_extension;
     use crate::{
         builder::{endo_sig, DFGBuilder, Dataflow, DataflowHugr},
         utils::test_quantum_extension::cx_gate,

hugr-core/src/extension/resolution.rs

@@ -1,5 +1,8 @@
 //! Utilities for resolving operations and types present in a HUGR, and updating
-//! the list of used extensions. See [`crate::Hugr::resolve_extension_defs`].
+//! the list of used extensions. The functionalities of this module can be
+//! called from the type methods [`crate::Hugr::resolve_extension_defs`],
+//! [`crate::ops::OpType::used_extensions`], and
+//! [`crate::types::Signature::used_extensions`].
 //!
 //! When listing "used extensions" we only care about _definitional_ extension
 //! requirements, i.e., the operations and types that are required to define the
@@ -13,21 +16,23 @@
 //! Note: These procedures are only temporary until `hugr-model` is stabilized.
 //! Once that happens, hugrs will no longer be directly deserialized using serde
 //! but instead will be created by the methods in `crate::import`. As these
-//! (will) automatically resolve extensions as the operations are created,
-//! we will no longer require this post-facto resolution step.
+//! (will) automatically resolve extensions as the operations are created, we
+//! will no longer require this post-facto resolution step.
 
 mod ops;
 mod types;
+mod types_mut;
 
-pub(crate) use ops::update_op_extensions;
-pub(crate) use types::update_op_types_extensions;
+pub(crate) use ops::{collect_op_extensions, update_op_extensions};
+pub(crate) use types::{collect_op_types_extensions, collect_signature_exts};
+pub(crate) use types_mut::update_op_types_extensions;
 
 use derive_more::{Display, Error, From};
 
 use super::{Extension, ExtensionId, ExtensionRegistry};
 use crate::ops::custom::OpaqueOpError;
 use crate::ops::{NamedOp, OpName, OpType};
-use crate::types::TypeName;
+use crate::types::{FuncTypeBase, MaybeRV, TypeName};
 use crate::Node;
 
 /// Errors that can occur during extension resolution.
@@ -101,3 +106,60 @@ impl ExtensionResolutionError {
         }
     }
 }
+
+/// Errors that can occur when collecting extension requirements.
+#[derive(Debug, Display, Clone, Error, From, PartialEq)]
+#[non_exhaustive]
+pub enum ExtensionCollectionError {
+    /// An operation requires an extension that is not in the given registry.
+    #[display(
+        "{op}{} contains custom types for which have lost the reference to their defining extensions. Dropped extensions: {}",
+        if let Some(node) = node { format!(" ({})", node) } else { "".to_string() },
+        missing_extensions.join(", ")
+    )]
+    DroppedOpExtensions {
+        /// The node that is missing extensions.
+        node: Option<Node>,
+        /// The operation that is missing extensions.
+        op: OpName,
+        /// The missing extensions.
+        missing_extensions: Vec<ExtensionId>,
+    },
+    /// A signature requires an extension that is not in the given registry.
+    #[display(
+        "Signature {signature} contains custom types for which have lost the reference to their defining extensions. Dropped extensions: {}",
+        missing_extensions.join(", ")
+    )]
+    DroppedSignatureExtensions {
+        /// The signature that is missing extensions.
+        signature: String,
+        /// The missing extensions.
+        missing_extensions: Vec<ExtensionId>,
+    },
+}
+
+impl ExtensionCollectionError {
+    /// Create a new error when operation extensions have been dropped.
+    pub fn dropped_op_extension(
+        node: Option<Node>,
+        op: &OpType,
+        missing_extension: impl IntoIterator<Item = ExtensionId>,
+    ) -> Self {
+        Self::DroppedOpExtensions {
+            node,
+            op: NamedOp::name(op),
+            missing_extensions: missing_extension.into_iter().collect(),
+        }
+    }
+
+    /// Create a new error when signature extensions have been dropped.
+    pub fn dropped_signature<RV: MaybeRV>(
+        signature: &FuncTypeBase<RV>,
+        missing_extension: impl IntoIterator<Item = ExtensionId>,
+    ) -> Self {
+        Self::DroppedSignatureExtensions {
+            signature: format!("{signature}"),
+            missing_extensions: missing_extension.into_iter().collect(),
+        }
+    }
+}

hugr-core/src/extension/resolution/ops.rs

@@ -1,12 +1,48 @@
-//! Resolve `OpaqueOp`s into `ExtensionOp`s and return an operation's required extension.
+//! Resolve `OpaqueOp`s into `ExtensionOp`s and return an operation's required
+//! extension.
+//!
+//! Contains both mutable ([`update_op_extensions`]) and immutable
+//! ([`collect_operation_extension`]) methods to resolve operations and collect
+//! the required extensions respectively.
 
 use std::sync::Arc;
 
-use super::{Extension, ExtensionRegistry, ExtensionResolutionError};
+use super::{Extension, ExtensionCollectionError, ExtensionRegistry, ExtensionResolutionError};
 use crate::ops::custom::OpaqueOpError;
 use crate::ops::{DataflowOpTrait, ExtensionOp, NamedOp, OpType};
 use crate::Node;
 
+/// Returns the extension in the registry required by the operation.
+///
+/// If the operation does not require an extension, returns `None`.
+///
+/// [`ExtensionOp`]s store a [`Weak`] reference to their extension, which can be
+/// invalidated if the original `Arc<Extension>` is dropped. On such cases, we
+/// return an error with the missing extension names.
+///
+/// # Attributes
+///
+/// - `node`: The node where the operation is located, if available. This is
+///   used to provide context in the error message.
+/// - `op`: The operation to collect the extensions from.
+pub(crate) fn collect_op_extensions(
+    node: Option<Node>,
+    op: &OpType,
+) -> Result<Option<Arc<Extension>>, ExtensionCollectionError> {
+    let OpType::ExtensionOp(ext_op) = op else {
+        return Ok(None);
+    };
+    let ext = ext_op.def().extension();
+    match ext.upgrade() {
+        Some(e) => Ok(Some(e)),
+        None => Err(ExtensionCollectionError::dropped_op_extension(
+            node,
+            op,
+            [ext_op.def().extension_id().clone()],
+        )),
+    }
+}
+
 /// Compute the required extension for an operation.
 ///
 /// If the op is a [`OpType::OpaqueOp`], replace it with a resolved