fix: Allocate ports on root nodes

CQCL · Oct 16, 2024 · 7915c3a · 7915c3a
1 parent 2412599
commit 7915c3a
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 32 deletions.
diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs
@@ -176,9 +176,7 @@ impl Hugr {
         let mut graph = MultiPortGraph::with_capacity(nodes, ports);
         let hierarchy = Hierarchy::new();
         let mut op_types = UnmanagedDenseMap::with_capacity(nodes);
-        let root = graph.add_node(0, 0);
-        // TODO: These extensions should be open in principle, but lets wait
-        // until extensions can be inferred for open sets until changing this
+        let root = graph.add_node(root_node.input_count(), root_node.output_count());
         op_types[root] = root_node;
 
         Self {

diff --git a/hugr-core/src/hugr/validate.rs b/hugr-core/src/hugr/validate.rs
@@ -166,11 +166,25 @@ impl<'a, 'b> ValidationContext<'a, 'b> {
             ))?;
         }
         // The Hugr can have only one root node.
+
+        for dir in Direction::BOTH {
+            // Check that we have the correct amount of ports and edges.
+            let num_ports = self.hugr.graph.num_ports(node.pg_index(), dir);
+            if num_ports != op_type.port_count(dir) {
+                return Err(ValidationError::WrongNumberOfPorts {
+                    node,
+                    optype: op_type.clone(),
+                    actual: num_ports,
+                    expected: op_type.port_count(dir),
+                    dir,
+                });
+            }
+        }
+
         if node == self.hugr.root() {
-            // The root node has no edges.
-            if self.hugr.graph.num_outputs(node.pg_index())
-                + self.hugr.graph.num_inputs(node.pg_index())
-                != 0
+            // The root node cannot have connected edges
+            if self.hugr.all_linked_inputs(node).next().is_some()
+                || self.hugr.all_linked_outputs(node).next().is_some()
             {
                 return Err(ValidationError::RootWithEdges { node });
             }
@@ -190,20 +204,6 @@ impl<'a, 'b> ValidationContext<'a, 'b> {
                     allowed_children,
                 });
             }
-
-            for dir in Direction::BOTH {
-                // Check that we have the correct amount of ports and edges.
-                let num_ports = self.hugr.graph.num_ports(node.pg_index(), dir);
-                if num_ports != op_type.port_count(dir) {
-                    return Err(ValidationError::WrongNumberOfPorts {
-                        node,
-                        optype: op_type.clone(),
-                        actual: num_ports,
-                        expected: op_type.port_count(dir),
-                        dir,
-                    });
-                }
-            }
         }
 
         // Thirdly that the node has correct children
@@ -602,10 +602,14 @@ impl<'a, 'b> ValidationContext<'a, 'b> {
         }
 
         // Check port connections.
