fix: breaking changes from extension weak references

tket2-hseries/src/extension/futures.rs

@@ -3,7 +3,7 @@
 //! `Future<t>` is a linear type representing a value that will be available in
 //! the future.  It can be consumed by `Read`, returning a `t`.  It can be
 //! duplicated by `Dup`, and discarded with `Free`.
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use hugr::{
     builder::{BuildError, Dataflow},
@@ -32,11 +32,11 @@ pub const EXTENSION_VERSION: Version = Version::new(0, 1, 0);
 lazy_static! {
     /// The "tket2.futures" extension.
     pub static ref EXTENSION: Arc<Extension>  = {
-        let mut ext = Extension::new(EXTENSION_ID, EXTENSION_VERSION);
-        let _ = add_future_type_def(&mut ext).unwrap();
+        Extension::new_arc(EXTENSION_ID, EXTENSION_VERSION, |ext, ext_ref| {
+            let _ = add_future_type_def(ext, ext_ref.clone()).unwrap();
 
-        FutureOpDef::load_all_ops(&mut ext).unwrap();
-        Arc::new(ext)
+            FutureOpDef::load_all_ops( ext, ext_ref).unwrap();
+        })
     };
 
     /// Extension registry including the "tket2.futures" extension.
@@ -48,12 +48,16 @@ lazy_static! {
     pub static ref FUTURE_TYPE_NAME: SmolStr = SmolStr::new_inline("Future");
 }
 
-fn add_future_type_def(ext: &mut Extension) -> Result<&TypeDef, ExtensionBuildError> {
+fn add_future_type_def(
+    ext: &mut Extension,
+    extension_ref: Weak<Extension>,
+) -> Result<&TypeDef, ExtensionBuildError> {
     ext.add_type(
         FUTURE_TYPE_NAME.to_owned(),
         vec![TypeBound::Any.into()],
         "A value that is computed asynchronously".into(),
         TypeBound::Any.into(),
+        &extension_ref,
     )
 }
 
@@ -121,7 +125,7 @@ impl MakeOpDef for FutureOpDef {
     }
 
     fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
-        try_from_name(op_def.name(), op_def.extension())
+        try_from_name(op_def.name(), op_def.extension_id())
     }
 
     fn description(&self) -> String {

tket2-hseries/src/extension/hseries.rs

@@ -42,13 +42,14 @@ pub const EXTENSION_VERSION: Version = Version::new(0, 1, 0);
 lazy_static! {
     /// The "tket2.hseries" extension.
     pub static ref EXTENSION: Arc<Extension> = {
-        let mut ext = Extension::new(EXTENSION_ID, EXTENSION_VERSION).with_reqs(ExtensionSet::from_iter([
-            futures::EXTENSION.name(),
-            PRELUDE.name(),
-            FLOAT_TYPES.name(),
-        ].into_iter().cloned()));
-        HSeriesOp::load_all_ops(&mut ext).unwrap();
-        Arc::new(ext)
+         Extension::new_arc(EXTENSION_ID, EXTENSION_VERSION, |ext, ext_ref| {
+            ext.add_requirements(ExtensionSet::from_iter([
+                futures::EXTENSION.name(),
+                PRELUDE.name(),
+                FLOAT_TYPES.name(),
+            ].into_iter().cloned()));
+            HSeriesOp::load_all_ops( ext, ext_ref).unwrap();
+        })
     };
 
     /// Extension registry including the "tket2.hseries" extension and
@@ -110,7 +111,7 @@ impl MakeOpDef for HSeriesOp {
     }
 
     fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
-        try_from_name(op_def.name(), op_def.extension())
+        try_from_name(op_def.name(), op_def.extension_id())
     }
 
     fn extension(&self) -> ExtensionId {

tket2-hseries/src/extension/hseries/lower.rs

@@ -156,7 +156,7 @@ pub fn check_lowered(hugr: &impl HugrView) -> Result<(), Vec<Node>> {
         .filter_map(|node| {
             let optype = hugr.get_optype(node);
             optype.as_extension_op().and_then(|ext| {
-                (ext.def().extension() == &tket2::extension::TKET2_EXTENSION_ID).then_some(node)
+                (ext.def().extension_id() == &tket2::extension::TKET2_EXTENSION_ID).then_some(node)
             })
         })
         .collect();

tket2-hseries/src/extension/result.rs

@@ -38,9 +38,10 @@ pub const EXTENSION_VERSION: Version = Version::new(0, 1, 0);
 lazy_static! {
     /// The "tket2.result" extension.
     pub static ref EXTENSION: Arc<Extension> = {
-        let mut ext = Extension::new(EXTENSION_ID, EXTENSION_VERSION).with_reqs(ExtensionSet::from_iter([INT_EXTENSION_ID, FLOAT_EXTENSION_ID]));
-        ResultOpDef::load_all_ops(&mut ext).unwrap();
-        Arc::new(ext)
+        Extension::new_arc(EXTENSION_ID, EXTENSION_VERSION, |ext, ext_ref| {
+            ext.add_requirements(ExtensionSet::from_iter([INT_EXTENSION_ID, FLOAT_EXTENSION_ID]));
+            ResultOpDef::load_all_ops(ext, ext_ref).unwrap();
+        })
     };
 
     /// Extension registry including the "tket2.result" extension and
@@ -196,7 +197,7 @@ impl MakeOpDef for ResultOpDef {
     }
 
     fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
-        try_from_name(op_def.name(), op_def.extension())
+        try_from_name(op_def.name(), op_def.extension_id())
     }
 
     fn extension(&self) -> ExtensionId {

tket2/src/extension.rs

@@ -43,16 +43,14 @@ pub static ref TKET1_OP_PAYLOAD : CustomType =
 
 /// The TKET1 extension, containing the opaque TKET1 operations.
 pub static ref TKET1_EXTENSION: Arc<Extension>  = {
-    let mut res = Extension::new(TKET1_EXTENSION_ID, TKET1_EXTENSION_VERSION);
-
-    let tket1_op_payload = TypeParam::String;
-    res.add_op(
-        TKET1_OP_NAME,
-        "An opaque TKET1 operation.".into(),
-        Tk1Signature([tket1_op_payload])
-    ).unwrap();
-
-    Arc::new(res)
+    Extension::new_arc(TKET1_EXTENSION_ID, TKET1_EXTENSION_VERSION, |res, ext_ref| {
+        res.add_op(
+            TKET1_OP_NAME,
+            "An opaque TKET1 operation.".into(),
+            Tk1Signature([TypeParam::String]),
+            ext_ref
+        ).unwrap();
+    })
 };
 
 /// Extension registry including the prelude, std, TKET1, and Tk2Ops extensions.
@@ -95,11 +93,11 @@ pub const TKET2_EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("tket2.qu
 pub const TKET2_EXTENSION_VERSION: Version = Version::new(0, 1, 0);
 
 lazy_static! {
-/// The extension definition for TKET2 ops and types.
-pub static ref TKET2_EXTENSION: Arc<Extension> = {
-    let mut e = Extension::new(TKET2_EXTENSION_ID, TKET2_EXTENSION_VERSION);
-    Tk2Op::load_all_ops(&mut e).expect("add fail");
-    SympyOpDef.add_to_extension(&mut e).unwrap();
-   Arc::new(e)
-};
+    /// The extension definition for TKET2 ops and types.
+    pub static ref TKET2_EXTENSION: Arc<Extension> = {
+        Extension::new_arc(TKET2_EXTENSION_ID, TKET2_EXTENSION_VERSION, |res, ext_ref| {
+            Tk2Op::load_all_ops(res, ext_ref).expect("add_fail");
+            SympyOpDef.add_to_extension(res, ext_ref).unwrap();
+        })
+    };
 }

tket2/src/extension/rotation.rs

@@ -10,7 +10,7 @@ use hugr::{
 };
 use smol_str::SmolStr;
 use std::f64::consts::PI;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use strum::{EnumIter, EnumString, IntoStaticStr};
 
 use lazy_static::lazy_static;
@@ -23,11 +23,11 @@ pub const ROTATION_EXTENSION_VERSION: Version = Version::new(0, 1, 0);
 
 lazy_static! {
     /// The extension definition for TKET2 rotation type and ops.
-    pub static ref ROTATION_EXTENSION: Arc<Extension> = {
-        let mut e = Extension::new(ROTATION_EXTENSION_ID, ROTATION_EXTENSION_VERSION);
-        add_to_extension(&mut e);
-        Arc::new(e)
-    };
+    pub static ref ROTATION_EXTENSION: Arc<Extension> =  {
+            Extension::new_arc(ROTATION_EXTENSION_ID, ROTATION_EXTENSION_VERSION, |e, extension_ref| {
+                add_to_extension(e, extension_ref);
+            }
+    )};
 }
 
 /// Identifier for the rotation type.
@@ -132,7 +132,7 @@ impl MakeOpDef for RotationOp {
     where
         Self: Sized,
     {
-        hugr::extension::simple_op::try_from_name(op_def.name(), op_def.extension())
+        hugr::extension::simple_op::try_from_name(op_def.name(), op_def.extension_id())
     }
 
     fn signature(&self) -> hugr::extension::SignatureFunc {
@@ -189,17 +189,18 @@ impl MakeRegisteredOp for RotationOp {
     }
 }
 
-pub(super) fn add_to_extension(extension: &mut Extension) {
+pub(super) fn add_to_extension(extension: &mut Extension, extension_ref: &Weak<Extension>) {
     extension
         .add_type(
             ROTATION_TYPE_ID,
             vec![],
             "rotation type expressed as number of half turns".to_owned(),
             TypeBound::Copyable.into(),
+            extension_ref,
         )
         .unwrap();
 
-    RotationOp::load_all_ops(extension).expect("add fail");
+    RotationOp::load_all_ops(extension, extension_ref).expect("add fail");
 }
 
 /// An extension trait for [Dataflow] providing methods to add

tket2/src/extension/sympy.rs

@@ -64,7 +64,7 @@ impl MakeOpDef for SympyOpDef {
     where
         Self: Sized,
     {
-        try_from_name(op_def.name(), op_def.extension())
+        try_from_name(op_def.name(), op_def.extension_id())
     }
 
     fn signature(&self) -> hugr::extension::SignatureFunc {

tket2/src/ops.rs

@@ -133,7 +133,7 @@ impl MakeOpDef for Tk2Op {
     }
 
     fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
-        try_from_name(op_def.name(), op_def.extension())
+        try_from_name(op_def.name(), op_def.extension_id())
     }
 }
 
@@ -190,7 +190,7 @@ pub(crate) fn match_symb_const_op(op: &OpType) -> Option<String> {
     };
 
     if let OpType::ExtensionOp(e) = op {
-        if e.def().name() == &SYM_OP_ID && e.def().extension() == &EXTENSION_ID {
+        if e.def().name() == &SYM_OP_ID && e.def().extension_id() == &EXTENSION_ID {
             Some(symbol_from_typeargs(e.args()))
         } else {
             None
@@ -276,7 +276,7 @@ pub(crate) mod test {
 
             let ext_op = op.into_extension_op();
             assert_eq!(ext_op.args(), &[]);
-            assert_eq!(ext_op.def().extension(), &EXTENSION_ID);
+            assert_eq!(ext_op.def().extension_id(), &EXTENSION_ID);
             let name = ext_op.def().name();
             assert_eq!(Tk2Op::from_str(name), Ok(op));
         }