test: check envelope roundtrips rather than json in HugrView::verify (#2186)

`verify` was calling an internal serialization test helpers to check
that roundtrip encoding was correct. With the latest changes, this
wasn't checking the actual user-facing envelope encoding.

~This PR currently fails some checks due to #2185~

Author: aborgna-q

.github/workflows/ci-rs.yml

@@ -19,7 +19,6 @@ env:
   CI: true # insta snapshots behave differently on ci
   SCCACHE_GHA_ENABLED: "true"
   RUSTC_WRAPPER: "sccache"
-  HUGR_TEST_SCHEMA: "1"
   # different strings for install action and feature name
   # adapted from https://github.com/TheDan64/inkwell/blob/master/.github/workflows/test.yml
   LLVM_VERSION: "14.0"

hugr-core/src/envelope.rs

@@ -319,14 +319,55 @@ fn encode_model<'h>(
 }
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod test {
     use super::*;
     use cool_asserts::assert_matches;
     use rstest::rstest;
+    use std::borrow::Cow;
     use std::io::BufReader;
 
     use crate::builder::test::{multi_module_package, simple_package};
     use crate::extension::PRELUDE_REGISTRY;
+    use crate::hugr::test::check_hugr_equality;
+    use crate::std_extensions::STD_REG;
+    use crate::HugrView;
+
+    /// Returns an `ExtensionRegistry` with the extensions from both
+    /// sets. Avoids cloning if the first one already contains all
+    /// extensions from the second one.
+    fn join_extensions<'a>(
+        extensions: &'a ExtensionRegistry,
+        other: &ExtensionRegistry,
+    ) -> Cow<'a, ExtensionRegistry> {
+        if other.iter().all(|e| extensions.contains(e.name())) {
+            Cow::Borrowed(extensions)
+        } else {
+            let mut extensions = extensions.clone();
+            extensions.extend(other);
+            Cow::Owned(extensions)
+        }
+    }
+
+    /// Serialize and deserialize a HUGR into an envelope with the given config,
+    /// and check that the result is the same as the original.
+    ///
+    /// We do not compare the before and after `Hugr`s for equality directly,
+    /// because impls of `CustomConst` are not required to implement equality
+    /// checking.
+    ///
+    /// Returns the deserialized HUGR.
+    pub(crate) fn check_hugr_roundtrip(hugr: &Hugr, config: EnvelopeConfig) -> Hugr {
+        let mut buffer = Vec::new();
+        hugr.store(&mut buffer, config).unwrap();
+
+        let extensions = join_extensions(&STD_REG, hugr.extensions());
+
+        let reader = BufReader::new(buffer.as_slice());
+        let extracted = Hugr::load(reader, Some(&extensions)).unwrap();
+
+        check_hugr_equality(&extracted, hugr);
+        extracted
+    }
 
     #[rstest]
     fn errors() {

hugr-core/src/envelope/header.rs

@@ -116,7 +116,7 @@ impl EnvelopeConfig {
     /// If the `zstd` feature is enabled, this will use zstd compression.
     pub const fn binary() -> Self {
         Self {
-            format: EnvelopeFormat::Model,
+            format: EnvelopeFormat::ModelWithExtensions,
             zstd: None,
         }
     }

hugr-core/src/envelope/package_json.rs

@@ -4,10 +4,10 @@ use itertools::Itertools;
 use std::io;
 
 use crate::extension::resolution::ExtensionResolutionError;
-use crate::extension::{ExtensionRegistry, PRELUDE_REGISTRY};
+use crate::extension::ExtensionRegistry;
 use crate::hugr::ExtensionError;
 use crate::package::Package;
-use crate::{Extension, Hugr, HugrView};
+use crate::{Extension, Hugr};
 
 /// Read a Package in json format from an io reader.
 pub(super) fn from_json_reader(
@@ -22,31 +22,9 @@ pub(super) fn from_json_reader(
     } = serde_json::from_value::<PackageDeser>(val.clone())?;
     let mut modules = modules.into_iter().map(|h| h.0).collect_vec();
 
-    // TODO: We don't currently store transitive extension dependencies in the
-    // package's extensions. For example, if we use a `collections.list` const
-    // value but don't use anything in `prelude` we would not include `prelude` in
-    // the package's extensions. But this would then fail when loading the
-    // extensions, as we _need_ the prelude to load the `collections.list` op
-    // definitions here.
-    //
-    // The current fix is to always include the prelude when decoding, but this
-    // only works for transitive `prelude` dependencies.
-    //
-    // Chains of custom extensions will cause this to fail.
-    let extension_registry = if PRELUDE_REGISTRY
-        .iter()
-        .any(|e| !extension_registry.contains(&e.name))
-    {
-        let mut reg_with_prelude = extension_registry.clone();
-        reg_with_prelude.extend(PRELUDE_REGISTRY.iter().cloned());
-        reg_with_prelude
-    } else {
-        extension_registry.clone()
-    };
-
-    let mut pkg_extensions = ExtensionRegistry::new_with_extension_resolution(
+    let pkg_extensions = ExtensionRegistry::new_with_extension_resolution(
         pkg_extensions,
-        &(&extension_registry).into(),
+        &extension_registry.into(),
     )?;
 
     // Resolve the operations in the modules using the defined registries.
@@ -55,7 +33,6 @@ pub(super) fn from_json_reader(
 
     for module in &mut modules {
         module.resolve_extension_defs(&combined_registry)?;
-        pkg_extensions.extend(module.extensions());
     }
 
     Ok(Package {

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

@@ -23,6 +23,7 @@ use crate::std_extensions::arithmetic::float_types::{self, float64_type, ConstF6
 use crate::std_extensions::arithmetic::int_ops;
 use crate::std_extensions::arithmetic::int_types::{self, int_type};
 use crate::std_extensions::collections::list::ListValue;
+use crate::std_extensions::std_reg;
 use crate::types::type_param::TypeParam;
 use crate::types::{PolyFuncType, Signature, Type, TypeArg, TypeBound};
 use crate::{type_row, Extension, Hugr, HugrView};
@@ -117,8 +118,13 @@ fn make_extension_self_referencing(name: &str, op_name: &str, type_name: &str) -
 /// Check that the extensions added during building coincide with read-only collected extensions
 /// and that they survive a serialization roundtrip.
 fn check_extension_resolution(mut hugr: Hugr) {
+    // Extensions used by the hugr, used to check that the roundtrip preserves them.
     let build_extensions = hugr.extensions().clone();
 
+    // Extensions used for resolution.
+    let mut resolution_extensions = std_reg();
+    resolution_extensions.extend(&build_extensions);
+
     // Check that the read-only methods collect the same extensions.
     let collected_exts = ExtensionRegistry::new(hugr.nodes().flat_map(|node| {
         hugr.get_optype(node)
@@ -132,7 +138,7 @@ fn check_extension_resolution(mut hugr: Hugr) {
     );
 
     // Check that the mutable methods collect the same extensions.
-    hugr.resolve_extension_defs(&build_extensions).unwrap();
+    hugr.resolve_extension_defs(&resolution_extensions).unwrap();
     assert_eq!(
         hugr.extensions(),
         &build_extensions,
@@ -143,7 +149,7 @@ fn check_extension_resolution(mut hugr: Hugr) {
     // Roundtrip serialize so all weak references are dropped.
     let ser = hugr.store_str(EnvelopeConfig::text()).unwrap();
 
-    let deser_hugr = Hugr::load_str(&ser, Some(&build_extensions)).unwrap();
+    let deser_hugr = Hugr::load_str(&ser, Some(&resolution_extensions)).unwrap();
 
     assert_eq!(
         deser_hugr.extensions(),

hugr-core/src/hugr.rs

@@ -489,14 +489,63 @@ fn make_module_hugr(root_op: OpType, nodes: usize, ports: usize) -> Option<Hugr>
 }
 
 #[cfg(test)]
-mod test {
+pub(crate) mod test {
     use std::{fs::File, io::BufReader};
 
-    use super::{Hugr, HugrView};
+    use super::*;
 
     use crate::envelope::{EnvelopeError, PackageEncodingError};
+    use crate::ops::OpaqueOp;
     use crate::test_file;
     use cool_asserts::assert_matches;
+    use portgraph::LinkView;
+
+    /// Check that two HUGRs are equivalent, up to node renumbering.
+    pub(crate) fn check_hugr_equality(lhs: &Hugr, rhs: &Hugr) {
+        // Original HUGR, with canonicalized node indices
+        //
+        // The internal port indices may still be different.
+        let mut lhs = lhs.clone();
+        lhs.canonicalize_nodes(|_, _| {});
+        let mut rhs = rhs.clone();
+        rhs.canonicalize_nodes(|_, _| {});
+
+        assert_eq!(rhs.module_root(), lhs.module_root());
+        assert_eq!(rhs.entrypoint(), lhs.entrypoint());
+        assert_eq!(rhs.hierarchy, lhs.hierarchy);
+        assert_eq!(rhs.metadata, lhs.metadata);
+
+        // Extension operations may have been downgraded to opaque operations.
+        for node in rhs.nodes() {
+            let new_op = rhs.get_optype(node);
+            let old_op = lhs.get_optype(node);
+            if !new_op.is_const() {
+                match (new_op, old_op) {
+                    (OpType::ExtensionOp(ext), OpType::OpaqueOp(opaque))
+                    | (OpType::OpaqueOp(opaque), OpType::ExtensionOp(ext)) => {
+                        let ext_opaque: OpaqueOp = ext.clone().into();
+                        assert_eq!(ext_opaque, opaque.clone());
+                    }
+                    _ => assert_eq!(new_op, old_op),
+                }
+            }
+        }
+
+        // Check that the graphs are equivalent up to port renumbering.
+        let new_graph = &rhs.graph;
+        let old_graph = &lhs.graph;
+        assert_eq!(new_graph.node_count(), old_graph.node_count());
+        assert_eq!(new_graph.port_count(), old_graph.port_count());
+        assert_eq!(new_graph.link_count(), old_graph.link_count());
+        for n in old_graph.nodes_iter() {
+            assert_eq!(new_graph.num_inputs(n), old_graph.num_inputs(n));
+            assert_eq!(new_graph.num_outputs(n), old_graph.num_outputs(n));
+            assert_eq!(
+                new_graph.output_neighbours(n).collect_vec(),
+                old_graph.output_neighbours(n).collect_vec()
+            );
+        }
+    }
 
     #[test]
     fn impls_send_and_sync() {

hugr-core/src/hugr/serialize/test.rs

@@ -11,6 +11,7 @@ use crate::extension::simple_op::MakeRegisteredOp;
 use crate::extension::test::SimpleOpDef;
 use crate::extension::ExtensionRegistry;
 use crate::hugr::internal::HugrMutInternals;
+use crate::hugr::test::check_hugr_equality;
 use crate::hugr::validate::ValidationError;
 use crate::hugr::views::ExtractionResult;
 use crate::ops::custom::{ExtensionOp, OpaqueOp, OpaqueOpError};
@@ -28,7 +29,6 @@ use crate::{type_row, OutgoingPort};
 use itertools::Itertools;
 use jsonschema::{Draft, Validator};
 use lazy_static::lazy_static;
-use portgraph::LinkView;
 use portgraph::{multiportgraph::MultiPortGraph, Hierarchy, LinkMut, PortMut, UnmanagedDenseMap};
 use rstest::rstest;
 
@@ -154,7 +154,7 @@ impl From<Type> for SerTestingLatest {
 
 #[test]
 fn empty_hugr_serialize() {
-    check_hugr_roundtrip(&Hugr::default(), true);
+    check_hugr_json_roundtrip(&Hugr::default(), true);
 }
 
 fn ser_deserialize_check_schema<T: serde::de::DeserializeOwned>(
@@ -184,71 +184,26 @@ fn ser_roundtrip_check_schema<TSer: Serialize, TDeser: serde::de::DeserializeOwn
 /// equality checking.
 ///
 /// Returns the deserialized HUGR.
-pub fn check_hugr_roundtrip(hugr: &impl HugrView, check_schema: bool) -> Hugr {
+fn check_hugr_json_roundtrip(hugr: &impl HugrView, check_schema: bool) -> Hugr {
     // Transform the whole view into a HUGR.
     let (mut base, extract_map) = hugr.extract_hugr(hugr.module_root());
     base.set_entrypoint(extract_map.extracted_node(hugr.entrypoint()));
 
     let new_hugr: HugrDeser =
         ser_roundtrip_check_schema(&HugrSer(&base), get_schemas(check_schema));
 
-    check_hugr(&base, &new_hugr.0);
+    check_hugr_equality(&base, &new_hugr.0);
     new_hugr.0
 }
 
 /// Deserialize a HUGR json, ensuring that it is valid against the schema.
 pub fn check_hugr_deserialize(hugr: &Hugr, value: serde_json::Value, check_schema: bool) -> Hugr {
     let new_hugr: HugrDeser = ser_deserialize_check_schema(value, get_schemas(check_schema));
 
-    check_hugr(hugr, &new_hugr.0);
+    check_hugr_equality(hugr, &new_hugr.0);
     new_hugr.0
 }
 
-/// Check that two HUGRs are equivalent, up to node renumbering.
-pub fn check_hugr(lhs: &Hugr, rhs: &Hugr) {
-    // Original HUGR, with canonicalized node indices
-    //
-    // The internal port indices may still be different.
-    let mut h_canon = lhs.clone();
-    h_canon.canonicalize_nodes(|_, _| {});
-
-    assert_eq!(rhs.module_root(), h_canon.module_root());
-    assert_eq!(rhs.entrypoint(), h_canon.entrypoint());
-    assert_eq!(rhs.hierarchy, h_canon.hierarchy);
-    assert_eq!(rhs.metadata, h_canon.metadata);
-
-    // Extension operations may have been downgraded to opaque operations.
-    for node in rhs.nodes() {
-        let new_op = rhs.get_optype(node);
-        let old_op = h_canon.get_optype(node);
-        if !new_op.is_const() {
-            match (new_op, old_op) {
-                (OpType::ExtensionOp(ext), OpType::OpaqueOp(opaque))
-                | (OpType::OpaqueOp(opaque), OpType::ExtensionOp(ext)) => {
-                    let ext_opaque: OpaqueOp = ext.clone().into();
-                    assert_eq!(ext_opaque, opaque.clone());
-                }
-                _ => assert_eq!(new_op, old_op),
-            }
-        }
-    }
-
-    // Check that the graphs are equivalent up to port renumbering.
-    let new_graph = &rhs.graph;
-    let old_graph = &h_canon.graph;
-    assert_eq!(new_graph.node_count(), old_graph.node_count());
-    assert_eq!(new_graph.port_count(), old_graph.port_count());
-    assert_eq!(new_graph.link_count(), old_graph.link_count());
-    for n in old_graph.nodes_iter() {
-        assert_eq!(new_graph.num_inputs(n), old_graph.num_inputs(n));
-        assert_eq!(new_graph.num_outputs(n), old_graph.num_outputs(n));
-        assert_eq!(
-            new_graph.output_neighbours(n).collect_vec(),
-            old_graph.output_neighbours(n).collect_vec()
-        );
-    }
-}
-
 fn check_testing_roundtrip(t: impl Into<SerTestingLatest>) {
     let before = Versioned::new_latest(t.into());
     let after = ser_roundtrip_check_schema(&before, get_testing_schemas(true));
@@ -306,7 +261,7 @@ fn simpleser() {
         extensions: ExtensionRegistry::default(),
     };
 
-    check_hugr_roundtrip(&hugr, true);
+    check_hugr_json_roundtrip(&hugr, true);
 }
 
 #[test]
@@ -335,7 +290,7 @@ fn weighted_hugr_ser() {
         module_builder.finish_hugr().unwrap()
     };
 
-    check_hugr_roundtrip(&hugr, true);
+    check_hugr_json_roundtrip(&hugr, true);
 }
 
 #[test]
@@ -351,7 +306,7 @@ fn dfg_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
     }
     let hugr = dfg.finish_hugr_with_outputs(params)?;
 
-    check_hugr_roundtrip(&hugr, true);
+    check_hugr_json_roundtrip(&hugr, true);
     Ok(())
 }
 
@@ -370,7 +325,7 @@ fn extension_ops() -> Result<(), Box<dyn std::error::Error>> {
 
     let hugr = dfg.finish_hugr_with_outputs([wire])?;
 
-    check_hugr_roundtrip(&hugr, true);
+    check_hugr_json_roundtrip(&hugr, true);
     Ok(())
 }
 
@@ -412,7 +367,7 @@ fn function_type() -> Result<(), Box<dyn std::error::Error>> {
     let op = bldr.add_dataflow_op(Noop(fn_ty), bldr.input_wires())?;
     let h = bldr.finish_hugr_with_outputs(op.outputs())?;
 
-    check_hugr_roundtrip(&h, true);
+    check_hugr_json_roundtrip(&h, true);
     Ok(())
 }
 
@@ -430,7 +385,7 @@ fn hierarchy_order() -> Result<(), Box<dyn std::error::Error>> {
     hugr.remove_node(old_in);
     hugr.validate()?;
 
-    let rhs: Hugr = check_hugr_roundtrip(&hugr, true);
+    let rhs: Hugr = check_hugr_json_roundtrip(&hugr, true);
     rhs.validate()?;
     Ok(())
 }