Non-panicking Hugr::new and Hugr::new_with_entrypoint

Author: aborgna-q

hugr-core/src/builder.rs

@@ -328,9 +328,10 @@ pub(crate) mod test {
     /// only (no wires), given a function type with extension delta.
     // TODO consider taking two type rows and using TO_BE_INFERRED
     pub(crate) fn closed_dfg_root_hugr(signature: Signature) -> Hugr {
-        let mut hugr = Hugr::new(ops::DFG {
+        let mut hugr = Hugr::new_with_entrypoint(ops::DFG {
             signature: signature.clone(),
-        });
+        })
+        .unwrap();
         hugr.add_node_with_parent(
             hugr.entrypoint(),
             ops::Input {

hugr-core/src/builder/cfg.rs

@@ -148,7 +148,7 @@ impl CFGBuilder<Hugr> {
             signature: signature.clone(),
         };
 
-        let base = Hugr::new(cfg_op);
+        let base = Hugr::new_with_entrypoint(cfg_op).expect("CFG entrypoints be valid");
         let cfg_node = base.entrypoint();
         CFGBuilder::create(base, cfg_node, signature.input, signature.output)
     }

hugr-core/src/builder/conditional.rs

@@ -2,7 +2,7 @@ use crate::hugr::views::HugrView;
 use crate::types::{Signature, TypeRow};
 
 use crate::ops;
-use crate::ops::handle::CaseID;
+use crate::ops::handle::{CaseID, NodeHandle};
 
 use super::build_traits::SubContainer;
 use super::handle::BuildHandle;
@@ -153,7 +153,7 @@ impl ConditionalBuilder<Hugr> {
     ) -> Result<Self, BuildError> {
         let sum_rows: Vec<_> = sum_rows.into_iter().collect();
         let other_inputs = other_inputs.into();
-        let outputs = outputs.into();
+        let outputs: TypeRow = outputs.into();
 
         let n_out_wires = outputs.len();
         let n_cases = sum_rows.len();
@@ -163,7 +163,7 @@ impl ConditionalBuilder<Hugr> {
             other_inputs,
             outputs,
         };
-        let base = Hugr::new(op);
+        let base = Hugr::new_with_entrypoint(op).expect("Conditional entrypoint should be valid");
         let conditional_node = base.entrypoint();
 
         Ok(ConditionalBuilder {
@@ -178,13 +178,15 @@ impl ConditionalBuilder<Hugr> {
 impl CaseBuilder<Hugr> {
     /// Initialize a Case rooted HUGR
     pub fn new(signature: Signature) -> Result<Self, BuildError> {
-        let op = ops::Case {
-            signature: signature.clone(),
-        };
-        let base = Hugr::new(op);
-        let root = base.entrypoint();
-        let dfg_builder = DFGBuilder::create_with_io(base, root, signature)?;
-
+        // Start by building a conditional with a single case
+        let mut conditional =
+            ConditionalBuilder::new([signature.input.clone()], vec![], signature.output.clone())?;
+        let case = conditional.case_builder(0)?.finish_sub_container()?.node();
+
+        // Extract the half-finished hugr, and wrap it in an owned case builder
+        let mut base = std::mem::take(conditional.hugr_mut());
+        base.set_entrypoint(case);
+        let dfg_builder = DFGBuilder::create(base, case)?;
         Ok(CaseBuilder::from_dfg_builder(dfg_builder))
     }
 }

hugr-core/src/builder/dataflow.rs

@@ -91,7 +91,7 @@ impl DFGBuilder<Hugr> {
         let dfg_op = ops::DFG {
             signature: signature.clone(),
         };
-        let base = Hugr::new(dfg_op);
+        let base = Hugr::new_with_entrypoint(dfg_op).expect("DFG entrypoint should be valid");
         let root = base.entrypoint();
         DFGBuilder::create_with_io(base, root, signature)
     }
@@ -166,7 +166,7 @@ impl FunctionBuilder<Hugr> {
             name: name.into(),
         };
 
-        let base = Hugr::new(op);
+        let base = Hugr::new_with_entrypoint(op).expect("FuncDefn entrypoint should be valid");
         let root = base.entrypoint();
 
         let db = DFGBuilder::create_with_io(base, root, body)?;

hugr-core/src/builder/tail_loop.rs

@@ -81,7 +81,8 @@ impl TailLoopBuilder<Hugr> {
             just_outputs: just_outputs.into(),
             rest: inputs_outputs.into(),
         };
-        let base = Hugr::new(tail_loop.clone());
+        let base = Hugr::new_with_entrypoint(tail_loop.clone())
+            .expect("tail_loop entrypoint should be valid");
         let root = base.entrypoint();
         Self::create_with_io(base, root, &tail_loop)
     }

hugr-core/src/hugr.rs

@@ -30,7 +30,7 @@ use crate::extension::resolution::{
     WeakExtensionRegistry,
 };
 use crate::extension::{ExtensionRegistry, ExtensionSet};
-use crate::ops::{self, NamedOp, OpTag, OpTrait};
+use crate::ops::{self, Module, NamedOp, OpName, OpTag, OpTrait};
 pub use crate::ops::{OpType, DEFAULT_OPTYPE};
 use crate::{Direction, Node};
 
@@ -63,7 +63,7 @@ pub struct Hugr {
 
 impl Default for Hugr {
     fn default() -> Self {
-        Self::new(crate::ops::Module::new())
+        Self::new()
     }
 }
 
@@ -89,19 +89,50 @@ pub type NodeMetadataMap = serde_json::Map<String, NodeMetadata>;
 
 /// Public API for HUGRs.
 impl Hugr {
-    /// Create a new Hugr, with a single root node.
-    pub fn new(root: impl Into<OpType>) -> Self {
-        make_module_hugr(root.into(), 0, 0)
+    /// Create a new Hugr, with a single [`Module`] operation as the root node.
+    pub fn new() -> Self {
+        make_module_hugr(Module::new().into(), 0, 0).unwrap()
     }
 
-    /// Create a new Hugr, with a given root node and preallocated capacity.
+    /// Create a new Hugr, with a given entrypoint operation.
     ///
-    /// If the root optype is [`OpType::Module`], the HUGR module root will
-    /// match the root node. Otherwise, the root node will be a child of the a
-    /// module created at the node hierarchy root. The specific HUGR created depends
-    /// on the operation type, and whether it can be contained
-    pub fn with_capacity(root: impl Into<OpType>, nodes: usize, ports: usize) -> Self {
-        make_module_hugr(root.into(), nodes, ports)
+    /// If the optype is [`OpType::Module`], the HUGR module root will match the
+    /// entrypoint node. Otherwise, the entrypoint will be a child of the a
+    /// module initialized at the node hierarchy root. The specific HUGR created
+    /// depends on the operation type.
+    ///
+    /// # Error
+    ///
+    /// Returns [`HugrError::UnsupportedEntrypoint`] if the entrypoint operation
+    /// requires additional context to be defined. This is the case for
+    /// [`OpType::Case`], [`OpType::DataflowBlock`], and [`OpType::ExitBlock`]
+    /// since they are context-specific definitions.
+    pub fn new_with_entrypoint(entrypoint_op: impl Into<OpType>) -> Result<Self, HugrError> {
+        Self::with_capacity(entrypoint_op, 0, 0)
+    }
+
+    /// Create a new Hugr, with a given entrypoint operation and preallocated capacity.
+    ///
+    /// If the optype is [`OpType::Module`], the HUGR module root will match the
+    /// entrypoint node. Otherwise, the entrypoint will be a child of the a
+    /// module initialized at the node hierarchy root. The specific HUGR created
+    /// depends on the operation type.
+    ///
+    /// # Error
+    ///
+    /// Returns [`HugrError::UnsupportedEntrypoint`] if the entrypoint operation
+    /// requires additional context to be defined. This is the case for
+    /// [`OpType::Case`], [`OpType::DataflowBlock`], and [`OpType::ExitBlock`]
+    /// since they are context-specific definitions.
+    pub fn with_capacity(
+        entrypoint_op: impl Into<OpType>,
+        nodes: usize,
+        ports: usize,
+    ) -> Result<Self, HugrError> {
+        let entrypoint_op: OpType = entrypoint_op.into();
+        let op_name = entrypoint_op.name();
+        make_module_hugr(entrypoint_op, nodes, ports)
+            .ok_or(HugrError::UnsupportedEntrypoint { op: op_name })
     }
 
     /// Load a Hugr from a json reader.
@@ -289,6 +320,12 @@ pub enum HugrError {
     /// An invalid port was specified.
     #[error("Invalid port direction {0:?}.")]
     InvalidPortDirection(Direction),
+    /// Cannot initialize a HUGR with the given entrypoint operation type.
+    #[error("Cannot initialize a HUGR with entrypoint type {op}")]
+    UnsupportedEntrypoint {
+        /// The name of the unsupported operation.
+        op: OpName,
+    },
 }
 
 /// Errors that can occur while loading and validating a Hugr json.
@@ -321,7 +358,7 @@ pub enum LoadHugrError {
 /// Some operation types are not allowed and will result in a panic. This is the
 /// case for [`OpType::Case`], [`OpType::DataflowBlock`], and [`OpType::ExitBlock`]
 /// since they are context-specific operation.
-fn make_module_hugr(root_op: OpType, nodes: usize, ports: usize) -> Hugr {
+fn make_module_hugr(root_op: OpType, nodes: usize, ports: usize) -> Option<Hugr> {
     let mut graph = MultiPortGraph::with_capacity(nodes, ports);
     let hierarchy = Hierarchy::new();
     let mut op_types = UnmanagedDenseMap::with_capacity(nodes);
@@ -395,28 +432,18 @@ fn make_module_hugr(root_op: OpType, nodes: usize, ports: usize) -> Hugr {
     }
     // Other more exotic ops are unsupported, and will cause a panic.
     else {
-        let is_expected = matches!(
+        debug_assert!(matches!(
             root_op,
             OpType::Input(_)
                 | OpType::Output(_)
                 | OpType::DataflowBlock(_)
                 | OpType::ExitBlock(_)
                 | OpType::Case(_)
-        );
-        if is_expected {
-            panic!(
-                "Operation type {} is not allowed as a root type of a Hugr being initialized",
-                root_op.name()
-            )
-        } else {
-            panic!(
-                "Operation type {} should have been allowed in `make_module_hugr`",
-                root_op.name()
-            )
-        }
+        ));
+        return None;
     }
 
-    hugr
+    Some(hugr)
 }
 
 #[cfg(test)]

hugr-core/src/hugr/patch/replace.rs

@@ -603,9 +603,10 @@ mod test {
         // Replacement: one BB with two DFGs inside.
         // Use Hugr rather than Builder because it must be empty (not even
         // Input/Output).
-        let mut replacement = Hugr::new(ops::CFG {
+        let mut replacement = Hugr::new_with_entrypoint(ops::CFG {
             signature: Signature::new_endo(just_list.clone()),
-        });
+        })
+        .expect("CFG is a valid entrypoint");
         let r_bb = replacement.add_node_with_parent(
             replacement.entrypoint(),
             DataflowBlock {
@@ -761,7 +762,7 @@ mod test {
         let cond = cond.finish_sub_container().unwrap();
         let h = h.finish_hugr_with_outputs(cond.outputs()).unwrap();
 
-        let mut r_hugr = Hugr::new(h.get_optype(cond.node()).clone());
+        let mut r_hugr = Hugr::new_with_entrypoint(h.get_optype(cond.node()).clone()).unwrap();
         let r1 = r_hugr.add_node_with_parent(
             r_hugr.entrypoint(),
             Case {

hugr-core/src/hugr/serialize.rs

@@ -16,7 +16,7 @@ use serde::{Deserialize, Deserializer, Serialize};
 
 use self::upgrade::UpgradeError;
 
-use super::{HugrMut, HugrView, NodeMetadataMap};
+use super::{HugrError, HugrMut, HugrView, NodeMetadataMap};
 
 mod upgrade;
 
@@ -133,6 +133,9 @@ pub enum HUGRSerializationError {
     /// First node in node list must be the HUGR root.
     #[error("The first node in the node list has parent {0}, should be itself (index 0)")]
     FirstNodeNotRoot(Node),
+    /// Failed to deserialize the HUGR.
+    #[error(transparent)]
+    HugrError(#[from] HugrError),
 }
 
 impl Serialize for Hugr {
@@ -243,7 +246,7 @@ impl TryFrom<SerHugrLatest> for Hugr {
         }
         // if there are any unconnected ports or copy nodes the capacity will be
         // an underestimate
-        let mut hugr = Hugr::with_capacity(root_type, nodes.len(), edges.len() * 2);
+        let mut hugr = Hugr::with_capacity(root_type, nodes.len(), edges.len() * 2)?;
 
         // Since the new Hugr may add some nodes to contain the root (if the
         // encoded file did not have a module at the root), we need a function

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

@@ -64,7 +64,7 @@ fn add_df_children(b: &mut Hugr, parent: Node, copies: usize) -> (Node, Node, No
 
 #[test]
 fn invalid_root() {
-    let mut b = Hugr::new(LogicOp::Not);
+    let mut b = Hugr::new_with_entrypoint(LogicOp::Not).unwrap();
     let root = b.module_root();
     assert_eq!(b.validate(), Ok(()));
 
@@ -99,7 +99,7 @@ fn invalid_root() {
 fn leaf_root() {
     let leaf_op: OpType = Noop(usize_t()).into();
 
-    let b = Hugr::new(leaf_op);
+    let b = Hugr::new_with_entrypoint(leaf_op).unwrap();
     assert_eq!(b.validate(), Ok(()));
 }
 
@@ -110,7 +110,7 @@ fn dfg_root() {
     }
     .into();
 
-    let mut b = Hugr::new(dfg_op);
+    let mut b = Hugr::new_with_entrypoint(dfg_op).unwrap();
     let root = b.entrypoint();
     add_df_children(&mut b, root, 1);
     assert_eq!(b.validate(), Ok(()));

hugr-core/src/hugr/views.rs

@@ -681,7 +681,7 @@ impl HugrView for Hugr {
         }
 
         let new_entrypoint_op = self.entrypoint_optype().clone();
-        let mut extracted = Hugr::new(new_entrypoint_op);
+        let mut extracted = Hugr::new_with_entrypoint(new_entrypoint_op).unwrap(); // TODO: Handle error
 
         let old_entrypoint = extracted.entrypoint();
         let old_parent = extracted.get_parent(old_entrypoint);