1use crate::{
9 mir::{Function, Immediate, InstKind, Terminator, Value, ValueId},
10 pass::FunctionPass,
11 utils::evm_word,
12};
13use alloy_primitives::U256;
14use solar_data_structures::map::FxHashMap;
15
16const DEFAULT_FUEL: usize = 10_000;
17
18#[derive(Clone, Debug, Default)]
20pub struct PureEvalStats {
21 pub functions_folded: usize,
23}
24
25#[derive(Debug)]
27pub struct PureEvaluator {
28 fuel: usize,
29 stats: PureEvalStats,
30}
31
32pub struct PureEvalPass;
34
35impl FunctionPass for PureEvalPass {
36 fn name(&self) -> &str {
37 "pure-eval"
38 }
39
40 fn run_on_function(&mut self, func: &mut Function) -> bool {
41 let changed = PureEvaluator::new().run(func).functions_folded != 0;
42 let repaired = crate::mir::utils::repair_reachability_phis(func);
43 changed || repaired
44 }
45}
46
47impl Default for PureEvaluator {
48 fn default() -> Self {
49 Self { fuel: DEFAULT_FUEL, stats: PureEvalStats::default() }
50 }
51}
52
53impl PureEvaluator {
54 #[must_use]
56 pub fn new() -> Self {
57 Self::default()
58 }
59
60 #[must_use]
62 pub const fn stats(&self) -> &PureEvalStats {
63 &self.stats
64 }
65
66 pub fn run(&mut self, func: &mut Function) -> &PureEvalStats {
68 self.stats = PureEvalStats::default();
69 if !func.params.is_empty() || !self.is_side_effect_free(func) {
70 return &self.stats;
71 }
72
73 let Some(values) = self.evaluate(func) else {
74 return &self.stats;
75 };
76 if self.is_already_folded(func, &values) {
77 return &self.stats;
78 }
79 self.rewrite_to_return(func, &values);
80 self.stats.functions_folded = 1;
81 &self.stats
82 }
83
84 fn is_already_folded(&self, func: &Function, values: &[U256]) -> bool {
88 let entry = func.entry_block;
89 for (block_id, block) in func.blocks.iter_enumerated() {
90 if !block.instructions.is_empty() {
91 return false;
92 }
93 if block_id != entry && !matches!(block.terminator, Some(Terminator::Invalid)) {
94 return false;
95 }
96 }
97 let Some(Terminator::Return { values: ret }) = &func.blocks[entry].terminator else {
98 return false;
99 };
100 ret.len() == values.len()
101 && ret.iter().zip(values).all(|(&ret_value, expected)| {
102 matches!(
103 func.value(ret_value),
104 Value::Immediate(imm) if imm.as_u256() == Some(*expected)
105 )
106 })
107 }
108
109 fn is_side_effect_free(&self, func: &Function) -> bool {
110 for block in &func.blocks {
111 for &inst_id in &block.instructions {
112 if func.instructions[inst_id].kind.has_side_effects() {
113 return false;
114 }
115 }
116 }
117 true
118 }
119
120 fn evaluate(&self, func: &Function) -> Option<Vec<U256>> {
121 let mut env = FxHashMap::default();
122 for (value_id, value) in func.values.iter_enumerated() {
123 if let Value::Immediate(imm) = value
124 && let Some(value) = imm.as_u256()
125 {
126 env.insert(value_id, value);
127 }
128 }
129
130 let mut current = func.entry_block;
131 let mut predecessor = None;
132 let mut fuel = self.fuel;
133 while fuel != 0 {
134 fuel -= 1;
135 let block = &func.blocks[current];
136
137 for &inst_id in &block.instructions {
138 let inst = &func.instructions[inst_id];
139 let result = match &inst.kind {
140 InstKind::Phi(incoming) => {
141 let pred = predecessor?;
142 let (_, value) = incoming.iter().find(|(block, _)| *block == pred)?;
143 self.value_const(&env, *value)?
144 }
145 kind => self.eval_inst(kind, &env)?,
146 };
147 if let Some(value_id) = func.inst_result_value(inst_id) {
148 env.insert(value_id, result);
149 }
150 }
151
152 match block.terminator.as_ref()? {
153 Terminator::Jump(target) => {
154 predecessor = Some(current);
155 current = *target;
156 }
157 Terminator::Branch { condition, then_block, else_block } => {
158 let condition = self.value_const(&env, *condition)?;
159 predecessor = Some(current);
160 current = if condition.is_zero() { *else_block } else { *then_block };
161 }
162 Terminator::Switch { value, default, cases } => {
163 let value = self.value_const(&env, *value)?;
164 predecessor = Some(current);
165 current = cases
166 .iter()
167 .find_map(|(case, target)| {
168 (self.value_const(&env, *case)? == value).then_some(*target)
169 })
170 .unwrap_or(*default);
171 }
172 Terminator::Return { values } => {
173 return values
174 .iter()
175 .map(|&value| self.value_const(&env, value))
176 .collect::<Option<Vec<_>>>();
177 }
178 Terminator::ReturnData { .. }
179 | Terminator::Revert { .. }
180 | Terminator::Stop
181 | Terminator::SelfDestruct { .. }
182 | Terminator::Invalid => return None,
183 }
184 }
185 None
186 }
187
188 fn value_const(&self, env: &FxHashMap<ValueId, U256>, value: ValueId) -> Option<U256> {
189 env.get(&value).copied()
190 }
191
192 fn eval_inst(&self, kind: &InstKind, env: &FxHashMap<ValueId, U256>) -> Option<U256> {
193 let get = |value| self.value_const(env, value);
194 Some(match *kind {
195 InstKind::Add(a, b) => get(a)?.wrapping_add(get(b)?),
196 InstKind::Sub(a, b) => get(a)?.wrapping_sub(get(b)?),
197 InstKind::Mul(a, b) => get(a)?.wrapping_mul(get(b)?),
198 InstKind::Div(a, b) => {
199 let b = get(b)?;
200 if b.is_zero() { U256::ZERO } else { get(a)? / b }
201 }
202 InstKind::Mod(a, b) => {
203 let b = get(b)?;
204 if b.is_zero() { U256::ZERO } else { get(a)? % b }
205 }
206 InstKind::Exp(a, b) => get(a)?.wrapping_pow(get(b)?),
207 InstKind::And(a, b) => get(a)? & get(b)?,
208 InstKind::Or(a, b) => get(a)? | get(b)?,
209 InstKind::Xor(a, b) => get(a)? ^ get(b)?,
210 InstKind::Not(a) => !get(a)?,
211 InstKind::Shl(shift, value) => {
212 let shift = get(shift)?;
213 if shift >= U256::from(256) {
214 U256::ZERO
215 } else {
216 get(value)? << shift.to::<usize>()
217 }
218 }
219 InstKind::Shr(shift, value) => {
220 let shift = get(shift)?;
221 if shift >= U256::from(256) {
222 U256::ZERO
223 } else {
224 get(value)? >> shift.to::<usize>()
225 }
226 }
227 InstKind::Sar(shift, value) => evm_word::sar(get(value)?, get(shift)?),
228 InstKind::Byte(index, value) => evm_word::byte(get(index)?, get(value)?),
229 InstKind::SignExtend(size, value) => evm_word::signextend(get(size)?, get(value)?),
230 InstKind::Lt(a, b) => U256::from(get(a)? < get(b)?),
231 InstKind::Gt(a, b) => U256::from(get(a)? > get(b)?),
232 InstKind::Eq(a, b) => U256::from(get(a)? == get(b)?),
233 InstKind::IsZero(a) => U256::from(get(a)?.is_zero()),
234 InstKind::Select(condition, then_value, else_value) => {
235 if get(condition)?.is_zero() {
236 get(else_value)?
237 } else {
238 get(then_value)?
239 }
240 }
241 InstKind::Phi(_) => unreachable!("phis are handled by the block interpreter"),
242 _ => return None,
243 })
244 }
245
246 fn rewrite_to_return(&self, func: &mut Function, values: &[U256]) {
247 let entry = func.entry_block;
248 let block_ids: Vec<_> = func.blocks.indices().collect();
249 for block_id in block_ids {
250 let block = &mut func.blocks[block_id];
251 block.instructions.clear();
252 if block_id == entry {
253 block.predecessors.clear();
254 } else {
255 block.predecessors.clear();
256 block.terminator = Some(Terminator::Invalid);
257 }
258 }
259
260 let values = values
261 .iter()
262 .map(|&value| func.alloc_value(Value::Immediate(Immediate::uint256(value))))
263 .collect();
264 func.blocks[entry].terminator = Some(Terminator::Return { values });
265 }
266}