perf: Pre-compute hugr capacity when decoding a hugr-model

hugr-core/src/import.rs

@@ -21,8 +21,8 @@ use crate::{
     hugr::HugrMut,
     ops::{
         AliasDecl, AliasDefn, CFG, Call, CallIndirect, Case, Conditional, Const, DFG,
-        DataflowBlock, ExitBlock, FuncDecl, FuncDefn, Input, LoadConstant, LoadFunction, OpType,
-        OpaqueOp, Output, Tag, TailLoop, Value,
+        DataflowBlock, ExitBlock, FuncDecl, FuncDefn, Input, LoadConstant, LoadFunction,
+        Module as ModuleOp, OpType, OpaqueOp, Output, Tag, TailLoop, Value,
         constant::{CustomConst, CustomSerialized, OpaqueValue},
     },
     package::Package,
@@ -322,18 +322,152 @@ struct Context<'a> {
     description: ModuleDesc,
 }
 
+/// Initial capacities for importer-owned graph and lookup storage.
+///
+/// These estimates are derived only from the model tables and the small set of
+/// nodes and ports synthesized by the importer. Correctness does not depend on
+/// them: malformed or future model constructs can still grow the collections.
+#[derive(Debug, Clone, Copy)]
+struct ImportCapacity {
+    hugr_nodes: usize,
+    hugr_ports: usize,
+    node_map: usize,
+    link_map: usize,
+    static_edges: usize,
+}
+
+impl ImportCapacity {
+    /// Estimate the storage needed while importing `module`.
+    fn for_module(module: &table::Module<'_>) -> Self {
+        let mut hugr_nodes = 1usize.saturating_add(module.nodes.len());
+        let mut hugr_ports = 0usize;
+        let mut static_edges = 0usize;
+
+        for node in &module.nodes {
+            hugr_ports = hugr_ports
+                .saturating_add(node.inputs.len())
+                .saturating_add(node.outputs.len());
+
+            match &node.operation {
+                table::Operation::DefineFunc(_) | table::Operation::DeclareFunc(_) => {
+                    // Function definitions and declarations have a static output.
+                    hugr_ports = hugr_ports.saturating_add(1);
+                }
+                table::Operation::Dfg
+                | table::Operation::Cfg
+                | table::Operation::TailLoop
+                | table::Operation::Conditional => {
+                    // Dataflow operations have an order port in each direction.
+                    hugr_ports = hugr_ports.saturating_add(2);
+                }
+                table::Operation::Custom(operation) => {
+                    // Imported custom operations are dataflow operations.
+                    hugr_ports = hugr_ports.saturating_add(2);
+
+                    let Some((name, args)) = applied_symbol(module, *operation) else {
+                        continue;
+                    };
+                    if name == model::CORE_CALL {
+                        hugr_ports = hugr_ports.saturating_add(1);
+                        static_edges = static_edges.saturating_add(1);
+                    } else if name == model::CORE_LOAD_CONST {
+                        let loads_function = args
+                            .last()
+                            .and_then(|value| match module.get_term(*value)? {
+                                table::Term::Apply(symbol, _) => module.get_node(*symbol),
+                                _ => None,
+                            })
+                            .is_some_and(|node| {
+                                matches!(
+                                    node.operation,
+                                    table::Operation::DefineFunc(_)
+                                        | table::Operation::DeclareFunc(_)
+                                )
+                            });
+
+                        if loads_function {
+                            // LoadFunction has one static input.
+                            hugr_ports = hugr_ports.saturating_add(1);
+                            static_edges = static_edges.saturating_add(1);
+                        } else {
+                            // A value load synthesizes a Const node and a static edge.
+                            hugr_nodes = hugr_nodes.saturating_add(1);
+                            hugr_ports = hugr_ports.saturating_add(2);
+                        }
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        let mut link_map = 0usize;
+        for region in &module.regions {
+            if let Some(scope) = &region.scope {
+                link_map = link_map.saturating_add(scope.links as usize);
+            }
+
+            match region.kind {
+                model::RegionKind::DataFlow => {
+                    // Input and Output boundary nodes each have an order port.
+                    hugr_nodes = hugr_nodes.saturating_add(2);
+                    hugr_ports = hugr_ports
+                        .saturating_add(region.sources.len())
+                        .saturating_add(region.targets.len())
+                        .saturating_add(2);
+                }
+                model::RegionKind::ControlFlow => {
+                    // Control-flow regions synthesize one ExitBlock.
+                    hugr_nodes = hugr_nodes.saturating_add(1);
+                    hugr_ports = hugr_ports.saturating_add(region.targets.len());
+                }
+                model::RegionKind::Module => {}
+            }
+        }
+
+        // Each conditional region is wrapped in an otherwise portless Case node.
+        let case_nodes = module
+            .nodes
+            .iter()
+            .filter(|node| matches!(node.operation, table::Operation::Conditional))
+            .map(|node| node.regions.len())
+            .fold(0usize, usize::saturating_add);
+        hugr_nodes = hugr_nodes.saturating_add(case_nodes);
+
+        Self {
+            hugr_nodes,
+            hugr_ports,
+            node_map: module.nodes.len(),
+            link_map,
+            static_edges,
+        }
+    }
+}
+
+/// Return the referenced symbol name and arguments of an application term.
+fn applied_symbol<'a>(
+    module: &table::Module<'a>,
+    term: table::TermId,
+) -> Option<(&'a str, &'a [table::TermId])> {
+    let table::Term::Apply(symbol, args) = module.get_term(term)? else {
+        return None;
+    };
+    Some((module.get_node(*symbol)?.operation.symbol()?, args))
+}
+
 impl<'a> Context<'a> {
     /// Creates the state used to import one model module.
     fn new(module: &'a table::Module<'a>, extensions: &'a ExtensionRegistry) -> Self {
         // TODO: Module should know about the number of edges, so that we can use a vector here.
         // For now we use a hashmap, which will be slower.
+        let capacity = ImportCapacity::for_module(module);
         Self {
             module,
-            hugr: Hugr::new(),
-            link_ports: FxHashMap::default(),
-            static_edges: Vec::new(),
+            hugr: Hugr::with_capacity(ModuleOp::new(), capacity.hugr_nodes, capacity.hugr_ports)
+                .expect("a module is always a supported HUGR root"),
+            link_ports: FxHashMap::with_capacity_and_hasher(capacity.link_map, Default::default()),
+            static_edges: Vec::with_capacity(capacity.static_edges),
             extensions,
-            nodes: FxHashMap::default(),
+            nodes: FxHashMap::with_capacity_and_hasher(capacity.node_map, Default::default()),
             local_vars: FxHashMap::default(),
             custom_name_cache: FxHashMap::default(),
             type_cache: FxHashMap::default(),
@@ -2262,21 +2396,53 @@ impl LocalVar {
 mod test {
     use std::str::FromStr;
 
-    use crate::{Hugr, std_extensions::std_reg};
+    use crate::{Direction, Hugr, HugrView, std_extensions::std_reg};
     use hugr_model::v0::{self as model, table};
     use rstest::rstest;
     use std::path::PathBuf;
 
-    use super::Context;
+    use super::{Context, ImportCapacity};
+
+    /// Parse the small model fixture used by importer unit tests.
+    fn model_add_ast() -> model::ast::Package {
+        model::ast::Package::from_str(include_str!(
+            "../../hugr-model/tests/fixtures/model-add.edn"
+        ))
+        .unwrap()
+    }
+
+    #[rstest]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-add.edn"))]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-call.edn"))]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-cfg.edn"))]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-cond.edn"))]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-const.edn"))]
+    #[case(include_str!("../../hugr-model/tests/fixtures/model-loop.edn"))]
+    fn import_capacity_covers_generated_graph_storage(#[case] source: &str) {
+        let bump = model::bumpalo::Bump::new();
+        let ast = model::ast::Package::from_str(source).unwrap();
+        let package = ast.resolve(&bump).unwrap();
+        let module = &package.modules[0];
+        let capacity = ImportCapacity::for_module(module);
+        let hugr = super::import_hugr(module, &std_reg()).unwrap();
+        let port_count = hugr
+            .nodes()
+            .map(|node| {
+                hugr.node_ports(node, Direction::Incoming).count()
+                    + hugr.node_ports(node, Direction::Outgoing).count()
+            })
+            .sum::<usize>();
+
+        assert!(capacity.hugr_nodes >= hugr.num_nodes());
+        assert!(capacity.hugr_ports >= port_count);
+        assert!(capacity.node_map >= module.nodes.len());
+    }
 
     /// Check that type imports of the same term share the same underlying storage.
     #[rstest]
     fn repeated_runtime_type_imports_share_storage() {
-        let ast = model::ast::Package::from_str(include_str!(
-            "../../hugr-model/tests/fixtures/model-add.edn"
-        ))
-        .unwrap();
         let bump = model::bumpalo::Bump::new();
+        let ast = model_add_ast();
         let package = ast.resolve(&bump).unwrap();
         let module = &package.modules[0];
         let registry = std_reg();