fe_mir/ir/
function.rs

1use fe_analyzer::namespace::items as analyzer_items;
2use fe_analyzer::namespace::types as analyzer_types;
3use fe_common::impl_intern_key;
4use fxhash::FxHashMap;
5use id_arena::Arena;
6use num_bigint::BigInt;
7use smol_str::SmolStr;
8use std::collections::BTreeMap;
9
10use super::{
11    basic_block::BasicBlock,
12    body_order::BodyOrder,
13    inst::{BranchInfo, Inst, InstId, InstKind},
14    types::TypeId,
15    value::{AssignableValue, Local, Value, ValueId},
16    BasicBlockId, SourceInfo,
17};
18
19/// Represents function signature.
20#[derive(Debug, Clone, PartialEq, Eq, Hash)]
21pub struct FunctionSignature {
22    pub params: Vec<FunctionParam>,
23    pub resolved_generics: BTreeMap<SmolStr, analyzer_types::TypeId>,
24    pub return_type: Option<TypeId>,
25    pub module_id: analyzer_items::ModuleId,
26    pub analyzer_func_id: analyzer_items::FunctionId,
27    pub linkage: Linkage,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct FunctionParam {
32    pub name: SmolStr,
33    pub ty: TypeId,
34    pub source: SourceInfo,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
38pub struct FunctionId(pub u32);
39impl_intern_key!(FunctionId);
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
42pub enum Linkage {
43    /// A function can only be called within the same module.
44    Private,
45
46    /// A function can be called from other modules, but can NOT be called from
47    /// other accounts and transactions.
48    Public,
49
50    /// A function can be called from other modules, and also can be called from
51    /// other accounts and transactions.
52    Export,
53}
54
55impl Linkage {
56    pub fn is_exported(self) -> bool {
57        self == Linkage::Export
58    }
59}
60
61/// A function body, which is not stored in salsa db to enable in-place
62/// transformation.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct FunctionBody {
65    pub fid: FunctionId,
66
67    pub store: BodyDataStore,
68
69    /// Tracks order of basic blocks and instructions in a function body.
70    pub order: BodyOrder,
71
72    pub source: SourceInfo,
73}
74
75impl FunctionBody {
76    pub fn new(fid: FunctionId, source: SourceInfo) -> Self {
77        let mut store = BodyDataStore::default();
78        let entry_bb = store.store_block(BasicBlock {});
79        Self {
80            fid,
81            store,
82            order: BodyOrder::new(entry_bb),
83            source,
84        }
85    }
86}
87
88/// A collection of basic block, instructions and values appear in a function
89/// body.
90#[derive(Default, Debug, Clone, PartialEq, Eq)]
91pub struct BodyDataStore {
92    /// Instructions appear in a function body.
93    insts: Arena<Inst>,
94
95    /// All values in a function.
96    values: Arena<Value>,
97
98    blocks: Arena<BasicBlock>,
99
100    /// Maps an immediate to a value to ensure the same immediate results in the
101    /// same value.
102    immediates: FxHashMap<(BigInt, TypeId), ValueId>,
103
104    unit_value: Option<ValueId>,
105
106    /// Maps an instruction to a value.
107    inst_results: FxHashMap<InstId, AssignableValue>,
108
109    /// All declared local variables in a function.
110    locals: Vec<ValueId>,
111}
112
113impl BodyDataStore {
114    pub fn store_inst(&mut self, inst: Inst) -> InstId {
115        self.insts.alloc(inst)
116    }
117
118    pub fn inst_data(&self, inst: InstId) -> &Inst {
119        &self.insts[inst]
120    }
121
122    pub fn inst_data_mut(&mut self, inst: InstId) -> &mut Inst {
123        &mut self.insts[inst]
124    }
125
126    pub fn replace_inst(&mut self, inst: InstId, new: Inst) -> Inst {
127        let old = &mut self.insts[inst];
128        std::mem::replace(old, new)
129    }
130
131    pub fn store_value(&mut self, value: Value) -> ValueId {
132        match value {
133            Value::Immediate { imm, ty } => self.store_immediate(imm, ty),
134
135            Value::Unit { .. } => {
136                if let Some(unit_value) = self.unit_value {
137                    unit_value
138                } else {
139                    let unit_value = self.values.alloc(value);
140                    self.unit_value = Some(unit_value);
141                    unit_value
142                }
143            }
144
145            Value::Local(ref local) => {
146                let is_user_defined = !local.is_tmp;
147                let value_id = self.values.alloc(value);
148                if is_user_defined {
149                    self.locals.push(value_id);
150                }
151                value_id
152            }
153
154            _ => self.values.alloc(value),
155        }
156    }
157
158    pub fn is_nop(&self, inst: InstId) -> bool {
159        matches!(&self.inst_data(inst).kind, InstKind::Nop)
160    }
161
162    pub fn is_terminator(&self, inst: InstId) -> bool {
163        self.inst_data(inst).is_terminator()
164    }
165
166    pub fn branch_info(&self, inst: InstId) -> BranchInfo {
167        self.inst_data(inst).branch_info()
168    }
169
170    pub fn value_data(&self, value: ValueId) -> &Value {
171        &self.values[value]
172    }
173
174    pub fn value_data_mut(&mut self, value: ValueId) -> &mut Value {
175        &mut self.values[value]
176    }
177
178    pub fn values(&self) -> impl Iterator<Item = &Value> {
179        self.values.iter().map(|(_, value_data)| value_data)
180    }
181
182    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut Value> {
183        self.values.iter_mut().map(|(_, value_data)| value_data)
184    }
185
186    pub fn store_block(&mut self, block: BasicBlock) -> BasicBlockId {
187        self.blocks.alloc(block)
188    }
189
190    /// Returns an instruction result
191    pub fn inst_result(&self, inst: InstId) -> Option<&AssignableValue> {
192        self.inst_results.get(&inst)
193    }
194
195    pub fn map_result(&mut self, inst: InstId, result: AssignableValue) {
196        self.inst_results.insert(inst, result);
197    }
198
199    pub fn remove_inst_result(&mut self, inst: InstId) -> Option<AssignableValue> {
200        self.inst_results.remove(&inst)
201    }
202
203    pub fn rewrite_branch_dest(&mut self, inst: InstId, from: BasicBlockId, to: BasicBlockId) {
204        match &mut self.inst_data_mut(inst).kind {
205            InstKind::Jump { dest } => {
206                if *dest == from {
207                    *dest = to;
208                }
209            }
210            InstKind::Branch { then, else_, .. } => {
211                if *then == from {
212                    *then = to;
213                }
214                if *else_ == from {
215                    *else_ = to;
216                }
217            }
218            _ => unreachable!("inst is not a branch"),
219        }
220    }
221
222    pub fn value_ty(&self, vid: ValueId) -> TypeId {
223        self.values[vid].ty()
224    }
225
226    pub fn locals(&self) -> &[ValueId] {
227        &self.locals
228    }
229
230    pub fn locals_mut(&mut self) -> &[ValueId] {
231        &mut self.locals
232    }
233
234    pub fn func_args(&self) -> impl Iterator<Item = ValueId> + '_ {
235        self.locals()
236            .iter()
237            .filter(|value| match self.value_data(**value) {
238                Value::Local(local) => local.is_arg,
239                _ => unreachable!(),
240            })
241            .copied()
242    }
243
244    pub fn func_args_mut(&mut self) -> impl Iterator<Item = &mut Value> {
245        self.values_mut().filter(|value| match value {
246            Value::Local(local) => local.is_arg,
247            _ => false,
248        })
249    }
250
251    /// Returns Some(`local_name`) if value is `Value::Local`.
252    pub fn local_name(&self, value: ValueId) -> Option<&str> {
253        match self.value_data(value) {
254            Value::Local(Local { name, .. }) => Some(name),
255            _ => None,
256        }
257    }
258
259    pub fn replace_value(&mut self, value: ValueId, to: Value) -> Value {
260        std::mem::replace(&mut self.values[value], to)
261    }
262
263    fn store_immediate(&mut self, imm: BigInt, ty: TypeId) -> ValueId {
264        if let Some(value) = self.immediates.get(&(imm.clone(), ty)) {
265            *value
266        } else {
267            let id = self.values.alloc(Value::Immediate {
268                imm: imm.clone(),
269                ty,
270            });
271            self.immediates.insert((imm, ty), id);
272            id
273        }
274    }
275}