vicis_core/ir/function/layout/
mod.rs

1use crate::ir::function::{basic_block::BasicBlockId, instruction::InstructionId};
2use rustc_hash::FxHashMap;
3
4pub struct Layout {
5    basic_blocks: FxHashMap<BasicBlockId, BasicBlockNode>,
6    instructions: FxHashMap<InstructionId, InstructionNode>,
7    pub first_block: Option<BasicBlockId>,
8    pub last_block: Option<BasicBlockId>,
9}
10
11#[derive(Debug)]
12pub struct BasicBlockNode {
13    prev: Option<BasicBlockId>,
14    next: Option<BasicBlockId>,
15    first_inst: Option<InstructionId>,
16    last_inst: Option<InstructionId>,
17}
18
19#[derive(Debug)]
20pub struct InstructionNode {
21    block: Option<BasicBlockId>,
22    prev: Option<InstructionId>,
23    next: Option<InstructionId>,
24}
25
26pub struct BasicBlockIter<'a> {
27    layout: &'a Layout,
28    cur: Option<BasicBlockId>,
29}
30
31pub struct InstructionIter<'a> {
32    layout: &'a Layout,
33    block: BasicBlockId,
34    head: Option<InstructionId>,
35    tail: Option<InstructionId>,
36}
37
38impl Default for Layout {
39    fn default() -> Self {
40        Self {
41            basic_blocks: FxHashMap::default(),
42            instructions: FxHashMap::default(),
43            first_block: None,
44            last_block: None,
45        }
46    }
47}
48
49impl Layout {
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    pub fn is_empty(&self) -> bool {
55        self.basic_blocks.is_empty()
56    }
57
58    pub fn block_node(&self, id: BasicBlockId) -> &BasicBlockNode {
59        &self.basic_blocks[&id]
60    }
61
62    pub fn block_iter(&self) -> BasicBlockIter {
63        BasicBlockIter {
64            layout: self,
65            cur: self.first_block,
66        }
67    }
68
69    pub fn inst_iter(&self, block: BasicBlockId) -> InstructionIter {
70        InstructionIter {
71            layout: self,
72            block,
73            head: self.basic_blocks[&block].first_inst,
74            tail: self.basic_blocks[&block].last_inst,
75        }
76    }
77
78    pub fn get_entry_block(&self) -> Option<BasicBlockId> {
79        self.first_block
80    }
81
82    pub fn next_block_of(&self, block: BasicBlockId) -> Option<BasicBlockId> {
83        self.basic_blocks[&block].next
84    }
85
86    pub fn append_block(&mut self, block: BasicBlockId) {
87        self.basic_blocks.entry(block).or_insert(BasicBlockNode {
88            prev: self.last_block,
89            next: None,
90            first_inst: None,
91            last_inst: None,
92        });
93
94        if let Some(last_block) = self.last_block {
95            self.basic_blocks.get_mut(&last_block).unwrap().next = Some(block);
96            self.basic_blocks.get_mut(&block).unwrap().prev = Some(last_block);
97        }
98
99        self.last_block = Some(block);
100
101        if self.first_block.is_none() {
102            self.first_block = Some(block)
103        }
104    }
105
106    pub fn append_inst(&mut self, inst: InstructionId, block: BasicBlockId) {
107        self.instructions
108            .entry(inst)
109            .or_insert(InstructionNode {
110                prev: self.basic_blocks[&block].last_inst,
111                next: None,
112                block: Some(block),
113            })
114            .block = Some(block);
115
116        if let Some(last_inst) = self.basic_blocks[&block].last_inst {
117            self.instructions.get_mut(&last_inst).unwrap().next = Some(inst);
118            self.instructions.get_mut(&inst).unwrap().prev = Some(last_inst);
119        }
120
121        self.basic_blocks.get_mut(&block).unwrap().last_inst = Some(inst);
122
123        if self.basic_blocks[&block].first_inst.is_none() {
124            self.basic_blocks.get_mut(&block).unwrap().first_inst = Some(inst);
125        }
126    }
127
128    pub fn insert_inst_at_start(&mut self, inst: InstructionId, block: BasicBlockId) {
129        self.instructions
130            .entry(inst)
131            .or_insert(InstructionNode {
132                prev: None,
133                next: None,
134                block: Some(block),
135            })
136            .block = Some(block);
137
138        if let Some(first_inst) = self.basic_blocks[&block].first_inst {
139            self.instructions.get_mut(&first_inst).unwrap().prev = Some(inst);
140            self.instructions.get_mut(&inst).unwrap().next = Some(first_inst);
141        }
142
143        self.basic_blocks.get_mut(&block).unwrap().first_inst = Some(inst);
144
145        if self.basic_blocks[&block].last_inst.is_none() {
146            self.basic_blocks.get_mut(&block).unwrap().last_inst = Some(inst);
147        }
148    }
149
150    pub fn remove_inst(&mut self, inst: InstructionId) -> Option<()> {
151        let block = self.instructions[&inst].block?;
152        let prev;
153        let next;
154        {
155            let inst_node = &mut self.instructions.get_mut(&inst)?;
156            prev = inst_node.prev;
157            next = inst_node.next;
158            inst_node.block = None;
159            inst_node.prev = None;
160            inst_node.next = None;
161        }
162        match prev {
163            Some(prev) => self.instructions.get_mut(&prev)?.next = next,
164            None => self.basic_blocks.get_mut(&block)?.first_inst = next,
165        }
166        match next {
167            Some(next) => self.instructions.get_mut(&next)?.prev = prev,
168            None => self.basic_blocks.get_mut(&block)?.last_inst = prev,
169        }
170        Some(())
171    }
172}
173
174impl BasicBlockNode {
175    pub fn last_inst(&self) -> &Option<InstructionId> {
176        &self.last_inst
177    }
178}
179
180impl<'a> Iterator for BasicBlockIter<'a> {
181    type Item = BasicBlockId;
182
183    fn next(&mut self) -> Option<Self::Item> {
184        let cur = self.cur?;
185        self.cur = self.layout.basic_blocks[&cur].next;
186        Some(cur)
187    }
188}
189
190impl<'a> Iterator for InstructionIter<'a> {
191    type Item = InstructionId;
192
193    fn next(&mut self) -> Option<Self::Item> {
194        let cur = self.head?;
195        if Some(cur) == self.layout.basic_blocks[&self.block].last_inst {
196            self.head = None;
197        } else {
198            self.head = self.layout.instructions[&cur].next;
199        }
200        Some(cur)
201    }
202}
203
204impl<'a> DoubleEndedIterator for InstructionIter<'a> {
205    fn next_back(&mut self) -> Option<Self::Item> {
206        let cur = self.tail?;
207        if self.head == self.tail {
208            self.head = None;
209            self.tail = None;
210        } else {
211            self.tail = self.layout.instructions[&cur].prev;
212        }
213        Some(cur)
214    }
215}
vicis_core/ir/function/layout/mod.rs

vicis_core/ir/function/layout/
mod.rs
