feat: Arbitrary types in the yaml extension definition (#839)

The current yaml extension definitions only support hard-coded types in
the operation signatures.

This PR leaves the existing `Q` and `USize` hardcoded values, but tries
to dereference the type name from the extensions in scopes otherwise.

drive-by: Extend the module doc with an example on how to load the
extension.

examples/extension/declarative.yaml

@@ -7,7 +7,7 @@ extensions:
     types:
       - # Types must have a name.
         # Parameters are not currently supported.
-        name: Copyable type
+        name: CopyableType
         description: A simple type with no parameters
         # Types may have a "Eq", "Copyable", or "Any" bound.
         # This field is optional and defaults to "Any".
@@ -29,3 +29,8 @@ extensions:
               [Q, 1]
             - # Or as a description, followed by the type and a number of repetitions.
               [Control, Q, 2]
+      - name: YetAnotherOperation
+        description: An operation that uses the declared Custom type
+        signature:
+          inputs: [CopyableType]
+          outputs: [CopyableType, CopyableType]

src/extension/declarative.rs

@@ -12,6 +12,16 @@
 #![doc = include_str!("../../examples/extension/declarative.yaml")]
 //! ```
 //!
+//! The definition can be loaded into a registry using the [`load_extensions`] or [`load_extensions_file`] functions.
+//!
+//! ```rust
+//! # const DECLARATIVE_YAML: &str = include_str!("../../examples/extension/declarative.yaml");
+//! # use hugr::extension::declarative::load_extensions;
+//! // Required extensions must already be present in the registry.
+//! let mut reg = hugr::std_extensions::logic::LOGIC_REG.clone();
+//! load_extensions(DECLARATIVE_YAML, &mut reg).unwrap();
+//! ```
+//!
 //! [specification]: https://github.com/CQCL/hugr/blob/main/specification/hugr.md#declarative-format
 
 mod ops;
@@ -290,8 +300,8 @@ extensions:
   - name: AnotherOperation
     description: An operation from 3 qubits to 3 qubits
     signature:
-        inputs: [Q, Q, Q]
-        outputs: [[Q, 1], [Control, Q, 2]]
+        inputs: [MyType, Q, Q]
+        outputs: [[MyType, 1], [Control, Q, 2]]
 "#;
 
     /// A yaml extension with unsupported features.
@@ -347,7 +357,7 @@ extensions:
 
     #[cfg_attr(miri, ignore)] // Opening files is not supported in (isolated) miri
     #[rstest]
-    #[case(EXAMPLE_YAML_FILE, 1, 1, 2, &std_extensions::logic::LOGIC_REG)]
+    #[case(EXAMPLE_YAML_FILE, 1, 1, 3, &std_extensions::logic::LOGIC_REG)]
     fn test_decode_file(
         #[case] yaml_file: &str,
         #[case] num_declarations: usize,

src/extension/declarative/signature.rs

@@ -174,30 +174,53 @@ impl TypeDeclaration {
         ctx: DeclarationContext<'_>,
         _op_params: &[TypeParam],
     ) -> Result<CustomType, ExtensionDeclarationError> {
+        let Some(type_def) = self.resolve_type(ext, ctx) else {
+            return Err(ExtensionDeclarationError::UnknownType {
+                ext: ext.name().clone(),
+                ty: self.0.clone(),
+            });
+        };
+
+        // The hard-coded types are not parametric.
+        assert!(type_def.params().is_empty());
+        let op = type_def.instantiate(&[]).unwrap();
+
+        Ok(op)
+    }
+
+    /// Resolve a type name to a type definition.
+    fn resolve_type<'a>(
+        &'a self,
+        ext: &'a Extension,
+        ctx: DeclarationContext<'a>,
+    ) -> Option<&'a TypeDef> {
         // The prelude is always in scope.
         debug_assert!(ctx.scope.contains(&PRELUDE_ID));
 
-        // Only hard-coded prelude types are supported for now.
+        // Some hard-coded prelude types are supported.
         let prelude = ctx.registry.get(&PRELUDE_ID).unwrap();
-        let op_def: &TypeDef = match self.0.as_str() {
-            "USize" => prelude.get_type("usize"),
-            "Q" => prelude.get_type("qubit"),
-            _ => {
-                return Err(ExtensionDeclarationError::UnknownType {
-                    ext: ext.name().clone(),
-                    ty: self.0.clone(),
-                })
-            }
+        match self.0.as_str() {
+            "USize" => return prelude.get_type("usize"),
+            "Q" => return prelude.get_type("qubit"),
+            _ => {}
         }
-        .ok_or(ExtensionDeclarationError::UnknownType {
-            ext: ext.name().clone(),
-            ty: self.0.clone(),
-        })?;
 
-        // The hard-coded types are not parametric.
-        assert!(op_def.params().is_empty());
-        let op = op_def.instantiate(&[]).unwrap();
+        // Try to resolve the type in the current extension.
+        if let Some(ty) = ext.get_type(self.0.as_str()) {
+            return Some(ty);
+        }
 
-        Ok(op)
+        // Try to resolve the type in the other extensions in scope.
+        for ext in ctx.scope.iter() {
+            if let Some(ty) = ctx
+                .registry
+                .get(ext)
+                .and_then(|ext| ext.get_type(self.0.as_str()))
+            {
+                return Some(ty);
+            }
+        }
+
+        None
     }
 }