Skip to main content

solar_codegen/mir/
function.rs

1//! MIR functions.
2
3use super::{
4    BasicBlock, BlockId, InstId, InstKind, Instruction, MemoryRegion, MirType, StorageAlias, Value,
5    ValueId,
6};
7use alloy_primitives::U256;
8use solar_data_structures::{
9    fmt::{self, FmtIteratorExt},
10    index::IndexVec,
11    map::{FxHashMap, FxHashSet},
12};
13use solar_interface::Ident;
14use solar_sema::hir::{StateMutability, Visibility};
15
16/// A function in the MIR.
17#[derive(Clone, Debug)]
18pub struct Function {
19    /// Function name.
20    pub name: Ident,
21    /// Function selector (4 bytes, for external functions).
22    pub selector: Option<[u8; 4]>,
23    /// Function attributes.
24    pub attributes: FunctionAttributes,
25    /// Parameter types.
26    pub params: Vec<MirType>,
27    /// Return types.
28    pub returns: Vec<MirType>,
29    /// Bytes reserved for lowered local memory slots.
30    ///
31    /// Internal-call functions place these in the internal frame; external entries
32    /// reserve the same space in their low-memory scratch layout.
33    pub internal_frame_size: u64,
34    /// Bytes reserved for the low-memory external ABI return buffer.
35    pub external_static_return_size: u64,
36    /// All values in this function.
37    pub values: IndexVec<ValueId, Value>,
38    /// All instructions in this function.
39    pub instructions: IndexVec<InstId, Instruction>,
40    /// All basic blocks in this function.
41    pub blocks: IndexVec<BlockId, BasicBlock>,
42    /// The entry block.
43    pub entry_block: BlockId,
44}
45
46impl Function {
47    /// Creates a new function.
48    #[must_use]
49    pub fn new(name: Ident) -> Self {
50        let mut blocks = IndexVec::new();
51        let entry_block = blocks.push(BasicBlock::new());
52
53        Self {
54            name,
55            selector: None,
56            attributes: FunctionAttributes::default(),
57            params: Vec::new(),
58            returns: Vec::new(),
59            internal_frame_size: 0,
60            external_static_return_size: 0,
61            values: IndexVec::new(),
62            instructions: IndexVec::new(),
63            blocks,
64            entry_block,
65        }
66    }
67
68    /// Returns the value for the given ID.
69    #[must_use]
70    pub fn value(&self, id: ValueId) -> &Value {
71        &self.values[id]
72    }
73
74    /// Returns an immediate value as U256.
75    #[must_use]
76    pub fn value_u256(&self, id: ValueId) -> Option<U256> {
77        self.value(id).as_immediate()?.as_u256()
78    }
79
80    /// Returns an immediate value as u64 when lossless.
81    #[must_use]
82    pub fn value_u64(&self, id: ValueId) -> Option<u64> {
83        self.value_u256(id).and_then(super::utils::u256_to_u64)
84    }
85
86    /// Returns a possibly replaced immediate value as U256.
87    #[must_use]
88    pub fn value_u256_after_replacements(
89        &self,
90        id: ValueId,
91        replacements: &FxHashMap<ValueId, ValueId>,
92    ) -> Option<U256> {
93        self.value_u256(super::utils::resolve_replacement(id, replacements))
94    }
95
96    /// Returns the statically known memory region for an address value.
97    #[must_use]
98    pub fn memory_region_for_addr(&self, addr: ValueId) -> MemoryRegion {
99        match self.value(addr) {
100            Value::Immediate(imm)
101                if imm.as_u256().is_some_and(|value| value < U256::from(0x80)) =>
102            {
103                MemoryRegion::Scratch
104            }
105            _ => MemoryRegion::Unknown,
106        }
107    }
108
109    /// Returns the instruction for the given ID.
110    #[must_use]
111    pub fn instruction(&self, id: InstId) -> &Instruction {
112        &self.instructions[id]
113    }
114
115    /// Returns the value produced by the given instruction, if it has one.
116    #[must_use]
117    pub fn inst_result_value(&self, id: InstId) -> Option<ValueId> {
118        self.values
119            .iter_enumerated()
120            .find(|(_, value)| matches!(value, Value::Inst(inst) if *inst == id))
121            .map(|(value_id, _)| value_id)
122    }
123
124    /// Returns a map from each instruction to its result value.
125    #[must_use]
126    pub fn inst_results(&self) -> FxHashMap<InstId, ValueId> {
127        let mut results =
128            FxHashMap::with_capacity_and_hasher(self.instructions.len(), Default::default());
129        for (value_id, value) in self.values.iter_enumerated() {
130            if let Value::Inst(inst_id) = value {
131                results.insert(*inst_id, value_id);
132            }
133        }
134        results
135    }
136
137    /// Returns a map from each instruction to the block containing it.
138    #[must_use]
139    pub fn inst_blocks(&self) -> FxHashMap<InstId, BlockId> {
140        let mut inst_blocks =
141            FxHashMap::with_capacity_and_hasher(self.instructions.len(), Default::default());
142        for (block_id, block) in self.blocks.iter_enumerated() {
143            for &inst_id in &block.instructions {
144                inst_blocks.insert(inst_id, block_id);
145            }
146        }
147        inst_blocks
148    }
149
150    /// Returns predecessors with duplicate CFG edges collapsed.
151    #[must_use]
152    pub fn unique_predecessors(&self, block: BlockId) -> Vec<BlockId> {
153        let mut predecessors = Vec::new();
154        for &pred in &self.blocks[block].predecessors {
155            if !predecessors.contains(&pred) {
156                predecessors.push(pred);
157            }
158        }
159        predecessors
160    }
161
162    /// Returns true if the block contains any phi instruction.
163    #[must_use]
164    pub fn block_has_phi(&self, block: BlockId) -> bool {
165        self.blocks[block]
166            .instructions
167            .iter()
168            .any(|&inst_id| matches!(self.instructions[inst_id].kind, InstKind::Phi(_)))
169    }
170
171    /// Returns true if every instruction in the block is a phi instruction.
172    #[must_use]
173    pub fn block_has_only_phis(&self, block: BlockId) -> bool {
174        self.blocks[block]
175            .instructions
176            .iter()
177            .all(|&inst_id| matches!(self.instructions[inst_id].kind, InstKind::Phi(_)))
178    }
179
180    /// Returns the result values produced by phi instructions in the block.
181    #[must_use]
182    pub fn block_phi_results(&self, block: BlockId) -> FxHashSet<ValueId> {
183        self.blocks[block]
184            .instructions
185            .iter()
186            .filter_map(|&inst_id| {
187                matches!(self.instructions[inst_id].kind, InstKind::Phi(_))
188                    .then(|| self.inst_result_value(inst_id))
189                    .flatten()
190            })
191            .collect()
192    }
193
194    /// Returns the basic block for the given ID.
195    #[must_use]
196    pub fn block(&self, id: BlockId) -> &BasicBlock {
197        &self.blocks[id]
198    }
199
200    /// Returns a mutable reference to the basic block.
201    pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock {
202        &mut self.blocks[id]
203    }
204
205    /// Returns the entry block.
206    #[must_use]
207    pub fn entry(&self) -> &BasicBlock {
208        &self.blocks[self.entry_block]
209    }
210
211    /// Allocates a new value.
212    pub fn alloc_value(&mut self, value: Value) -> ValueId {
213        self.values.push(value)
214    }
215
216    /// Allocates a new instruction.
217    pub fn alloc_inst(&mut self, inst: Instruction) -> InstId {
218        self.instructions.push(inst)
219    }
220
221    /// Allocates a new basic block.
222    pub fn alloc_block(&mut self) -> BlockId {
223        self.blocks.push(BasicBlock::new())
224    }
225
226    /// Replaces all value uses according to a one-step replacement map.
227    pub fn replace_uses(&mut self, replacements: &FxHashMap<ValueId, ValueId>) {
228        if replacements.is_empty() {
229            return;
230        }
231
232        for inst in self.instructions.iter_mut() {
233            super::utils::replace_inst_uses(&mut inst.kind, replacements);
234        }
235        for block in self.blocks.iter_mut() {
236            if let Some(term) = &mut block.terminator {
237                super::utils::replace_terminator_uses(term, replacements);
238            }
239        }
240    }
241
242    /// Replaces all value uses according to a canonicalized replacement map.
243    pub fn replace_uses_canonicalized(&mut self, replacements: &FxHashMap<ValueId, ValueId>) {
244        if replacements.is_empty() {
245            return;
246        }
247
248        for inst in self.instructions.iter_mut() {
249            super::utils::replace_inst_uses_canonicalized(&mut inst.kind, replacements);
250        }
251        for block in self.blocks.iter_mut() {
252            if let Some(term) = &mut block.terminator {
253                super::utils::replace_terminator_uses_canonicalized(term, replacements);
254            }
255        }
256    }
257
258    /// Annotates storage-alias metadata for state-access instructions.
259    pub(crate) fn annotate_storage_aliases(&mut self, scope: super::utils::StorageAliasScope) {
260        let inst_ids: Vec<_> =
261            self.instructions.iter_enumerated().map(|(inst_id, _)| inst_id).collect();
262        for inst_id in inst_ids {
263            let slot = match self.instructions[inst_id].kind {
264                InstKind::SLoad(slot) | InstKind::SStore(slot, _) => Some(slot),
265                InstKind::TLoad(slot) | InstKind::TStore(slot, _)
266                    if scope == super::utils::StorageAliasScope::StorageAndTransient =>
267                {
268                    Some(slot)
269                }
270                _ => None,
271            };
272            let alias = slot.map(|slot| StorageAlias::for_value(self, slot));
273            self.instructions[inst_id].metadata.set_storage_alias(alias);
274        }
275    }
276
277    /// Returns stored storage-alias metadata, or computes a conservative alias key.
278    #[must_use]
279    pub fn storage_alias(&self, inst_id: InstId, slot: ValueId) -> StorageAlias {
280        self.instructions[inst_id]
281            .metadata
282            .storage_alias()
283            .unwrap_or_else(|| StorageAlias::for_value(self, slot))
284    }
285
286    /// Returns storage-alias metadata after applying value replacements.
287    #[must_use]
288    pub fn storage_alias_after_replacements(
289        &self,
290        inst_id: InstId,
291        slot: ValueId,
292        replacements: &FxHashMap<ValueId, ValueId>,
293    ) -> StorageAlias {
294        let original_slot = slot;
295        let slot = super::utils::resolve_replacement(slot, replacements);
296        if slot == original_slot {
297            self.storage_alias(inst_id, slot)
298        } else {
299            StorageAlias::for_value(self, slot)
300        }
301    }
302
303    /// Returns true if this function is public or external.
304    #[must_use]
305    pub fn is_public(&self) -> bool {
306        matches!(self.attributes.visibility, Visibility::Public | Visibility::External)
307    }
308
309    /// Returns the function selector as a hex string.
310    #[must_use]
311    pub fn selector_hex(&self) -> Option<String> {
312        self.selector.map(alloy_primitives::hex::encode)
313    }
314
315    /// Returns the human-readable textual MIR representation of this function.
316    pub fn to_text(&self) -> impl fmt::Display + '_ {
317        super::display::display_function_text(self)
318    }
319
320    /// Returns this function's DOT-format CFG.
321    pub fn to_dot(&self) -> impl fmt::Display + '_ {
322        super::display::display_function_dot(self)
323    }
324}
325
326/// Function attributes.
327#[derive(Clone, Debug)]
328pub struct FunctionAttributes {
329    /// Visibility modifier.
330    pub visibility: Visibility,
331    /// State mutability.
332    pub state_mutability: StateMutability,
333    /// Whether this is a constructor.
334    pub is_constructor: bool,
335    /// Whether this is a fallback function.
336    pub is_fallback: bool,
337    /// Whether this is a receive function.
338    pub is_receive: bool,
339}
340
341impl Default for FunctionAttributes {
342    fn default() -> Self {
343        Self {
344            visibility: Visibility::Internal,
345            state_mutability: StateMutability::NonPayable,
346            is_constructor: false,
347            is_fallback: false,
348            is_receive: false,
349        }
350    }
351}
352
353impl fmt::Display for Function {
354    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
355        write!(f, "fn {}({})", self.name, self.params.iter().format(", "))?;
356
357        if !self.returns.is_empty() {
358            write!(f, " -> ({})", self.returns.iter().format(", "))?;
359        }
360
361        Ok(())
362    }
363}