1use std::collections::{HashMap, HashSet};
62
63use rucc_base::Symbol;
64use rucc_ir::{Datum, Def, Extra, Flags, Func, FuncId, Inst, Linkage, Module, Opcode, Type, Value};
65
66use crate::discharge::{Fact, about, alive, covers, derives, normal};
67use crate::extents::extents;
68
69pub fn annotate(module: &mut Module) -> usize {
73 let reachable = reachable(module);
74 let closed: Vec<FuncId> = module
75 .funcs()
76 .filter(|&id| {
77 let func = &module[id];
78 !func.is_declaration()
79 && func.linkage == Linkage::Internal
80 && !reachable.contains(&func.name)
81 })
82 .collect();
83 if closed.is_empty() {
84 return 0;
85 }
86 let globals = extents(module);
87 let handed = handed(module, &closed, &globals);
88 if handed.is_empty() {
89 return 0;
90 }
91 let mut marked = 0;
92 for id in closed {
93 let Some(sizes) = handed.get(&id) else { continue };
94 let func = &module[id];
95 let Some(entry) = func.entry() else { continue };
96 let object = |base: Value| -> Option<Fact> {
97 let Def::Param { block, index } = func[base].def else { return None };
98 if block != entry {
99 return None;
100 }
101 Some(Fact::whole(base, i128::from(*sizes.get(&index)?)))
102 };
103 let marks: Vec<Inst> = func
104 .blocks()
105 .flat_map(|block| func.insts(block).collect::<Vec<_>>())
106 .filter(|&inst| !func[inst].flags.contains(Flags::HANDED))
107 .filter(|&inst| inside(func, inst, &object))
108 .collect();
109 marked += marks.len();
110 let func = &mut module[id];
111 for inst in marks {
112 func[inst].flags |= Flags::HANDED;
113 }
114 }
115 marked
116}
117
118fn handed(
124 module: &Module,
125 closed: &[FuncId],
126 globals: &HashMap<Symbol, u64>,
127) -> HashMap<FuncId, HashMap<u32, u64>> {
128 let mut where_defined: HashMap<_, FuncId> = HashMap::new();
129 for &id in closed {
130 where_defined.insert(module[id].name, id);
131 }
132 let sites = sites(module, &where_defined);
133 let mut known: HashMap<FuncId, HashMap<u32, u64>> = HashMap::new();
134 loop {
135 let mut settled = true;
136 for &id in closed {
137 let Some(calls) = sites.get(&id) else { continue };
138 let count = module[id].signature().params.len();
139 let mut sizes = HashMap::new();
140 for index in 0..count {
141 if module[id].signature().params[index].ty != Type::PTR {
142 continue;
143 }
144 let Some(least) = least(module, calls, index, globals, &known) else { continue };
145 sizes.insert(u32::try_from(index).unwrap_or(u32::MAX), least);
146 }
147 if known.get(&id) != Some(&sizes) {
148 known.insert(id, sizes);
149 settled = false;
150 }
151 }
152 if settled {
153 known.retain(|_, sizes| !sizes.is_empty());
154 return known;
155 }
156 }
157}
158
159fn least(
166 module: &Module,
167 calls: &[(FuncId, Inst)],
168 index: usize,
169 globals: &HashMap<Symbol, u64>,
170 known: &HashMap<FuncId, HashMap<u32, u64>>,
171) -> Option<u64> {
172 let mut least = None;
173 for &(caller, inst) in calls {
174 let func = &module[caller];
175 let &value = func[func[inst].args].get(index)?;
176 let left = passed(caller, func, value, globals, known)?;
177 least = Some(least.map_or(left, |so_far: u64| so_far.min(left)));
178 }
179 least
180}
181
182fn passed(
188 caller: FuncId,
189 func: &Func,
190 value: Value,
191 globals: &HashMap<Symbol, u64>,
192 known: &HashMap<FuncId, HashMap<u32, u64>>,
193) -> Option<u64> {
194 let (base, offset) = normal(func, value);
195 let whole = i128::from(object(caller, func, base, globals, known)?);
196 if offset < 0 || offset > whole {
197 return None;
198 }
199 u64::try_from(whole - offset).ok()
200}
201
202fn object(
204 caller: FuncId,
205 func: &Func,
206 base: Value,
207 globals: &HashMap<Symbol, u64>,
208 known: &HashMap<FuncId, HashMap<u32, u64>>,
209) -> Option<u64> {
210 match func[base].def {
211 Def::Param { block, index } => {
215 if func.entry() != Some(block) {
216 return None;
217 }
218 known.get(&caller)?.get(&index).copied()
219 }
220 Def::Result { inst, .. } => match func[inst].opcode {
221 Opcode::Alloca if func[func[inst].args].is_empty() => {
222 let Extra::Mem(info) = func[inst].extra else { return None };
223 Some(func[info].size)
224 }
225 Opcode::GlobalAddr => {
226 let Extra::Symbol(name) = func[inst].extra else { return None };
227 globals.get(&name).copied()
228 }
229 _ => None,
230 },
231 }
232}
233
234fn sites(
241 module: &Module,
242 where_defined: &HashMap<Symbol, FuncId>,
243) -> HashMap<FuncId, Vec<(FuncId, Inst)>> {
244 let mut sites: HashMap<FuncId, Vec<(FuncId, Inst)>> = HashMap::new();
245 for id in module.funcs() {
246 let func = &module[id];
247 if func.is_declaration() {
248 continue;
249 }
250 for block in func.blocks() {
251 for inst in func.insts(block) {
252 if !matches!(func[inst].opcode, Opcode::Call | Opcode::TailCall) {
253 continue;
254 }
255 let Extra::Call(at) = func[inst].extra else { continue };
256 let Some(callee) = func[at].callee else { continue };
257 let Some(&target) = where_defined.get(&callee) else { continue };
258 let signature = module[target].signature();
259 if signature.variadic || signature.params.len() != func[func[inst].args].len() {
260 continue;
261 }
262 sites.entry(target).or_default().push((id, inst));
263 }
264 }
265 }
266 sites
267}
268
269fn reachable(module: &Module) -> HashSet<Symbol> {
275 let mut taken = HashSet::new();
276 for id in module.funcs() {
277 let func = &module[id];
278 if func.is_declaration() {
279 continue;
280 }
281 for block in func.blocks() {
282 for inst in func.insts(block) {
283 if func[inst].opcode != Opcode::GlobalAddr {
284 continue;
285 }
286 if let Extra::Symbol(name) = func[inst].extra {
287 taken.insert(name);
288 }
289 }
290 }
291 }
292 for id in module.globals() {
293 let Some(init) = module[id].init else { continue };
294 for &datum in &module[init] {
295 if let Datum::Addr(at) = datum {
296 taken.insert(module[at].symbol);
297 }
298 }
299 }
300 for id in module.aliases() {
301 taken.insert(module[id].target);
302 }
303 taken
304}
305
306fn inside(func: &Func, inst: Inst, object: &impl Fn(Value) -> Option<Fact>) -> bool {
311 match func[inst].opcode {
312 Opcode::CheckBounds => {
313 if func[func[inst].args].len() > 2 {
314 return false;
315 }
316 let Some(asked) = about(func, inst) else { return false };
317 object(asked.base).is_some_and(|whole| covers(&whole, &asked))
318 }
319 Opcode::CheckLive => {
320 let Some(asked) = alive(func, inst) else { return false };
321 object(asked.base).is_some_and(|whole| covers(&whole, &asked))
322 }
323 Opcode::CheckDeriv => {
324 let Some((from, to)) = derives(func, inst) else { return false };
325 object(from.base).is_some_and(|whole| covers(&whole, &from) && covers(&whole, &to))
326 }
327 _ => false,
328 }
329}
330
331#[cfg(test)]
332mod tests {
333 use rucc_base::Interner;
334 use rucc_ir::{
335 Builder, Extra, Func, Global, InstData, Linkage, MemInfo, MemOrder, Module, Opcode,
336 Restrict, Signature, Type, Value,
337 };
338 use rucc_target::{TargetInfo, Triple};
339
340 use super::annotate;
341
342 fn module(names: &mut Interner) -> Module {
344 let target = TargetInfo::new("x86_64-unknown-linux-gnu".parse::<Triple>().unwrap());
345 Module::new(names.intern("t.c"), &target)
346 }
347
348 fn callee(names: &mut Interner, module: &mut Module, size: u64) {
351 let name = names.intern("g");
352 let mut func = Func::new(name, Signature::new().with_params(&[Type::PTR]));
353 func.linkage = Linkage::Internal;
354 let block = func.create_block();
355 let pointer = func.append_param(block, Type::PTR);
356 let mut build = Builder::new(&mut func, block);
357 check(&mut build, pointer, size);
358 live(&mut build, pointer);
359 build.ret(&[]);
360 module.add_func(func);
361 }
362
363 fn caller(
365 names: &mut Interner,
366 module: &mut Module,
367 name: &str,
368 argument: impl FnOnce(&mut Builder<'_>) -> Value,
369 ) {
370 let at = names.intern(name);
371 let called = names.intern("g");
372 let mut func = Func::new(at, Signature::new());
373 let block = func.create_block();
374 let mut build = Builder::new(&mut func, block);
375 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
376 let value = argument(&mut build);
377 build.call(called, signature, &[value]);
378 build.ret(&[]);
379 module.add_func(func);
380 }
381
382 fn local(build: &mut Builder<'_>, size: u64) -> Value {
384 let info = MemInfo {
385 size,
386 align: 8,
387 order: MemOrder::NotAtomic,
388 tbaa: None,
389 restrict: Restrict::NONE,
390 };
391 let extra = Extra::Mem(build.func().add_mem(info));
392 build.value(InstData { extra, ..InstData::new(Opcode::Alloca) }, Type::PTR)
393 }
394
395 fn check(build: &mut Builder<'_>, pointer: Value, size: u64) {
397 let args = build.func().push_values(&[pointer]);
398 let capability = build.value(InstData { args, ..InstData::new(Opcode::CapOf) }, Type::CAP);
399 let info = MemInfo {
400 size,
401 align: 1,
402 order: MemOrder::NotAtomic,
403 tbaa: None,
404 restrict: Restrict::NONE,
405 };
406 let args = build.func().push_values(&[capability, pointer]);
407 let extra = Extra::Mem(build.func().add_mem(info));
408 build.inst(InstData { args, extra, ..InstData::new(Opcode::CheckBounds) }, &[]);
409 }
410
411 fn live(build: &mut Builder<'_>, pointer: Value) {
413 let args = build.func().push_values(&[pointer]);
414 let capability = build.value(InstData { args, ..InstData::new(Opcode::CapOf) }, Type::CAP);
415 let args = build.func().push_values(&[capability, pointer]);
416 build.inst(InstData { args, ..InstData::new(Opcode::CheckLive) }, &[]);
417 }
418
419 fn past(build: &mut Builder<'_>, pointer: Value, bytes: i128) -> Value {
421 let offset = build.iconst(Type::int(64), bytes);
422 let args = build.func().push_values(&[pointer, offset]);
423 build.value(InstData { args, ..InstData::new(Opcode::PtrAdd) }, Type::PTR)
424 }
425
426 #[test]
427 fn a_check_on_a_parameter_every_call_hands_a_slot_big_enough_is_marked() {
428 let mut names = Interner::new();
429 let mut module = module(&mut names);
430 callee(&mut names, &mut module, 16);
431 caller(&mut names, &mut module, "f", |build| local(build, 32));
432 assert_eq!(annotate(&mut module), 2, "the bounds check and the lifetime one");
433 }
434
435 #[test]
436 fn a_call_handing_a_slot_too_small_marks_only_the_lifetime_check() {
437 let mut names = Interner::new();
441 let mut module = module(&mut names);
442 callee(&mut names, &mut module, 16);
443 caller(&mut names, &mut module, "f", |build| local(build, 8));
444 assert_eq!(annotate(&mut module), 1);
445 }
446
447 #[test]
448 fn the_fewest_bytes_any_call_hands_is_what_the_parameter_gets() {
449 let mut names = Interner::new();
453 let mut module = module(&mut names);
454 callee(&mut names, &mut module, 16);
455 caller(&mut names, &mut module, "f", |build| local(build, 32));
456 caller(&mut names, &mut module, "h", |build| local(build, 8));
457 assert_eq!(annotate(&mut module), 1, "the lifetime check, which eight bytes settle");
458 }
459
460 #[test]
461 fn a_call_handing_a_field_of_a_slot_leaves_what_is_past_the_field() {
462 let mut names = Interner::new();
465 let mut module = module(&mut names);
466 callee(&mut names, &mut module, 16);
467 caller(&mut names, &mut module, "f", |build| {
468 let slot = local(build, 32);
469 past(build, slot, 16)
470 });
471 assert_eq!(annotate(&mut module), 2);
472 }
473
474 #[test]
475 fn a_call_handing_a_field_that_leaves_too_little_marks_only_the_lifetime_check() {
476 let mut names = Interner::new();
479 let mut module = module(&mut names);
480 callee(&mut names, &mut module, 16);
481 caller(&mut names, &mut module, "f", |build| {
482 let slot = local(build, 32);
483 past(build, slot, 20)
484 });
485 assert_eq!(annotate(&mut module), 1);
486 }
487
488 #[test]
489 fn a_callee_anything_can_reach_is_left_alone() {
490 let mut names = Interner::new();
493 let mut module = module(&mut names);
494 callee(&mut names, &mut module, 16);
495 let id = module.funcs().next().expect("the callee");
496 module[id].linkage = Linkage::External;
497 caller(&mut names, &mut module, "f", |build| local(build, 32));
498 assert_eq!(annotate(&mut module), 0);
499 }
500
501 #[test]
502 fn a_callee_whose_address_is_taken_is_left_alone() {
503 let mut names = Interner::new();
506 let mut module = module(&mut names);
507 callee(&mut names, &mut module, 16);
508 caller(&mut names, &mut module, "f", |build| local(build, 32));
509 let called = names.intern("g");
510 caller(&mut names, &mut module, "h", |build| {
511 let extra = Extra::Symbol(called);
512 build.value(InstData { extra, ..InstData::new(Opcode::GlobalAddr) }, Type::PTR);
513 local(build, 32)
514 });
515 assert_eq!(annotate(&mut module), 0);
516 }
517
518 #[test]
519 fn a_callee_named_by_a_globals_image_is_left_alone() {
520 let mut names = Interner::new();
521 let mut module = module(&mut names);
522 callee(&mut names, &mut module, 16);
523 caller(&mut names, &mut module, "f", |build| local(build, 32));
524 let at = module.add_reloc(rucc_ir::Reloc { symbol: names.intern("g"), addend: 0, size: 8 });
525 let init = module.push_data(&[rucc_ir::Datum::Addr(at)]);
526 let mut global = Global::new(names.intern("table"), 8, 8);
527 global.init = Some(init);
528 module.add_global(global);
529 assert_eq!(annotate(&mut module), 0);
530 }
531
532 #[test]
533 fn a_callee_nothing_in_the_module_calls_is_left_alone() {
534 let mut names = Interner::new();
537 let mut module = module(&mut names);
538 callee(&mut names, &mut module, 16);
539 assert_eq!(annotate(&mut module), 0);
540 }
541
542 #[test]
543 fn a_chain_of_static_helpers_reaches_the_slot_at_the_top() {
544 let mut names = Interner::new();
548 let mut module = module(&mut names);
549 callee(&mut names, &mut module, 16);
550 let name = names.intern("h");
551 let called = names.intern("g");
552 let mut func = Func::new(name, Signature::new().with_params(&[Type::PTR]));
553 func.linkage = Linkage::Internal;
554 let block = func.create_block();
555 let pointer = func.append_param(block, Type::PTR);
556 let mut build = Builder::new(&mut func, block);
557 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
558 build.call(called, signature, &[pointer]);
559 build.ret(&[]);
560 module.add_func(func);
561 let at = names.intern("f");
562 let called = names.intern("h");
563 let mut func = Func::new(at, Signature::new());
564 let block = func.create_block();
565 let mut build = Builder::new(&mut func, block);
566 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
567 let slot = local(&mut build, 32);
568 build.call(called, signature, &[slot]);
569 build.ret(&[]);
570 module.add_func(func);
571 assert_eq!(annotate(&mut module), 2);
572 }
573
574 #[test]
575 fn two_functions_handing_each_other_their_own_parameter_hold_nothing_up() {
576 let mut names = Interner::new();
579 let mut module = module(&mut names);
580 relay(&mut names, &mut module, "g", "h", 16);
581 relay(&mut names, &mut module, "h", "g", 16);
582 assert_eq!(annotate(&mut module), 0);
583 }
584
585 fn relay(names: &mut Interner, module: &mut Module, name: &str, on: &str, size: u64) {
587 let at = names.intern(name);
588 let called = names.intern(on);
589 let mut func = Func::new(at, Signature::new().with_params(&[Type::PTR]));
590 func.linkage = Linkage::Internal;
591 let block = func.create_block();
592 let pointer = func.append_param(block, Type::PTR);
593 let mut build = Builder::new(&mut func, block);
594 check(&mut build, pointer, size);
595 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
596 build.call(called, signature, &[pointer]);
597 build.ret(&[]);
598 module.add_func(func);
599 }
600}