Add qk_dag_successors and qk_dag_predecessors to the C API

crates/cext/cbindgen.toml

@@ -79,3 +79,4 @@ prefix = "Qk"
 "CDelayUnit" = "DelayUnit"
 "CNeighbors" = "Neighbors"
 "COperationKind" = "OperationKind"
+CDagNeighbors = "DagNeighbors"

crates/cext/src/dag.rs

@@ -706,6 +706,146 @@ pub unsafe extern "C" fn qk_dag_op_node_kind(dag: *const DAGCircuit, node: u32)
     }
 }
 
+/// A struct for storing successors and predecessors information
+/// retrieved from `qk_dag_successors` and `qk_dag_predecessors`, respectively.
+#[repr(C)]
+pub struct CDagNeighbors {
+    /// Array of size `num_neighbors` of node indices.
+    neighbors: *const u32,
+    /// The length of the `neighbors` array.
+    num_neighbors: usize,
+}
+
+/// @ingroup QkDag
+/// Retrieve the successors of the specified node.
+///
+/// The successors array and its length are stored in a `QkDagNeighbors` struct through
+/// the `successors` output parameter. Each element in the array is a DAG node index.
+/// You must call the `qk_dag_neighbors_clear`
+/// function when done to free the memory allocated for the struct.
+///
+/// @param dag A pointer to the DAG.
+/// @param node The node to get the successors of.
+/// @param successors A pointer to a `QkDagNeighbors` instance where the successors information will be written.
+///
+/// # Example
+/// ```c
+/// QkDag *dag = qk_dag_new();
+/// QkQuantumRegister *qr = qk_quantum_register_new(2, "qr");
+/// qk_dag_add_quantum_register(dag, qr);
+///
+/// uint32_t node_cx = qk_dag_apply_gate(dag, QkGate_CX, (uint32_t[]){0, 1}, NULL, false);
+///
+/// QkDagNeighbors successors;
+/// qk_dag_successors(dag, node_cx, &successors);
+///
+/// qk_dag_neighbors_clear(&successors);
+/// qk_dag_free(dag);
+/// ```
+///
+/// # Safety
+///
+/// Behavior is undefined if ``dag`` is not a valid, non-null pointer to a ``QkDag``.
+#[unsafe(no_mangle)]
+#[cfg(feature = "cbinding")]
+pub unsafe extern "C" fn qk_dag_successors(
+    dag: *const DAGCircuit,
+    node: u32,
+    successors: *mut CDagNeighbors,
+) {
+    // SAFETY: Per documentation, the pointers are to valid data.
+    let dag = unsafe { const_ptr_as_ref(dag) };
+    let dag_neighbors = unsafe { mut_ptr_as_ref(successors) };
+
+    let successors: Vec<u32> = dag
+        .successors(NodeIndex::new(node as usize))
+        .map(|node| node.index() as u32)
+        .collect();
+    dag_neighbors.num_neighbors = successors.len();
+    dag_neighbors.neighbors = Box::into_raw(successors.into_boxed_slice()) as *const u32;
+}
+
+/// @ingroup QkDag
+/// Retrieve the predecessors of the specified node.
+///
+/// The predecessors array and its length are returned in a `QkDagNeighbors` struct through
+/// the `predecessors` output parameter. Each element in the array is a DAG node index.
+/// You must call the `qk_dag_neighbors_clear`
+/// function when done to free the memory allocated for the struct.
+///
+/// @param dag A pointer to the DAG.
+/// @param node The node to get the predecessors of.
+/// @param predecessors A pointer to a `QkDagNeighbors` instance where the predecessors information will be written.
+///
+/// # Example
+/// ```c
+/// QkDag *dag = qk_dag_new();
+/// QkQuantumRegister *qr = qk_quantum_register_new(2, "qr");
+/// qk_dag_add_quantum_register(dag, qr);
+///
+/// uint32_t node_cx = qk_dag_apply_gate(dag, QkGate_CX, (uint32_t[]){0, 1}, NULL, false);
+///
+/// QkDagNeighbors predecessors;
+/// qk_dag_predecessors(dag, node_cx, &predecessors);
+///
+/// qk_dag_neighbors_clear(&predecessors);
+/// qk_dag_free(dag);
+/// ```
+///
+/// # Safety
+///
+/// Behavior is undefined if ``dag`` is not a valid, non-null pointer to a ``QkDag``.
+#[unsafe(no_mangle)]
+#[cfg(feature = "cbinding")]
+pub unsafe extern "C" fn qk_dag_predecessors(
+    dag: *const DAGCircuit,
+    node: u32,
+    predecessors: *mut CDagNeighbors,
+) {
+    // SAFETY: Per documentation, the pointers are to valid data.
+    let dag = unsafe { const_ptr_as_ref(dag) };
+    let dag_neighbors = unsafe { mut_ptr_as_ref(predecessors) };
+
+    let predecessors: Vec<u32> = dag
+        .predecessors(NodeIndex::new(node as usize))
+        .map(|node| node.index() as u32)
+        .collect();
+    dag_neighbors.num_neighbors = predecessors.len();
+    dag_neighbors.neighbors = Box::into_raw(predecessors.into_boxed_slice()) as *const u32;
+}
+
+/// @ingroup QkDag
+/// Clear the fields of the input `QkDagNeighbors` struct.
+///
+/// The function deallocates the memory pointed to by the `neighbors` field and sets it to NULL.
+/// It also sets the `num_neighbors` field to 0.
+///
+/// @param neighbors A pointer to a `QkDagNeighbors` object.
+///
+/// # Safety
+///
+/// Behavior is undefined if ``neighbors`` is not a valid, non-null pointer to a QkDagNeighbors
+/// object populated with either ``qk_dag_successors`` or ``qk_dag_predecessors``.
+#[unsafe(no_mangle)]
+#[cfg(feature = "cbinding")]
+pub unsafe extern "C" fn qk_dag_neighbors_clear(neighbors: *mut CDagNeighbors) {
+    // SAFETY: Per documentation, the pointer is to a valid data.
+    let neighbors = unsafe { mut_ptr_as_ref(neighbors) };
+
+    if neighbors.num_neighbors > 0 {
+        let slice = std::ptr::slice_from_raw_parts_mut(
+            neighbors.neighbors as *mut u32,
+            neighbors.num_neighbors,
+        );
+        unsafe {
+            let _ = Box::from_raw(slice);
+        }
+    }
+
+    neighbors.num_neighbors = 0;
+    neighbors.neighbors = std::ptr::null();
+}
+
 /// @ingroup QkDag
 /// Free the DAG.
 ///

