diff --git a/src/builder/build_traits.rs b/src/builder/build_traits.rs index 22c7267af..d76b5727b 100644 --- a/src/builder/build_traits.rs +++ b/src/builder/build_traits.rs @@ -292,10 +292,7 @@ pub trait Dataflow: Container { let op = ops::DFG { signature: signature.clone(), }; - let nodetype = match &input_extensions { - None => NodeType::open_extensions(op), - Some(rs) => NodeType::new(op, rs.clone()), - }; + let nodetype = NodeType::new(op, input_extensions.clone()); let (dfg_n, _) = add_node_with_wires(self, nodetype, input_wires.into_iter().collect())?; DFGBuilder::create_with_io(self.hugr_mut(), dfg_n, signature, input_extensions) diff --git a/src/builder/dataflow.rs b/src/builder/dataflow.rs index 89ac001ac..07c1e2f97 100644 --- a/src/builder/dataflow.rs +++ b/src/builder/dataflow.rs @@ -49,20 +49,14 @@ impl + AsRef> DFGBuilder { let output = ops::Output { types: signature.output().clone(), }; + base.as_mut() + .add_node_with_parent(parent, NodeType::new(input, input_extensions.clone()))?; base.as_mut().add_node_with_parent( parent, - match &input_extensions { - None => NodeType::open_extensions(input), - Some(rs) => NodeType::new(input, rs.clone()), - }, - )?; - base.as_mut().add_node_with_parent( - parent, - match input_extensions.map(|inp| inp.union(&signature.extension_reqs)) { - // TODO: Make this NodeType::open_extensions - None => NodeType::open_extensions(output), - Some(rs) => NodeType::new(output, rs), - }, + NodeType::new( + output, + input_extensions.map(|inp| inp.union(&signature.extension_reqs)), + ), )?; Ok(Self { diff --git a/src/hugr.rs b/src/hugr.rs index 8e5a30c45..f697c63b7 100644 --- a/src/hugr.rs +++ b/src/hugr.rs @@ -66,10 +66,10 @@ pub struct NodeType { impl NodeType { /// Create a new optype with some ExtensionSet - pub fn new(op: impl Into, input_extensions: ExtensionSet) -> Self { + pub fn new(op: impl Into, input_extensions: impl Into>) -> Self { NodeType { op: op.into(), - input_extensions: Some(input_extensions), + input_extensions: input_extensions.into(), } } diff --git a/src/hugr/serialize.rs b/src/hugr/serialize.rs index 953253e8e..5870da98c 100644 --- a/src/hugr/serialize.rs +++ b/src/hugr/serialize.rs @@ -210,10 +210,7 @@ impl TryFrom for Hugr { // if there are any unconnected ports or copy nodes the capacity will be // an underestimate let mut hugr = Hugr::with_capacity( - match input_extensions { - None => NodeType::open_extensions(root_type), - Some(rs) => NodeType::new(root_type, rs), - }, + NodeType::new(root_type, input_extensions), nodes.len(), edges.len() * 2, ); diff --git a/src/ops/custom.rs b/src/ops/custom.rs index 85dacb559..c5262ee43 100644 --- a/src/ops/custom.rs +++ b/src/ops/custom.rs @@ -252,10 +252,7 @@ pub fn resolve_extension_ops( // Only now can we perform the replacements as the 'for' loop was borrowing 'h' preventing use from using it mutably for (n, op) in replacements { let leaf: LeafOp = op.into(); - let node_type = match h.get_nodetype(n).input_extensions() { - None => NodeType::open_extensions(leaf), - Some(exts) => NodeType::new(leaf, exts.clone()), - }; + let node_type = NodeType::new(leaf, h.get_nodetype(n).input_extensions().cloned()); h.replace_op(n, node_type); } Ok(())