feat: include generator in package reading errors

Author: ss2165

hugr-core/src/envelope.rs

@@ -47,12 +47,13 @@ pub mod serde_with;
 pub use header::{EnvelopeConfig, EnvelopeFormat, MAGIC_NUMBERS, ZstdConfig};
 pub use package_json::PackageEncodingError;
 
-use crate::Hugr;
+use crate::{Hugr, HugrView};
 use crate::{
     extension::{ExtensionRegistry, Version},
     package::Package,
 };
 use header::EnvelopeHeader;
+use std::collections::HashSet;
 use std::io::BufRead;
 use std::io::Write;
 use std::str::FromStr;
@@ -64,6 +65,49 @@ use itertools::Itertools as _;
 use crate::import::ImportError;
 use crate::{Extension, import::import_package};
 
+/// Key used to store the name of the generator that produced the envelope.
+pub const GENERATOR_KEY: &str = "__generator";
+
+/// Get the name of the generator from the metadata of the HUGR modules.
+/// If multiple modules have different generators, only the first one is returned.
+fn get_generator<H: HugrView>(modules: &[H]) -> Option<String> {
+    let generators: HashSet<String> = modules
+        .iter()
+        .filter_map(|hugr| hugr.get_metadata(hugr.module_root(), GENERATOR_KEY))
+        .map(|v| v.to_string())
+        .collect();
+    debug_assert!(
+        generators.len() <= 1,
+        "Multiple generators found in the package metadata: {generators:?}"
+    );
+    generators.into_iter().next()
+}
+
+fn gen_str(generator: &Option<String>) -> String {
+    match generator {
+        Some(g) => format!("\ngenerated by {g}"),
+        None => String::new(),
+    }
+}
+
+/// Wrap an error with a generator string.
+#[derive(Error, Debug)]
+#[error("{inner}{}", gen_str(&self.generator))]
+pub struct WithGenerator<E: std::fmt::Display> {
+    inner: E,
+    /// The name of the generator that produced the envelope, if any.
+    generator: Option<String>,
+}
+
+impl<E: std::fmt::Display> WithGenerator<E> {
+    fn new(err: E, modules: &[impl HugrView]) -> Self {
+        Self {
+            inner: err,
+            generator: get_generator(modules),
+        }
+    }
+}
+
 /// Read a HUGR envelope from a reader.
 ///
 /// Returns the deserialized package and the configuration used to encode it.
@@ -216,6 +260,7 @@ pub enum EnvelopeError {
         /// The source error.
         #[from]
         source: ImportError,
+        // TODO add generator to model import errors
     },
     /// Error reading a HUGR model payload.
     #[error(transparent)]
@@ -454,7 +499,7 @@ fn check_breaking_extensions(
     hugr: impl crate::HugrView,
     registry: &ExtensionRegistry,
 ) -> Result<(), ExtensionBreakingError> {
-    let Some(exts) = hugr.get_metadata(hugr.module_root(), &USED_EXTENSIONS_KEY) else {
+    let Some(exts) = hugr.get_metadata(hugr.module_root(), USED_EXTENSIONS_KEY) else {
         return Ok(()); // No used extensions metadata, nothing to check
     };
     let used_exts: Vec<UsedExtension> = serde_json::from_value(exts.clone())?; // TODO handle errors properly
@@ -690,4 +735,28 @@ pub(crate) mod test {
             Err(ExtensionBreakingError::Deserialization(_))
         );
     }
+
+    #[test]
+    fn test_with_generator_error_message() {
+        let test_ext = Extension::new(ExtensionId::new_unchecked("test"), Version::new(1, 0, 0));
+        let registry = ExtensionRegistry::new([Arc::new(test_ext)]);
+
+        let mut hugr = simple_package().modules.remove(0);
+
+        // Set a generator name in the metadata
+        let generator_name = json!({ "name": "TestGenerator", "version": "1.2.3" });
+        hugr.set_metadata(hugr.module_root(), GENERATOR_KEY, generator_name.clone());
+
+        // Set incompatible extension version in metadata
+        let used_exts = json!([{ "name": "test", "version": "2.0.0" }]);
+        hugr.set_metadata(hugr.module_root(), USED_EXTENSIONS_KEY, used_exts);
+
+        // Create the error and wrap it with WithGenerator
+        let err = check_breaking_extensions(&hugr, &registry).unwrap_err();
+        let with_gen = WithGenerator::new(err, &[&hugr]);
+
+        let err_msg = with_gen.to_string();
+        assert!(err_msg.contains("Extension 'test' version mismatch"));
+        assert!(err_msg.contains(generator_name.to_string().as_str()));
+    }
 }

hugr-core/src/envelope/package_json.rs

@@ -3,6 +3,7 @@ use derive_more::{Display, Error, From};
 use itertools::Itertools;
 use std::io;
 
+use super::{ExtensionBreakingError, WithGenerator, check_breaking_extensions};
 use crate::extension::ExtensionRegistry;
 use crate::extension::resolution::ExtensionResolutionError;
 use crate::hugr::ExtensionError;
@@ -21,20 +22,24 @@ pub(super) fn from_json_reader(
         extensions: pkg_extensions,
     } = serde_json::from_value::<PackageDeser>(val.clone())?;
     let mut modules = modules.into_iter().map(|h| h.0).collect_vec();
-
     let pkg_extensions = ExtensionRegistry::new_with_extension_resolution(
         pkg_extensions,
         &extension_registry.into(),
-    )?;
+    )
+    .map_err(|err| WithGenerator::new(err, &modules))?;
 
     // Resolve the operations in the modules using the defined registries.
     let mut combined_registry = extension_registry.clone();
     combined_registry.extend(&pkg_extensions);
 
-    for module in &mut modules {
-        super::check_breaking_extensions(&module, &combined_registry)?;
-        module.resolve_extension_defs(&combined_registry)?;
+    for module in &modules {
+        check_breaking_extensions(module, &combined_registry)
+            .map_err(|err| WithGenerator::new(err, &modules))?;
     }
+    modules
+        .iter_mut()
+        .try_for_each(|module| module.resolve_extension_defs(&combined_registry))
+        .map_err(|err| WithGenerator::new(err, &modules))?;
 
     Ok(Package {
         modules,
@@ -65,9 +70,9 @@ pub enum PackageEncodingError {
     /// Error raised while reading from a file.
     IOError(io::Error),
     /// Could not resolve the extension needed to encode the hugr.
-    ExtensionResolution(ExtensionResolutionError),
+    ExtensionResolution(WithGenerator<ExtensionResolutionError>),
     /// Error raised while checking for breaking extension version mismatch.
-    ExtensionVersion(super::ExtensionBreakingError),
+    ExtensionVersion(WithGenerator<ExtensionBreakingError>),
     /// Could not resolve the runtime extensions for the hugr.
     RuntimeExtensionResolution(ExtensionError),
 }