Skip to main content

solar_codegen/transform/
pure_eval.rs

1//! Bounded evaluator for closed, pure MIR functions.
2//!
3//! This pass executes no-argument functions whose reachable instructions are pure and whose control
4//! flow becomes deterministic under the evaluator. It is intentionally fuel-limited and only
5//! rewrites functions that end in a raw `Return`, so ABI-returning external entries are left to the
6//! normal encoder path.
7
8use 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/// Statistics from bounded pure evaluation.
19#[derive(Clone, Debug, Default)]
20pub struct PureEvalStats {
21    /// Number of functions folded to constant returns.
22    pub functions_folded: usize,
23}
24
25/// Bounded pure MIR evaluator.
26#[derive(Debug)]
27pub struct PureEvaluator {
28    fuel: usize,
29    stats: PureEvalStats,
30}
31
32/// Function pass for bounded pure MIR evaluation.
33pub 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    /// Creates a new evaluator with the default fuel.
55    #[must_use]
56    pub fn new() -> Self {
57        Self::default()
58    }
59
60    /// Returns statistics for the most recent run.
61    #[must_use]
62    pub const fn stats(&self) -> &PureEvalStats {
63        &self.stats
64    }
65
66    /// Runs the evaluator on one function.
67    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    /// Returns true when the function is already in the exact shape
85    /// [`Self::rewrite_to_return`] would produce, so rewriting again would
86    /// report a change (and allocate fresh immediates) without progress.
87    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}