1use rucc_mir::{Constraint, Func, Operand, Reg, Role};
72use rucc_target::{PhysReg, RegClass};
73
74use crate::live::{Live, Range};
75use crate::order::{Order, Point};
76
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum Place {
80 Reg(PhysReg),
82 Slot(u32),
85}
86
87#[derive(Debug, Clone, Default)]
97pub struct Env {
98 classes: Vec<Class>,
99}
100
101#[derive(Debug, Clone, Default)]
103struct Class {
104 order: Vec<PhysReg>,
105 scratch: Vec<PhysReg>,
106}
107
108impl Env {
109 #[must_use]
111 pub fn new() -> Self {
112 Self::default()
113 }
114
115 #[must_use]
117 pub fn with(mut self, class: RegClass, order: &[PhysReg], scratch: &[PhysReg]) -> Self {
118 let index = usize::from(class.number());
119 if self.classes.len() <= index {
120 self.classes.resize(index + 1, Class::default());
121 }
122 self.classes[index] = Class { order: order.to_vec(), scratch: scratch.to_vec() };
123 self
124 }
125
126 #[must_use]
128 pub fn order(&self, class: RegClass) -> &[PhysReg] {
129 self.classes.get(usize::from(class.number())).map_or(&[], |class| &class.order)
130 }
131
132 #[must_use]
134 pub fn scratch(&self, class: RegClass) -> &[PhysReg] {
135 self.classes.get(usize::from(class.number())).map_or(&[], |class| &class.scratch)
136 }
137}
138
139#[derive(Debug, Clone)]
141pub struct Assignment {
142 places: Vec<Option<Place>>,
143 slots: Vec<RegClass>,
144}
145
146impl Assignment {
147 #[must_use]
154 pub fn empty(vregs: usize) -> Self {
155 Self { places: vec![None; vregs], slots: Vec::new() }
156 }
157
158 pub fn put(&mut self, reg: Reg, place: Place) {
165 self.places[index(reg)] = Some(place);
166 }
167
168 pub fn take_slot(&mut self, class: RegClass) -> u32 {
174 let slot = u32::try_from(self.slots.len()).expect("too many spilled values");
175 self.slots.push(class);
176 slot
177 }
178
179 #[must_use]
182 pub fn place(&self, reg: Reg) -> Option<Place> {
183 self.places.get(usize::try_from(reg.number()?).ok()?).copied().flatten()
184 }
185
186 #[must_use]
188 pub fn slots(&self) -> &[RegClass] {
189 &self.slots
190 }
191
192 #[must_use]
194 pub fn spilled(&self) -> usize {
195 self.slots.len()
196 }
197
198 fn spill(&mut self, reg: Reg, class: RegClass) {
200 let slot = self.take_slot(class);
201 self.put(reg, Place::Slot(slot));
202 }
203}
204
205#[derive(Debug, Clone, Copy)]
207struct Interval {
208 reg: Reg,
209 class: RegClass,
210 range: Range,
211}
212
213#[derive(Debug, Clone, Copy)]
215struct Held {
216 reg: Reg,
217 class: RegClass,
218 range: Range,
219 at: PhysReg,
220}
221
222#[derive(Debug, Clone, Copy)]
224struct Blocked {
225 class: RegClass,
226 at: PhysReg,
227 point: Point,
231 by: Option<Reg>,
236}
237
238#[derive(Debug, Clone, Copy)]
240struct Reuse {
241 source: Reg,
243 at: Point,
245}
246
247#[must_use]
254pub fn assign(func: &Func, order: &Order, live: &Live, env: &Env) -> Assignment {
255 let blocked = blocked(func, order);
256 let forced = forced(func);
257 let reuses = reuses(func, order);
258 let hints = hints(func);
259
260 let mut intervals = Vec::with_capacity(func.vregs());
261 for (number, reuse) in reuses.iter().enumerate() {
262 let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
263 let (Some(mut range), Some(class)) = (live.range(reg), func.class_of(reg)) else {
264 continue;
265 };
266 if let Some(reuse) = reuse {
267 range.start = range.start.min(reuse.at);
268 }
269 intervals.push(Interval { reg, class, range });
270 }
271 intervals.sort_by_key(|interval| (interval.range.start, interval.reg));
272
273 let mut assignment = Assignment::empty(func.vregs());
274 let mut active: Vec<Held> = Vec::new();
275 for interval in intervals {
276 active.retain(|held| held.range.end >= interval.range.start);
277 if forced.contains(&interval.reg) {
278 assignment.spill(interval.reg, interval.class);
279 continue;
280 }
281 assert!(
286 !env.order(interval.class).is_empty(),
287 "a value in class {}, which the target hands out no registers from",
288 interval.class.number()
289 );
290 let two_address = reuses[index(interval.reg)]
291 .and_then(|reuse| coalesce(&assignment, &active, &blocked, interval, reuse));
292 let hinted = hints[index(interval.reg)].filter(|&at| {
296 env.order(interval.class).contains(&at)
297 && available(&active, &blocked, interval, at, None)
298 });
299 let chosen = two_address.or(hinted).or_else(|| {
300 env.order(interval.class)
301 .iter()
302 .copied()
303 .find(|&at| available(&active, &blocked, interval, at, None))
304 });
305 match chosen {
306 Some(at) => {
307 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
308 let reg = interval.reg;
309 active.push(Held { reg, class: interval.class, range: interval.range, at });
310 }
311 None => spill_one(&mut assignment, &mut active, &blocked, interval),
312 }
313 }
314 assignment
315}
316
317fn available(
322 active: &[Held],
323 blocked: &[Blocked],
324 interval: Interval,
325 at: PhysReg,
326 except: Option<Reg>,
327) -> bool {
328 let taken = active
329 .iter()
330 .any(|held| held.at == at && held.class == interval.class && Some(held.reg) != except);
331 let insisted = blocked.iter().any(|one| {
332 one.at == at
333 && one.class == interval.class
334 && one.by != Some(interval.reg)
335 && interval.range.covers(one.point)
336 });
337 !taken && !insisted
338}
339
340fn coalesce(
343 assignment: &Assignment,
344 active: &[Held],
345 blocked: &[Blocked],
346 interval: Interval,
347 reuse: Reuse,
348) -> Option<PhysReg> {
349 let Some(Place::Reg(at)) = assignment.place(reuse.source) else { return None };
350 let source = active.iter().find(|held| held.reg == reuse.source)?;
351 let dies = source.range.end == reuse.at;
354 let begins = interval.range.start == reuse.at;
361 (dies && begins && available(active, blocked, interval, at, Some(reuse.source))).then_some(at)
362}
363
364fn spill_one(
367 assignment: &mut Assignment,
368 active: &mut Vec<Held>,
369 blocked: &[Blocked],
370 interval: Interval,
371) {
372 let victim = active
375 .iter()
376 .enumerate()
377 .filter(|(_, held)| held.class == interval.class)
378 .filter(|(_, held)| available(&[], blocked, interval, held.at, None))
379 .max_by_key(|(_, held)| held.range.end)
380 .map(|(at, held)| (at, held.at, held.range.end));
381 match victim {
382 Some((victim, at, end)) if end > interval.range.end => {
383 let held = active.remove(victim);
384 assignment.spill(held.reg, held.class);
385 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
386 let reg = interval.reg;
387 active.push(Held { reg, class: interval.class, range: interval.range, at });
388 }
389 _ => assignment.spill(interval.reg, interval.class),
390 }
391}
392
393fn blocked(func: &Func, order: &Order) -> Vec<Blocked> {
399 let mut blocked = Vec::new();
400 let mut claimed: Vec<(RegClass, PhysReg)> = Vec::new();
401 for block in func.blocks() {
402 for inst in func.insts(block) {
403 let operands = &func[func[inst].operands];
404 claimed.clear();
405 for operand in operands {
406 if let Some(at) = insisted(operand) {
407 let key = (operand.class, at);
408 if !claimed.contains(&key) {
409 claimed.push(key);
410 }
411 }
412 }
413 for &(class, at) in &claimed {
414 for (point, role) in [(order.early(inst), Role::Use), (order.late(inst), Role::Def)]
419 {
420 let mut named = false;
421 for operand in operands {
422 let mine = insisted(operand) == Some(at) && operand.class == class;
423 if !mine || !(operand.role == role || operand.role == Role::EarlyDef) {
424 continue;
425 }
426 named = true;
427 let by = operand.reg.is_virtual().then_some(operand.reg);
428 blocked.push(Blocked { class, at, point, by });
429 }
430 if !named {
431 blocked.push(Blocked { class, at, point, by: None });
432 }
433 }
434 }
435 }
436 }
437 blocked
438}
439
440fn insisted(operand: &Operand) -> Option<PhysReg> {
443 match operand.constraint {
444 Constraint::Fixed(at) => Some(at),
445 _ => operand.reg.phys(),
446 }
447}
448
449fn hints(func: &Func) -> Vec<Option<PhysReg>> {
456 let mut hints = vec![None; func.vregs()];
457 for block in func.blocks() {
458 for inst in func.insts(block) {
459 for operand in &func[func[inst].operands] {
460 let Constraint::Fixed(at) = operand.constraint else { continue };
461 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
462 let Some(number) = number else { continue };
463 if func.class_of(operand.reg) == Some(operand.class) && hints[number].is_none() {
464 hints[number] = Some(at);
465 }
466 }
467 }
468 }
469 hints
470}
471
472fn forced(func: &Func) -> Vec<Reg> {
474 let mut forced = Vec::new();
475 for block in func.blocks() {
476 for inst in func.insts(block) {
477 for operand in &func[func[inst].operands] {
478 if operand.constraint == Constraint::Stack
479 && operand.reg.is_virtual()
480 && !forced.contains(&operand.reg)
481 {
482 forced.push(operand.reg);
483 }
484 }
485 }
486 }
487 forced
488}
489
490fn reuses(func: &Func, order: &Order) -> Vec<Option<Reuse>> {
492 let mut reuses = vec![None; func.vregs()];
493 for block in func.blocks() {
494 for inst in func.insts(block) {
495 let operands = &func[func[inst].operands];
496 for operand in operands {
497 let Constraint::Reuse(other) = operand.constraint else { continue };
498 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
499 let Some(number) = number else { continue };
500 let source = operands[usize::from(other)].reg;
501 reuses[number] = Some(Reuse { source, at: order.early(inst) });
502 }
503 }
504 }
505 reuses
506}
507
508fn index(reg: Reg) -> usize {
510 usize::try_from(reg.number().expect("a virtual register")).expect("a register number")
511}
512
513#[cfg(test)]
514mod tests {
515 use rucc_base::Interner;
516 use rucc_mir::{BlockCall, Opcode, Operand};
517 use rucc_target::x86_64::{GPR, R13, R14, R15, RAX, RCX, RDX, REGS, SYSV};
518
519 use super::*;
520
521 fn env() -> Env {
523 let (order, scratch) = SYSV.int_order.split_at(SYSV.int_order.len() - 3);
524 Env::new().with(GPR, order, scratch)
525 }
526
527 fn narrow(count: usize) -> Env {
530 Env::new().with(GPR, &SYSV.int_order[..count], &SYSV.int_order[count..count + 1])
531 }
532
533 fn named(place: Option<Place>) -> String {
535 match place {
536 Some(Place::Reg(reg)) => REGS.name(GPR, reg).expect("a register").to_string(),
537 Some(Place::Slot(slot)) => format!("slot {slot}"),
538 None => "nowhere".to_string(),
539 }
540 }
541
542 fn places(func: &Func, env: &Env) -> Vec<String> {
544 let order = Order::of(func);
545 let live = Live::of(func, &order);
546 let assignment = assign(func, &order, &live, env);
547 (0..func.vregs())
548 .map(|number| {
549 let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
550 named(assignment.place(reg))
551 })
552 .collect()
553 }
554
555 #[test]
556 fn two_values_that_are_never_both_wanted_share_a_register() {
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).uses(first, GPR).finish();
565 func.build(block, opcode).def(second, GPR).finish();
566 func.build(block, opcode).uses(second, GPR).finish();
567
568 assert_eq!(places(&func, &env()), ["rax", "rax"]);
571 }
572
573 #[test]
574 fn two_values_that_are_both_wanted_do_not() {
575 let mut names = Interner::new();
576 let mut func = Func::new(names.intern("f"));
577 let opcode = Opcode::new(names.intern("x64.nop"));
578 let block = func.create_block();
579 let first = func.new_vreg(GPR);
580 let second = func.new_vreg(GPR);
581 func.build(block, opcode).def(first, GPR).finish();
582 func.build(block, opcode).def(second, GPR).finish();
583 func.build(block, opcode).uses(first, GPR).finish();
584 func.build(block, opcode).uses(second, GPR).finish();
585
586 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
587 }
588
589 #[test]
590 fn a_value_written_early_that_nothing_reads_still_holds_its_register() {
591 let mut names = Interner::new();
592 let mut func = Func::new(names.intern("f"));
593 let opcode = Opcode::new(names.intern("x64.nop"));
594 let block = func.create_block();
595 let wanted = func.new_vreg(GPR);
596 let spare = func.new_vreg(GPR);
597 func.build(block, opcode)
600 .def(wanted, GPR)
601 .operand(Operand::write_early(spare, GPR))
602 .finish();
603 func.build(block, opcode).uses(wanted, GPR).finish();
604
605 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
611 }
612
613 #[test]
614 fn the_value_wanted_longest_is_the_one_that_goes_to_the_stack() {
615 let mut names = Interner::new();
616 let mut func = Func::new(names.intern("f"));
617 let opcode = Opcode::new(names.intern("x64.nop"));
618 let block = func.create_block();
619 let long = func.new_vreg(GPR);
620 let short = func.new_vreg(GPR);
621 let third = func.new_vreg(GPR);
622 func.build(block, opcode).def(long, GPR).finish();
623 func.build(block, opcode).def(short, GPR).finish();
624 func.build(block, opcode).def(third, GPR).finish();
625 func.build(block, opcode).uses(short, GPR).finish();
626 func.build(block, opcode).uses(third, GPR).finish();
627 func.build(block, opcode).uses(long, GPR).finish();
628
629 assert_eq!(places(&func, &narrow(2)), ["slot 0", "rcx", "rax"]);
632 }
633
634 #[test]
635 fn a_register_an_instruction_insists_on_goes_to_the_values_that_asked_for_it() {
636 let mut names = Interner::new();
637 let mut func = Func::new(names.intern("f"));
638 let opcode = Opcode::new(names.intern("x64.nop"));
639 let block = func.create_block();
640 let across = func.new_vreg(GPR);
641 let dividend = func.new_vreg(GPR);
642 let quotient = func.new_vreg(GPR);
643 let remainder = func.new_vreg(GPR);
644 func.build(block, opcode).def(across, GPR).finish();
645 func.build(block, opcode).def(dividend, GPR).finish();
646 func.build(block, opcode)
647 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
648 .operand(Operand::write_early(remainder, GPR).with(Constraint::Fixed(RDX)))
649 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
650 .finish();
651 func.build(block, opcode).uses(across, GPR).finish();
652
653 assert_eq!(places(&func, &env()), ["rcx", "rax", "rax", "rdx"]);
658 }
659
660 #[test]
661 fn a_value_wanted_after_the_instruction_that_insists_does_not_get_that_register() {
662 let mut names = Interner::new();
663 let mut func = Func::new(names.intern("f"));
664 let opcode = Opcode::new(names.intern("x64.nop"));
665 let block = func.create_block();
666 let dividend = func.new_vreg(GPR);
667 let quotient = func.new_vreg(GPR);
668 func.build(block, opcode).def(dividend, GPR).finish();
669 func.build(block, opcode)
670 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
671 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
672 .finish();
673 func.build(block, opcode).uses(dividend, GPR).finish();
674
675 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
679 }
680
681 #[test]
682 fn a_value_an_instruction_can_only_read_from_memory_is_on_the_stack() {
683 let mut names = Interner::new();
684 let mut func = Func::new(names.intern("f"));
685 let opcode = Opcode::new(names.intern("x64.nop"));
686 let block = func.create_block();
687 let value = func.new_vreg(GPR);
688 func.build(block, opcode).def(value, GPR).finish();
689 func.build(block, opcode)
690 .operand(Operand::read(value, GPR).with(Constraint::Stack))
691 .finish();
692
693 assert_eq!(places(&func, &env()), ["slot 0"]);
694 }
695
696 #[test]
697 fn a_two_address_instruction_writes_the_register_it_read_when_it_can() {
698 let mut names = Interner::new();
699 let mut func = Func::new(names.intern("f"));
700 let opcode = Opcode::new(names.intern("x64.nop"));
701 let block = func.create_block();
702 let left = func.new_vreg(GPR);
703 let right = func.new_vreg(GPR);
704 let sum = func.new_vreg(GPR);
705 func.build(block, opcode).def(left, GPR).finish();
706 func.build(block, opcode).def(right, GPR).finish();
707 func.build(block, opcode)
708 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
709 .uses(left, GPR)
710 .uses(right, GPR)
711 .finish();
712 func.build(block, opcode).uses(right, GPR).finish();
713
714 assert_eq!(places(&func, &env()), ["rax", "rcx", "rax"]);
717 }
718
719 #[test]
720 fn a_two_address_instruction_that_cannot_gets_a_register_nothing_it_reads_is_in() {
721 let mut names = Interner::new();
722 let mut func = Func::new(names.intern("f"));
723 let opcode = Opcode::new(names.intern("x64.nop"));
724 let block = func.create_block();
725 let left = func.new_vreg(GPR);
726 let right = func.new_vreg(GPR);
727 let sum = func.new_vreg(GPR);
728 func.build(block, opcode).def(left, GPR).finish();
729 func.build(block, opcode).def(right, GPR).finish();
730 func.build(block, opcode)
731 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
732 .uses(left, GPR)
733 .uses(right, GPR)
734 .finish();
735 func.build(block, opcode).uses(left, GPR).finish();
736
737 assert_eq!(places(&func, &env()), ["rax", "rcx", "rdx"]);
741 }
742
743 #[test]
744 fn a_value_live_across_a_whole_loop_holds_its_register_over_all_of_it() {
745 let mut names = Interner::new();
746 let mut func = Func::new(names.intern("f"));
747 let opcode = Opcode::new(names.intern("x64.nop"));
748 let head = func.create_block();
749 let body = func.create_block();
750 let carried = func.new_vreg(GPR);
751 let inside = func.new_vreg(GPR);
752 func.build(head, opcode).def(carried, GPR).finish();
753 *func.succs_mut(head) = vec![BlockCall::to(body)];
754 func.build(body, opcode).def(inside, GPR).finish();
755 func.build(body, opcode).uses(inside, GPR).uses(carried, GPR).finish();
756 *func.succs_mut(body) = vec![BlockCall::to(body)];
757
758 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
761 }
762
763 #[test]
764 fn a_two_address_answer_already_live_does_not_take_the_register_it_read() {
765 let mut names = Interner::new();
766 let mut func = Func::new(names.intern("f"));
767 let opcode = Opcode::new(names.intern("x64.nop"));
768 let head = func.create_block();
769 let latch = func.create_block();
770 let out = func.create_block();
771 let source = func.new_vreg(GPR);
772 let carried = func.new_vreg(GPR);
773 func.build(head, opcode).def(source, GPR).finish();
774 func.build(head, opcode).def(carried, GPR).finish();
775 *func.succs_mut(head) = vec![BlockCall::to(latch)];
776 func.build(latch, opcode)
779 .operand(Operand::write(carried, GPR).with(Constraint::Reuse(1)))
780 .uses(source, GPR)
781 .uses(carried, GPR)
782 .finish();
783 *func.succs_mut(latch) = vec![BlockCall::to(head), BlockCall::to(out)];
784 func.build(out, opcode).uses(carried, GPR).finish();
785
786 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
791
792 let order = Order::of(&func);
795 let live = Live::of(&func, &order);
796 let assignment = assign(&func, &order, &live, &env());
797 assert!(crate::check::check(&func, &order, &live, &assignment).is_empty());
798 }
799
800 #[test]
801 fn a_frame_says_what_each_of_its_slots_is_for() {
802 let mut names = Interner::new();
803 let mut func = Func::new(names.intern("f"));
804 let opcode = Opcode::new(names.intern("x64.nop"));
805 let block = func.create_block();
806 let first = func.new_vreg(GPR);
807 let second = func.new_vreg(GPR);
808 func.build(block, opcode).def(first, GPR).finish();
809 func.build(block, opcode).def(second, GPR).finish();
810 func.build(block, opcode).uses(first, GPR).uses(second, GPR).finish();
811
812 let order = Order::of(&func);
813 let live = Live::of(&func, &order);
814 let assignment = assign(&func, &order, &live, &narrow(1));
815 assert_eq!(assignment.spilled(), 1);
816 assert_eq!(assignment.slots(), [GPR]);
817 assert_eq!(assignment.place(Reg::physical(RCX)), None);
820 assert_eq!(env().scratch(GPR), [R13, R14, R15]);
821 }
822}