solar_codegen/analysis/
scalar_evolution.rs

1//! Scalar-evolution-style affine analysis for MIR loops.
2//!
3//! This is intentionally small: it recognizes expressions of the form
4//! `base + c + sum(iv * scale)` inside one natural loop. Optimization passes can use this to avoid
5//! ad hoc pattern matching when reasoning about memory/storage addresses derived from loop indices.
6//!
7//! Analysis contract:
8//! - loop-invariant values are represented as a single optional base
9//! - constants and induction scales use checked signed arithmetic
10//! - unrecognized or overflowing expressions are omitted instead of guessed
11
12use crate::{
13    analysis::Loop,
14    mir::{Function, InstKind, Value, ValueId},
15};
16use alloy_primitives::U256;
17use smallvec::SmallVec;
18use solar_data_structures::map::FxHashMap;
19
20/// One affine induction-variable term.
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub struct AffineTerm {
23    /// Loop induction variable.
24    pub value: ValueId,
25    /// Signed scale applied to the induction variable.
26    pub scale: i128,
27}
28
29/// An affine expression in one loop.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct AffineExpr {
32    /// Optional loop-invariant base value.
33    pub base: Option<ValueId>,
34    /// Signed constant offset.
35    pub constant: i128,
36    /// Induction-variable terms.
37    pub terms: SmallVec<[AffineTerm; 2]>,
38}
39
40impl AffineExpr {
41    fn constant(constant: i128) -> Self {
42        Self { base: None, constant, terms: SmallVec::new() }
43    }
44
45    fn base(base: ValueId) -> Self {
46        Self { base: Some(base), constant: 0, terms: SmallVec::new() }
47    }
48
49    fn induction(value: ValueId) -> Self {
50        Self { base: None, constant: 0, terms: smallvec::smallvec![AffineTerm { value, scale: 1 }] }
51    }
52
53    fn add(mut self, other: Self) -> Option<Self> {
54        self.constant = self.constant.checked_add(other.constant)?;
55        self.base = match (self.base, other.base) {
56            (None, base) | (base, None) => base,
57            (Some(_), Some(_)) => return None,
58        };
59        for term in other.terms {
60            self.add_term(term.value, term.scale)?;
61        }
62        Some(self)
63    }
64
65    fn sub(mut self, other: Self) -> Option<Self> {
66        self.constant = self.constant.checked_sub(other.constant)?;
67        if other.base.is_some() {
68            return None;
69        }
70        for term in other.terms {
71            self.add_term(term.value, term.scale.checked_neg()?)?;
72        }
73        Some(self)
74    }
75
76    fn mul_const(mut self, scale: i128) -> Option<Self> {
77        if self.base.is_some() && scale != 1 {
78            return None;
79        }
80        self.constant = self.constant.checked_mul(scale)?;
81        for term in &mut self.terms {
82            term.scale = term.scale.checked_mul(scale)?;
83        }
84        Some(self)
85    }
86
87    fn add_term(&mut self, value: ValueId, scale: i128) -> Option<()> {
88        if scale == 0 {
89            return Some(());
90        }
91        if let Some(term) = self.terms.iter_mut().find(|term| term.value == value) {
92            term.scale = term.scale.checked_add(scale)?;
93            if term.scale == 0 {
94                self.terms.retain(|term| term.value != value);
95            }
96        } else {
97            self.terms.push(AffineTerm { value, scale });
98        }
99        Some(())
100    }
101}
102
103/// Affine expressions recognized for one loop.
104#[derive(Clone, Debug, Default)]
105pub struct ScalarEvolution {
106    expressions: FxHashMap<ValueId, AffineExpr>,
107}
108
109impl ScalarEvolution {
110    /// Computes affine expressions for values used by `loop_data`.
111    #[must_use]
112    pub fn analyze(func: &Function, loop_data: &Loop) -> Self {
113        let mut analysis = Self::default();
114        for value in func.values.indices() {
115            let _ = analysis.affine_expr(func, loop_data, value);
116        }
117        analysis
118    }
119
120    /// Returns the affine expression for a value, if recognized.
121    #[must_use]
122    pub fn get(&self, value: ValueId) -> Option<&AffineExpr> {
123        self.expressions.get(&value)
124    }
125
126    fn affine_expr(
127        &mut self,
128        func: &Function,
129        loop_data: &Loop,
130        value: ValueId,
131    ) -> Option<AffineExpr> {
132        if let Some(expr) = self.expressions.get(&value) {
133            return Some(expr.clone());
134        }
135
136        let expr = match func.value(value) {
137            Value::Immediate(imm) => AffineExpr::constant(u256_to_i128(imm.as_u256()?)?),
138            Value::Arg { .. } => AffineExpr::base(value),
139            Value::Undef(_) | Value::Error(_) => return None,
140            Value::Inst(_) if !value_defined_in_loop(func, value, loop_data) => {
141                AffineExpr::base(value)
142            }
143            Value::Inst(inst_id) => {
144                if loop_data.induction_vars.iter().any(|iv| iv.value == value) {
145                    AffineExpr::induction(value)
146                } else {
147                    match func.instructions[*inst_id].kind {
148                        InstKind::Add(a, b) => {
149                            let a = self.affine_expr(func, loop_data, a)?;
150                            let b = self.affine_expr(func, loop_data, b)?;
151                            a.add(b)?
152                        }
153                        InstKind::Sub(a, b) => {
154                            let a = self.affine_expr(func, loop_data, a)?;
155                            let b = self.affine_expr(func, loop_data, b)?;
156                            a.sub(b)?
157                        }
158                        InstKind::Mul(a, b) => {
159                            let a_expr = self.affine_expr(func, loop_data, a);
160                            let b_expr = self.affine_expr(func, loop_data, b);
161                            match (a_expr, b_expr) {
162                                (Some(expr), Some(scale))
163                                    if scale.base.is_none() && scale.terms.is_empty() =>
164                                {
165                                    expr.mul_const(scale.constant)?
166                                }
167                                (Some(scale), Some(expr))
168                                    if scale.base.is_none() && scale.terms.is_empty() =>
169                                {
170                                    expr.mul_const(scale.constant)?
171                                }
172                                _ => return None,
173                            }
174                        }
175                        InstKind::Shl(shift, value) => {
176                            let shift = self.affine_expr(func, loop_data, shift)?;
177                            if shift.base.is_some() || !shift.terms.is_empty() {
178                                return None;
179                            }
180                            let shift = u32::try_from(shift.constant).ok()?;
181                            let scale = 1i128.checked_shl(shift)?;
182                            self.affine_expr(func, loop_data, value)?.mul_const(scale)?
183                        }
184                        _ => return None,
185                    }
186                }
187            }
188        };
189
190        self.expressions.insert(value, expr.clone());
191        Some(expr)
192    }
193}
194
195fn value_defined_in_loop(func: &Function, value: ValueId, loop_data: &Loop) -> bool {
196    match func.value(value) {
197        Value::Inst(inst_id) => loop_data
198            .blocks
199            .iter()
200            .any(|&block_id| func.blocks[block_id].instructions.contains(inst_id)),
201        Value::Undef(_) | Value::Error(_) => true,
202        Value::Arg { .. } | Value::Immediate(_) => false,
203    }
204}
205
206fn u256_to_i128(value: U256) -> Option<i128> {
207    if value <= U256::from(i128::MAX as u128) { Some(value.to::<u128>() as i128) } else { None }
208}
solar_codegen/analysis/scalar_evolution.rs

solar_codegen/analysis/
scalar_evolution.rs
