dfir_lang/graph/
eliminate_extra_unions_tees.rs

1#![warn(missing_docs)]
2
3use super::ops::tee::TEE;
4use super::ops::union::UNION;
5use super::{DfirGraph, GraphNodeId};
6
7fn find_unary_ops<'a>(
8    graph: &'a DfirGraph,
9    op_name: &'static str,
10) -> impl 'a + Iterator<Item = GraphNodeId> {
11    graph
12        .node_ids()
13        .filter(move |&node_id| {
14            graph
15                .node_op_inst(node_id)
16                .is_some_and(|op_inst| op_name == op_inst.op_constraints.name)
17        })
18        .filter(|&node_id| {
19            1 == graph.node_degree_in(node_id) && 1 == graph.node_degree_out(node_id)
20        })
21}
22
23/// Removes missing unions and tees. Must be applied BEFORE subgraph partitioning, i.e. on a flat
24/// graph.
25pub fn eliminate_extra_unions_tees(graph: &mut DfirGraph) {
26    let extra_ops = find_unary_ops(graph, UNION.name)
27        .chain(find_unary_ops(graph, TEE.name))
28        .collect::<Vec<_>>();
29    for extra_op in extra_ops {
30        graph.remove_intermediate_node(extra_op);
31    }
32}