llvm_transforms/
const_prop.rs

1//! Constant propagation pass.
2//!
3//! Walks every instruction in a function in **reverse post-order (RPO)** so
4//! that each block's predecessors are generally visited before the block
5//! itself.  When all operands of an instruction are constants it folds the
6//! instruction to a constant (via `try_fold`), records the substitution, and
7//! rewrites all downstream uses of that instruction result.  Folded
8//! instructions are then dropped from block bodies.
9//!
10//! RPO traversal propagates constants through straight-line code and through
11//! forward edges of loops in a single pass.  Back-edges (loop-carried
12//! constants) require a second pass; use `PassManager::run_until_fixed_point`
13//! when that is needed.
14
15use crate::constant_fold::try_fold;
16use crate::pass::FunctionPass;
17use llvm_ir::{Context, Function, InstrId, InstrKind, ValueRef};
18use std::collections::{HashMap, HashSet};
19
20/// Constant propagation / constant folding pass.
21pub struct ConstProp;
22
23impl FunctionPass for ConstProp {
24    fn name(&self) -> &'static str {
25        "const-prop"
26    }
27
28    fn run_on_function(&mut self, ctx: &mut Context, func: &mut Function) -> bool {
29        if func.blocks.is_empty() {
30            return false;
31        }
32
33        // Map InstrId → its constant replacement (ValueRef::Constant).
34        let mut subst: HashMap<InstrId, ValueRef> = HashMap::new();
35
36        // Process blocks in RPO so each block's predecessors generally come first.
37        for bi in rpo(func) {
38            let body: Vec<InstrId> = func.blocks[bi].body.clone();
39            for iid in body {
40                // Apply pending substitutions to this instruction's operands.
41                if !subst.is_empty() {
42                    let new_kind = subst_kind(func.instr(iid).kind.clone(), &subst);
43                    func.instr_mut(iid).kind = new_kind;
44                }
45                // Try to fold the (possibly updated) instruction.
46                let kind = func.instr(iid).kind.clone();
47                if let Some(cid) = try_fold(ctx, &kind) {
48                    subst.insert(iid, ValueRef::Constant(cid));
49                }
50            }
51            // Also propagate into the terminator.
52            if let Some(tid) = func.blocks[bi].terminator {
53                if !subst.is_empty() {
54                    let new_kind = subst_kind(func.instr(tid).kind.clone(), &subst);
55                    func.instr_mut(tid).kind = new_kind;
56                }
57            }
58        }
59
60        if subst.is_empty() {
61            return false;
62        }
63
64        // Remove folded instructions from block bodies (they are now dead).
65        for bb in &mut func.blocks {
66            bb.body.retain(|id| !subst.contains_key(id));
67        }
68        true
69    }
70}
71
72// ---------------------------------------------------------------------------
73// Value substitution helper — replaces InstrId refs appearing in `subst`
74// ---------------------------------------------------------------------------
75
76/// Substitute ValueRef operands throughout `kind`, replacing every
77/// `ValueRef::Instruction(id)` that appears in `subst` with `subst[id]`.
78///
79/// This function is `pub(crate)` so that `mem2reg` can reuse it.
80pub(crate) fn subst_kind(kind: InstrKind, subst: &HashMap<InstrId, ValueRef>) -> InstrKind {
81    let s = |v: ValueRef| -> ValueRef {
82        if let ValueRef::Instruction(id) = v {
83            subst.get(&id).copied().unwrap_or(v)
84        } else {
85            v
86        }
87    };
88
89    match kind {
90        // --- Integer arithmetic ---
91        InstrKind::Add { flags, lhs, rhs } => InstrKind::Add {
92            flags,
93            lhs: s(lhs),
94            rhs: s(rhs),
95        },
96        InstrKind::Sub { flags, lhs, rhs } => InstrKind::Sub {
97            flags,
98            lhs: s(lhs),
99            rhs: s(rhs),
100        },
101        InstrKind::Mul { flags, lhs, rhs } => InstrKind::Mul {
102            flags,
103            lhs: s(lhs),
104            rhs: s(rhs),
105        },
106        InstrKind::UDiv { exact, lhs, rhs } => InstrKind::UDiv {
107            exact,
108            lhs: s(lhs),
109            rhs: s(rhs),
110        },
111        InstrKind::SDiv { exact, lhs, rhs } => InstrKind::SDiv {
112            exact,
113            lhs: s(lhs),
114            rhs: s(rhs),
115        },
116        InstrKind::URem { lhs, rhs } => InstrKind::URem {
117            lhs: s(lhs),
118            rhs: s(rhs),
119        },
120        InstrKind::SRem { lhs, rhs } => InstrKind::SRem {
121            lhs: s(lhs),
122            rhs: s(rhs),
123        },
124        // --- Bitwise ---
125        InstrKind::And { lhs, rhs } => InstrKind::And {
126            lhs: s(lhs),
127            rhs: s(rhs),
128        },
129        InstrKind::Or { lhs, rhs } => InstrKind::Or {
130            lhs: s(lhs),
131            rhs: s(rhs),
132        },
133        InstrKind::Xor { lhs, rhs } => InstrKind::Xor {
134            lhs: s(lhs),
135            rhs: s(rhs),
136        },
137        InstrKind::Shl { flags, lhs, rhs } => InstrKind::Shl {
138            flags,
139            lhs: s(lhs),
140            rhs: s(rhs),
141        },
142        InstrKind::LShr { exact, lhs, rhs } => InstrKind::LShr {
143            exact,
144            lhs: s(lhs),
145            rhs: s(rhs),
146        },
147        InstrKind::AShr { exact, lhs, rhs } => InstrKind::AShr {
148            exact,
149            lhs: s(lhs),
150            rhs: s(rhs),
151        },
152        // --- Floating-point ---
153        InstrKind::FAdd { flags, lhs, rhs } => InstrKind::FAdd {
154            flags,
155            lhs: s(lhs),
156            rhs: s(rhs),
157        },
158        InstrKind::FSub { flags, lhs, rhs } => InstrKind::FSub {
159            flags,
160            lhs: s(lhs),
161            rhs: s(rhs),
162        },
163        InstrKind::FMul { flags, lhs, rhs } => InstrKind::FMul {
164            flags,
165            lhs: s(lhs),
166            rhs: s(rhs),
167        },
168        InstrKind::FDiv { flags, lhs, rhs } => InstrKind::FDiv {
169            flags,
170            lhs: s(lhs),
171            rhs: s(rhs),
172        },
173        InstrKind::FRem { flags, lhs, rhs } => InstrKind::FRem {
174            flags,
175            lhs: s(lhs),
176            rhs: s(rhs),
177        },
178        InstrKind::FNeg { flags, operand } => InstrKind::FNeg {
179            flags,
180            operand: s(operand),
181        },
182        // --- Comparisons ---
183        InstrKind::ICmp { pred, lhs, rhs } => InstrKind::ICmp {
184            pred,
185            lhs: s(lhs),
186            rhs: s(rhs),
187        },
188        InstrKind::FCmp {
189            flags,
190            pred,
191            lhs,
192            rhs,
193        } => InstrKind::FCmp {
194            flags,
195            pred,
196            lhs: s(lhs),
197            rhs: s(rhs),
198        },
199        // --- Memory ---
200        InstrKind::Alloca {
201            alloc_ty,
202            num_elements,
203            align,
204        } => InstrKind::Alloca {
205            alloc_ty,
206            num_elements: num_elements.map(s),
207            align,
208        },
209        InstrKind::Load {
210            ty,
211            ptr,
212            align,
213            volatile,
214        } => InstrKind::Load {
215            ty,
216            ptr: s(ptr),
217            align,
218            volatile,
219        },
220        InstrKind::Store {
221            val,
222            ptr,
223            align,
224            volatile,
225        } => InstrKind::Store {
226            val: s(val),
227            ptr: s(ptr),
228            align,
229            volatile,
230        },
231        InstrKind::GetElementPtr {
232            inbounds,
233            base_ty,
234            ptr,
235            indices,
236        } => InstrKind::GetElementPtr {
237            inbounds,
238            base_ty,
239            ptr: s(ptr),
240            indices: indices.into_iter().map(s).collect(),
241        },
242        // --- Casts ---
243        InstrKind::Trunc { val, to } => InstrKind::Trunc { val: s(val), to },
244        InstrKind::ZExt { val, to } => InstrKind::ZExt { val: s(val), to },
245        InstrKind::SExt { val, to } => InstrKind::SExt { val: s(val), to },
246        InstrKind::FPTrunc { val, to } => InstrKind::FPTrunc { val: s(val), to },
247        InstrKind::FPExt { val, to } => InstrKind::FPExt { val: s(val), to },
248        InstrKind::FPToUI { val, to } => InstrKind::FPToUI { val: s(val), to },
249        InstrKind::FPToSI { val, to } => InstrKind::FPToSI { val: s(val), to },
250        InstrKind::UIToFP { val, to } => InstrKind::UIToFP { val: s(val), to },
251        InstrKind::SIToFP { val, to } => InstrKind::SIToFP { val: s(val), to },
252        InstrKind::PtrToInt { val, to } => InstrKind::PtrToInt { val: s(val), to },
253        InstrKind::IntToPtr { val, to } => InstrKind::IntToPtr { val: s(val), to },
254        InstrKind::BitCast { val, to } => InstrKind::BitCast { val: s(val), to },
255        InstrKind::AddrSpaceCast { val, to } => InstrKind::AddrSpaceCast { val: s(val), to },
256        InstrKind::Freeze { val } => InstrKind::Freeze { val: s(val) },
257        // --- Misc ---
258        InstrKind::Select {
259            cond,
260            then_val,
261            else_val,
262        } => InstrKind::Select {
263            cond: s(cond),
264            then_val: s(then_val),
265            else_val: s(else_val),
266        },
267        InstrKind::Phi { ty, incoming } => InstrKind::Phi {
268            ty,
269            incoming: incoming.into_iter().map(|(v, b)| (s(v), b)).collect(),
270        },
271        InstrKind::ExtractValue { aggregate, indices } => InstrKind::ExtractValue {
272            aggregate: s(aggregate),
273            indices,
274        },
275        InstrKind::InsertValue {
276            aggregate,
277            val,
278            indices,
279        } => InstrKind::InsertValue {
280            aggregate: s(aggregate),
281            val: s(val),
282            indices,
283        },
284        InstrKind::ExtractElement { vec, idx } => InstrKind::ExtractElement {
285            vec: s(vec),
286            idx: s(idx),
287        },
288        InstrKind::InsertElement { vec, val, idx } => InstrKind::InsertElement {
289            vec: s(vec),
290            val: s(val),
291            idx: s(idx),
292        },
293        InstrKind::ShuffleVector { v1, v2, mask } => InstrKind::ShuffleVector {
294            v1: s(v1),
295            v2: s(v2),
296            mask,
297        },
298        // --- Call ---
299        InstrKind::Call {
300            tail,
301            callee_ty,
302            callee,
303            args,
304        } => InstrKind::Call {
305            tail,
306            callee_ty,
307            callee: s(callee),
308            args: args.into_iter().map(s).collect(),
309        },
310        // --- Terminators ---
311        InstrKind::Ret { val } => InstrKind::Ret { val: val.map(s) },
312        InstrKind::Br { dest } => InstrKind::Br { dest },
313        InstrKind::CondBr {
314            cond,
315            then_dest,
316            else_dest,
317        } => InstrKind::CondBr {
318            cond: s(cond),
319            then_dest,
320            else_dest,
321        },
322        InstrKind::Switch {
323            val,
324            default,
325            cases,
326        } => InstrKind::Switch {
327            val: s(val),
328            default,
329            cases: cases.into_iter().map(|(v, b)| (s(v), b)).collect(),
330        },
331        InstrKind::Unreachable => InstrKind::Unreachable,
332    }
333}
334
335// ---------------------------------------------------------------------------
336// RPO helper — produces block indices in reverse post-order
337// ---------------------------------------------------------------------------
338
339/// Returns block indices of `func` in reverse post-order (RPO) from the
340/// entry block (index 0).
341///
342/// RPO ensures that for any non-back-edge A→B in the CFG, A appears before B
343/// in the returned sequence.  This means constant values defined in a block
344/// are available for substitution in successor blocks in the same pass,
345/// maximising the number of folds performed per iteration.
346///
347/// Unreachable blocks are appended at the end (in their stored order) so that
348/// the pass still processes them for correctness.
349pub(crate) fn rpo(func: &Function) -> Vec<usize> {
350    let n = func.blocks.len();
351    let mut visited: HashSet<usize> = HashSet::with_capacity(n);
352    let mut post_order: Vec<usize> = Vec::with_capacity(n);
353
354    // Iterative DFS to avoid stack overflow on deep CFGs.
355    let mut stack: Vec<(usize, bool)> = vec![(0, false)]; // (block_idx, post-visit?)
356    while let Some((bi, post)) = stack.pop() {
357        if post {
358            post_order.push(bi);
359            continue;
360        }
361        if visited.contains(&bi) {
362            continue;
363        }
364        visited.insert(bi);
365        // Push post-visit marker first, then successors in reverse so we
366        // process them left-to-right.
367        stack.push((bi, true));
368        if let Some(tid) = func.blocks[bi].terminator {
369            let succs = func.instr(tid).successors();
370            for &succ in succs.iter().rev() {
371                let si = succ.0 as usize;
372                if si < n && !visited.contains(&si) {
373                    stack.push((si, false));
374                }
375            }
376        }
377    }
378
379    // Reverse post-order.
380    post_order.reverse();
381
382    // Append any unreachable blocks not visited by the DFS.
383    for bi in 0..n {
384        if !visited.contains(&bi) {
385            post_order.push(bi);
386        }
387    }
388
389    post_order
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395    use crate::pass::FunctionPass;
396    use llvm_ir::{Builder, Context, Linkage, Module, ValueRef};
397
398    // Build: f() -> i32 { ret (2 + 3) }  — the add folds to 5
399    fn make_const_add_fn() -> (Context, Module) {
400        let mut ctx = Context::new();
401        let mut module = Module::new("test");
402        let mut b = Builder::new(&mut ctx, &mut module);
403        b.add_function("f", b.ctx.i32_ty, vec![], vec![], false, Linkage::External);
404        let entry = b.add_block("entry");
405        b.position_at_end(entry);
406        let c2 = b.const_int(b.ctx.i32_ty, 2);
407        let c3 = b.const_int(b.ctx.i32_ty, 3);
408        let sum = b.build_add("sum", c2, c3);
409        b.build_ret(sum);
410        (ctx, module)
411    }
412
413    #[test]
414    fn const_prop_folds_add() {
415        let (mut ctx, mut module) = make_const_add_fn();
416        // Before: body has 'sum = add 2, 3'.
417        assert_eq!(module.functions[0].blocks[0].body.len(), 1);
418
419        let mut pass = ConstProp;
420        let changed = pass.run_on_function(&mut ctx, &mut module.functions[0]);
421        assert!(changed);
422
423        // After: 'sum' folded away; body is empty.
424        assert_eq!(
425            module.functions[0].blocks[0].body.len(),
426            0,
427            "folded add must be removed from body"
428        );
429
430        // The ret terminator should now reference the constant 5 directly.
431        let func = &module.functions[0];
432        let tid = func.blocks[0].terminator.unwrap();
433        if let llvm_ir::InstrKind::Ret { val: Some(v) } = &func.instr(tid).kind {
434            if let ValueRef::Constant(cid) = v {
435                assert_eq!(
436                    ctx.get_const(*cid),
437                    &llvm_ir::ConstantData::Int {
438                        ty: ctx.i32_ty,
439                        val: 5
440                    }
441                );
442            } else {
443                panic!("ret operand should be a constant");
444            }
445        } else {
446            panic!("terminator should be ret with value");
447        }
448    }
449
450    #[test]
451    fn const_prop_chain() {
452        // f() -> i32 { a = 2+3; b = a*10; ret b }  → 50
453        let mut ctx = Context::new();
454        let mut module = Module::new("test");
455        let mut b = Builder::new(&mut ctx, &mut module);
456        b.add_function("f", b.ctx.i32_ty, vec![], vec![], false, Linkage::External);
457        let entry = b.add_block("entry");
458        b.position_at_end(entry);
459        let c2 = b.const_int(b.ctx.i32_ty, 2);
460        let c3 = b.const_int(b.ctx.i32_ty, 3);
461        let c10 = b.const_int(b.ctx.i32_ty, 10);
462        let sum = b.build_add("sum", c2, c3);
463        let prod = b.build_mul("prod", sum, c10);
464        b.build_ret(prod);
465
466        let mut pass = ConstProp;
467        pass.run_on_function(&mut ctx, &mut module.functions[0]);
468
469        let func = &module.functions[0];
470        assert_eq!(
471            func.blocks[0].body.len(),
472            0,
473            "both sum and prod should be folded"
474        );
475        let tid = func.blocks[0].terminator.unwrap();
476        if let llvm_ir::InstrKind::Ret {
477            val: Some(ValueRef::Constant(cid)),
478        } = &func.instr(tid).kind
479        {
480            assert_eq!(
481                ctx.get_const(*cid),
482                &llvm_ir::ConstantData::Int {
483                    ty: ctx.i32_ty,
484                    val: 50
485                }
486            );
487        } else {
488            panic!("expected ret with constant 50");
489        }
490    }
491
492    #[test]
493    fn const_prop_across_blocks_rpo() {
494        // Two blocks: entry defines a constant, then branches to exit which uses it.
495        //   entry: a = add 2, 3; br exit
496        //   exit:  b = add a, 10; ret b
497        // With RPO ordering, entry is processed before exit so `a` is folded (=5)
498        // and then `b = add 5, 10` is also folded (=15).
499        let mut ctx = Context::new();
500        let mut module = Module::new("test");
501        let mut b = Builder::new(&mut ctx, &mut module);
502        b.add_function("f", b.ctx.i32_ty, vec![], vec![], false, Linkage::External);
503
504        let entry = b.add_block("entry");
505        let exit = b.add_block("exit");
506
507        b.position_at_end(entry);
508        let c2 = b.const_int(b.ctx.i32_ty, 2);
509        let c3 = b.const_int(b.ctx.i32_ty, 3);
510        let a = b.build_add("a", c2, c3);
511        b.build_br(exit);
512
513        b.position_at_end(exit);
514        let c10 = b.const_int(b.ctx.i32_ty, 10);
515        let bv = b.build_add("b", a, c10);
516        b.build_ret(bv);
517
518        let mut pass = ConstProp;
519        pass.run_on_function(&mut ctx, &mut module.functions[0]);
520
521        let func = &module.functions[0];
522        // Both `a` and `b` should be folded away.
523        assert_eq!(
524            func.blocks[0].body.len(),
525            0,
526            "`a` should be folded in entry"
527        );
528        assert_eq!(func.blocks[1].body.len(), 0, "`b` should be folded in exit");
529        // ret in exit block should reference constant 15.
530        let tid = func.blocks[1].terminator.unwrap();
531        if let llvm_ir::InstrKind::Ret {
532            val: Some(ValueRef::Constant(cid)),
533        } = &func.instr(tid).kind
534        {
535            assert_eq!(
536                ctx.get_const(*cid),
537                &llvm_ir::ConstantData::Int {
538                    ty: ctx.i32_ty,
539                    val: 15
540                }
541            );
542        } else {
543            panic!("expected ret with constant 15");
544        }
545    }
546
547    #[test]
548    fn rpo_order_entry_before_successor() {
549        // Verify that rpo() returns entry block (0) before its successor.
550        let mut ctx = Context::new();
551        let mut module = Module::new("test");
552        let mut b = Builder::new(&mut ctx, &mut module);
553        b.add_function("f", b.ctx.i32_ty, vec![], vec![], false, Linkage::External);
554        let entry = b.add_block("entry");
555        let exit = b.add_block("exit");
556        b.position_at_end(entry);
557        b.build_br(exit);
558        b.position_at_end(exit);
559        let c = b.const_int(b.ctx.i32_ty, 0);
560        b.build_ret(c);
561
562        let func = &module.functions[0];
563        let order = rpo(func);
564        let entry_pos = order.iter().position(|&i| i == entry.0 as usize).unwrap();
565        let exit_pos = order.iter().position(|&i| i == exit.0 as usize).unwrap();
566        assert!(entry_pos < exit_pos, "entry must come before exit in RPO");
567    }
568}
llvm_transforms/const_prop.rs

llvm_transforms/
const_prop.rs
