diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index 7bc80061e055..3e27af126346 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -39,35 +39,43 @@ use std::f64::consts::PI; use std::ffi::c_double; use std::hash::{Hash, Hasher}; -trait IntoUniqueNodes { +trait IntoUnique { type Output; - fn into_unique(self) -> Self::Output; + fn unique(self) -> Self::Output; } -struct UniqueNeighbors<'a, E: 'a, Ix: 'a> { - neighbors: Neighbors<'a, E, Ix>, - seen: HashSet, +struct UniqueIterator { + neighbors: I, + seen: HashSet, } -impl<'a, E: 'a, Ix: 'a> IntoUniqueNodes for Neighbors<'a, E, Ix> { - type Output = UniqueNeighbors<'a, E, Ix>; +impl IntoUnique for I +where + I: Iterator, + I::Item: Hash + Eq + Clone, +{ + type Output = UniqueIterator; - fn into_unique(self) -> Self::Output { - UniqueNeighbors { + fn unique(self) -> Self::Output { + UniqueIterator { neighbors: self, seen: HashSet::new(), } } } -impl<'a, E> Iterator for UniqueNeighbors<'a, E, DefaultIx> { - type Item = NodeIndex; +impl Iterator for UniqueIterator +where + I: Iterator, + I::Item: Hash + Eq + Clone, +{ + type Item = I::Item; - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { // First any outgoing edges while let Some(node) = self.neighbors.next() { if !self.seen.contains(&node) { - self.seen.insert(node); + self.seen.insert(node.clone()); return Some(node); } } @@ -745,7 +753,7 @@ impl DAGCircuit { let child = self .dag .neighbors_directed(input_node, Outgoing) - .into_unique() + .unique() .next() .ok_or_else(|| { DAGCircuitError::new_err(format!( @@ -1983,25 +1991,46 @@ impl DAGCircuit { // for source, dest, edge in raw_nodes: // yield (self._multi_graph[source], self._multi_graph[dest], edge) // - // def op_nodes(self, op=None, include_directives=True): - // """Get the list of "op" nodes in the dag. - // - // Args: - // op (Type): :class:`qiskit.circuit.Operation` subclass op nodes to - // return. If None, return all op nodes. - // include_directives (bool): include `barrier`, `snapshot` etc. - // - // Returns: - // list[DAGOpNode]: the list of node ids containing the given op. - // """ - // nodes = [] - // for node in self._multi_graph.nodes(): - // if isinstance(node, DAGOpNode): - // if not include_directives and getattr(node.op, "_directive", False): - // continue - // if op is None or isinstance(node.op, op): - // nodes.append(node) - // return nodes + /// Get the list of "op" nodes in the dag. + /// + /// Args: + /// op (Type): :class:`qiskit.circuit.Operation` subclass op nodes to + /// return. If None, return all op nodes. + /// include_directives (bool): include `barrier`, `snapshot` etc. + /// + /// Returns: + /// list[DAGOpNode]: the list of node ids containing the given op. + #[pyo3(signature=(op=None, include_directives=true))] + fn op_nodes( + &self, + py: Python, + op: Option<&Bound>, + include_directives: bool, + ) -> PyResult>> { + let mut nodes = Vec::new(); + for (node, weight) in self.dag.node_references() { + if let NodeType::Operation(ref packed) = weight { + if !include_directives { + let directive = if packed.op.bind(py).hasattr(intern!(py, "_directive"))? { + packed + .op + .bind(py) + .getattr(intern!(py, "_directive"))? + .extract()? + } else { + false + }; + if directive { + continue; + } + } + if op.is_none() || packed.op.bind(py).is_instance(op.unwrap())? { + nodes.push(node); + } + } + } + nodes.into_iter().map(|i| self.get_node(py, i)).collect() + } // // def gate_nodes(self): // """Get the list of gate nodes in the dag. @@ -2055,7 +2084,7 @@ impl DAGCircuit { let successors: PyResult> = self .dag .neighbors_directed(NodeIndex::new(node._node_id as usize), Outgoing) - .into_unique() + .unique() .map(|i| self.get_node(py, i)) .collect(); Ok(PyTuple::new_bound(py, successors?) @@ -2070,7 +2099,7 @@ impl DAGCircuit { let predecessors: PyResult> = self .dag .neighbors_directed(NodeIndex::new(node._node_id as usize), Incoming) - .into_unique() + .unique() .map(|i| self.get_node(py, i)) .collect(); Ok(PyTuple::new_bound(py, predecessors?) @@ -2099,57 +2128,97 @@ impl DAGCircuit { ) .is_some() } - // - // def quantum_predecessors(self, node): - // """Returns iterator of the predecessors of a node that are - // connected by a quantum edge as DAGOpNodes and DAGInNodes.""" - // return iter( - // self._multi_graph.find_predecessors_by_edge( - // node._node_id, lambda edge_data: isinstance(edge_data, Qubit) - // ) - // ) - // - // def classical_predecessors(self, node): - // """Returns iterator of the predecessors of a node that are - // connected by a classical edge as DAGOpNodes and DAGInNodes.""" - // return iter( - // self._multi_graph.find_predecessors_by_edge( - // node._node_id, lambda edge_data: isinstance(edge_data, Clbit) - // ) - // ) - // - // def ancestors(self, node): - // """Returns set of the ancestors of a node as DAGOpNodes and DAGInNodes.""" - // return {self._multi_graph[x] for x in rx.ancestors(self._multi_graph, node._node_id)} - // - // def descendants(self, node): - // """Returns set of the descendants of a node as DAGOpNodes and DAGOutNodes.""" - // return {self._multi_graph[x] for x in rx.descendants(self._multi_graph, node._node_id)} - // - // def bfs_successors(self, node): - // """ - // Returns an iterator of tuples of (DAGNode, [DAGNodes]) where the DAGNode is the current node - // and [DAGNode] is its successors in BFS order. - // """ - // return iter(rx.bfs_successors(self._multi_graph, node._node_id)) - // - // def quantum_successors(self, node): - // """Returns iterator of the successors of a node that are - // connected by a quantum edge as Opnodes and DAGOutNodes.""" - // return iter( - // self._multi_graph.find_successors_by_edge( - // node._node_id, lambda edge_data: isinstance(edge_data, Qubit) - // ) - // ) - // - // def classical_successors(self, node): - // """Returns iterator of the successors of a node that are - // connected by a classical edge as DAGOpNodes and DAGInNodes.""" - // return iter( - // self._multi_graph.find_successors_by_edge( - // node._node_id, lambda edge_data: isinstance(edge_data, Clbit) - // ) - // ) + + /// Returns iterator of the predecessors of a node that are + /// connected by a quantum edge as DAGOpNodes and DAGInNodes. + fn quantum_predecessors(&self, py: Python, node: &DAGNode) -> PyResult> { + let node = NodeIndex::new(node._node_id as usize); + let edges = self.dag.edges_directed(node, Incoming); + let filtered = edges.filter_map(|e| match e.weight() { + Wire::Qubit(_) => Some(e.source()), + _ => None, + }); + let predecessors: PyResult> = + filtered.unique().map(|i| self.get_node(py, i)).collect(); + Ok(PyTuple::new_bound(py, predecessors?) + .into_any() + .iter() + .unwrap() + .unbind()) + } + + /// Returns iterator of the predecessors of a node that are + /// connected by a classical edge as DAGOpNodes and DAGInNodes. + fn classical_predecessors(&self, py: Python, node: &DAGNode) -> PyResult> { + let node = NodeIndex::new(node._node_id as usize); + let edges = self.dag.edges_directed(node, Incoming); + let filtered = edges.filter_map(|e| match e.weight() { + Wire::Clbit(_) => Some(e.source()), + _ => None, + }); + let predecessors: PyResult> = + filtered.unique().map(|i| self.get_node(py, i)).collect(); + Ok(PyTuple::new_bound(py, predecessors?) + .into_any() + .iter() + .unwrap() + .unbind()) + } + + /// Returns set of the ancestors of a node as DAGOpNodes and DAGInNodes. + fn ancestors(&self, py: Python, node: &DAGNode) -> PyResult> { + // return {self._multi_graph[x] for x in rx.ancestors(self._multi_graph, node._node_id)} + todo!() + } + + /// Returns set of the descendants of a node as DAGOpNodes and DAGOutNodes. + fn descendants(&self, py: Python, node: &DAGNode) -> PyResult> { + // return {self._multi_graph[x] for x in rx.descendants(self._multi_graph, node._node_id)} + todo!() + } + + /// Returns an iterator of tuples of (DAGNode, [DAGNodes]) where the DAGNode is the current node + /// and [DAGNode] is its successors in BFS order. + fn bfs_successors(&self, py: Python, node: &DAGNode) -> PyResult> { + // return iter(rx.bfs_successors(self._multi_graph, node._node_id)) + todo!() + } + + /// Returns iterator of the successors of a node that are + /// connected by a quantum edge as DAGOpNodes and DAGOutNodes. + fn quantum_successors(&self, py: Python, node: &DAGNode) -> PyResult> { + let node = NodeIndex::new(node._node_id as usize); + let edges = self.dag.edges_directed(node, Outgoing); + let filtered = edges.filter_map(|e| match e.weight() { + Wire::Qubit(_) => Some(e.target()), + _ => None, + }); + let predecessors: PyResult> = + filtered.unique().map(|i| self.get_node(py, i)).collect(); + Ok(PyTuple::new_bound(py, predecessors?) + .into_any() + .iter() + .unwrap() + .unbind()) + } + + /// Returns iterator of the successors of a node that are + /// connected by a classical edge as DAGOpNodes and DAGOutNodes. + fn classical_successors(&self, py: Python, node: &DAGNode) -> PyResult> { + let node = NodeIndex::new(node._node_id as usize); + let edges = self.dag.edges_directed(node, Incoming); + let filtered = edges.filter_map(|e| match e.weight() { + Wire::Clbit(_) => Some(e.target()), + _ => None, + }); + let predecessors: PyResult> = + filtered.unique().map(|i| self.get_node(py, i)).collect(); + Ok(PyTuple::new_bound(py, predecessors?) + .into_any() + .iter() + .unwrap() + .unbind()) + } // // def remove_op_node(self, node): // """Remove an operation node n.