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 owns: 0,
392 restrict: Restrict::NONE,
393 };
394 let extra = Extra::Mem(build.func().add_mem(info));
395 build.value(InstData { extra, ..InstData::new(Opcode::Alloca) }, Type::PTR)
396 }
397
398 fn check(build: &mut Builder<'_>, pointer: Value, size: u64) {
400 let args = build.func().push_values(&[pointer]);
401 let capability = build.value(InstData { args, ..InstData::new(Opcode::CapOf) }, Type::CAP);
402 let info = MemInfo {
403 size,
404 align: 1,
405 order: MemOrder::NotAtomic,
406 tbaa: None,
407 owns: 0,
408 restrict: Restrict::NONE,
409 };
410 let args = build.func().push_values(&[capability, pointer]);
411 let extra = Extra::Mem(build.func().add_mem(info));
412 build.inst(InstData { args, extra, ..InstData::new(Opcode::CheckBounds) }, &[]);
413 }
414
415 fn live(build: &mut Builder<'_>, pointer: Value) {
417 let args = build.func().push_values(&[pointer]);
418 let capability = build.value(InstData { args, ..InstData::new(Opcode::CapOf) }, Type::CAP);
419 let args = build.func().push_values(&[capability, pointer]);
420 build.inst(InstData { args, ..InstData::new(Opcode::CheckLive) }, &[]);
421 }
422
423 fn past(build: &mut Builder<'_>, pointer: Value, bytes: i128) -> Value {
425 let offset = build.iconst(Type::int(64), bytes);
426 let args = build.func().push_values(&[pointer, offset]);
427 build.value(InstData { args, ..InstData::new(Opcode::PtrAdd) }, Type::PTR)
428 }
429
430 #[test]
431 fn a_check_on_a_parameter_every_call_hands_a_slot_big_enough_is_marked() {
432 let mut names = Interner::new();
433 let mut module = module(&mut names);
434 callee(&mut names, &mut module, 16);
435 caller(&mut names, &mut module, "f", |build| local(build, 32));
436 assert_eq!(
437 annotate(&mut module, Pic::Executable),
438 2,
439 "the bounds check and the lifetime one"
440 );
441 }
442
443 #[test]
444 fn a_call_handing_a_slot_too_small_marks_only_the_lifetime_check() {
445 let mut names = Interner::new();
449 let mut module = module(&mut names);
450 callee(&mut names, &mut module, 16);
451 caller(&mut names, &mut module, "f", |build| local(build, 8));
452 assert_eq!(annotate(&mut module, Pic::Executable), 1);
453 }
454
455 #[test]
456 fn the_fewest_bytes_any_call_hands_is_what_the_parameter_gets() {
457 let mut names = Interner::new();
461 let mut module = module(&mut names);
462 callee(&mut names, &mut module, 16);
463 caller(&mut names, &mut module, "f", |build| local(build, 32));
464 caller(&mut names, &mut module, "h", |build| local(build, 8));
465 assert_eq!(
466 annotate(&mut module, Pic::Executable),
467 1,
468 "the lifetime check, which eight bytes settle"
469 );
470 }
471
472 #[test]
473 fn a_call_handing_a_field_of_a_slot_leaves_what_is_past_the_field() {
474 let mut names = Interner::new();
477 let mut module = module(&mut names);
478 callee(&mut names, &mut module, 16);
479 caller(&mut names, &mut module, "f", |build| {
480 let slot = local(build, 32);
481 past(build, slot, 16)
482 });
483 assert_eq!(annotate(&mut module, Pic::Executable), 2);
484 }
485
486 #[test]
487 fn a_call_handing_a_field_that_leaves_too_little_marks_only_the_lifetime_check() {
488 let mut names = Interner::new();
491 let mut module = module(&mut names);
492 callee(&mut names, &mut module, 16);
493 caller(&mut names, &mut module, "f", |build| {
494 let slot = local(build, 32);
495 past(build, slot, 20)
496 });
497 assert_eq!(annotate(&mut module, Pic::Executable), 1);
498 }
499
500 #[test]
501 fn a_callee_anything_can_reach_is_left_alone() {
502 let mut names = Interner::new();
505 let mut module = module(&mut names);
506 callee(&mut names, &mut module, 16);
507 let id = module.funcs().next().expect("the callee");
508 module[id].linkage = Linkage::External;
509 caller(&mut names, &mut module, "f", |build| local(build, 32));
510 assert_eq!(annotate(&mut module, Pic::Executable), 0);
511 }
512
513 #[test]
514 fn a_callee_whose_address_is_taken_is_left_alone() {
515 let mut names = Interner::new();
518 let mut module = module(&mut names);
519 callee(&mut names, &mut module, 16);
520 caller(&mut names, &mut module, "f", |build| local(build, 32));
521 let called = names.intern("g");
522 caller(&mut names, &mut module, "h", |build| {
523 let extra = Extra::Symbol(called);
524 build.value(InstData { extra, ..InstData::new(Opcode::GlobalAddr) }, Type::PTR);
525 local(build, 32)
526 });
527 assert_eq!(annotate(&mut module, Pic::Executable), 0);
528 }
529
530 #[test]
531 fn a_callee_named_by_a_globals_image_is_left_alone() {
532 let mut names = Interner::new();
533 let mut module = module(&mut names);
534 callee(&mut names, &mut module, 16);
535 caller(&mut names, &mut module, "f", |build| local(build, 32));
536 let at = module.add_reloc(rucc_ir::Reloc { symbol: names.intern("g"), addend: 0, size: 8 });
537 let init = module.push_data(&[rucc_ir::Datum::Addr(at)]);
538 let mut global = Global::new(names.intern("table"), 8, 8);
539 global.init = Some(init);
540 module.add_global(global);
541 assert_eq!(annotate(&mut module, Pic::Executable), 0);
542 }
543
544 #[test]
545 fn a_callee_nothing_in_the_module_calls_is_left_alone() {
546 let mut names = Interner::new();
549 let mut module = module(&mut names);
550 callee(&mut names, &mut module, 16);
551 assert_eq!(annotate(&mut module, Pic::Executable), 0);
552 }
553
554 #[test]
555 fn a_chain_of_static_helpers_reaches_the_slot_at_the_top() {
556 let mut names = Interner::new();
560 let mut module = module(&mut names);
561 callee(&mut names, &mut module, 16);
562 let name = names.intern("h");
563 let called = names.intern("g");
564 let mut func = Func::new(name, Signature::new().with_params(&[Type::PTR]));
565 func.linkage = Linkage::Internal;
566 let block = func.create_block();
567 let pointer = func.append_param(block, Type::PTR);
568 let mut build = Builder::new(&mut func, block);
569 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
570 build.call(called, signature, &[pointer]);
571 build.ret(&[]);
572 module.add_func(func);
573 let at = names.intern("f");
574 let called = names.intern("h");
575 let mut func = Func::new(at, Signature::new());
576 let block = func.create_block();
577 let mut build = Builder::new(&mut func, block);
578 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
579 let slot = local(&mut build, 32);
580 build.call(called, signature, &[slot]);
581 build.ret(&[]);
582 module.add_func(func);
583 assert_eq!(annotate(&mut module, Pic::Executable), 2);
584 }
585
586 #[test]
587 fn two_functions_handing_each_other_their_own_parameter_hold_nothing_up() {
588 let mut names = Interner::new();
591 let mut module = module(&mut names);
592 relay(&mut names, &mut module, "g", "h", 16);
593 relay(&mut names, &mut module, "h", "g", 16);
594 assert_eq!(annotate(&mut module, Pic::Executable), 0);
595 }
596
597 fn relay(names: &mut Interner, module: &mut Module, name: &str, on: &str, size: u64) {
599 let at = names.intern(name);
600 let called = names.intern(on);
601 let mut func = Func::new(at, Signature::new().with_params(&[Type::PTR]));
602 func.linkage = Linkage::Internal;
603 let block = func.create_block();
604 let pointer = func.append_param(block, Type::PTR);
605 let mut build = Builder::new(&mut func, block);
606 check(&mut build, pointer, size);
607 let signature = build.func().add_signature(Signature::new().with_params(&[Type::PTR]));
608 build.call(called, signature, &[pointer]);
609 build.ret(&[]);
610 module.add_func(func);
611 }
612}