fe_mir/ir/
body_cursor.rs

1//! This module provides a collection of structs to modify function body
2//! in-place.
3// The design used here is greatly inspired by [`cranelift`](https://crates.io/crates/cranelift)
4
5use super::{
6    value::AssignableValue, BasicBlock, BasicBlockId, FunctionBody, Inst, InstId, ValueId,
7};
8
9/// Specify a current location of [`BodyCursor`]
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum CursorLocation {
12    Inst(InstId),
13    BlockTop(BasicBlockId),
14    BlockBottom(BasicBlockId),
15    NoWhere,
16}
17
18pub struct BodyCursor<'a> {
19    body: &'a mut FunctionBody,
20    loc: CursorLocation,
21}
22
23impl<'a> BodyCursor<'a> {
24    pub fn new(body: &'a mut FunctionBody, loc: CursorLocation) -> Self {
25        Self { body, loc }
26    }
27
28    pub fn new_at_entry(body: &'a mut FunctionBody) -> Self {
29        let entry = body.order.entry();
30        Self {
31            body,
32            loc: CursorLocation::BlockTop(entry),
33        }
34    }
35    pub fn set_loc(&mut self, loc: CursorLocation) {
36        self.loc = loc;
37    }
38
39    pub fn loc(&self) -> CursorLocation {
40        self.loc
41    }
42
43    pub fn next_loc(&self) -> CursorLocation {
44        match self.loc() {
45            CursorLocation::Inst(inst) => self.body.order.next_inst(inst).map_or_else(
46                || CursorLocation::BlockBottom(self.body.order.inst_block(inst)),
47                CursorLocation::Inst,
48            ),
49            CursorLocation::BlockTop(block) => self
50                .body
51                .order
52                .first_inst(block)
53                .map_or_else(|| CursorLocation::BlockBottom(block), CursorLocation::Inst),
54            CursorLocation::BlockBottom(block) => self
55                .body()
56                .order
57                .next_block(block)
58                .map_or(CursorLocation::NoWhere, |next_block| {
59                    CursorLocation::BlockTop(next_block)
60                }),
61            CursorLocation::NoWhere => CursorLocation::NoWhere,
62        }
63    }
64
65    pub fn prev_loc(&self) -> CursorLocation {
66        match self.loc() {
67            CursorLocation::Inst(inst) => self.body.order.prev_inst(inst).map_or_else(
68                || CursorLocation::BlockTop(self.body.order.inst_block(inst)),
69                CursorLocation::Inst,
70            ),
71            CursorLocation::BlockTop(block) => self
72                .body
73                .order
74                .prev_block(block)
75                .map_or(CursorLocation::NoWhere, |prev_block| {
76                    CursorLocation::BlockBottom(prev_block)
77                }),
78            CursorLocation::BlockBottom(block) => self
79                .body
80                .order
81                .last_inst(block)
82                .map_or_else(|| CursorLocation::BlockTop(block), CursorLocation::Inst),
83            CursorLocation::NoWhere => CursorLocation::NoWhere,
84        }
85    }
86
87    pub fn next_block(&self) -> Option<BasicBlockId> {
88        let block = self.expect_block();
89        self.body.order.next_block(block)
90    }
91
92    pub fn prev_block(&self) -> Option<BasicBlockId> {
93        let block = self.expect_block();
94        self.body.order.prev_block(block)
95    }
96
97    pub fn proceed(&mut self) {
98        self.set_loc(self.next_loc())
99    }
100
101    pub fn back(&mut self) {
102        self.set_loc(self.prev_loc());
103    }
104
105    pub fn body(&self) -> &FunctionBody {
106        self.body
107    }
108
109    pub fn body_mut(&mut self) -> &mut FunctionBody {
110        self.body
111    }
112
113    /// Sets a cursor to an entry block.
114    pub fn set_to_entry(&mut self) {
115        let entry_bb = self.body().order.entry();
116        let loc = CursorLocation::BlockTop(entry_bb);
117        self.set_loc(loc);
118    }
119
120    /// Insert [`InstId`] to a location where a cursor points.
121    /// If you need to store and insert [`Inst`], use [`store_and_insert_inst`].
122    ///
123    /// # Panics
124    /// Panics if a cursor points [`CursorLocation::NoWhere`].
125    pub fn insert_inst(&mut self, inst: InstId) {
126        match self.loc() {
127            CursorLocation::Inst(at) => self.body.order.insert_inst_after(inst, at),
128            CursorLocation::BlockTop(block) => self.body.order.prepend_inst(inst, block),
129            CursorLocation::BlockBottom(block) => self.body.order.append_inst(inst, block),
130            CursorLocation::NoWhere => panic!("cursor loc points to `NoWhere`"),
131        }
132    }
133
134    pub fn store_and_insert_inst(&mut self, data: Inst) -> InstId {
135        let inst = self.body.store.store_inst(data);
136        self.insert_inst(inst);
137        inst
138    }
139
140    /// Remove a current pointed [`Inst`] from a function body. A cursor
141    /// proceeds to a next inst.
142    ///
143    /// # Panics
144    /// Panics if a cursor doesn't point [`CursorLocation::Inst`].
145    pub fn remove_inst(&mut self) {
146        let inst = self.expect_inst();
147        let next_loc = self.next_loc();
148        self.body.order.remove_inst(inst);
149        self.set_loc(next_loc);
150    }
151
152    /// Remove a current pointed `block` and contained insts from a function
153    /// body. A cursor proceeds to a next block.
154    ///
155    /// # Panics
156    /// Panics if a cursor doesn't point [`CursorLocation::Inst`].
157    pub fn remove_block(&mut self) {
158        let block = match self.loc() {
159            CursorLocation::Inst(inst) => self.body.order.inst_block(inst),
160            CursorLocation::BlockTop(block) | CursorLocation::BlockBottom(block) => block,
161            CursorLocation::NoWhere => panic!("cursor loc points `NoWhere`"),
162        };
163
164        // Store next block of the current block for later use.
165        let next_block = self.body.order.next_block(block);
166
167        // Remove all insts in the current block.
168        if let Some(first_inst) = self.body.order.first_inst(block) {
169            self.set_loc(CursorLocation::Inst(first_inst));
170            while matches!(self.loc(), CursorLocation::Inst(..)) {
171                self.remove_inst();
172            }
173        }
174        // Remove current block.
175        self.body.order.remove_block(block);
176
177        // Set cursor location to next block if exists.
178        if let Some(next_block) = next_block {
179            self.set_loc(CursorLocation::BlockTop(next_block))
180        } else {
181            self.set_loc(CursorLocation::NoWhere)
182        }
183    }
184
185    /// Insert [`BasicBlockId`] to a location where a cursor points.
186    /// If you need to store and insert [`BasicBlock`], use
187    /// [`store_and_insert_block`].
188    ///
189    /// # Panics
190    /// Panics if a cursor points [`CursorLocation::NoWhere`].
191    pub fn insert_block(&mut self, block: BasicBlockId) {
192        let current = self.expect_block();
193        self.body.order.insert_block_after_block(block, current)
194    }
195
196    pub fn store_and_insert_block(&mut self, block: BasicBlock) -> BasicBlockId {
197        let block_id = self.body.store.store_block(block);
198        self.insert_block(block_id);
199        block_id
200    }
201
202    pub fn map_result(&mut self, result: AssignableValue) -> Option<ValueId> {
203        let inst = self.expect_inst();
204        let result_value = result.value_id();
205        self.body.store.map_result(inst, result);
206        result_value
207    }
208
209    /// Returns current inst that cursor points.
210    ///
211    /// # Panics
212    /// Panics if a cursor doesn't point [`CursorLocation::Inst`].
213    pub fn expect_inst(&self) -> InstId {
214        match self.loc {
215            CursorLocation::Inst(inst) => inst,
216            _ => panic!("Cursor doesn't point any inst."),
217        }
218    }
219
220    /// Returns current block that cursor points.
221    ///
222    /// # Panics
223    /// Panics if a cursor points [`CursorLocation::NoWhere`].
224    pub fn expect_block(&self) -> BasicBlockId {
225        match self.loc {
226            CursorLocation::Inst(inst) => self.body.order.inst_block(inst),
227            CursorLocation::BlockTop(block) | CursorLocation::BlockBottom(block) => block,
228            CursorLocation::NoWhere => panic!("cursor loc points `NoWhere`"),
229        }
230    }
231}
fe_mir/ir/body_cursor.rs

fe_mir/ir/
body_cursor.rs
