Skip to main content

solar_codegen/transform/
cse.rs

1//! Common Subexpression Elimination (CSE) optimization pass.
2//!
3//! This pass identifies and eliminates redundant computations within basic blocks.
4//! When the same expression is computed multiple times with the same operands,
5//! only the first computation is kept and subsequent uses reference the cached result.
6//!
7//! ## Example
8//!
9//! Before CSE:
10//! ```text
11//! v1 = add v0, 42
12//! v2 = mul v1, 2
13//! v3 = add v0, 42  // redundant - same as v1
14//! v4 = mul v3, 3
15//! ```
16//!
17//! After CSE:
18//! ```text
19//! v1 = add v0, 42
20//! v2 = mul v1, 2
21//! // v3 removed, uses of v3 replaced with v1
22//! v4 = mul v1, 3
23//! ```
24//!
25//! The pass performs dominator-tree CSE with path-local invalidation for
26//! alias-sensitive memory/storage reads, then runs a local cleanup pass.
27//!
28//! Safety contract:
29//! - cache only pure expressions, classified memory reads, and exact storage or transient-storage
30//!   reads
31//! - invalidate memory reads by overlapping memory writes and unknown memory effects
32//! - invalidate storage reads by possibly-aliasing writes or calls that may re-enter and mutate the
33//!   current contract
34//! - when inheriting a cache across a dominator-tree edge, also invalidate state-dependent reads by
35//!   clobbers in every block that can lie on a CFG path between the dominator and its child
36//!   (diamond arms, loop bodies), including the child itself when it sits on a cycle
37
38use crate::{
39    analysis::{CfgInfo, DominatorTree},
40    mir::{
41        BlockId, Function, Immediate, InstId, InstKind, Instruction, MemoryRegion, MirType,
42        StorageAlias, Value, ValueId, utils as mir_utils,
43    },
44    pass::FunctionPass,
45};
46use alloy_primitives::U256;
47use solar_data_structures::map::{FxHashMap, FxHashSet};
48use std::cmp::Ordering;
49
50/// Common Subexpression Elimination pass.
51#[derive(Debug, Default)]
52pub struct CommonSubexprEliminator {
53    /// Number of instructions eliminated.
54    pub eliminated_count: usize,
55}
56
57/// Function pass for local common subexpression elimination.
58pub struct CsePass;
59
60impl FunctionPass for CsePass {
61    fn name(&self) -> &str {
62        "cse"
63    }
64
65    fn run_on_function(&mut self, func: &mut Function) -> bool {
66        CommonSubexprEliminator::new().run_to_fixpoint(func) != 0
67    }
68}
69
70/// A normalized expression key for CSE lookup.
71/// Expressions are normalized so that equivalent computations map to the same key.
72#[derive(Clone, Debug, PartialEq, Eq, Hash)]
73enum ExprKey {
74    Add(OperandKey, OperandKey),
75    Offset(OperandKey, U256),
76    Sub(OperandKey, OperandKey),
77    Mul(OperandKey, OperandKey),
78    Div(OperandKey, OperandKey),
79    SDiv(OperandKey, OperandKey),
80    Mod(OperandKey, OperandKey),
81    SMod(OperandKey, OperandKey),
82    Exp(OperandKey, OperandKey),
83    AddMod(OperandKey, OperandKey, OperandKey),
84    MulMod(OperandKey, OperandKey, OperandKey),
85    And(OperandKey, OperandKey),
86    Or(OperandKey, OperandKey),
87    Xor(OperandKey, OperandKey),
88    Shl(OperandKey, OperandKey),
89    Shr(OperandKey, OperandKey),
90    Sar(OperandKey, OperandKey),
91    Byte(OperandKey, OperandKey),
92    /// Also keys `Gt(a, b)`, normalized as `Lt(b, a)`.
93    Lt(OperandKey, OperandKey),
94    /// Also keys `SGt(a, b)`, normalized as `SLt(b, a)`.
95    SLt(OperandKey, OperandKey),
96    Eq(OperandKey, OperandKey),
97    IsZero(OperandKey),
98    Not(OperandKey),
99    SignExtend(OperandKey, OperandKey),
100    Select(OperandKey, OperandKey, OperandKey),
101    MLoad(MemRangeKey),
102    Keccak256(MemRangeKey),
103    SLoad(StorageAlias),
104    TLoad(StorageAlias),
105    CalldataLoad(OperandKey),
106    ExtCodeSize(OperandKey),
107    ExtCodeHash(OperandKey),
108    BlockHash(OperandKey),
109    Balance(OperandKey),
110    SelfBalance,
111    BlobHash(OperandKey),
112    LoadImmutable(u32),
113}
114
115#[derive(Clone, Debug, PartialEq, Eq, Hash)]
116enum OperandKey {
117    Value(ValueId),
118    Immediate(Immediate),
119}
120
121#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
122struct MemRangeKey {
123    region: MemoryRegion,
124    base: Option<ValueId>,
125    offset: Option<u64>,
126    size: Option<u64>,
127    /// The canonical size operand when `size` is not a known constant.
128    ///
129    /// Participates only in key equality so that reads with different dynamic sizes never
130    /// unify while reads with the same dynamic size operand still do. Aliasing checks ignore
131    /// it: `memory_ranges_may_alias` stays conservative whenever `size` is `None`, so write
132    /// keys can always leave it unset.
133    dyn_size: Option<ValueId>,
134}
135
136struct GlobalCseContext<'a> {
137    dom_tree: &'a DominatorTree,
138    inst_results: &'a FxHashMap<InstId, ValueId>,
139    block_clobbers: &'a FxHashMap<BlockId, Vec<Clobber>>,
140    reachability: &'a FxHashMap<BlockId, FxHashSet<BlockId>>,
141    replacements: &'a mut FxHashMap<ValueId, ValueId>,
142    dead: &'a mut FxHashSet<InstId>,
143}
144
145/// A single cache-invalidating effect of a side-effecting instruction.
146#[derive(Clone, Copy, Debug)]
147enum Clobber {
148    /// A memory write; `None` clobbers all of memory.
149    Memory(Option<MemRangeKey>),
150    /// A storage write to a possibly-aliasing slot.
151    Storage(StorageAlias),
152    /// An effect that may mutate any storage slot (e.g. a re-entering call).
153    AllStorage,
154    /// A transient-storage write to a possibly-aliasing slot.
155    Transient(StorageAlias),
156    /// An effect that may mutate any transient-storage slot.
157    AllTransient,
158    /// An effect that may change account balances or deployed code.
159    AccountEnvironment,
160}
161
162struct PhiExpressionCandidate {
163    block_id: BlockId,
164    phi_inst: InstId,
165    phi_result: ValueId,
166    kind: InstKind,
167    result_ty: MirType,
168    incoming: Vec<(ValueId, InstId)>,
169}
170
171struct PhiSinkContext<'a> {
172    dominators: &'a DominatorTree,
173    inst_blocks: &'a FxHashMap<InstId, BlockId>,
174    inst_results: &'a FxHashMap<InstId, ValueId>,
175    replacements: &'a FxHashMap<ValueId, ValueId>,
176}
177
178impl CommonSubexprEliminator {
179    /// Creates a new CSE pass.
180    pub fn new() -> Self {
181        Self::default()
182    }
183
184    /// Runs CSE on a function.
185    /// Returns the number of expressions eliminated.
186    pub fn run(&mut self, func: &mut Function) -> usize {
187        self.eliminated_count = 0;
188
189        self.sink_redundant_phi_expressions(func);
190
191        // Neither the global nor the local pass allocates values, so the map stays valid.
192        let inst_results = func.inst_results();
193        self.process_global_pure(func, &inst_results);
194
195        // Process each block independently (local CSE)
196        let block_ids: Vec<BlockId> = func.blocks.indices().collect();
197        for block_id in block_ids {
198            self.process_block(func, block_id, &inst_results);
199        }
200
201        self.eliminated_count
202    }
203
204    /// Runs CSE iteratively until no more changes.
205    pub fn run_to_fixpoint(&mut self, func: &mut Function) -> usize {
206        let mut total = 0;
207        loop {
208            let eliminated = self.run(func);
209            if eliminated == 0 {
210                break;
211            }
212            total += eliminated;
213        }
214        total
215    }
216
217    fn process_global_pure(
218        &mut self,
219        func: &mut Function,
220        inst_results: &FxHashMap<InstId, ValueId>,
221    ) {
222        let mut cfg = CfgInfo::new(func);
223        let block_clobbers = self.block_clobber_summaries(func);
224        let reachability = if block_clobbers.is_empty() {
225            FxHashMap::default()
226        } else {
227            cfg.transitive_reachability().clone()
228        };
229        let mut replacements = FxHashMap::default();
230        let mut dead = FxHashSet::default();
231        let mut cache = FxHashMap::default();
232        let mut ctx = GlobalCseContext {
233            dom_tree: cfg.dominators(),
234            inst_results,
235            block_clobbers: &block_clobbers,
236            reachability: &reachability,
237            replacements: &mut replacements,
238            dead: &mut dead,
239        };
240
241        self.process_global_block(func, func.entry_block, &mut cache, &mut ctx);
242
243        if !replacements.is_empty() {
244            self.apply_replacements_to_all_blocks(func, &replacements);
245        }
246        if !dead.is_empty() {
247            for block in func.blocks.iter_mut() {
248                block.instructions.retain(|id| !dead.contains(id));
249            }
250        }
251    }
252
253    fn sink_redundant_phi_expressions(&mut self, func: &mut Function) {
254        let cfg = CfgInfo::new(func);
255        let inst_results = func.inst_results();
256        let inst_blocks = func.inst_blocks();
257        let use_counts = Self::value_use_counts(func);
258        let replacements = FxHashMap::default();
259        let ctx = PhiSinkContext {
260            dominators: cfg.dominators(),
261            inst_blocks: &inst_blocks,
262            inst_results: &inst_results,
263            replacements: &replacements,
264        };
265        let mut candidates = Vec::new();
266
267        for block_id in func.blocks.indices() {
268            let phi_insts: Vec<_> = func.blocks[block_id]
269                .instructions
270                .iter()
271                .copied()
272                .take_while(|&inst_id| matches!(func.instructions[inst_id].kind, InstKind::Phi(_)))
273                .collect();
274            for phi_inst in phi_insts {
275                if let Some(candidate) =
276                    self.phi_expression_candidate(func, block_id, phi_inst, &ctx)
277                {
278                    candidates.push(candidate);
279                }
280            }
281        }
282
283        if candidates.is_empty() {
284            return;
285        }
286
287        let mut dead = FxHashSet::default();
288        let mut replacements = FxHashMap::default();
289        let mut inserted_by_block: FxHashMap<BlockId, usize> = FxHashMap::default();
290
291        for candidate in candidates {
292            let new_inst =
293                func.alloc_inst(Instruction::new(candidate.kind, Some(candidate.result_ty)));
294            let new_value = func.alloc_value(Value::Inst(new_inst));
295
296            let phi_count = func.blocks[candidate.block_id]
297                .instructions
298                .iter()
299                .take_while(|&&inst_id| matches!(func.instructions[inst_id].kind, InstKind::Phi(_)))
300                .count();
301            let inserted = inserted_by_block.entry(candidate.block_id).or_default();
302            func.blocks[candidate.block_id].instructions.insert(phi_count + *inserted, new_inst);
303            *inserted += 1;
304
305            replacements.insert(candidate.phi_result, new_value);
306            dead.insert(candidate.phi_inst);
307            for (value, inst_id) in candidate.incoming {
308                if use_counts.get(&value).copied().unwrap_or_default() == 1 {
309                    dead.insert(inst_id);
310                }
311            }
312            self.eliminated_count += 1;
313        }
314
315        self.apply_replacements_to_all_blocks(func, &replacements);
316        for block in func.blocks.iter_mut() {
317            block.instructions.retain(|id| !dead.contains(id));
318        }
319    }
320
321    fn phi_expression_candidate(
322        &self,
323        func: &Function,
324        block_id: BlockId,
325        phi_inst: InstId,
326        ctx: &PhiSinkContext<'_>,
327    ) -> Option<PhiExpressionCandidate> {
328        let inst = &func.instructions[phi_inst];
329        let result_ty = inst.result_ty?;
330        let phi_result = *ctx.inst_results.get(&phi_inst)?;
331        let InstKind::Phi(incoming) = &inst.kind else { return None };
332        if incoming.len() < 2 {
333            return None;
334        }
335
336        let mut expected_key = None;
337        let mut candidate_kind = None;
338        let mut incoming_insts = Vec::with_capacity(incoming.len());
339
340        for &(_, value) in incoming {
341            let Value::Inst(inst_id) = func.value(value) else { return None };
342            let source_inst = &func.instructions[*inst_id];
343            if source_inst.kind.has_side_effects()
344                || !Self::operands_dominate_block(
345                    func,
346                    &source_inst.kind,
347                    block_id,
348                    ctx.inst_blocks,
349                    ctx.dominators,
350                )
351            {
352                return None;
353            }
354
355            let key = self.make_expr_key(func, *inst_id, &source_inst.kind, ctx.replacements)?;
356            if !Self::is_sinkable_pure_expr(&key) {
357                return None;
358            }
359            if expected_key.as_ref().is_some_and(|expected| expected != &key) {
360                return None;
361            }
362            expected_key = Some(key);
363            candidate_kind.get_or_insert_with(|| source_inst.kind.clone());
364            incoming_insts.push((value, *inst_id));
365        }
366
367        Some(PhiExpressionCandidate {
368            block_id,
369            phi_inst,
370            phi_result,
371            kind: candidate_kind?,
372            result_ty,
373            incoming: incoming_insts,
374        })
375    }
376
377    fn process_global_block(
378        &mut self,
379        func: &Function,
380        block_id: BlockId,
381        cache: &mut FxHashMap<ExprKey, ValueId>,
382        ctx: &mut GlobalCseContext<'_>,
383    ) {
384        let inst_ids = func.blocks[block_id].instructions.clone();
385        for inst_id in inst_ids {
386            let kind = func.instructions[inst_id].kind.clone();
387            if kind.has_side_effects() {
388                self.invalidate_for_side_effect(func, inst_id, &kind, ctx.replacements, cache);
389                continue;
390            }
391
392            let Some(key) = self.make_expr_key(func, inst_id, &kind, ctx.replacements) else {
393                continue;
394            };
395
396            let Some(&result) = ctx.inst_results.get(&inst_id) else {
397                continue;
398            };
399            if let Some(cached) = cache.get(&key) {
400                ctx.replacements.insert(result, *cached);
401                ctx.dead.insert(inst_id);
402                self.eliminated_count += 1;
403            } else {
404                cache.insert(key, result);
405            }
406        }
407
408        for &child in ctx.dom_tree.children(block_id) {
409            let mut child_cache = cache.clone();
410            self.filter_inherited_cache(block_id, child, &mut child_cache, ctx);
411            self.process_global_block(func, child, &mut child_cache, ctx);
412        }
413    }
414
415    /// Invalidates state-dependent cache entries inherited across the dominator-tree edge
416    /// `parent -> child`.
417    ///
418    /// Dominance alone is sound only for pure expressions: memory, storage, transient-storage, and
419    /// account-environment reads must also survive every CFG path from `parent` to `child`, which
420    /// may pass through blocks that are not on the dominator-tree path (diamond arms, loop bodies).
421    /// Applies the clobber summary of every such intermediate block, including `child` itself when
422    /// it lies on a cycle (clobbers wrap around the backedge to the child's entry).
423    fn filter_inherited_cache(
424        &self,
425        parent: BlockId,
426        child: BlockId,
427        cache: &mut FxHashMap<ExprKey, ValueId>,
428        ctx: &GlobalCseContext<'_>,
429    ) {
430        if ctx.block_clobbers.is_empty() || !cache.keys().any(Self::is_path_sensitive_expr) {
431            return;
432        }
433        let Some(reachable_from_parent) = ctx.reachability.get(&parent) else { return };
434        for (&mid, clobbers) in ctx.block_clobbers {
435            // Clobbers in `parent` itself were already applied while processing it sequentially.
436            if mid == parent || !reachable_from_parent.contains(&mid) {
437                continue;
438            }
439            if !ctx.reachability.get(&mid).is_some_and(|reachable| reachable.contains(&child)) {
440                continue;
441            }
442            for clobber in clobbers {
443                self.apply_clobber(cache, clobber);
444            }
445        }
446    }
447
448    /// Returns the per-block invalidation summaries for blocks with clobbering effects.
449    fn block_clobber_summaries(&self, func: &Function) -> FxHashMap<BlockId, Vec<Clobber>> {
450        let no_replacements = FxHashMap::default();
451        let mut summaries = FxHashMap::default();
452        for (block_id, block) in func.blocks.iter_enumerated() {
453            let mut clobbers = Vec::new();
454            for &inst_id in &block.instructions {
455                let kind = &func.instructions[inst_id].kind;
456                if kind.has_side_effects() {
457                    self.side_effect_clobbers(func, inst_id, kind, &no_replacements, &mut clobbers);
458                }
459            }
460            if !clobbers.is_empty() {
461                summaries.insert(block_id, clobbers);
462            }
463        }
464        summaries
465    }
466
467    fn is_path_sensitive_expr(key: &ExprKey) -> bool {
468        Self::is_memory_expr(key)
469            || Self::is_account_environment_expr(key)
470            || matches!(key, ExprKey::SLoad(_) | ExprKey::TLoad(_))
471    }
472
473    /// Processes a single basic block.
474    fn process_block(
475        &mut self,
476        func: &mut Function,
477        block_id: BlockId,
478        inst_results: &FxHashMap<InstId, ValueId>,
479    ) {
480        // Map from expression key to the ValueId that computed it
481        let mut expr_cache: FxHashMap<ExprKey, ValueId> = FxHashMap::default();
482
483        // Map from ValueId to its replacement ValueId
484        let mut replacements: FxHashMap<ValueId, ValueId> = FxHashMap::default();
485
486        // Instructions to remove
487        let mut to_remove: FxHashSet<InstId> = FxHashSet::default();
488
489        // Get instruction list for this block
490        let block = func.block(block_id);
491        let inst_ids: Vec<InstId> = block.instructions.clone();
492
493        for inst_id in inst_ids {
494            let inst = &func.instructions[inst_id];
495            let kind = inst.kind.clone();
496
497            if kind.has_side_effects() {
498                self.invalidate_for_side_effect(
499                    func,
500                    inst_id,
501                    &kind,
502                    &replacements,
503                    &mut expr_cache,
504                );
505                continue;
506            }
507
508            // Try to create an expression key
509            if let Some(key) = self.make_expr_key(func, inst_id, &kind, &replacements)
510                && let Some(&result) = inst_results.get(&inst_id)
511            {
512                if let Some(&cached_value) = expr_cache.get(&key) {
513                    // This expression was already computed - mark for elimination
514                    replacements.insert(result, cached_value);
515                    to_remove.insert(inst_id);
516                    self.eliminated_count += 1;
517                } else {
518                    // First occurrence - cache it
519                    expr_cache.insert(key, result);
520                }
521            }
522        }
523
524        // Apply replacements everywhere: the eliminated result may be used in dominated blocks.
525        if !replacements.is_empty() {
526            self.apply_replacements_to_all_blocks(func, &replacements);
527        }
528
529        // Remove eliminated instructions
530        let block = func.block_mut(block_id);
531        block.instructions.retain(|id| !to_remove.contains(id));
532    }
533
534    /// Creates a normalized expression key for an instruction.
535    /// Returns None for instructions that shouldn't be cached.
536    fn make_expr_key(
537        &self,
538        func: &Function,
539        inst_id: InstId,
540        kind: &InstKind,
541        replacements: &FxHashMap<ValueId, ValueId>,
542    ) -> Option<ExprKey> {
543        // Helper to get canonical operands after in-block replacements.
544        let operand = |v: ValueId| Self::operand_key(func, v, replacements);
545        let value = |v: ValueId| mir_utils::resolve_replacement(v, replacements);
546
547        match kind {
548            // Commutative operations - normalize operand order
549            InstKind::Add(a, b) => {
550                if let Some((base, offset)) = Self::offset_expr_for_add(func, *a, *b, replacements)
551                {
552                    Some(ExprKey::Offset(base, offset))
553                } else {
554                    let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
555                    Some(ExprKey::Add(a, b))
556                }
557            }
558            InstKind::Mul(a, b) => {
559                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
560                Some(ExprKey::Mul(a, b))
561            }
562            InstKind::And(a, b) => {
563                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
564                Some(ExprKey::And(a, b))
565            }
566            InstKind::Or(a, b) => {
567                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
568                Some(ExprKey::Or(a, b))
569            }
570            InstKind::Xor(a, b) => {
571                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
572                Some(ExprKey::Xor(a, b))
573            }
574            InstKind::Eq(a, b) => {
575                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
576                Some(ExprKey::Eq(a, b))
577            }
578
579            // Non-commutative operations - preserve order
580            InstKind::Sub(a, b) => {
581                if let Some((base, offset)) = Self::offset_expr_for_sub(func, *a, *b, replacements)
582                {
583                    Some(ExprKey::Offset(base, offset))
584                } else {
585                    Some(ExprKey::Sub(operand(*a), operand(*b)))
586                }
587            }
588            InstKind::Div(a, b) => Some(ExprKey::Div(operand(*a), operand(*b))),
589            InstKind::SDiv(a, b) => Some(ExprKey::SDiv(operand(*a), operand(*b))),
590            InstKind::Mod(a, b) => Some(ExprKey::Mod(operand(*a), operand(*b))),
591            InstKind::SMod(a, b) => Some(ExprKey::SMod(operand(*a), operand(*b))),
592            InstKind::Exp(a, b) => Some(ExprKey::Exp(operand(*a), operand(*b))),
593            InstKind::AddMod(a, b, n) => {
594                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
595                Some(ExprKey::AddMod(a, b, operand(*n)))
596            }
597            InstKind::MulMod(a, b, n) => {
598                let (a, b) = Self::ordered_pair(operand(*a), operand(*b));
599                Some(ExprKey::MulMod(a, b, operand(*n)))
600            }
601            InstKind::Shl(a, b) => Some(ExprKey::Shl(operand(*a), operand(*b))),
602            InstKind::Shr(a, b) => Some(ExprKey::Shr(operand(*a), operand(*b))),
603            InstKind::Sar(a, b) => Some(ExprKey::Sar(operand(*a), operand(*b))),
604            InstKind::Byte(a, b) => Some(ExprKey::Byte(operand(*a), operand(*b))),
605            // Swapped comparisons - canonicalize `a > b` as `b < a` so they unify
606            InstKind::Lt(a, b) => Some(ExprKey::Lt(operand(*a), operand(*b))),
607            InstKind::Gt(a, b) => Some(ExprKey::Lt(operand(*b), operand(*a))),
608            InstKind::SLt(a, b) => Some(ExprKey::SLt(operand(*a), operand(*b))),
609            InstKind::SGt(a, b) => Some(ExprKey::SLt(operand(*b), operand(*a))),
610            InstKind::SignExtend(a, b) => Some(ExprKey::SignExtend(operand(*a), operand(*b))),
611
612            // Unary operations
613            InstKind::IsZero(a) => Some(ExprKey::IsZero(operand(*a))),
614            InstKind::Not(a) => Some(ExprKey::Not(operand(*a))),
615            InstKind::CalldataLoad(a) => Some(ExprKey::CalldataLoad(operand(*a))),
616            InstKind::ExtCodeSize(a) => Some(ExprKey::ExtCodeSize(operand(*a))),
617            InstKind::ExtCodeHash(a) => Some(ExprKey::ExtCodeHash(operand(*a))),
618            InstKind::Balance(a) => Some(ExprKey::Balance(operand(*a))),
619            InstKind::BlockHash(a) => Some(ExprKey::BlockHash(operand(*a))),
620            InstKind::BlobHash(a) => Some(ExprKey::BlobHash(operand(*a))),
621            // Immutable reads are constant once the runtime code is patched.
622            InstKind::LoadImmutable(offset) => Some(ExprKey::LoadImmutable(*offset)),
623
624            InstKind::Select(condition, then_value, else_value) => Some(ExprKey::Select(
625                operand(*condition),
626                operand(*then_value),
627                operand(*else_value),
628            )),
629
630            InstKind::MLoad(addr) => {
631                let key = self.memory_range_key(func, inst_id, value(*addr), Some(32))?;
632                Some(ExprKey::MLoad(key))
633            }
634            InstKind::Keccak256(offset, size) => {
635                let size = value(*size);
636                let const_size = func.value_u64(size);
637                let mut key = self.memory_range_key(func, inst_id, value(*offset), const_size)?;
638                if const_size.is_none() {
639                    // Key the dynamic size operand so reads of different lengths never unify.
640                    key.dyn_size = Some(size);
641                }
642                Some(ExprKey::Keccak256(key))
643            }
644
645            InstKind::SLoad(slot) => Some(ExprKey::SLoad(func.storage_alias_after_replacements(
646                inst_id,
647                *slot,
648                replacements,
649            ))),
650            InstKind::TLoad(slot) => Some(ExprKey::TLoad(func.storage_alias_after_replacements(
651                inst_id,
652                *slot,
653                replacements,
654            ))),
655
656            InstKind::SelfBalance => Some(ExprKey::SelfBalance),
657
658            // Don't cache these:
659            // - Cheap nullary reads usually cost less than their extra stack lifetime
660            // - Memory size/gas/returndata-size reads can change inside a block
661            // - Storage writes - side effects
662            // - Phi nodes - not expressions
663            // - Calls - side effects
664            _ => None,
665        }
666    }
667
668    fn invalidate_for_side_effect(
669        &self,
670        func: &Function,
671        inst_id: InstId,
672        kind: &InstKind,
673        replacements: &FxHashMap<ValueId, ValueId>,
674        expr_cache: &mut FxHashMap<ExprKey, ValueId>,
675    ) {
676        let mut clobbers = Vec::new();
677        self.side_effect_clobbers(func, inst_id, kind, replacements, &mut clobbers);
678        for clobber in &clobbers {
679            self.apply_clobber(expr_cache, clobber);
680        }
681    }
682
683    /// Collects the cache-invalidating effects of a side-effecting instruction.
684    fn side_effect_clobbers(
685        &self,
686        func: &Function,
687        inst_id: InstId,
688        kind: &InstKind,
689        replacements: &FxHashMap<ValueId, ValueId>,
690        clobbers: &mut Vec<Clobber>,
691    ) {
692        match *kind {
693            InstKind::MStore(addr, _) => {
694                let addr = mir_utils::resolve_replacement(addr, replacements);
695                clobbers.push(Clobber::Memory(self.memory_range_key(
696                    func,
697                    inst_id,
698                    addr,
699                    Some(32),
700                )));
701            }
702            InstKind::MStore8(addr, _) => {
703                let addr = mir_utils::resolve_replacement(addr, replacements);
704                clobbers.push(Clobber::Memory(self.memory_range_key(func, inst_id, addr, Some(1))));
705            }
706            InstKind::MCopy(dest, _, size)
707            | InstKind::CalldataCopy(dest, _, size)
708            | InstKind::CodeCopy(dest, _, size)
709            | InstKind::ReturnDataCopy(dest, _, size)
710            | InstKind::ExtCodeCopy(_, dest, _, size) => {
711                let dest = mir_utils::resolve_replacement(dest, replacements);
712                let size = func.value_u64(mir_utils::resolve_replacement(size, replacements));
713                clobbers.push(Clobber::Memory(self.memory_range_key(func, inst_id, dest, size)));
714            }
715            InstKind::SStore(slot, _) => {
716                clobbers.push(Clobber::Storage(func.storage_alias_after_replacements(
717                    inst_id,
718                    slot,
719                    replacements,
720                )));
721            }
722            InstKind::TStore(slot, _) => {
723                clobbers.push(Clobber::Transient(func.storage_alias_after_replacements(
724                    inst_id,
725                    slot,
726                    replacements,
727                )));
728            }
729            _ if kind.may_mutate_memory() => {
730                clobbers.push(Clobber::Memory(None));
731                if Self::may_change_account_environment(kind) {
732                    clobbers.push(Clobber::AccountEnvironment);
733                }
734                if kind.may_mutate_storage() {
735                    clobbers.push(Clobber::AllStorage);
736                }
737                if kind.may_mutate_transient_storage() {
738                    clobbers.push(Clobber::AllTransient);
739                }
740            }
741            _ if kind.may_mutate_storage() => clobbers.push(Clobber::AllStorage),
742            _ if kind.may_mutate_transient_storage() => clobbers.push(Clobber::AllTransient),
743            _ => {}
744        }
745    }
746
747    /// Removes cache entries invalidated by a single clobbering effect.
748    fn apply_clobber(&self, expr_cache: &mut FxHashMap<ExprKey, ValueId>, clobber: &Clobber) {
749        match *clobber {
750            Clobber::Memory(write) => self.invalidate_memory(expr_cache, write),
751            Clobber::Storage(alias) => {
752                expr_cache.retain(|key, _| match key {
753                    ExprKey::SLoad(cached) => !cached.may_alias(alias),
754                    _ => true,
755                });
756            }
757            Clobber::AllStorage => {
758                expr_cache.retain(|key, _| !matches!(key, ExprKey::SLoad(_)));
759            }
760            Clobber::Transient(alias) => {
761                expr_cache.retain(|key, _| match key {
762                    ExprKey::TLoad(cached) => !cached.may_alias(alias),
763                    _ => true,
764                });
765            }
766            Clobber::AllTransient => {
767                expr_cache.retain(|key, _| !matches!(key, ExprKey::TLoad(_)));
768            }
769            Clobber::AccountEnvironment => {
770                expr_cache.retain(|key, _| !Self::is_account_environment_expr(key));
771            }
772        }
773    }
774
775    fn invalidate_memory(
776        &self,
777        expr_cache: &mut FxHashMap<ExprKey, ValueId>,
778        write: Option<MemRangeKey>,
779    ) {
780        expr_cache.retain(|key, _| match key {
781            ExprKey::MLoad(read) | ExprKey::Keccak256(read) => {
782                write.is_some_and(|write| !Self::memory_ranges_may_alias(*read, write))
783            }
784            _ => true,
785        });
786    }
787
788    fn is_memory_expr(key: &ExprKey) -> bool {
789        matches!(key, ExprKey::MLoad(_) | ExprKey::Keccak256(_))
790    }
791
792    fn is_account_environment_expr(key: &ExprKey) -> bool {
793        matches!(
794            key,
795            ExprKey::ExtCodeSize(_)
796                | ExprKey::ExtCodeHash(_)
797                | ExprKey::Balance(_)
798                | ExprKey::SelfBalance
799        )
800    }
801
802    /// STATICCALL is excluded: the whole static context forbids value transfers, `SSTORE`,
803    /// `CREATE`, and `SELFDESTRUCT`, so balances and deployed code cannot change. Its memory
804    /// clobber (the return buffer write) is handled separately via `may_mutate_memory`.
805    fn may_change_account_environment(kind: &InstKind) -> bool {
806        matches!(
807            kind,
808            InstKind::Call { .. }
809                | InstKind::DelegateCall { .. }
810                | InstKind::InternalCall { .. }
811                | InstKind::Create(_, _, _)
812                | InstKind::Create2(_, _, _, _)
813        )
814    }
815
816    fn is_sinkable_pure_expr(key: &ExprKey) -> bool {
817        !matches!(
818            key,
819            ExprKey::MLoad(_)
820                | ExprKey::Keccak256(_)
821                | ExprKey::SLoad(_)
822                | ExprKey::TLoad(_)
823                | ExprKey::CalldataLoad(_)
824                | ExprKey::ExtCodeSize(_)
825                | ExprKey::ExtCodeHash(_)
826                | ExprKey::BlockHash(_)
827                | ExprKey::Balance(_)
828                | ExprKey::SelfBalance
829                | ExprKey::BlobHash(_)
830        )
831    }
832
833    fn operands_dominate_block(
834        func: &Function,
835        kind: &InstKind,
836        block_id: BlockId,
837        inst_blocks: &FxHashMap<InstId, BlockId>,
838        dominators: &DominatorTree,
839    ) -> bool {
840        kind.operands().into_iter().all(|value| {
841            Self::value_dominates_block(func, value, block_id, inst_blocks, dominators)
842        })
843    }
844
845    fn value_dominates_block(
846        func: &Function,
847        value: ValueId,
848        block_id: BlockId,
849        inst_blocks: &FxHashMap<InstId, BlockId>,
850        dominators: &DominatorTree,
851    ) -> bool {
852        match func.value(value) {
853            Value::Immediate(_) | Value::Arg { .. } | Value::Undef(_) | Value::Error(_) => true,
854            Value::Inst(inst_id) => inst_blocks
855                .get(inst_id)
856                .is_some_and(|&def_block| dominators.dominates(def_block, block_id)),
857        }
858    }
859
860    fn memory_range_key(
861        &self,
862        func: &Function,
863        inst_id: InstId,
864        addr: ValueId,
865        size: Option<u64>,
866    ) -> Option<MemRangeKey> {
867        let region = func.instructions[inst_id]
868            .metadata
869            .memory_region()
870            .unwrap_or_else(|| func.memory_region_for_addr(addr));
871        let (base, offset) = Self::memory_addr_base_offset(func, addr);
872        Some(MemRangeKey { region, base, offset, size, dyn_size: None })
873    }
874
875    fn memory_addr_base_offset(func: &Function, addr: ValueId) -> (Option<ValueId>, Option<u64>) {
876        if let Some((base, offset)) = Self::offset_value(func, addr, &FxHashMap::default(), 0) {
877            if let (OperandKey::Value(base), Some(offset)) = (base, mir_utils::u256_to_u64(offset))
878            {
879                return (Some(base), Some(offset));
880            }
881            return (Some(addr), Some(0));
882        }
883        match func.value(addr) {
884            Value::Immediate(imm) => (None, imm.as_u256().and_then(mir_utils::u256_to_u64)),
885            Value::Arg { .. } | Value::Inst(_) | Value::Undef(_) | Value::Error(_) => {
886                (Some(addr), Some(0))
887            }
888        }
889    }
890
891    fn memory_ranges_may_alias(read: MemRangeKey, write: MemRangeKey) -> bool {
892        if read.region != MemoryRegion::Unknown
893            && write.region != MemoryRegion::Unknown
894            && read.region != write.region
895        {
896            return false;
897        }
898        if read.base != write.base {
899            return true;
900        }
901        let (Some(read_offset), Some(read_size), Some(write_offset), Some(write_size)) =
902            (read.offset, read.size, write.offset, write.size)
903        else {
904            return true;
905        };
906        mir_utils::ranges_overlap(read_offset, read_size, write_offset, write_size)
907    }
908
909    fn offset_expr_for_add(
910        func: &Function,
911        a: ValueId,
912        b: ValueId,
913        replacements: &FxHashMap<ValueId, ValueId>,
914    ) -> Option<(OperandKey, U256)> {
915        if let Some(offset) = func.value_u256_after_replacements(b, replacements) {
916            let (base, existing) = Self::offset_value(func, a, replacements, 0)?;
917            Some((base, existing.wrapping_add(offset)))
918        } else if let Some(offset) = func.value_u256_after_replacements(a, replacements) {
919            let (base, existing) = Self::offset_value(func, b, replacements, 0)?;
920            Some((base, existing.wrapping_add(offset)))
921        } else {
922            None
923        }
924    }
925
926    fn offset_expr_for_sub(
927        func: &Function,
928        a: ValueId,
929        b: ValueId,
930        replacements: &FxHashMap<ValueId, ValueId>,
931    ) -> Option<(OperandKey, U256)> {
932        let offset = func.value_u256_after_replacements(b, replacements)?;
933        let (base, existing) = Self::offset_value(func, a, replacements, 0)?;
934        Some((base, existing.wrapping_sub(offset)))
935    }
936
937    fn offset_value(
938        func: &Function,
939        value: ValueId,
940        replacements: &FxHashMap<ValueId, ValueId>,
941        depth: usize,
942    ) -> Option<(OperandKey, U256)> {
943        if depth >= 4 {
944            return None;
945        }
946
947        let value = mir_utils::resolve_replacement(value, replacements);
948        match func.value(value) {
949            Value::Immediate(_) => None,
950            Value::Arg { .. } | Value::Undef(_) | Value::Error(_) => {
951                Some((OperandKey::Value(value), U256::ZERO))
952            }
953            Value::Inst(inst_id) => match func.instructions[*inst_id].kind {
954                InstKind::Add(a, b) => {
955                    if let Some(offset) = func.value_u256_after_replacements(b, replacements) {
956                        let (base, existing) =
957                            Self::offset_value(func, a, replacements, depth + 1)?;
958                        Some((base, existing.wrapping_add(offset)))
959                    } else if let Some(offset) = func.value_u256_after_replacements(a, replacements)
960                    {
961                        let (base, existing) =
962                            Self::offset_value(func, b, replacements, depth + 1)?;
963                        Some((base, existing.wrapping_add(offset)))
964                    } else {
965                        Some((OperandKey::Value(value), U256::ZERO))
966                    }
967                }
968                InstKind::Sub(a, b) => {
969                    let offset = func.value_u256_after_replacements(b, replacements)?;
970                    let (base, existing) = Self::offset_value(func, a, replacements, depth + 1)?;
971                    Some((base, existing.wrapping_sub(offset)))
972                }
973                _ => Some((OperandKey::Value(value), U256::ZERO)),
974            },
975        }
976    }
977
978    fn operand_key(
979        func: &Function,
980        value: ValueId,
981        replacements: &FxHashMap<ValueId, ValueId>,
982    ) -> OperandKey {
983        let value = mir_utils::resolve_replacement(value, replacements);
984        match func.value(value) {
985            Value::Immediate(imm) => OperandKey::Immediate(imm.clone()),
986            _ => OperandKey::Value(value),
987        }
988    }
989
990    fn ordered_pair(a: OperandKey, b: OperandKey) -> (OperandKey, OperandKey) {
991        if Self::cmp_operand_key(&a, &b).is_gt() { (b, a) } else { (a, b) }
992    }
993
994    fn cmp_operand_key(a: &OperandKey, b: &OperandKey) -> Ordering {
995        match (a, b) {
996            (OperandKey::Immediate(a), OperandKey::Immediate(b)) => Self::cmp_immediate(a, b),
997            (OperandKey::Immediate(_), OperandKey::Value(_)) => Ordering::Less,
998            (OperandKey::Value(_), OperandKey::Immediate(_)) => Ordering::Greater,
999            (OperandKey::Value(a), OperandKey::Value(b)) => a.index().cmp(&b.index()),
1000        }
1001    }
1002
1003    fn cmp_immediate(a: &Immediate, b: &Immediate) -> Ordering {
1004        let rank = |imm: &Immediate| match imm {
1005            Immediate::Bool(_) => 0,
1006            Immediate::UInt(_, _) => 1,
1007            Immediate::Int(_, _) => 2,
1008            Immediate::Address(_) => 3,
1009            Immediate::FixedBytes(_, _) => 4,
1010        };
1011        rank(a).cmp(&rank(b)).then_with(|| match (a, b) {
1012            (Immediate::Bool(a), Immediate::Bool(b)) => a.cmp(b),
1013            (Immediate::UInt(a_value, a_bits), Immediate::UInt(b_value, b_bits))
1014            | (Immediate::Int(a_value, a_bits), Immediate::Int(b_value, b_bits)) => {
1015                a_bits.cmp(b_bits).then_with(|| a_value.cmp(b_value))
1016            }
1017            (Immediate::Address(a), Immediate::Address(b)) => a.cmp(b),
1018            (Immediate::FixedBytes(a_value, a_len), Immediate::FixedBytes(b_value, b_len)) => {
1019                a_len.cmp(b_len).then_with(|| a_value.cmp(b_value))
1020            }
1021            _ => Ordering::Equal,
1022        })
1023    }
1024
1025    fn value_use_counts(func: &Function) -> FxHashMap<ValueId, usize> {
1026        let mut counts = FxHashMap::default();
1027        for inst in func.instructions.iter() {
1028            for value in inst.operands() {
1029                *counts.entry(value).or_insert(0) += 1;
1030            }
1031        }
1032        for block in func.blocks.iter() {
1033            if let Some(term) = &block.terminator {
1034                Self::count_terminator_uses(term, &mut counts);
1035            }
1036        }
1037        counts
1038    }
1039
1040    fn count_terminator_uses(
1041        term: &crate::mir::Terminator,
1042        counts: &mut FxHashMap<ValueId, usize>,
1043    ) {
1044        use crate::mir::Terminator;
1045
1046        let mut count = |value| {
1047            *counts.entry(value).or_insert(0) += 1;
1048        };
1049
1050        match term {
1051            Terminator::Jump(_) | Terminator::Stop | Terminator::Invalid => {}
1052            Terminator::Branch { condition, .. } => count(*condition),
1053            Terminator::Switch { value, cases, .. } => {
1054                count(*value);
1055                for (case, _) in cases {
1056                    count(*case);
1057                }
1058            }
1059            Terminator::Return { values } => {
1060                for &value in values {
1061                    count(value);
1062                }
1063            }
1064            Terminator::Revert { offset, size } | Terminator::ReturnData { offset, size } => {
1065                count(*offset);
1066                count(*size);
1067            }
1068            Terminator::SelfDestruct { recipient } => count(*recipient),
1069        }
1070    }
1071
1072    fn apply_replacements_to_all_blocks(
1073        &self,
1074        func: &mut Function,
1075        replacements: &FxHashMap<ValueId, ValueId>,
1076    ) {
1077        let block_ids: Vec<_> = func.blocks.indices().collect();
1078        for block_id in block_ids {
1079            self.apply_replacements(func, block_id, replacements);
1080        }
1081    }
1082
1083    /// Applies value replacements to all instructions in a block.
1084    fn apply_replacements(
1085        &self,
1086        func: &mut Function,
1087        block_id: BlockId,
1088        replacements: &FxHashMap<ValueId, ValueId>,
1089    ) {
1090        let block = func.block(block_id);
1091        let inst_ids: Vec<InstId> = block.instructions.clone();
1092
1093        for inst_id in inst_ids {
1094            let inst = &mut func.instructions[inst_id];
1095            if mir_utils::replace_inst_uses_canonicalized(&mut inst.kind, replacements) != 0 {
1096                if mir_utils::is_memory_inst(&inst.kind) {
1097                    inst.metadata.set_memory_region(None);
1098                }
1099                if matches!(
1100                    inst.kind,
1101                    InstKind::SLoad(_)
1102                        | InstKind::SStore(_, _)
1103                        | InstKind::TLoad(_)
1104                        | InstKind::TStore(_, _)
1105                ) {
1106                    inst.metadata.set_storage_alias(None);
1107                }
1108            }
1109        }
1110
1111        // Also update terminator if present
1112        let block = func.block_mut(block_id);
1113        if let Some(term) = &mut block.terminator {
1114            mir_utils::replace_terminator_uses_canonicalized(term, replacements);
1115        }
1116    }
1117}