Skip to main content

jetro_core/
ssa.rs

1//! SSA-style numbering + data-flow graph for v2 programs.
2//!
3//! v2's stack machine has no named intermediate values.  To express a
4//! data-flow graph, we assign each stack push a fresh `ValueId` and
5//! record which earlier values each opcode consumes.
6//!
7//! Limitations:
8//! - Sub-programs (branches) are opaque to this pass — they are walked
9//!   separately but their values live in their own namespace.
10//! - Phi nodes are not synthesised; there are no merge points within
11//!   a single block.
12//!
13//! Useful for: def-use queries, live-range analysis, value-numbering CSE.
14
15use std::sync::Arc;
16use std::collections::HashMap;
17use super::vm::{Program, Opcode};
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub struct ValueId(pub u32);
21
22#[derive(Debug, Clone)]
23pub struct SsaInstr {
24    pub id:   ValueId,
25    pub op:   Opcode,
26    /// Value ids consumed from the stack (pops) in order.
27    pub uses: Vec<ValueId>,
28}
29
30/// A phi node: at a merge point, value takes one of several incoming
31/// values depending on which predecessor path was taken.
32#[derive(Debug, Clone)]
33pub struct Phi {
34    pub id:       ValueId,
35    /// Incoming (predecessor label, value) pairs.
36    pub incoming: Vec<(PhiEdge, ValueId)>,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
40pub enum PhiEdge { AndLhs, AndRhs, OrLhs, OrRhs, CoalesceLhs, CoalesceRhs }
41
42#[derive(Debug, Clone, Default)]
43pub struct SsaGraph {
44    pub instrs: Vec<SsaInstr>,
45    pub phis:   Vec<Phi>,
46    /// Last pushed value id — the program's result.
47    pub result: Option<ValueId>,
48}
49
50impl SsaGraph {
51    pub fn build(program: &Program) -> SsaGraph {
52        let mut g = SsaGraph::default();
53        let mut stack: Vec<ValueId> = Vec::new();
54        for op in program.ops.iter() {
55            let arity = op_arity(op);
56            let mut uses = Vec::with_capacity(arity.pops);
57            for _ in 0..arity.pops {
58                if let Some(v) = stack.pop() { uses.push(v); }
59            }
60            let id = ValueId(g.instrs.len() as u32);
61            // Synthesise phi at short-circuit / coalesce merge points.
62            // At an AndOp, the lhs is on the stack; rhs sub-program produces
63            // a new value conditionally.  Result is a phi(lhs, rhs).
64            let phi_edge = match op {
65                Opcode::AndOp(_)      => Some((PhiEdge::AndLhs, PhiEdge::AndRhs)),
66                Opcode::OrOp(_)       => Some((PhiEdge::OrLhs, PhiEdge::OrRhs)),
67                Opcode::CoalesceOp(_) => Some((PhiEdge::CoalesceLhs, PhiEdge::CoalesceRhs)),
68                _ => None,
69            };
70            if let Some((lhs_edge, rhs_edge)) = phi_edge {
71                // uses[0] is the lhs consumed; the "rhs" value is synthetic — we
72                // model it as the instruction's own id (sub-program result).
73                let lhs = uses.first().copied().unwrap_or(id);
74                g.phis.push(Phi {
75                    id,
76                    incoming: vec![(lhs_edge, lhs), (rhs_edge, id)],
77                });
78            }
79            g.instrs.push(SsaInstr { id, op: op.clone(), uses });
80            if arity.pushes { stack.push(id); }
81        }
82        g.result = stack.pop();
83        g
84    }
85
86    /// Use-list (consumers) for each value id.
87    pub fn use_list(&self) -> HashMap<ValueId, Vec<ValueId>> {
88        let mut m: HashMap<ValueId, Vec<ValueId>> = HashMap::new();
89        for instr in &self.instrs {
90            for u in &instr.uses {
91                m.entry(*u).or_default().push(instr.id);
92            }
93        }
94        m
95    }
96
97    /// Dead values: pushed but never consumed and not the final result.
98    pub fn dead_values(&self) -> Vec<ValueId> {
99        let uses = self.use_list();
100        self.instrs.iter()
101            .filter(|i| {
102                op_arity(&i.op).pushes
103                && !uses.contains_key(&i.id)
104                && Some(i.id) != self.result
105            })
106            .map(|i| i.id)
107            .collect()
108    }
109}
110
111#[derive(Debug, Clone, Copy)]
112struct Arity { pops: usize, pushes: bool }
113
114fn op_arity(op: &Opcode) -> Arity {
115    match op {
116        Opcode::PushNull | Opcode::PushBool(_) | Opcode::PushInt(_)
117            | Opcode::PushFloat(_) | Opcode::PushStr(_)
118            | Opcode::PushRoot | Opcode::PushCurrent | Opcode::LoadIdent(_)
119            | Opcode::RootChain(_) | Opcode::GetPointer(_)
120            | Opcode::MakeObj(_) | Opcode::MakeArr(_) | Opcode::FString(_)
121            | Opcode::ListComp(_) | Opcode::DictComp(_) | Opcode::SetComp(_)
122            | Opcode::PatchEval(_) =>
123            Arity { pops: 0, pushes: true },
124
125        Opcode::GetField(_) | Opcode::OptField(_) | Opcode::GetIndex(_)
126            | Opcode::GetSlice(..) | Opcode::Descendant(_) | Opcode::DescendAll
127            | Opcode::DynIndex(_) | Opcode::InlineFilter(_)
128            | Opcode::Quantifier(_) | Opcode::FilterCount(_)
129            | Opcode::FindFirst(_) | Opcode::FindOne(_)
130            | Opcode::FilterMap { .. } | Opcode::FilterFilter { .. }
131            | Opcode::MapMap { .. } | Opcode::MapSum(_) | Opcode::MapAvg(_)
132            | Opcode::MapFlatten(_) | Opcode::FilterTakeWhile { .. }
133            | Opcode::FilterDropWhile { .. } | Opcode::MapUnique(_)
134            | Opcode::EquiJoin { .. }
135            | Opcode::TopN { .. } | Opcode::KindCheck { .. }
136            | Opcode::Not | Opcode::Neg
137            | Opcode::CallMethod(_) | Opcode::CallOptMethod(_)
138            | Opcode::AndOp(_) | Opcode::OrOp(_) | Opcode::CoalesceOp(_)
139            | Opcode::CastOp(_) =>
140            Arity { pops: 1, pushes: true },
141
142        Opcode::Add | Opcode::Sub | Opcode::Mul | Opcode::Div | Opcode::Mod
143            | Opcode::Eq | Opcode::Neq | Opcode::Lt | Opcode::Lte
144            | Opcode::Gt | Opcode::Gte | Opcode::Fuzzy =>
145            Arity { pops: 2, pushes: true },
146
147        Opcode::StoreVar(_) => Arity { pops: 1, pushes: false },
148        Opcode::SetCurrent | Opcode::BindVar(_)
149            | Opcode::BindObjDestructure(_) | Opcode::BindArrDestructure(_)
150            | Opcode::LetExpr { .. } => Arity { pops: 0, pushes: false },
151    }
152}
153
154/// Value-numbering CSE on top of SSA: two instructions with identical
155/// opcode and matching `uses` are mapped to the same canonical id.
156pub fn value_number(g: &SsaGraph) -> HashMap<ValueId, ValueId> {
157    let mut canon: HashMap<ValueId, ValueId> = HashMap::new();
158    let mut seen: HashMap<(u64, Vec<ValueId>), ValueId> = HashMap::new();
159    for instr in &g.instrs {
160        let canon_uses: Vec<ValueId> = instr.uses.iter()
161            .map(|u| *canon.get(u).unwrap_or(u)).collect();
162        let key = (op_hash(&instr.op), canon_uses);
163        match seen.get(&key) {
164            Some(&existing) => { canon.insert(instr.id, existing); }
165            None            => { seen.insert(key, instr.id); canon.insert(instr.id, instr.id); }
166        }
167    }
168    canon
169}
170
171fn op_hash(op: &Opcode) -> u64 {
172    use std::collections::hash_map::DefaultHasher;
173    use std::hash::{Hash, Hasher};
174    let mut h = DefaultHasher::new();
175    std::mem::discriminant(op).hash(&mut h);
176    match op {
177        Opcode::PushInt(n) => n.hash(&mut h),
178        Opcode::PushStr(s) => s.as_bytes().hash(&mut h),
179        Opcode::PushBool(b) => b.hash(&mut h),
180        Opcode::GetField(k) | Opcode::OptField(k) | Opcode::LoadIdent(k) =>
181            k.as_bytes().hash(&mut h),
182        Opcode::GetIndex(i) => i.hash(&mut h),
183        _ => {}
184    }
185    h.finish()
186}
187
188// Silence unused-Arc warning.
189#[allow(dead_code)]
190fn _use_arc<T>(_: Arc<T>) {}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::vm::Compiler;
196
197    #[test]
198    fn ssa_builds_graph() {
199        let p = Compiler::compile_str("$.a + $.b").unwrap();
200        let g = SsaGraph::build(&p);
201        assert!(g.result.is_some());
202        let add = g.instrs.last().unwrap();
203        assert_eq!(add.uses.len(), 2);
204    }
205
206    #[test]
207    fn ssa_use_list() {
208        let p = Compiler::compile_str("$.a + $.b").unwrap();
209        let g = SsaGraph::build(&p);
210        let uses = g.use_list();
211        assert_eq!(uses.values().map(|v| v.len()).sum::<usize>(), 2);
212    }
213
214    #[test]
215    fn value_numbering_dedups_identical() {
216        // Use GetField sequences that don't const-fold.
217        let p = Compiler::compile_str("[$.a, $.a]").unwrap();
218        let g = SsaGraph::build(&p);
219        let canon = value_number(&g);
220        // Both root-chain loads share canonical id.
221        let load_ids: Vec<ValueId> = g.instrs.iter()
222            .filter(|i| matches!(i.op, crate::vm::Opcode::RootChain(_)))
223            .map(|i| canon[&i.id]).collect();
224        if load_ids.len() >= 2 {
225            assert_eq!(load_ids[0], load_ids[1]);
226        }
227    }
228}