use std::collections::HashSet;
use rucc_base::Interner;
use rucc_ir::{Block, Def, Extra, Flags, Func, FuncId, IntPred, Module, Opcode, Value};
use crate::cfg::Cfg;
use crate::discharge::{Fact, constant, operand_of};
const MAKES: &[&str] = &["aligned_alloc", "calloc", "malloc", "realloc", "reallocf", "valloc"];
const LARGEST: i128 = 4 * 1024 * 1024 * 1024;
pub fn annotate(module: &mut Module, names: &Interner) -> usize {
let defined: HashSet<&str> = module
.funcs()
.filter(|&id| !module[id].is_declaration())
.map(|id| names.resolve(module[id].name))
.collect();
let mut marked = 0;
let ids: Vec<FuncId> = module.funcs().collect();
for id in ids {
if module[id].is_declaration() {
continue;
}
let func = &module[id];
let marks: Vec<_> = func
.blocks()
.flat_map(|block| func.insts(block).collect::<Vec<_>>())
.filter(|&inst| func[inst].opcode == Opcode::Call)
.filter(|&inst| !func[inst].flags.contains(Flags::HEAP))
.filter(|&inst| {
let Extra::Call(at) = func[inst].extra else { return false };
let Some(callee) = func[at].callee else { return false };
let name = names.resolve(callee);
!defined.contains(name) && MAKES.binary_search(&name).is_ok()
})
.collect();
marked += marks.len();
let func = &mut module[id];
for inst in marks {
func[inst].flags |= Flags::HEAP;
}
}
marked
}
pub(crate) fn allocates(func: &Func) -> bool {
func.blocks().any(|block| {
func.insts(block)
.any(|inst| func[inst].opcode == Opcode::Call && func[inst].flags.contains(Flags::HEAP))
})
}
pub(crate) fn made(func: &Func, base: Value) -> Option<Fact> {
let Def::Result { inst, index: 0 } = func[base].def else { return None };
if !func[inst].flags.contains(Flags::HEAP) {
return None;
}
let &last = func[func[inst].args].last()?;
let size = constant(func, last)?;
(size > 0 && size <= LARGEST).then(|| Fact::whole(base, size))
}
pub(crate) fn tested(func: &Func, cfg: &Cfg, pointer: Value) -> HashSet<Block> {
let Some(entry) = cfg.entry() else { return HashSet::new() };
let order: Vec<Block> = cfg.reverse_postorder().collect();
let mut known: HashSet<Block> = order.iter().copied().filter(|&block| block != entry).collect();
loop {
let mut settled = true;
for &block in &order {
if !known.contains(&block) {
continue;
}
let preds = cfg.predecessors(block);
let holds = !preds.is_empty()
&& preds
.iter()
.all(|&pred| known.contains(&pred) || proves(func, pred, block, pointer));
if !holds {
known.remove(&block);
settled = false;
}
}
if settled {
return known;
}
}
}
fn proves(func: &Func, pred: Block, into: Block, pointer: Value) -> bool {
let Some(term) = func.terminator(pred) else { return false };
if func[term].opcode != Opcode::BrIf {
return false;
}
let Extra::Targets(targets) = func[term].extra else { return false };
let [first, second] = func[targets] else { return false };
if first.block == second.block {
return false;
}
let Some(&condition) = func[func[term].args].first() else { return false };
match against_null(func, condition, pointer) {
Some(IntPred::Ne) => first.block == into,
Some(IntPred::Eq) => second.block == into,
_ => false,
}
}
fn against_null(func: &Func, condition: Value, pointer: Value) -> Option<IntPred> {
let Def::Result { inst, .. } = func[condition].def else { return None };
if func[inst].opcode != Opcode::ICmp {
return None;
}
let Extra::IntPred(pred) = func[inst].extra else { return None };
if !matches!(pred, IntPred::Eq | IntPred::Ne) {
return None;
}
let args = &func[func[inst].args];
let (&left, &right) = (args.first()?, args.get(1)?);
let matched = (left == pointer && null(func, right)) || (right == pointer && null(func, left));
matched.then_some(pred)
}
fn null(func: &Func, value: Value) -> bool {
if constant(func, value) == Some(0) {
return true;
}
operand_of(func, value, Opcode::IntToPtr, 0)
.is_some_and(|inner| constant(func, inner) == Some(0))
}
#[cfg(test)]
mod tests {
use rucc_base::Interner;
use rucc_ir::{Builder, Func, Linkage, Signature, Type};
use rucc_target::{TargetInfo, Triple};
use super::{
Fact, Flags, IntPred, LARGEST, MAKES, Module, Opcode, Value, allocates, made, tested,
};
use crate::cfg::Cfg;
fn module(names: &mut Interner) -> Module {
let target = TargetInfo::new("x86_64-unknown-linux-gnu".parse::<Triple>().unwrap());
Module::new(names.intern("t.c"), &target)
}
fn shape() -> Signature {
Signature::new().with_params(&[Type::int(64)]).with_returns(&[Type::PTR])
}
fn calls(names: &mut Interner, module: &mut Module, at: &str, name: &str, size: i128) {
let at = names.intern(at);
let called = names.intern(name);
let mut func = Func::new(at, Signature::new());
let block = func.create_block();
let mut build = Builder::new(&mut func, block);
let signature = build.func().add_signature(shape());
let bytes = build.iconst(Type::int(64), size);
build.call(called, signature, &[bytes]);
build.ret(&[]);
module.add_func(func);
}
fn marked(func: &Func) -> bool {
allocates(func)
}
fn allocation(names: &mut Interner, build: &mut Builder<'_>, size: i128) -> Value {
let called = names.intern("malloc");
let signature = build.func().add_signature(shape());
let bytes = build.iconst(Type::int(64), size);
let call = build.call(called, signature, &[bytes]);
let func = build.func();
func[call].flags |= Flags::HEAP;
func[call].results().next().expect("a call that gives back a pointer")
}
#[test]
fn the_names_are_sorted_and_each_one_is_written_once() {
let mut sorted = MAKES.to_vec();
sorted.sort_unstable();
sorted.dedup();
assert_eq!(sorted, MAKES);
}
#[test]
fn a_call_to_an_allocator_is_marked_and_a_call_to_anything_else_is_not() {
let mut names = Interner::new();
let mut module = module(&mut names);
calls(&mut names, &mut module, "f", "malloc", 16);
calls(&mut names, &mut module, "g", "reallocf", 16);
calls(&mut names, &mut module, "h", "memcpy", 16);
assert_eq!(super::annotate(&mut module, &names), 2);
let found: Vec<bool> = module.funcs().map(|id| marked(&module[id])).collect();
assert_eq!(found, [true, true, false]);
}
#[test]
fn a_module_that_writes_its_own_allocator_is_left_alone() {
let mut names = Interner::new();
let mut module = module(&mut names);
calls(&mut names, &mut module, "f", "malloc", 16);
let name = names.intern("malloc");
let mut mine = Func::new(name, shape());
mine.linkage = Linkage::External;
let block = mine.create_block();
let mut build = Builder::new(&mut mine, block);
build.ret(&[]);
module.add_func(mine);
assert_eq!(super::annotate(&mut module, &names), 0);
}
#[test]
fn a_declaration_is_walked_over_rather_than_into() {
let mut names = Interner::new();
let mut module = module(&mut names);
let name = names.intern("malloc");
module.add_func(Func::new(name, shape()));
calls(&mut names, &mut module, "f", "malloc", 16);
assert_eq!(super::annotate(&mut module, &names), 1);
}
#[test]
fn the_size_is_read_off_the_last_argument_and_has_to_be_a_number_in_range() {
for (size, works) in
[(16, true), (LARGEST, true), (0, false), (-1, false), (LARGEST + 1, false)]
{
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new());
let block = func.create_block();
let mut build = Builder::new(&mut func, block);
let pointer = allocation(&mut names, &mut build, size);
build.ret(&[]);
let wanted = works.then(|| Fact::whole(pointer, size));
assert_eq!(made(&func, pointer), wanted, "{size}");
}
}
#[test]
fn a_pointer_from_a_call_nobody_marked_is_not_an_allocation() {
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new());
let block = func.create_block();
let mut build = Builder::new(&mut func, block);
let called = names.intern("mine");
let signature = build.func().add_signature(shape());
let bytes = build.iconst(Type::int(64), 16);
let call = build.call(called, signature, &[bytes]);
let pointer = build.func()[call].results().next().expect("a pointer");
build.ret(&[]);
assert_eq!(made(&func, pointer), None);
}
#[test]
fn the_tested_arm_is_the_one_the_comparison_says_it_is() {
for pred in [IntPred::Ne, IntPred::Eq] {
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new());
let entry = func.create_block();
let first = func.create_block();
let second = func.create_block();
let mut build = Builder::new(&mut func, entry);
let pointer = allocation(&mut names, &mut build, 16);
let zero = build.iconst(Type::int(64), 0);
let null = build.unary(Opcode::IntToPtr, zero, Type::PTR);
let condition = build.icmp(pred, pointer, null);
build.br_if(condition, first, &[], second, &[]);
for block in [first, second] {
let mut build = Builder::new(&mut func, block);
build.ret(&[]);
}
let cfg = Cfg::new(&func);
let known = tested(&func, &cfg, pointer);
let good = if pred == IntPred::Ne { first } else { second };
let bad = if pred == IntPred::Ne { second } else { first };
assert!(known.contains(&good), "{pred:?}");
assert!(!known.contains(&bad), "{pred:?}");
assert!(!known.contains(&entry), "nothing is known before the test runs");
}
}
#[test]
fn a_join_of_two_paths_that_have_both_been_through_the_test_is_still_tested() {
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new().with_params(&[Type::int(32)]));
let entry = func.create_block();
let inside = func.create_block();
let arm = func.create_block();
let join = func.create_block();
let outside = func.create_block();
let mut build = Builder::new(&mut func, entry);
let value = build.func().append_param(entry, Type::int(32));
let pointer = allocation(&mut names, &mut build, 16);
let zero = build.iconst(Type::int(64), 0);
let null = build.unary(Opcode::IntToPtr, zero, Type::PTR);
let condition = build.icmp(IntPred::Ne, pointer, null);
build.br_if(condition, inside, &[], outside, &[]);
let mut build = Builder::new(&mut func, inside);
let none = build.iconst(Type::int(32), 0);
let again = build.icmp(IntPred::Ne, value, none);
build.br_if(again, arm, &[], join, &[]);
let mut build = Builder::new(&mut func, arm);
build.jump(join, &[]);
for block in [join, outside] {
let mut build = Builder::new(&mut func, block);
build.ret(&[]);
}
let cfg = Cfg::new(&func);
let known = tested(&func, &cfg, pointer);
assert!(known.contains(&inside));
assert!(known.contains(&arm));
assert!(known.contains(&join), "both ways in have been through the test");
assert!(!known.contains(&outside));
assert!(!known.contains(&entry));
}
#[test]
fn a_block_a_path_reaches_without_the_test_is_not_tested() {
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new());
let entry = func.create_block();
let sorry = func.create_block();
let after = func.create_block();
let mut build = Builder::new(&mut func, entry);
let pointer = allocation(&mut names, &mut build, 16);
let zero = build.iconst(Type::int(64), 0);
let null = build.unary(Opcode::IntToPtr, zero, Type::PTR);
let condition = build.icmp(IntPred::Eq, pointer, null);
build.br_if(condition, sorry, &[], after, &[]);
let mut build = Builder::new(&mut func, sorry);
build.jump(after, &[]);
let mut build = Builder::new(&mut func, after);
build.ret(&[]);
let cfg = Cfg::new(&func);
let known = tested(&func, &cfg, pointer);
assert!(!known.contains(&after));
assert!(!known.contains(&sorry));
}
#[test]
fn a_loop_the_test_is_outside_of_is_tested_all_the_way_round() {
let mut names = Interner::new();
let name = names.intern("f");
let mut func = Func::new(name, Signature::new().with_params(&[Type::int(32)]));
let entry = func.create_block();
let header = func.create_block();
let body = func.create_block();
let done = func.create_block();
let mut build = Builder::new(&mut func, entry);
let value = build.func().append_param(entry, Type::int(32));
let pointer = allocation(&mut names, &mut build, 16);
let zero = build.iconst(Type::int(64), 0);
let null = build.unary(Opcode::IntToPtr, zero, Type::PTR);
let condition = build.icmp(IntPred::Ne, pointer, null);
build.br_if(condition, header, &[], done, &[]);
let mut build = Builder::new(&mut func, header);
let none = build.iconst(Type::int(32), 0);
let again = build.icmp(IntPred::Ne, value, none);
build.br_if(again, body, &[], done, &[]);
let mut build = Builder::new(&mut func, body);
build.jump(header, &[]);
let mut build = Builder::new(&mut func, done);
build.ret(&[]);
let cfg = Cfg::new(&func);
let known = tested(&func, &cfg, pointer);
assert!(known.contains(&header));
assert!(known.contains(&body));
assert!(!known.contains(&done));
}
}