Skip to main content

solar_codegen/transform/
storage_dse.rs

1//! Local dead storage-store elimination.
2//!
3//! This pass removes persistent `sstore` instructions inside a single basic
4//! block when a later store to the same definitely-known slot overwrites them
5//! before any storage observer can see the intermediate value. It also removes
6//! repeated equal stores when no intervening instruction can clobber storage.
7
8use crate::{
9    mir::{BlockId, Function, InstId, InstKind, StorageAlias, ValueId, utils as mir_utils},
10    pass::FunctionPass,
11};
12use solar_data_structures::map::{FxHashMap, FxHashSet};
13
14/// Local dead storage-store elimination pass.
15#[derive(Debug, Default)]
16pub struct StorageStoreEliminator {
17    /// Number of storage stores eliminated.
18    pub eliminated_count: usize,
19}
20
21/// Function pass for local dead storage-store elimination.
22pub struct StorageDsePass;
23
24impl FunctionPass for StorageDsePass {
25    fn name(&self) -> &str {
26        "storage-dse"
27    }
28
29    fn run_on_function(&mut self, func: &mut Function) -> bool {
30        StorageStoreEliminator::new().run_to_fixpoint(func) != 0
31    }
32}
33
34impl StorageStoreEliminator {
35    /// Creates a new storage-store eliminator.
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Runs local storage DSE on a function.
41    pub fn run(&mut self, func: &mut Function) -> usize {
42        self.eliminated_count = 0;
43        func.annotate_storage_aliases(mir_utils::StorageAliasScope::Storage);
44
45        let block_ids: Vec<BlockId> = func.blocks.indices().collect();
46        for block_id in block_ids {
47            self.remove_overwritten_stores(func, block_id);
48            self.remove_equal_stores(func, block_id);
49        }
50
51        self.eliminated_count
52    }
53
54    /// Runs local storage DSE to a fixed point.
55    pub fn run_to_fixpoint(&mut self, func: &mut Function) -> usize {
56        let mut total = 0;
57        loop {
58            let eliminated = self.run(func);
59            if eliminated == 0 {
60                break;
61            }
62            total += eliminated;
63        }
64        total
65    }
66
67    fn remove_overwritten_stores(&mut self, func: &mut Function, block_id: BlockId) {
68        let inst_ids = func.blocks[block_id].instructions.clone();
69        let mut later_writes: FxHashSet<StorageAlias> = FxHashSet::default();
70        let mut dead: FxHashSet<InstId> = FxHashSet::default();
71
72        for &inst_id in inst_ids.iter().rev() {
73            match &func.instructions[inst_id].kind {
74                InstKind::SStore(slot, _) => {
75                    let alias = func.storage_alias(inst_id, *slot);
76                    if later_writes.contains(&alias) {
77                        dead.insert(inst_id);
78                        self.eliminated_count += 1;
79                        continue;
80                    }
81
82                    Self::remove_aliasing_set(&mut later_writes, alias);
83                    later_writes.insert(alias);
84                }
85                InstKind::SLoad(slot) => {
86                    let alias = func.storage_alias(inst_id, *slot);
87                    Self::remove_aliasing_set(&mut later_writes, alias);
88                }
89                kind if Self::may_observe_or_mutate_storage(kind) => {
90                    later_writes.clear();
91                }
92                _ => {}
93            }
94        }
95
96        if dead.is_empty() {
97            return;
98        }
99
100        func.blocks[block_id].instructions.retain(|id| !dead.contains(id));
101    }
102
103    fn remove_equal_stores(&mut self, func: &mut Function, block_id: BlockId) {
104        let inst_ids = func.blocks[block_id].instructions.clone();
105        let mut stored_values: FxHashMap<StorageAlias, ValueId> = FxHashMap::default();
106        let mut dead: FxHashSet<InstId> = FxHashSet::default();
107
108        for &inst_id in &inst_ids {
109            match &func.instructions[inst_id].kind {
110                InstKind::SStore(slot, value) => {
111                    let alias = func.storage_alias(inst_id, *slot);
112                    if stored_values.get(&alias).is_some_and(|&stored| stored == *value) {
113                        dead.insert(inst_id);
114                        self.eliminated_count += 1;
115                        continue;
116                    }
117
118                    Self::remove_aliasing_map(&mut stored_values, alias);
119                    stored_values.insert(alias, *value);
120                }
121                kind if Self::may_observe_or_mutate_storage(kind) => {
122                    stored_values.clear();
123                }
124                _ => {}
125            }
126        }
127
128        if dead.is_empty() {
129            return;
130        }
131
132        func.blocks[block_id].instructions.retain(|id| !dead.contains(id));
133    }
134
135    fn remove_aliasing_set(aliases: &mut FxHashSet<StorageAlias>, alias: StorageAlias) {
136        aliases.retain(|cached| !cached.may_alias(alias));
137    }
138
139    fn remove_aliasing_map(values: &mut FxHashMap<StorageAlias, ValueId>, alias: StorageAlias) {
140        values.retain(|cached, _| !cached.may_alias(alias));
141    }
142
143    fn may_observe_or_mutate_storage(kind: &InstKind) -> bool {
144        matches!(
145            kind,
146            InstKind::Call { .. }
147                | InstKind::StaticCall { .. }
148                | InstKind::DelegateCall { .. }
149                | InstKind::InternalCall { .. }
150                | InstKind::Create(_, _, _)
151                | InstKind::Create2(_, _, _, _)
152        ) || kind.may_mutate_storage()
153    }
154}