Skip to main content

solar_codegen/transform/
storage_load_cse.rs

1//! Local storage-load forwarding.
2//!
3//! This pass removes redundant `sload` instructions on straight-line paths when
4//! no intervening storage write may alias the loaded slot.
5
6use crate::{
7    analysis::Liveness,
8    mir::{BlockId, Function, InstId, InstKind, StorageAlias, ValueId, utils as mir_utils},
9    pass::{AnalysisManager, FunctionPass, LivenessAnalysis},
10};
11use solar_data_structures::map::{FxHashMap, FxHashSet};
12
13/// Local storage load CSE pass.
14#[derive(Debug, Default)]
15pub struct StorageLoadCse {
16    /// Number of storage loads eliminated.
17    pub eliminated_count: usize,
18}
19
20/// Function pass for straight-line storage-load CSE.
21pub struct StorageLoadCsePass;
22
23impl FunctionPass for StorageLoadCsePass {
24    fn name(&self) -> &str {
25        "storage-load-cse"
26    }
27
28    fn run_on_function(&mut self, func: &mut Function) -> bool {
29        StorageLoadCse::new().run_to_fixpoint(func) != 0
30    }
31}
32
33impl StorageLoadCse {
34    /// Creates a new storage-load CSE pass.
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Runs storage-load CSE on a function.
40    pub fn run(&mut self, func: &mut Function) -> usize {
41        self.eliminated_count = 0;
42        func.annotate_storage_aliases(mir_utils::StorageAliasScope::Storage);
43
44        let mut analyses = AnalysisManager::new();
45        let liveness = analyses.get_or_compute(&LivenessAnalysis, func);
46        let inst_results = func.inst_results();
47        let block_ids: Vec<BlockId> = func.blocks.indices().collect();
48        let mut replacements = FxHashMap::default();
49        let mut dead = FxHashSet::default();
50
51        for block_id in block_ids {
52            self.process_block(
53                func,
54                block_id,
55                liveness,
56                &inst_results,
57                &mut replacements,
58                &mut dead,
59            );
60        }
61
62        if !replacements.is_empty() {
63            Self::replace_uses(func, &replacements);
64        }
65        if !dead.is_empty() {
66            for block in func.blocks.iter_mut() {
67                block.instructions.retain(|id| !dead.contains(id));
68            }
69        }
70
71        self.eliminated_count
72    }
73
74    /// Runs storage-load CSE to a fixed point.
75    pub fn run_to_fixpoint(&mut self, func: &mut Function) -> usize {
76        let mut total = 0;
77        loop {
78            let eliminated = self.run(func);
79            if eliminated == 0 {
80                break;
81            }
82            total += eliminated;
83        }
84        total
85    }
86
87    fn process_block(
88        &mut self,
89        func: &Function,
90        block_id: BlockId,
91        liveness: &Liveness,
92        inst_results: &FxHashMap<InstId, ValueId>,
93        replacements: &mut FxHashMap<ValueId, ValueId>,
94        dead: &mut FxHashSet<InstId>,
95    ) {
96        let mut cached_loads: FxHashMap<StorageAlias, ValueId> = FxHashMap::default();
97        let inst_ids = func.blocks[block_id].instructions.clone();
98
99        for (inst_idx, inst_id) in inst_ids.into_iter().enumerate() {
100            match &func.instructions[inst_id].kind {
101                InstKind::SLoad(slot) => {
102                    let alias = func.storage_alias_after_replacements(inst_id, *slot, replacements);
103                    let Some(&result) = inst_results.get(&inst_id) else {
104                        continue;
105                    };
106                    if let Some(&cached) = cached_loads.get(&alias) {
107                        if !liveness
108                            .live_at_inst(func, block_id, inst_idx)
109                            .live_before
110                            .contains(cached)
111                        {
112                            cached_loads.insert(alias, result);
113                            continue;
114                        }
115                        replacements.insert(result, cached);
116                        dead.insert(inst_id);
117                        self.eliminated_count += 1;
118                    } else {
119                        cached_loads.insert(alias, result);
120                    }
121                }
122                InstKind::SStore(slot, _) => {
123                    let alias = func.storage_alias_after_replacements(inst_id, *slot, replacements);
124                    cached_loads.retain(|cached_alias, _| !cached_alias.may_alias(alias));
125                }
126                kind if kind.may_mutate_storage() => cached_loads.clear(),
127                _ => {}
128            }
129        }
130    }
131
132    fn replace_uses(func: &mut Function, replacements: &FxHashMap<ValueId, ValueId>) {
133        if replacements.is_empty() {
134            return;
135        }
136
137        for inst in func.instructions.iter_mut() {
138            mir_utils::replace_inst_uses_canonicalized(&mut inst.kind, replacements);
139            if matches!(inst.kind, InstKind::SLoad(_) | InstKind::SStore(_, _)) {
140                inst.metadata.set_storage_alias(None);
141            }
142        }
143
144        for block in func.blocks.iter_mut() {
145            if let Some(term) = &mut block.terminator {
146                mir_utils::replace_terminator_uses_canonicalized(term, replacements);
147            }
148        }
149    }
150}