1use rucc_mir::{Constraint, Func, Operand, Reg, Role};
66use rucc_target::{PhysReg, RegClass};
67
68use crate::live::{Live, Range};
69use crate::order::{Order, Point};
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub enum Place {
74 Reg(PhysReg),
76 Slot(u32),
79}
80
81#[derive(Debug, Clone, Default)]
91pub struct Env {
92 classes: Vec<Class>,
93}
94
95#[derive(Debug, Clone, Default)]
97struct Class {
98 order: Vec<PhysReg>,
99 scratch: Vec<PhysReg>,
100}
101
102impl Env {
103 #[must_use]
105 pub fn new() -> Self {
106 Self::default()
107 }
108
109 #[must_use]
111 pub fn with(mut self, class: RegClass, order: &[PhysReg], scratch: &[PhysReg]) -> Self {
112 let index = usize::from(class.number());
113 if self.classes.len() <= index {
114 self.classes.resize(index + 1, Class::default());
115 }
116 self.classes[index] = Class { order: order.to_vec(), scratch: scratch.to_vec() };
117 self
118 }
119
120 #[must_use]
122 pub fn order(&self, class: RegClass) -> &[PhysReg] {
123 self.classes.get(usize::from(class.number())).map_or(&[], |class| &class.order)
124 }
125
126 #[must_use]
128 pub fn scratch(&self, class: RegClass) -> &[PhysReg] {
129 self.classes.get(usize::from(class.number())).map_or(&[], |class| &class.scratch)
130 }
131}
132
133#[derive(Debug, Clone)]
135pub struct Assignment {
136 places: Vec<Option<Place>>,
137 slots: Vec<RegClass>,
138}
139
140impl Assignment {
141 #[must_use]
148 pub fn empty(vregs: usize) -> Self {
149 Self { places: vec![None; vregs], slots: Vec::new() }
150 }
151
152 pub fn put(&mut self, reg: Reg, place: Place) {
159 self.places[index(reg)] = Some(place);
160 }
161
162 pub fn take_slot(&mut self, class: RegClass) -> u32 {
168 let slot = u32::try_from(self.slots.len()).expect("too many spilled values");
169 self.slots.push(class);
170 slot
171 }
172
173 #[must_use]
176 pub fn place(&self, reg: Reg) -> Option<Place> {
177 self.places.get(usize::try_from(reg.number()?).ok()?).copied().flatten()
178 }
179
180 #[must_use]
182 pub fn slots(&self) -> &[RegClass] {
183 &self.slots
184 }
185
186 #[must_use]
188 pub fn spilled(&self) -> usize {
189 self.slots.len()
190 }
191
192 fn spill(&mut self, reg: Reg, class: RegClass) {
194 let slot = self.take_slot(class);
195 self.put(reg, Place::Slot(slot));
196 }
197}
198
199#[derive(Debug, Clone, Copy)]
201struct Interval {
202 reg: Reg,
203 class: RegClass,
204 range: Range,
205}
206
207#[derive(Debug, Clone, Copy)]
209struct Held {
210 reg: Reg,
211 class: RegClass,
212 range: Range,
213 at: PhysReg,
214}
215
216#[derive(Debug, Clone, Copy)]
218struct Blocked {
219 class: RegClass,
220 at: PhysReg,
221 point: Point,
225 by: Option<Reg>,
230}
231
232#[derive(Debug, Clone, Copy)]
234struct Reuse {
235 source: Reg,
237 at: Point,
239}
240
241#[must_use]
248pub fn assign(func: &Func, order: &Order, live: &Live, env: &Env) -> Assignment {
249 let blocked = blocked(func, order);
250 let forced = forced(func);
251 let reuses = reuses(func, order);
252 let hints = hints(func);
253
254 let mut intervals = Vec::with_capacity(func.vregs());
255 for (number, reuse) in reuses.iter().enumerate() {
256 let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
257 let (Some(mut range), Some(class)) = (live.range(reg), func.class_of(reg)) else {
258 continue;
259 };
260 if let Some(reuse) = reuse {
261 range.start = range.start.min(reuse.at);
262 }
263 intervals.push(Interval { reg, class, range });
264 }
265 intervals.sort_by_key(|interval| (interval.range.start, interval.reg));
266
267 let mut assignment = Assignment::empty(func.vregs());
268 let mut active: Vec<Held> = Vec::new();
269 for interval in intervals {
270 active.retain(|held| held.range.end >= interval.range.start);
271 if forced.contains(&interval.reg) {
272 assignment.spill(interval.reg, interval.class);
273 continue;
274 }
275 assert!(
276 !env.order(interval.class).is_empty(),
277 "a value in a class the target hands out no registers from"
278 );
279 let two_address = reuses[index(interval.reg)]
280 .and_then(|reuse| coalesce(&assignment, &active, &blocked, interval, reuse));
281 let hinted = hints[index(interval.reg)].filter(|&at| {
285 env.order(interval.class).contains(&at)
286 && available(&active, &blocked, interval, at, None)
287 });
288 let chosen = two_address.or(hinted).or_else(|| {
289 env.order(interval.class)
290 .iter()
291 .copied()
292 .find(|&at| available(&active, &blocked, interval, at, None))
293 });
294 match chosen {
295 Some(at) => {
296 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
297 let reg = interval.reg;
298 active.push(Held { reg, class: interval.class, range: interval.range, at });
299 }
300 None => spill_one(&mut assignment, &mut active, &blocked, interval),
301 }
302 }
303 assignment
304}
305
306fn available(
311 active: &[Held],
312 blocked: &[Blocked],
313 interval: Interval,
314 at: PhysReg,
315 except: Option<Reg>,
316) -> bool {
317 let taken = active
318 .iter()
319 .any(|held| held.at == at && held.class == interval.class && Some(held.reg) != except);
320 let insisted = blocked.iter().any(|one| {
321 one.at == at
322 && one.class == interval.class
323 && one.by != Some(interval.reg)
324 && interval.range.covers(one.point)
325 });
326 !taken && !insisted
327}
328
329fn coalesce(
332 assignment: &Assignment,
333 active: &[Held],
334 blocked: &[Blocked],
335 interval: Interval,
336 reuse: Reuse,
337) -> Option<PhysReg> {
338 let Some(Place::Reg(at)) = assignment.place(reuse.source) else { return None };
339 let source = active.iter().find(|held| held.reg == reuse.source)?;
340 let dies = source.range.end == reuse.at;
343 (dies && available(active, blocked, interval, at, Some(reuse.source))).then_some(at)
344}
345
346fn spill_one(
349 assignment: &mut Assignment,
350 active: &mut Vec<Held>,
351 blocked: &[Blocked],
352 interval: Interval,
353) {
354 let victim = active
357 .iter()
358 .enumerate()
359 .filter(|(_, held)| held.class == interval.class)
360 .filter(|(_, held)| available(&[], blocked, interval, held.at, None))
361 .max_by_key(|(_, held)| held.range.end)
362 .map(|(at, held)| (at, held.at, held.range.end));
363 match victim {
364 Some((victim, at, end)) if end > interval.range.end => {
365 let held = active.remove(victim);
366 assignment.spill(held.reg, held.class);
367 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
368 let reg = interval.reg;
369 active.push(Held { reg, class: interval.class, range: interval.range, at });
370 }
371 _ => assignment.spill(interval.reg, interval.class),
372 }
373}
374
375fn blocked(func: &Func, order: &Order) -> Vec<Blocked> {
381 let mut blocked = Vec::new();
382 let mut claimed: Vec<(RegClass, PhysReg)> = Vec::new();
383 for block in func.blocks() {
384 for inst in func.insts(block) {
385 let operands = &func[func[inst].operands];
386 claimed.clear();
387 for operand in operands {
388 if let Some(at) = insisted(operand) {
389 let key = (operand.class, at);
390 if !claimed.contains(&key) {
391 claimed.push(key);
392 }
393 }
394 }
395 for &(class, at) in &claimed {
396 for (point, role) in [(order.early(inst), Role::Use), (order.late(inst), Role::Def)]
401 {
402 let mut named = false;
403 for operand in operands {
404 let mine = insisted(operand) == Some(at) && operand.class == class;
405 if !mine || !(operand.role == role || operand.role == Role::EarlyDef) {
406 continue;
407 }
408 named = true;
409 let by = operand.reg.is_virtual().then_some(operand.reg);
410 blocked.push(Blocked { class, at, point, by });
411 }
412 if !named {
413 blocked.push(Blocked { class, at, point, by: None });
414 }
415 }
416 }
417 }
418 }
419 blocked
420}
421
422fn insisted(operand: &Operand) -> Option<PhysReg> {
425 match operand.constraint {
426 Constraint::Fixed(at) => Some(at),
427 _ => operand.reg.phys(),
428 }
429}
430
431fn hints(func: &Func) -> Vec<Option<PhysReg>> {
438 let mut hints = vec![None; func.vregs()];
439 for block in func.blocks() {
440 for inst in func.insts(block) {
441 for operand in &func[func[inst].operands] {
442 let Constraint::Fixed(at) = operand.constraint else { continue };
443 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
444 let Some(number) = number else { continue };
445 if func.class_of(operand.reg) == Some(operand.class) && hints[number].is_none() {
446 hints[number] = Some(at);
447 }
448 }
449 }
450 }
451 hints
452}
453
454fn forced(func: &Func) -> Vec<Reg> {
456 let mut forced = Vec::new();
457 for block in func.blocks() {
458 for inst in func.insts(block) {
459 for operand in &func[func[inst].operands] {
460 if operand.constraint == Constraint::Stack
461 && operand.reg.is_virtual()
462 && !forced.contains(&operand.reg)
463 {
464 forced.push(operand.reg);
465 }
466 }
467 }
468 }
469 forced
470}
471
472fn reuses(func: &Func, order: &Order) -> Vec<Option<Reuse>> {
474 let mut reuses = vec![None; func.vregs()];
475 for block in func.blocks() {
476 for inst in func.insts(block) {
477 let operands = &func[func[inst].operands];
478 for operand in operands {
479 let Constraint::Reuse(other) = operand.constraint else { continue };
480 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
481 let Some(number) = number else { continue };
482 let source = operands[usize::from(other)].reg;
483 reuses[number] = Some(Reuse { source, at: order.early(inst) });
484 }
485 }
486 }
487 reuses
488}
489
490fn index(reg: Reg) -> usize {
492 usize::try_from(reg.number().expect("a virtual register")).expect("a register number")
493}
494
495#[cfg(test)]
496mod tests {
497 use rucc_base::Interner;
498 use rucc_mir::{BlockCall, Opcode, Operand};
499 use rucc_target::x86_64::{GPR, R13, R14, R15, RAX, RCX, RDX, REGS, SYSV};
500
501 use super::*;
502
503 fn env() -> Env {
505 let (order, scratch) = SYSV.int_order.split_at(SYSV.int_order.len() - 3);
506 Env::new().with(GPR, order, scratch)
507 }
508
509 fn narrow(count: usize) -> Env {
512 Env::new().with(GPR, &SYSV.int_order[..count], &SYSV.int_order[count..count + 1])
513 }
514
515 fn named(place: Option<Place>) -> String {
517 match place {
518 Some(Place::Reg(reg)) => REGS.name(GPR, reg).expect("a register").to_string(),
519 Some(Place::Slot(slot)) => format!("slot {slot}"),
520 None => "nowhere".to_string(),
521 }
522 }
523
524 fn places(func: &Func, env: &Env) -> Vec<String> {
526 let order = Order::of(func);
527 let live = Live::of(func, &order);
528 let assignment = assign(func, &order, &live, env);
529 (0..func.vregs())
530 .map(|number| {
531 let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
532 named(assignment.place(reg))
533 })
534 .collect()
535 }
536
537 #[test]
538 fn two_values_that_are_never_both_wanted_share_a_register() {
539 let mut names = Interner::new();
540 let mut func = Func::new(names.intern("f"));
541 let opcode = Opcode::new(names.intern("x64.nop"));
542 let block = func.create_block();
543 let first = func.new_vreg(GPR);
544 let second = func.new_vreg(GPR);
545 func.build(block, opcode).def(first, GPR).finish();
546 func.build(block, opcode).uses(first, GPR).finish();
547 func.build(block, opcode).def(second, GPR).finish();
548 func.build(block, opcode).uses(second, GPR).finish();
549
550 assert_eq!(places(&func, &env()), ["rax", "rax"]);
553 }
554
555 #[test]
556 fn two_values_that_are_both_wanted_do_not() {
557 let mut names = Interner::new();
558 let mut func = Func::new(names.intern("f"));
559 let opcode = Opcode::new(names.intern("x64.nop"));
560 let block = func.create_block();
561 let first = func.new_vreg(GPR);
562 let second = func.new_vreg(GPR);
563 func.build(block, opcode).def(first, GPR).finish();
564 func.build(block, opcode).def(second, GPR).finish();
565 func.build(block, opcode).uses(first, GPR).finish();
566 func.build(block, opcode).uses(second, GPR).finish();
567
568 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
569 }
570
571 #[test]
572 fn a_value_written_early_that_nothing_reads_still_holds_its_register() {
573 let mut names = Interner::new();
574 let mut func = Func::new(names.intern("f"));
575 let opcode = Opcode::new(names.intern("x64.nop"));
576 let block = func.create_block();
577 let wanted = func.new_vreg(GPR);
578 let spare = func.new_vreg(GPR);
579 func.build(block, opcode)
582 .def(wanted, GPR)
583 .operand(Operand::write_early(spare, GPR))
584 .finish();
585 func.build(block, opcode).uses(wanted, GPR).finish();
586
587 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
593 }
594
595 #[test]
596 fn the_value_wanted_longest_is_the_one_that_goes_to_the_stack() {
597 let mut names = Interner::new();
598 let mut func = Func::new(names.intern("f"));
599 let opcode = Opcode::new(names.intern("x64.nop"));
600 let block = func.create_block();
601 let long = func.new_vreg(GPR);
602 let short = func.new_vreg(GPR);
603 let third = func.new_vreg(GPR);
604 func.build(block, opcode).def(long, GPR).finish();
605 func.build(block, opcode).def(short, GPR).finish();
606 func.build(block, opcode).def(third, GPR).finish();
607 func.build(block, opcode).uses(short, GPR).finish();
608 func.build(block, opcode).uses(third, GPR).finish();
609 func.build(block, opcode).uses(long, GPR).finish();
610
611 assert_eq!(places(&func, &narrow(2)), ["slot 0", "rcx", "rax"]);
614 }
615
616 #[test]
617 fn a_register_an_instruction_insists_on_goes_to_the_values_that_asked_for_it() {
618 let mut names = Interner::new();
619 let mut func = Func::new(names.intern("f"));
620 let opcode = Opcode::new(names.intern("x64.nop"));
621 let block = func.create_block();
622 let across = func.new_vreg(GPR);
623 let dividend = func.new_vreg(GPR);
624 let quotient = func.new_vreg(GPR);
625 let remainder = func.new_vreg(GPR);
626 func.build(block, opcode).def(across, GPR).finish();
627 func.build(block, opcode).def(dividend, GPR).finish();
628 func.build(block, opcode)
629 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
630 .operand(Operand::write_early(remainder, GPR).with(Constraint::Fixed(RDX)))
631 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
632 .finish();
633 func.build(block, opcode).uses(across, GPR).finish();
634
635 assert_eq!(places(&func, &env()), ["rcx", "rax", "rax", "rdx"]);
640 }
641
642 #[test]
643 fn a_value_wanted_after_the_instruction_that_insists_does_not_get_that_register() {
644 let mut names = Interner::new();
645 let mut func = Func::new(names.intern("f"));
646 let opcode = Opcode::new(names.intern("x64.nop"));
647 let block = func.create_block();
648 let dividend = func.new_vreg(GPR);
649 let quotient = func.new_vreg(GPR);
650 func.build(block, opcode).def(dividend, GPR).finish();
651 func.build(block, opcode)
652 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
653 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
654 .finish();
655 func.build(block, opcode).uses(dividend, GPR).finish();
656
657 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
661 }
662
663 #[test]
664 fn a_value_an_instruction_can_only_read_from_memory_is_on_the_stack() {
665 let mut names = Interner::new();
666 let mut func = Func::new(names.intern("f"));
667 let opcode = Opcode::new(names.intern("x64.nop"));
668 let block = func.create_block();
669 let value = func.new_vreg(GPR);
670 func.build(block, opcode).def(value, GPR).finish();
671 func.build(block, opcode)
672 .operand(Operand::read(value, GPR).with(Constraint::Stack))
673 .finish();
674
675 assert_eq!(places(&func, &env()), ["slot 0"]);
676 }
677
678 #[test]
679 fn a_two_address_instruction_writes_the_register_it_read_when_it_can() {
680 let mut names = Interner::new();
681 let mut func = Func::new(names.intern("f"));
682 let opcode = Opcode::new(names.intern("x64.nop"));
683 let block = func.create_block();
684 let left = func.new_vreg(GPR);
685 let right = func.new_vreg(GPR);
686 let sum = func.new_vreg(GPR);
687 func.build(block, opcode).def(left, GPR).finish();
688 func.build(block, opcode).def(right, GPR).finish();
689 func.build(block, opcode)
690 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
691 .uses(left, GPR)
692 .uses(right, GPR)
693 .finish();
694 func.build(block, opcode).uses(right, GPR).finish();
695
696 assert_eq!(places(&func, &env()), ["rax", "rcx", "rax"]);
699 }
700
701 #[test]
702 fn a_two_address_instruction_that_cannot_gets_a_register_nothing_it_reads_is_in() {
703 let mut names = Interner::new();
704 let mut func = Func::new(names.intern("f"));
705 let opcode = Opcode::new(names.intern("x64.nop"));
706 let block = func.create_block();
707 let left = func.new_vreg(GPR);
708 let right = func.new_vreg(GPR);
709 let sum = func.new_vreg(GPR);
710 func.build(block, opcode).def(left, GPR).finish();
711 func.build(block, opcode).def(right, GPR).finish();
712 func.build(block, opcode)
713 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
714 .uses(left, GPR)
715 .uses(right, GPR)
716 .finish();
717 func.build(block, opcode).uses(left, GPR).finish();
718
719 assert_eq!(places(&func, &env()), ["rax", "rcx", "rdx"]);
723 }
724
725 #[test]
726 fn a_value_live_across_a_whole_loop_holds_its_register_over_all_of_it() {
727 let mut names = Interner::new();
728 let mut func = Func::new(names.intern("f"));
729 let opcode = Opcode::new(names.intern("x64.nop"));
730 let head = func.create_block();
731 let body = func.create_block();
732 let carried = func.new_vreg(GPR);
733 let inside = func.new_vreg(GPR);
734 func.build(head, opcode).def(carried, GPR).finish();
735 *func.succs_mut(head) = vec![BlockCall::to(body)];
736 func.build(body, opcode).def(inside, GPR).finish();
737 func.build(body, opcode).uses(inside, GPR).uses(carried, GPR).finish();
738 *func.succs_mut(body) = vec![BlockCall::to(body)];
739
740 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
743 }
744
745 #[test]
746 fn a_frame_says_what_each_of_its_slots_is_for() {
747 let mut names = Interner::new();
748 let mut func = Func::new(names.intern("f"));
749 let opcode = Opcode::new(names.intern("x64.nop"));
750 let block = func.create_block();
751 let first = func.new_vreg(GPR);
752 let second = func.new_vreg(GPR);
753 func.build(block, opcode).def(first, GPR).finish();
754 func.build(block, opcode).def(second, GPR).finish();
755 func.build(block, opcode).uses(first, GPR).uses(second, GPR).finish();
756
757 let order = Order::of(&func);
758 let live = Live::of(&func, &order);
759 let assignment = assign(&func, &order, &live, &narrow(1));
760 assert_eq!(assignment.spilled(), 1);
761 assert_eq!(assignment.slots(), [GPR]);
762 assert_eq!(assignment.place(Reg::physical(RCX)), None);
765 assert_eq!(env().scratch(GPR), [R13, R14, R15]);
766 }
767}