-        for dir in Direction::BOTH {
-            for (i, port_index) in self.hugr.graph.ports(node.pg_index(), dir).enumerate() {
-                let port = Port::new(dir, i);
-                self.validate_port(node, port, port_index, op_type, var_decls)?;
+        //
+        // Root nodes are ignored, as they cannot have connected edges.
+        if node != self.hugr.root() {
+            for dir in Direction::BOTH {
+                for (i, port_index) in self.hugr.graph.ports(node.pg_index(), dir).enumerate() {
+                    let port = Port::new(dir, i);
+                    self.validate_port(node, port, port_index, op_type, var_decls)?;
+                }
             }
         }
 

diff --git a/hugr-core/src/hugr/views.rs b/hugr-core/src/hugr/views.rs
@@ -24,7 +24,7 @@ use itertools::{Itertools, MapInto};
 use portgraph::render::{DotFormat, MermaidFormat};
 use portgraph::{multiportgraph, LinkView, PortView};
 
-use super::internal::{HugrInternals, HugrMutInternals};
+use super::internal::HugrInternals;
 use super::{
     Hugr, HugrError, HugrMut, NodeMetadata, NodeMetadataMap, ValidationError, DEFAULT_OPTYPE,
 };
@@ -509,7 +509,6 @@ pub trait ExtractHugr: HugrView + Sized {
         let old_root = hugr.root();
         let new_root = hugr.insert_from_view(old_root, &self).new_root;
         hugr.set_root(new_root);
-        hugr.set_num_ports(new_root, 0, 0);
         hugr.remove_node(old_root);
         hugr
     }

diff --git a/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_cfg.snap b/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_cfg.snap
@@ -3,7 +3,7 @@ source: hugr-core/src/hugr/views/tests.rs
 expression: h.dot_string()
 ---
 digraph {
-0 [shape=plain label=<<table border="1"><tr><td align="text" border="0" colspan="1">(0) CFG</td></tr></table>>]
+0 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="1" cellpadding="1" >0: usize</td></tr><tr><td align="text" border="0" colspan="1">(0) CFG</td></tr><tr><td port="out0" align="text" colspan="1" cellpadding="1" >0: usize</td></tr></table>>]
 1 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="1" cellpadding="1" border="0">0</td></tr><tr><td align="text" border="0" colspan="1">(1) ExitBlock</td></tr></table>>]
 2 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="2" cellpadding="1" border="0">0</td></tr><tr><td align="text" border="0" colspan="2">(2) DataflowBlock</td></tr><tr><td port="out0" align="text" colspan="1" cellpadding="1" border="0">0</td><td port="out1" align="text" colspan="1" cellpadding="1" border="0">1</td></tr></table>>]
 2:out0 -> 6:in0 [style="dashed"]
@@ -44,4 +44,3 @@ hier8 [shape=plain label="8"]
 hier9 [shape=plain label="9"]
 hier10 [shape=plain label="10"]
 }
-
diff --git a/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_dfg.snap b/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_dfg.snap
@@ -1,9 +1,9 @@
 ---
-source: hugr-core/views/tests.rs
+source: hugr-core/src/hugr/views/tests.rs
 expression: h.dot_string()
 ---
 digraph {
-0 [shape=plain label=<<table border="1"><tr><td align="text" border="0" colspan="1">(0) DFG</td></tr></table>>]
+0 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="2" cellpadding="1" >0: qubit</td><td port="in1" align="text" colspan="2" cellpadding="1" >1: qubit</td></tr><tr><td align="text" border="0" colspan="4">(0) DFG</td></tr><tr><td port="out0" align="text" colspan="2" cellpadding="1" >0: qubit</td><td port="out1" align="text" colspan="2" cellpadding="1" >1: qubit</td></tr></table>>]
 1 [shape=plain label=<<table border="1"><tr><td align="text" border="0" colspan="2">(1) Input</td></tr><tr><td port="out0" align="text" colspan="1" cellpadding="1" >0: qubit</td><td port="out1" align="text" colspan="1" cellpadding="1" >1: qubit</td></tr></table>>]
 1:out0 -> 3:in0 [style=""]
 1:out1 -> 3:in1 [style=""]

diff --git a/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_empty_dfg.snap b/hugr-core/src/hugr/views/snapshots/hugr_core__hugr__views__tests__dot_empty_dfg.snap
@@ -3,7 +3,7 @@ source: hugr-core/src/hugr/views/tests.rs
 expression: h.dot_string()
 ---
 digraph {
-0 [shape=plain label=<<table border="1"><tr><td align="text" border="0" colspan="1">(0) DFG</td></tr></table>>]
+0 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="1" cellpadding="1" >0: []+[]</td></tr><tr><td align="text" border="0" colspan="1">(0) DFG</td></tr><tr><td port="out0" align="text" colspan="1" cellpadding="1" >0: []+[]</td></tr></table>>]
 1 [shape=plain label=<<table border="1"><tr><td align="text" border="0" colspan="1">(1) Input</td></tr><tr><td port="out0" align="text" colspan="1" cellpadding="1" >0: []+[]</td></tr></table>>]
 1:out0 -> 2:in0 [style=""]
 2 [shape=plain label=<<table border="1"><tr><td port="in0" align="text" colspan="1" cellpadding="1" >0: []+[]</td></tr><tr><td align="text" border="0" colspan="1">(2) Output</td></tr></table>>]