1use 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#[derive(Clone, Debug)]
18pub struct Function {
19 pub name: Ident,
21 pub selector: Option<[u8; 4]>,
23 pub attributes: FunctionAttributes,
25 pub params: Vec<MirType>,
27 pub returns: Vec<MirType>,
29 pub internal_frame_size: u64,
34 pub external_static_return_size: u64,
36 pub values: IndexVec<ValueId, Value>,
38 pub instructions: IndexVec<InstId, Instruction>,
40 pub blocks: IndexVec<BlockId, BasicBlock>,
42 pub entry_block: BlockId,
44}
45
46impl Function {
47 #[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 #[must_use]
70 pub fn value(&self, id: ValueId) -> &Value {
71 &self.values[id]
72 }
73
74 #[must_use]
76 pub fn value_u256(&self, id: ValueId) -> Option<U256> {
77 self.value(id).as_immediate()?.as_u256()
78 }
79
80 #[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 #[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 #[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 #[must_use]
111 pub fn instruction(&self, id: InstId) -> &Instruction {
112 &self.instructions[id]
113 }
114
115 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[must_use]
196 pub fn block(&self, id: BlockId) -> &BasicBlock {
197 &self.blocks[id]
198 }
199
200 pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock {
202 &mut self.blocks[id]
203 }
204
205 #[must_use]
207 pub fn entry(&self) -> &BasicBlock {
208 &self.blocks[self.entry_block]
209 }
210
211 pub fn alloc_value(&mut self, value: Value) -> ValueId {
213 self.values.push(value)
214 }
215
216 pub fn alloc_inst(&mut self, inst: Instruction) -> InstId {
218 self.instructions.push(inst)
219 }
220
221 pub fn alloc_block(&mut self) -> BlockId {
223 self.blocks.push(BasicBlock::new())
224 }
225
226 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 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 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 #[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 #[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 #[must_use]
305 pub fn is_public(&self) -> bool {
306 matches!(self.attributes.visibility, Visibility::Public | Visibility::External)
307 }
308
309 #[must_use]
311 pub fn selector_hex(&self) -> Option<String> {
312 self.selector.map(alloy_primitives::hex::encode)
313 }
314
315 pub fn to_text(&self) -> impl fmt::Display + '_ {
317 super::display::display_function_text(self)
318 }
319
320 pub fn to_dot(&self) -> impl fmt::Display + '_ {
322 super::display::display_function_dot(self)
323 }
324}
325
326#[derive(Clone, Debug)]
328pub struct FunctionAttributes {
329 pub visibility: Visibility,
331 pub state_mutability: StateMutability,
333 pub is_constructor: bool,
335 pub is_fallback: bool,
337 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}