llvm_transforms/
dead_arg_elim.rs

1//! Inter-procedural dead argument elimination.
2
3use crate::pass::ModulePass;
4use llvm_ir::{
5    Context, FunctionId, GlobalId, InstrKind, Module, TypeData, ValueRef,
6};
7
8/// Removes trailing unused function parameters and rewrites direct callsites.
9///
10/// This intentionally targets only trailing dead args to avoid argument
11/// reindexing across function bodies.
12pub struct DeadArgElim;
13
14impl ModulePass for DeadArgElim {
15    fn name(&self) -> &'static str {
16        "dead-arg-elim"
17    }
18
19    fn run_on_module(&mut self, ctx: &mut Context, module: &mut Module) -> bool {
20        let mut changed = false;
21        let mut updates: Vec<(FunctionId, usize)> = Vec::new();
22
23        for (fi, f) in module.functions.iter().enumerate() {
24            if f.is_declaration || f.args.is_empty() {
25                continue;
26            }
27            let mut max_used: Option<usize> = None;
28            for instr in &f.instructions {
29                for op in instr.kind.operands() {
30                    if let ValueRef::Argument(aid) = op {
31                        let idx = aid.0 as usize;
32                        max_used = Some(max_used.map_or(idx, |m| m.max(idx)));
33                    }
34                }
35            }
36            let keep_len = max_used.map_or(0, |m| m + 1).min(f.args.len());
37            if keep_len < f.args.len() {
38                updates.push((FunctionId(fi as u32), keep_len));
39            }
40        }
41
42        for (fid, keep_len) in updates {
43            let f = &mut module.functions[fid.0 as usize];
44            let old_len = f.args.len();
45            if keep_len >= old_len {
46                continue;
47            }
48            f.args.truncate(keep_len);
49            f.arg_names.retain(|_, aid| (aid.0 as usize) < keep_len);
50
51            if let TypeData::Function(ft) = ctx.get_type(f.ty).clone() {
52                let mut params = ft.params;
53                params.truncate(keep_len);
54                f.ty = ctx.mk_fn_type(ft.ret, params, ft.variadic);
55            }
56            let new_ty = f.ty;
57
58            for caller in &mut module.functions {
59                for instr in &mut caller.instructions {
60                    let InstrKind::Call {
61                        callee,
62                        callee_ty,
63                        args,
64                        ..
65                    } = &mut instr.kind
66                    else {
67                        continue;
68                    };
69                    if *callee == ValueRef::Global(GlobalId(fid.0)) {
70                        args.truncate(keep_len);
71                        *callee_ty = new_ty;
72                    }
73                }
74            }
75            changed = true;
76        }
77
78        changed
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use llvm_ir::{Builder, Linkage};
86
87    #[test]
88    fn dead_arg_elim_removes_trailing_unused_arg_and_rewrites_calls() {
89        let mut ctx = Context::new();
90        let mut module = Module::new("m");
91        let mut b = Builder::new(&mut ctx, &mut module);
92
93        b.add_function(
94            "f",
95            b.ctx.i64_ty,
96            vec![b.ctx.i64_ty, b.ctx.i64_ty, b.ctx.i64_ty],
97            vec!["a".into(), "b".into(), "dead".into()],
98            false,
99            Linkage::External,
100        );
101        let f_entry = b.add_block("f.entry");
102        b.position_at_end(f_entry);
103        let a = b.get_arg(0);
104        let bv = b.get_arg(1);
105        let s = b.build_add("s", a, bv);
106        b.build_ret(s);
107
108        b.add_function(
109            "caller",
110            b.ctx.i64_ty,
111            vec![b.ctx.i64_ty],
112            vec!["x".into()],
113            false,
114            Linkage::External,
115        );
116        let c_entry = b.add_block("caller.entry");
117        b.position_at_end(c_entry);
118        let x = b.get_arg(0);
119        let c1 = b.const_int(b.ctx.i64_ty, 1);
120        let c2 = b.const_int(b.ctx.i64_ty, 2);
121        let call_ty = b
122            .ctx
123            .mk_fn_type(b.ctx.i64_ty, vec![b.ctx.i64_ty, b.ctx.i64_ty, b.ctx.i64_ty], false);
124        let r = b.build_call(
125            "r",
126            b.ctx.i64_ty,
127            call_ty,
128            ValueRef::Global(GlobalId(0)),
129            vec![x, c1, c2],
130        );
131        b.build_ret(r);
132
133        let mut pass = DeadArgElim;
134        let changed = pass.run_on_module(&mut ctx, &mut module);
135        assert!(changed);
136        assert_eq!(module.functions[0].args.len(), 2);
137        let call = module.functions[1]
138            .instructions
139            .iter()
140            .find(|i| matches!(i.kind, InstrKind::Call { .. }))
141            .expect("call exists");
142        if let InstrKind::Call { args, .. } = &call.kind {
143            assert_eq!(args.len(), 2);
144        } else {
145            panic!("expected call");
146        }
147    }
148}
llvm_transforms/dead_arg_elim.rs

llvm_transforms/
dead_arg_elim.rs
