From 5786561786b3c9f1b04645fc2a6a88f03223f06e Mon Sep 17 00:00:00 2001
From: Lukas Heidemann <lukas@heidemann.me>
Date: Tue, 21 Jan 2025 11:09:23 +0000
Subject: [PATCH] Bugfix: determine correct bounds of custom types

---
 hugr-core/src/import.rs                       | 24 +++++++++++++++----
 .../model__roundtrip_constraints.snap         | 19 ++++++++-------
 .../tests/fixtures/model-constraints.edn      |  7 +++---
 3 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/hugr-core/src/import.rs b/hugr-core/src/import.rs
index f542ac695..62de95146 100644
--- a/hugr-core/src/import.rs
+++ b/hugr-core/src/import.rs
@@ -18,7 +18,7 @@ use crate::{
     types::{
         type_param::TypeParam, type_row::TypeRowBase, CustomType, FuncTypeBase, MaybeRV,
         PolyFuncType, PolyFuncTypeBase, RowVariable, Signature, Type, TypeArg, TypeBase, TypeBound,
-        TypeEnum, TypeRow,
+        TypeEnum, TypeName, TypeRow,
     },
     Direction, Hugr, HugrView, Node, Port,
 };
@@ -52,6 +52,14 @@ pub enum ImportError {
         /// The available extensions in the registry.
         available: Vec<ExtensionId>,
     },
+    /// An extension type is missing.
+    #[error("Importing the hugr requires extension {ext} to have a type named {name}, but it was not found.")]
+    ExtensionType {
+        /// The extension that is missing the type.
+        ext: ExtensionId,
+        /// The name of the missing type.
+        name: TypeName,
+    },
     /// The model is not well-formed.
     #[error("validate error: {0}")]
     Model(#[from] model::ModelError),
@@ -1088,13 +1096,21 @@ impl<'a> Context<'a> {
                             available: self.extensions.ids().cloned().collect(),
                         })?;
 
+                let ext_type =
+                    extension_ref
+                        .get_type(&id)
+                        .ok_or_else(|| ImportError::ExtensionType {
+                            ext: extension.clone(),
+                            name: id.clone(),
+                        })?;
+
+                let bound = ext_type.bound(&args);
+
                 Ok(TypeBase::new_extension(CustomType::new(
                     id,
                     args,
                     extension,
-                    // As part of the migration from `TypeBound`s to constraints, we pretend that all
-                    // `TypeBound`s are copyable.
-                    TypeBound::Copyable,
+                    bound,
                     &Arc::downgrade(extension_ref),
                 )))
             }
diff --git a/hugr-core/tests/snapshots/model__roundtrip_constraints.snap b/hugr-core/tests/snapshots/model__roundtrip_constraints.snap
index a6c6a34f4..abac5878f 100644
--- a/hugr-core/tests/snapshots/model__roundtrip_constraints.snap
+++ b/hugr-core/tests/snapshots/model__roundtrip_constraints.snap
@@ -4,18 +4,21 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-cons
 ---
 (hugr 0)
 
-(import prelude.Array)
+(import collections.array.array)
 
 (declare-func array.replicate
-  (forall ?0 type)
-  (forall ?1 nat)
-  (where (nonlinear ?0))
-  [?0] [(@ prelude.Array ?0 ?1)] (ext))
+  (forall ?0 nat)
+  (forall ?1 type)
+  (where (nonlinear ?1))
+  [?1] [(@ collections.array.array ?0 ?1)] (ext))
 
 (declare-func array.copy
-  (forall ?0 type)
-  (where (nonlinear ?0))
-  [(@ prelude.Array ?0)] [(@ prelude.Array ?0) (@ prelude.Array ?0)] (ext))
+  (forall ?0 nat)
+  (forall ?1 type)
+  (where (nonlinear ?1))
+  [(@ collections.array.array ?0 ?1)]
+  [(@ collections.array.array ?0 ?1) (@ collections.array.array ?0 ?1)]
+  (ext))
 
 (define-func util.copy
   (forall ?0 type)
diff --git a/hugr-model/tests/fixtures/model-constraints.edn b/hugr-model/tests/fixtures/model-constraints.edn
index 8987e4c69..c2d58346b 100644
--- a/hugr-model/tests/fixtures/model-constraints.edn
+++ b/hugr-model/tests/fixtures/model-constraints.edn
@@ -1,16 +1,17 @@
 (hugr 0)
 
 (declare-func array.replicate
-  (forall ?t type)
   (forall ?n nat)
+  (forall ?t type)
   (where (nonlinear ?t))
-  [?t] [(@ prelude.Array ?t ?n)]
+  [?t] [(@ collections.array.array ?n ?t)]
   (ext))
 
 (declare-func array.copy
+  (forall ?n nat)
   (forall ?t type)
   (where (nonlinear ?t))
-  [(@ prelude.Array ?t)] [(@ prelude.Array ?t) (@ prelude.Array ?t)] (ext))
+  [(@ collections.array.array ?n ?t)] [(@ collections.array.array ?n ?t) (@ collections.array.array ?n ?t)] (ext))
 
 (define-func util.copy
   (forall ?t type)