Skip to main content

thalir_core/builder/
cursor.rs

1use super::{BlockBuilder, IRContext};
2use crate::{
3    block::{BasicBlock, BlockId, Terminator},
4    instructions::Instruction,
5    IrError, Result,
6};
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, Copy)]
10pub enum CursorPosition {
11    None,
12    BlockStart(BlockId),
13    BlockEnd(BlockId),
14    After(BlockId, usize),
15    Before(BlockId, usize),
16}
17
18pub struct FunctionCursor<'a> {
19    position: CursorPosition,
20    context: &'a mut IRContext,
21    blocks: &'a mut HashMap<BlockId, BasicBlock>,
22}
23
24impl<'a> FunctionCursor<'a> {
25    pub fn new(context: &'a mut IRContext, blocks: &'a mut HashMap<BlockId, BasicBlock>) -> Self {
26        Self {
27            position: CursorPosition::None,
28            context,
29            blocks,
30        }
31    }
32
33    pub fn goto_block_start(&mut self, block: BlockId) {
34        self.position = CursorPosition::BlockStart(block);
35        self.context.set_current_block(block);
36    }
37
38    pub fn goto_block_end(&mut self, block: BlockId) {
39        self.position = CursorPosition::BlockEnd(block);
40        self.context.set_current_block(block);
41    }
42
43    pub fn goto_after_inst(&mut self, block: BlockId, inst_index: usize) {
44        self.position = CursorPosition::After(block, inst_index);
45        self.context.set_current_block(block);
46    }
47
48    pub fn goto_before_inst(&mut self, block: BlockId, inst_index: usize) {
49        self.position = CursorPosition::Before(block, inst_index);
50        self.context.set_current_block(block);
51    }
52
53    pub fn current_block(&self) -> Option<BlockId> {
54        match self.position {
55            CursorPosition::None => None,
56            CursorPosition::BlockStart(b)
57            | CursorPosition::BlockEnd(b)
58            | CursorPosition::After(b, _)
59            | CursorPosition::Before(b, _) => Some(b),
60        }
61    }
62
63    pub fn insert_inst(&mut self, inst: Instruction) -> Result<()> {
64        let block_id = self
65            .current_block()
66            .ok_or_else(|| IrError::BuilderError("Cursor not positioned".to_string()))?;
67
68        let block = self
69            .blocks
70            .get_mut(&block_id)
71            .ok_or_else(|| IrError::BuilderError(format!("Block {:?} not found", block_id)))?;
72
73        match self.position {
74            CursorPosition::BlockStart(_) => {
75                block.instructions.insert(0, inst);
76                self.position = CursorPosition::After(block_id, 0);
77            }
78            CursorPosition::BlockEnd(_) => {
79                let index = block.instructions.len();
80                block.instructions.push(inst);
81                self.position = CursorPosition::After(block_id, index);
82            }
83            CursorPosition::After(_, index) => {
84                let insert_at = index + 1;
85                if insert_at <= block.instructions.len() {
86                    block.instructions.insert(insert_at, inst);
87                    self.position = CursorPosition::After(block_id, insert_at);
88                } else {
89                    block.instructions.push(inst);
90                    self.position = CursorPosition::After(block_id, block.instructions.len() - 1);
91                }
92            }
93            CursorPosition::Before(_, index) => {
94                block.instructions.insert(index, inst);
95                self.position = CursorPosition::After(block_id, index);
96            }
97            CursorPosition::None => {
98                return Err(IrError::BuilderError("Cursor not positioned".to_string()));
99            }
100        }
101
102        Ok(())
103    }
104
105    pub fn set_terminator(&mut self, term: Terminator) -> Result<()> {
106        let block_id = self
107            .current_block()
108            .ok_or_else(|| IrError::BuilderError("Cursor not positioned".to_string()))?;
109
110        let block = self
111            .blocks
112            .get_mut(&block_id)
113            .ok_or_else(|| IrError::BuilderError(format!("Block {:?} not found", block_id)))?;
114
115        block.terminator = term;
116        Ok(())
117    }
118
119    pub fn create_block(&mut self, _name: String) -> BlockId {
120        let block_id = BlockId(self.context.next_id() as u32);
121        let block = BasicBlock::new(block_id);
122
123        self.blocks.insert(block_id, block);
124        self.goto_block_start(block_id);
125
126        block_id
127    }
128
129    pub fn split_block(&mut self, _new_block_name: String) -> Result<BlockId> {
130        let (block_id, split_at) = match self.position {
131            CursorPosition::After(b, i) => (b, i + 1),
132            CursorPosition::Before(b, i) => (b, i),
133            _ => {
134                return Err(IrError::BuilderError(
135                    "Can only split block when cursor is at an instruction".to_string(),
136                ))
137            }
138        };
139
140        let new_block_id = BlockId(self.context.next_id() as u32);
141        let mut new_block = BasicBlock::new(new_block_id);
142
143        let old_block = self
144            .blocks
145            .get_mut(&block_id)
146            .ok_or_else(|| IrError::BuilderError(format!("Block {:?} not found", block_id)))?;
147
148        if split_at < old_block.instructions.len() {
149            new_block.instructions = old_block.instructions.split_off(split_at);
150        }
151
152        new_block.terminator = old_block.terminator.clone();
153        old_block.terminator = Terminator::Jump(new_block_id, Vec::new());
154
155        self.blocks.insert(new_block_id, new_block);
156
157        self.goto_block_start(new_block_id);
158
159        Ok(new_block_id)
160    }
161
162    pub fn is_terminated(&self) -> Result<bool> {
163        let block_id = self
164            .current_block()
165            .ok_or_else(|| IrError::BuilderError("Cursor not positioned".to_string()))?;
166
167        let block = self
168            .blocks
169            .get(&block_id)
170            .ok_or_else(|| IrError::BuilderError(format!("Block {:?} not found", block_id)))?;
171
172        Ok(!matches!(block.terminator, Terminator::Invalid))
173    }
174
175    pub fn next_inst_index(&self) -> Option<usize> {
176        match self.position {
177            CursorPosition::BlockStart(_) => Some(0),
178            CursorPosition::After(_, i) => Some(i + 1),
179            CursorPosition::Before(_, i) => Some(i),
180            _ => None,
181        }
182    }
183}
184
185impl<'a> BlockBuilder<'a> {
186    pub fn cursor_at_end(&mut self) -> CursorPosition {
187        CursorPosition::BlockEnd(self.block_id)
188    }
189
190    pub fn cursor_at_start(&mut self) -> CursorPosition {
191        CursorPosition::BlockStart(self.block_id)
192    }
193}