pulsar_ir/
control_flow_graph.rs

1// Copyright (C) 2024 Ethan Uppal. All rights reserved.
2use super::basic_block::{BasicBlock, BasicBlockCell};
3use pulsar_utils::digraph::Digraph;
4use std::fmt::Display;
5
6pub struct ControlFlowGraph {
7    in_graph: Digraph<BasicBlockCell, bool>,
8    out_graph: Digraph<BasicBlockCell, bool>,
9    entry: BasicBlockCell
10}
11
12impl ControlFlowGraph {
13    pub fn new() -> Self {
14        let entry = BasicBlockCell::new(BasicBlock::new());
15        let mut in_graph = Digraph::new();
16        let mut out_graph = Digraph::new();
17
18        in_graph.add_node(entry.clone());
19        out_graph.add_node(entry.clone());
20
21        Self {
22            in_graph,
23            out_graph,
24            entry
25        }
26    }
27
28    pub fn entry(&self) -> BasicBlockCell {
29        self.entry.clone()
30    }
31
32    pub fn new_block(&mut self) -> BasicBlockCell {
33        let block = BasicBlockCell::new(BasicBlock::new());
34        self.in_graph.add_node(block.clone());
35        self.out_graph.add_node(block.clone());
36        block
37    }
38
39    pub fn add_branch(
40        &mut self, block: BasicBlockCell, condition: bool, dest: BasicBlockCell
41    ) {
42        self.in_graph
43            .add_edge(dest.clone(), condition, block.clone());
44        self.out_graph.add_edge(block, condition, dest);
45    }
46
47    pub fn size(&self) -> usize {
48        self.out_graph.node_count()
49    }
50
51    pub fn blocks(&self) -> Vec<BasicBlockCell> {
52        let mut result = vec![];
53        self.out_graph
54            .dfs(|node| result.push(node), self.entry.clone());
55        result
56    }
57}
58
59impl Display for ControlFlowGraph {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        for (i, basic_block) in self.out_graph.nodes().iter().enumerate() {
62            if i > 0 {
63                writeln!(f)?;
64            }
65            write!(f, "{}", basic_block)?;
66        }
67        Ok(())
68    }
69}