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#[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 Private,
45
46 Public,
49
50 Export,
53}
54
55impl Linkage {
56 pub fn is_exported(self) -> bool {
57 self == Linkage::Export
58 }
59}
60
61#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct FunctionBody {
65 pub fid: FunctionId,
66
67 pub store: BodyDataStore,
68
69 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#[derive(Default, Debug, Clone, PartialEq, Eq)]
91pub struct BodyDataStore {
92 insts: Arena<Inst>,
94
95 values: Arena<Value>,
97
98 blocks: Arena<BasicBlock>,
99
100 immediates: FxHashMap<(BigInt, TypeId), ValueId>,
103
104 unit_value: Option<ValueId>,
105
106 inst_results: FxHashMap<InstId, AssignableValue>,
108
109 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 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 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}