Skip to content

Commit

Permalink
Add: front_layer method to DAGCircuit
Browse files Browse the repository at this point in the history
- Add rust-native counterpart for `front_layer`.
  • Loading branch information
raynelfss committed Jul 10, 2024
1 parent b71bdd8 commit 6355619
Showing 1 changed file with 26 additions and 2 deletions.
28 changes: 26 additions & 2 deletions crates/circuit/src/dag_circuit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2959,7 +2959,8 @@ def _format(operand):
}

/// Return a list of op nodes in the first layer of this dag.
fn front_layers(&self) -> Py<PyList> {
#[pyo3(name = "front_layer")]
fn py_front_layer(&self, py: Python) -> PyResult<Py<PyList>> {
// graph_layers = self.multigraph_layers()
// try:
// next(graph_layers) # Remove input nodes
Expand All @@ -2969,7 +2970,13 @@ def _format(operand):
// op_nodes = [node for node in next(graph_layers) if isinstance(node, DAGOpNode)]
//
// return op_nodes
todo!()
// todo!()
let native_front_layer = self.front_layer();
let front_layer_list = PyList::empty_bound(py);
for node in native_front_layer {
front_layer_list.append(self.get_node(py, node)?)?;
}
Ok(front_layer_list.into())
}

/// Yield a shallow view on a layer of this DAGCircuit for all d layers of this circuit.
Expand Down Expand Up @@ -3928,4 +3935,21 @@ impl DAGCircuit {
Err(_) => unreachable!("Not a DAG."),
})
}

pub fn front_layer<'a>(&'a self) -> Box<dyn Iterator<Item = NodeIndex> + 'a> {
let mut graph_layers = self.multigraph_layers();
graph_layers.next();

let next_layer = graph_layers.next();
match next_layer {
Some(layer) => Box::new(layer.into_iter().filter_map(|node| {
if matches!(self.dag.node_weight(node).unwrap(), NodeType::Operation(_)) {
Some(node)
} else {
None
}
})),
None => Box::new(vec![].into_iter()),
}
}
}

0 comments on commit 6355619

Please sign in to comment.