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