llvm_transforms/
constant_fold_pass.rs

1//! Dedicated constant-folding function pass.
2//!
3//! This pass evaluates foldable instructions whose operands are compile-time
4//! constants (via [`crate::constant_fold::try_fold`]) and rewrites downstream
5//! uses to direct constants.
6
7use crate::const_prop::{rpo, subst_kind};
8use crate::constant_fold::try_fold;
9use crate::pass::FunctionPass;
10use llvm_ir::{Context, Function, InstrId, ValueRef};
11use std::collections::HashMap;
12
13/// Function pass that folds compile-time constant expressions.
14pub struct ConstantFold;
15
16impl FunctionPass for ConstantFold {
17    fn name(&self) -> &'static str {
18        "constant-fold"
19    }
20
21    fn run_on_function(&mut self, ctx: &mut Context, func: &mut Function) -> bool {
22        if func.blocks.is_empty() {
23            return false;
24        }
25
26        // Map InstrId -> folded constant replacement.
27        let mut subst: HashMap<InstrId, ValueRef> = HashMap::new();
28
29        for bi in rpo(func) {
30            let body = func.blocks[bi].body.clone();
31            for iid in body {
32                if !subst.is_empty() {
33                    let new_kind = subst_kind(func.instr(iid).kind.clone(), &subst);
34                    func.instr_mut(iid).kind = new_kind;
35                }
36                let kind = func.instr(iid).kind.clone();
37                if let Some(cid) = try_fold(ctx, &kind) {
38                    subst.insert(iid, ValueRef::Constant(cid));
39                }
40            }
41            if let Some(tid) = func.blocks[bi].terminator {
42                if !subst.is_empty() {
43                    let new_kind = subst_kind(func.instr(tid).kind.clone(), &subst);
44                    func.instr_mut(tid).kind = new_kind;
45                }
46            }
47        }
48
49        if subst.is_empty() {
50            return false;
51        }
52
53        for bb in &mut func.blocks {
54            bb.body.retain(|id| !subst.contains_key(id));
55        }
56        true
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use llvm_ir::{Builder, Context, InstrKind, Linkage, Module, ValueRef};
64
65    fn make_const_add_fn() -> (Context, Module) {
66        let mut ctx = Context::new();
67        let mut module = Module::new("const_add");
68        let mut b = Builder::new(&mut ctx, &mut module);
69        b.add_function("f", b.ctx.i32_ty, vec![], vec![], false, Linkage::External);
70        let entry = b.add_block("entry");
71        b.position_at_end(entry);
72        let c2 = b.const_int(b.ctx.i32_ty, 2);
73        let sum = b.build_add("sum", c2, c2);
74        b.build_ret(sum);
75        (ctx, module)
76    }
77
78    fn make_non_const_add_fn() -> (Context, Module) {
79        let mut ctx = Context::new();
80        let mut module = Module::new("non_const_add");
81        let mut b = Builder::new(&mut ctx, &mut module);
82        b.add_function(
83            "f",
84            b.ctx.i32_ty,
85            vec![b.ctx.i32_ty],
86            vec!["x".into()],
87            false,
88            Linkage::External,
89        );
90        let entry = b.add_block("entry");
91        b.position_at_end(entry);
92        let x = b.get_arg(0);
93        let c2 = b.const_int(b.ctx.i32_ty, 2);
94        let sum = b.build_add("sum", x, c2);
95        b.build_ret(sum);
96        (ctx, module)
97    }
98
99    #[test]
100    fn folds_add_2_plus_2() {
101        let (mut ctx, mut module) = make_const_add_fn();
102        let mut pass = ConstantFold;
103        let changed = pass.run_on_function(&mut ctx, &mut module.functions[0]);
104        assert!(changed);
105        assert_eq!(module.functions[0].blocks[0].body.len(), 0);
106        let func = &module.functions[0];
107        let tid = func.blocks[0].terminator.expect("terminator");
108        match &func.instr(tid).kind {
109            InstrKind::Ret {
110                val: Some(ValueRef::Constant(cid)),
111            } => match ctx.get_const(*cid) {
112                llvm_ir::ConstantData::Int { val, .. } => assert_eq!(*val, 4),
113                other => panic!("unexpected ret constant: {other:?}"),
114            },
115            other => panic!("expected ret constant, got {other:?}"),
116        }
117    }
118
119    #[test]
120    fn does_not_fold_non_constant_expression() {
121        let (mut ctx, mut module) = make_non_const_add_fn();
122        let mut pass = ConstantFold;
123        let changed = pass.run_on_function(&mut ctx, &mut module.functions[0]);
124        assert!(!changed);
125        assert_eq!(module.functions[0].blocks[0].body.len(), 1);
126    }
127}
llvm_transforms/constant_fold_pass.rs

llvm_transforms/
constant_fold_pass.rs
