feat: Hugr::store_with_exts and auto-include in serde_as

hugr-core/src/envelope/serde_with.rs

@@ -16,7 +16,9 @@ use crate::std_extensions::STD_REG;
 /// storing it as a string.
 ///
 /// Note that only PRELUDE extensions are used to decode the package's content.
-/// Additional extensions should be included in the serialized envelope.
+/// When serializing a HUGR, any additional extensions required to load it are
+/// embedded in the envelope. Packages should manually add any required
+/// extensions before serializing.
 ///
 /// # Examples
 ///
@@ -186,8 +188,20 @@ macro_rules! impl_serde_as_string_envelope {
             where
                 S: serde::Serializer,
             {
+                // Include any additional extension required to load the HUGR in the envelope.
+                let extensions: &$crate::extension::ExtensionRegistry = $extension_reg;
+                let mut extra_extensions = $crate::extension::ExtensionRegistry::default();
+                for ext in $crate::hugr::views::HugrView::extensions(source).iter() {
+                    if !extensions.contains(ext.name()) {
+                        extra_extensions.register_updated(ext.clone());
+                    }
+                }
+
                 let str = source
-                    .store_str($crate::envelope::EnvelopeConfig::text())
+                    .store_str_with_exts(
+                        $crate::envelope::EnvelopeConfig::text(),
+                        &extra_extensions,
+                    )
                     .map_err(serde::ser::Error::custom)?;
                 serializer.collect_str(&str)
             }

hugr-core/src/hugr.rs

@@ -32,7 +32,7 @@ use crate::extension::resolution::{
     ExtensionResolutionError, WeakExtensionRegistry, resolve_op_extensions,
     resolve_op_types_extensions,
 };
-use crate::extension::{ExtensionRegistry, ExtensionSet};
+use crate::extension::{EMPTY_REG, ExtensionRegistry, ExtensionSet};
 use crate::ops::{self, Module, NamedOp, OpName, OpTag, OpTrait};
 pub use crate::ops::{DEFAULT_OPTYPE, OpType};
 use crate::package::Package;
@@ -150,7 +150,13 @@ impl Hugr {
         self.graph.reserve(nodes, ports);
     }
 
-    /// Read a Package from a HUGR envelope.
+    /// Read a HUGR from an Envelope.
+    ///
+    /// To load a HUGR, all the extensions used in its definition must be
+    /// available. The Envelope may include some of the extensions, but any
+    /// additional extensions must be provided in the `extensions` parameter. If
+    /// `extensions` is `None`, the default [`crate::std_extensions::STD_REG`]
+    /// is used.
     pub fn load(
         reader: impl io::BufRead,
         extensions: Option<&ExtensionRegistry>,
@@ -162,40 +168,87 @@ impl Hugr {
         }
     }
 
-    /// Read a Package from a HUGR envelope encoded in a string.
+    /// Read a HUGR from an Envelope encoded in a string.
+    ///
+    /// Note that not all Envelopes are valid strings. In the general case,
+    /// it is recommended to use [`Hugr::load`] with a bytearray instead.
     ///
-    /// Note that not all envelopes are valid strings. In the general case,
-    /// it is recommended to use `Package::load` with a bytearray instead.
+    /// To load a HUGR, all the extensions used in its definition must be
+    /// available. The Envelope may include some of the extensions, but any
+    /// additional extensions must be provided in the `extensions` parameter. If
+    /// `extensions` is `None`, the default [`crate::std_extensions::STD_REG`]
+    /// is used.
     pub fn load_str(
         envelope: impl AsRef<str>,
         extensions: Option<&ExtensionRegistry>,
     ) -> Result<Self, EnvelopeError> {
         Self::load(envelope.as_ref().as_bytes(), extensions)
     }
 
-    /// Store the Package in a HUGR envelope.
+    /// Store the HUGR in an Envelope.
+    ///
+    /// The Envelope will not include any extension definition, and will require
+    /// an adequate [`ExtensionRegistry`] to be loaded (see [`Hugr::load`]).
+    /// Use [`Hugr::store_with_exts`] to include additional extensions in the
+    /// Envelope.
     pub fn store(
         &self,
         writer: impl io::Write,
         config: EnvelopeConfig,
     ) -> Result<(), EnvelopeError> {
-        envelope::write_envelope_impl(writer, [self], &self.extensions, config)
+        self.store_with_exts(writer, config, &EMPTY_REG)
+    }
+
+    /// Store the HUGR in an Envelope.
+    ///
+    /// The Envelope will embed the definitions of the extensions in the
+    /// `extensions` registry. Any other extension used in the HUGR definition
+    /// must be passed to [`Hugr::load`] to load back the HUGR.
+    pub fn store_with_exts(
+        &self,
+        writer: impl io::Write,
+        config: EnvelopeConfig,
+        extensions: &ExtensionRegistry,
+    ) -> Result<(), EnvelopeError> {
+        envelope::write_envelope_impl(writer, [self], extensions, config)
     }
 
-    /// Store the Package in a HUGR envelope encoded in a string.
+    /// Store the HUGR in an Envelope encoded in a string.
     ///
-    /// Note that not all envelopes are valid strings. In the general case,
-    /// it is recommended to use `Package::store` with a bytearray instead.
+    /// Note that not all Envelopes are valid strings. In the general case,
+    /// it is recommended to use [`Hugr::store`] with a bytearray instead.
     /// See [`EnvelopeFormat::ascii_printable`][crate::envelope::EnvelopeFormat::ascii_printable].
+    ///
+    /// The Envelope will not include any extension definition, and will require
+    /// an adequate [`ExtensionRegistry`] to be loaded (see [`Hugr::load_str`]).
+    /// Use [`Hugr::store_str_with_exts`] to include additional extensions in the
+    /// Envelope.
     pub fn store_str(&self, config: EnvelopeConfig) -> Result<String, EnvelopeError> {
+        self.store_str_with_exts(config, &EMPTY_REG)
+    }
+
+    /// Store the HUGR in an Envelope encoded in a string.
+    ///
+    /// Note that not all Envelopes are valid strings. In the general case,
+    /// it is recommended to use [`Hugr::store_str`] with a bytearray instead.
+    /// See [`EnvelopeFormat::ascii_printable`][crate::envelope::EnvelopeFormat::ascii_printable].
+    ///
+    /// The Envelope will embed the definitions of the extensions in the
+    /// `extensions` registry. Any other extension used in the HUGR definition
+    /// must be passed to [`Hugr::load_str`] to load back the HUGR.
+    pub fn store_str_with_exts(
+        &self,
+        config: EnvelopeConfig,
+        extensions: &ExtensionRegistry,
+    ) -> Result<String, EnvelopeError> {
         if !config.format.ascii_printable() {
             return Err(EnvelopeError::NonASCIIFormat {
                 format: config.format,
             });
         }
 
         let mut buf = Vec::new();
-        self.store(&mut buf, config)?;
+        self.store_with_exts(&mut buf, config, extensions)?;
         Ok(String::from_utf8(buf).expect("Envelope is valid utf8"))
     }
 

hugr-core/src/package.rs

@@ -57,6 +57,12 @@ impl Package {
     }
 
     /// Read a Package from a HUGR envelope.
+    ///
+    /// To load a Package, all the extensions used in its definition must be
+    /// available. The Envelope may include some of the extensions, but any
+    /// additional extensions must be provided in the `extensions` parameter. If
+    /// `extensions` is `None`, the default [`crate::std_extensions::STD_REG`]
+    /// is used.
     pub fn load(
         reader: impl io::BufRead,
         extensions: Option<&ExtensionRegistry>,
@@ -70,6 +76,12 @@ impl Package {
     ///
     /// Note that not all envelopes are valid strings. In the general case,
     /// it is recommended to use `Package::load` with a bytearray instead.
+    ///
+    /// To load a Package, all the extensions used in its definition must be
+    /// available. The Envelope may include some of the extensions, but any
+    /// additional extensions must be provided in the `extensions` parameter. If
+    /// `extensions` is `None`, the default [`crate::std_extensions::STD_REG`]
+    /// is used.
     pub fn load_str(
         envelope: impl AsRef<str>,
         extensions: Option<&ExtensionRegistry>,
@@ -78,6 +90,10 @@ impl Package {
     }
 
     /// Store the Package in a HUGR envelope.
+    ///
+    /// The Envelope will embed the definitions of the extensions in
+    /// [`Package::extensions`]. Any other extension used in the definition must
+    /// be passed to [`Package::load`] to load back the package.
     pub fn store(
         &self,
         writer: impl io::Write,
@@ -91,6 +107,10 @@ impl Package {
     /// Note that not all envelopes are valid strings. In the general case,
     /// it is recommended to use `Package::store` with a bytearray instead.
     /// See [`EnvelopeFormat::ascii_printable`][crate::envelope::EnvelopeFormat::ascii_printable].
+    ///
+    /// The Envelope will embed the definitions of the extensions in
+    /// [`Package::extensions`]. Any other extension used in the definition must
+    /// be passed to [`Package::load_str`] to load back the package.
     pub fn store_str(&self, config: EnvelopeConfig) -> Result<String, EnvelopeError> {
         if !config.format.ascii_printable() {
             return Err(EnvelopeError::NonASCIIFormat {

hugr/benches/benchmarks/hugr.rs

@@ -9,7 +9,7 @@ use hugr::envelope::{EnvelopeConfig, EnvelopeFormat};
 use hugr::std_extensions::STD_REG;
 use std::hint::black_box;
 
-pub use examples::{circuit, simple_cfg_hugr, simple_dfg_hugr};
+pub use examples::{BENCH_EXTENSIONS, circuit, simple_cfg_hugr, simple_dfg_hugr};
 
 trait Serializer {
     fn serialize(&self, hugr: &Hugr) -> Vec<u8>;
@@ -28,7 +28,7 @@ impl Serializer for JsonSer {
         bytes
     }
     fn deserialize(&self, bytes: &[u8]) -> Hugr {
-        Hugr::load(bytes, None).unwrap()
+        Hugr::load(bytes, Some(&BENCH_EXTENSIONS)).unwrap()
     }
 }
 
@@ -46,7 +46,7 @@ impl Serializer for CapnpSer {
     }
 
     fn deserialize(&self, bytes: &[u8]) -> Hugr {
-        Hugr::load(bytes, None).unwrap()
+        Hugr::load(bytes, Some(&BENCH_EXTENSIONS)).unwrap()
     }
 }
 

hugr/benches/benchmarks/hugr/examples.rs

@@ -1,17 +1,18 @@
 //! Builders and utilities for benchmarks.
 
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use hugr::builder::{
     BuildError, CFGBuilder, Container, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer,
     HugrBuilder, ModuleBuilder,
 };
+use hugr::extension::ExtensionRegistry;
 use hugr::extension::prelude::{bool_t, qb_t, usize_t};
 use hugr::ops::OpName;
+use hugr::std_extensions::STD_REG;
 use hugr::std_extensions::arithmetic::float_types::{ConstF64, float64_type};
 use hugr::types::Signature;
 use hugr::{CircuitUnit, Extension, Hugr, Node, type_row};
-use lazy_static::lazy_static;
 
 pub fn simple_dfg_hugr() -> Hugr {
     let dfg_builder = DFGBuilder::new(Signature::new(vec![bool_t()], vec![bool_t()])).unwrap();
@@ -51,38 +52,42 @@ pub fn simple_cfg_hugr() -> Hugr {
     cfg_builder.finish_hugr().unwrap()
 }
 
-lazy_static! {
-    static ref QUANTUM_EXT: Arc<Extension> = {
-        Extension::new_arc(
-            "bench.quantum".try_into().unwrap(),
-            hugr::extension::Version::new(0, 0, 0),
-            |ext, extension_ref| {
-                ext.add_op(
-                    OpName::new_inline("H"),
-                    String::new(),
-                    Signature::new_endo(qb_t()),
-                    extension_ref,
-                )
-                .unwrap();
-                ext.add_op(
-                    OpName::new_inline("Rz"),
-                    String::new(),
-                    Signature::new(vec![qb_t(), float64_type()], vec![qb_t()]),
-                    extension_ref,
-                )
-                .unwrap();
-
-                ext.add_op(
-                    OpName::new_inline("CX"),
-                    String::new(),
-                    Signature::new_endo(vec![qb_t(), qb_t()]),
-                    extension_ref,
-                )
-                .unwrap();
-            },
-        )
-    };
-}
+pub static QUANTUM_EXT: LazyLock<Arc<Extension>> = LazyLock::new(|| {
+    Extension::new_arc(
+        "bench.quantum".try_into().unwrap(),
+        hugr::extension::Version::new(0, 0, 0),
+        |ext, extension_ref| {
+            ext.add_op(
+                OpName::new_inline("H"),
+                String::new(),
+                Signature::new_endo(qb_t()),
+                extension_ref,
+            )
+            .unwrap();
+            ext.add_op(
+                OpName::new_inline("Rz"),
+                String::new(),
+                Signature::new(vec![qb_t(), float64_type()], vec![qb_t()]),
+                extension_ref,
+            )
+            .unwrap();
+
+            ext.add_op(
+                OpName::new_inline("CX"),
+                String::new(),
+                Signature::new_endo(vec![qb_t(), qb_t()]),
+                extension_ref,
+            )
+            .unwrap();
+        },
+    )
+});
+
+pub static BENCH_EXTENSIONS: LazyLock<ExtensionRegistry> = LazyLock::new(|| {
+    let mut reg = STD_REG.clone();
+    reg.register_updated(QUANTUM_EXT.clone());
+    reg
+});
 
 /// The node ids for a layer generated by the [`circuit`] function.
 pub struct CircuitLayer {