crates/circuit/src/dag_circuit.rs

@@ -3404,11 +3404,10 @@ impl DAGCircuit {
 
     /// Returns iterator of the successors of a node as :class:`.DAGOpNode`\ s and
     /// :class:`.DAGOutNode`\ s.
-    fn successors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
+    #[pyo3(name = "successors")]
+    fn py_successors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
         let successors: PyResult<Vec<_>> = self
-            .dag
-            .neighbors_directed(node.node.unwrap(), Outgoing)
-            .unique()
+            .successors(node.node.unwrap())
             .map(|i| self.get_node(py, i))
             .collect();
         Ok(PyTuple::new(py, successors?)?
@@ -3420,11 +3419,10 @@ impl DAGCircuit {
 
     /// Returns iterator of the predecessors of a node as :class:`.DAGOpNode`\ s and
     /// :class:`.DAGInNode`\ s.
-    fn predecessors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
+    #[pyo3(name = "predecessors")]
+    fn py_predecessors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
         let predecessors: PyResult<Vec<_>> = self
-            .dag
-            .neighbors_directed(node.node.unwrap(), Incoming)
-            .unique()
+            .predecessors(node.node.unwrap())
             .map(|i| self.get_node(py, i))
             .collect();
         Ok(PyTuple::new(py, predecessors?)?
@@ -7769,6 +7767,16 @@ impl DAGCircuit {
     pub fn get_metadata(&self) -> Option<&Py<PyAny>> {
         self.metadata.as_ref()
     }
+
+    // Returns an iterator over the unique successors of the given node
+    pub fn successors(&self, node: NodeIndex) -> impl Iterator<Item = NodeIndex> {
+        self.dag.neighbors_directed(node, Outgoing).unique()
+    }
+
+    // Returns an iterator over the unique predecessors of the given node
+    pub fn predecessors(&self, node: NodeIndex) -> impl Iterator<Item = NodeIndex> {
+        self.dag.neighbors_directed(node, Incoming).unique()
+    }
 }
 
 pub struct DAGCircuitBuilder {

docs/cdoc/qk-dag.rst

@@ -27,6 +27,9 @@ Data Types
 
 .. doxygenenum:: QkOperationKind
 
+.. doxygenstruct:: QkDagNeighbors
+   :members:
+
 Functions
 =========
 

releasenotes/notes/c-api-dag-node-neighbors-a6ac08bce6d42e4a.yaml

@@ -0,0 +1,8 @@
+---
+features_c:
+  - |
+    Added the :c:func:`qk_dag_successors` and :c:func:`qk_dag_predecessors` functions to retrieve
+    the successors and predecessors of a DAG node, respectively. The neighbors information is returned
+    by these functions through an output parameter of the newly added :c:struct:`QkDagNeighbors` struct. 
+    The :c:func:`qk_dag_neighbors_clear` function was also added to free the memory and reset the fields 
+    of :c:struct:`QkDagNeighbors`.

test/c/test_dag.c

@@ -312,6 +312,77 @@ static int test_dag_endpoint_node_value(void) {
     return result;
 }
 
+/*
+ * Test qk_dag_successors and qk_dag_predecessors
+ */
+static int test_dag_node_neighbors(void) {
+    int result = Ok;
+    QkDag *dag = qk_dag_new();
+    QkQuantumRegister *qr = qk_quantum_register_new(3, "qr");
+    qk_dag_add_quantum_register(dag, qr);
+    qk_quantum_register_free(qr);
+
+    uint32_t node_h = qk_dag_apply_gate(dag, QkGate_H, (uint32_t[]){0}, NULL, false);
+    uint32_t node_ccx = qk_dag_apply_gate(dag, QkGate_CCX, (uint32_t[]){0, 1, 2}, NULL, false);
+    uint32_t node_cx = qk_dag_apply_gate(dag, QkGate_CX, (uint32_t[]){1, 2}, NULL, false);
+
+    QkDagNeighbors successors, predecessors;
+
+    // H node
+    qk_dag_successors(dag, node_h, &successors);
+    qk_dag_predecessors(dag, node_h, &predecessors);
+    if (successors.num_neighbors != 1 || successors.neighbors[0] != node_ccx ||
+        predecessors.num_neighbors != 1 ||
+        qk_dag_node_type(dag, predecessors.neighbors[0]) != QkDagNodeType_QubitIn) {
+        printf("Incorrect neighbors information for the H node!\n");
+        result = EqualityError;
+        goto cleanup;
+    }
+    qk_dag_neighbors_clear(&successors);
+    qk_dag_neighbors_clear(&predecessors);
+
+    if (successors.neighbors != NULL || successors.num_neighbors != 0) {
+        printf("qk_dag_neighbors_clear didn't work!\n");
+        result = RuntimeError;
+        goto cleanup;
+    }
+
+    // CCX node
+    qk_dag_successors(dag, node_ccx, &successors);
+    qk_dag_predecessors(dag, node_ccx, &predecessors);
+    if (successors.num_neighbors != 2 || // CX is counted as a unique successor
+        successors.neighbors[0] != node_cx ||
+        qk_dag_node_type(dag, successors.neighbors[1]) != QkDagNodeType_QubitOut ||
+        predecessors.num_neighbors != 3 ||
+        qk_dag_node_type(dag, predecessors.neighbors[0]) != QkDagNodeType_QubitIn ||
+        qk_dag_node_type(dag, predecessors.neighbors[1]) != QkDagNodeType_QubitIn ||
+        predecessors.neighbors[2] != node_h) {
+        printf("Incorrect neighbors information for the CCX node!\n");
+        result = EqualityError;
+        goto cleanup;
+    }
+    qk_dag_neighbors_clear(&successors);
+    qk_dag_neighbors_clear(&predecessors);
+
+    // CX node
+    qk_dag_successors(dag, node_cx, &successors);
+    qk_dag_predecessors(dag, node_cx, &predecessors);
+    if (successors.num_neighbors != 2 ||
+        qk_dag_node_type(dag, successors.neighbors[0]) != QkDagNodeType_QubitOut ||
+        qk_dag_node_type(dag, successors.neighbors[1]) != QkDagNodeType_QubitOut ||
+        predecessors.num_neighbors != 1 || // CCX is counted as a unique predecessor
+        predecessors.neighbors[0] != node_ccx) {
+        printf("Incorrect neighbors information for the CX node!\n");
+        result = EqualityError;
+    }
+
+cleanup:
+    qk_dag_neighbors_clear(&successors);
+    qk_dag_neighbors_clear(&predecessors);
+    qk_dag_free(dag);
+    return result;
+}
+
 int test_dag(void) {
     int num_failed = 0;
     num_failed += RUN_TEST(test_empty);
@@ -321,6 +392,7 @@ int test_dag(void) {
     num_failed += RUN_TEST(test_dag_node_type);
     num_failed += RUN_TEST(test_dag_endpoint_node_value);
     num_failed += RUN_TEST(test_op_node_bits_explicit);
+    num_failed += RUN_TEST(test_dag_node_neighbors);
 
     fflush(stderr);
     fprintf(stderr, "=== Number of failed subtests: %i\n", num_failed);