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!(
280 !env.order(interval.class).is_empty(),
281 "a value in class {}, which the target hands out no registers from",
282 interval.class.number()
283 );
284 let two_address = reuses[index(interval.reg)]
285 .and_then(|reuse| coalesce(&assignment, &active, &blocked, interval, reuse));
286 let hinted = hints[index(interval.reg)].filter(|&at| {
290 env.order(interval.class).contains(&at)
291 && available(&active, &blocked, interval, at, None)
292 });
293 let chosen = two_address.or(hinted).or_else(|| {
294 env.order(interval.class)
295 .iter()
296 .copied()
297 .find(|&at| available(&active, &blocked, interval, at, None))
298 });
299 match chosen {
300 Some(at) => {
301 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
302 let reg = interval.reg;
303 active.push(Held { reg, class: interval.class, range: interval.range, at });
304 }
305 None => spill_one(&mut assignment, &mut active, &blocked, interval),
306 }
307 }
308 assignment
309}
310
311fn available(
316 active: &[Held],
317 blocked: &[Blocked],
318 interval: Interval,
319 at: PhysReg,
320 except: Option<Reg>,
321) -> bool {
322 let taken = active
323 .iter()
324 .any(|held| held.at == at && held.class == interval.class && Some(held.reg) != except);
325 let insisted = blocked.iter().any(|one| {
326 one.at == at
327 && one.class == interval.class
328 && one.by != Some(interval.reg)
329 && interval.range.covers(one.point)
330 });
331 !taken && !insisted
332}
333
334fn coalesce(
337 assignment: &Assignment,
338 active: &[Held],
339 blocked: &[Blocked],
340 interval: Interval,
341 reuse: Reuse,
342) -> Option<PhysReg> {
343 let Some(Place::Reg(at)) = assignment.place(reuse.source) else { return None };
344 let source = active.iter().find(|held| held.reg == reuse.source)?;
345 let dies = source.range.end == reuse.at;
348 (dies && available(active, blocked, interval, at, Some(reuse.source))).then_some(at)
349}
350
351fn spill_one(
354 assignment: &mut Assignment,
355 active: &mut Vec<Held>,
356 blocked: &[Blocked],
357 interval: Interval,
358) {
359 let victim = active
362 .iter()
363 .enumerate()
364 .filter(|(_, held)| held.class == interval.class)
365 .filter(|(_, held)| available(&[], blocked, interval, held.at, None))
366 .max_by_key(|(_, held)| held.range.end)
367 .map(|(at, held)| (at, held.at, held.range.end));
368 match victim {
369 Some((victim, at, end)) if end > interval.range.end => {
370 let held = active.remove(victim);
371 assignment.spill(held.reg, held.class);
372 assignment.places[index(interval.reg)] = Some(Place::Reg(at));
373 let reg = interval.reg;
374 active.push(Held { reg, class: interval.class, range: interval.range, at });
375 }
376 _ => assignment.spill(interval.reg, interval.class),
377 }
378}
379
380fn blocked(func: &Func, order: &Order) -> Vec<Blocked> {
386 let mut blocked = Vec::new();
387 let mut claimed: Vec<(RegClass, PhysReg)> = Vec::new();
388 for block in func.blocks() {
389 for inst in func.insts(block) {
390 let operands = &func[func[inst].operands];
391 claimed.clear();
392 for operand in operands {
393 if let Some(at) = insisted(operand) {
394 let key = (operand.class, at);
395 if !claimed.contains(&key) {
396 claimed.push(key);
397 }
398 }
399 }
400 for &(class, at) in &claimed {
401 for (point, role) in [(order.early(inst), Role::Use), (order.late(inst), Role::Def)]
406 {
407 let mut named = false;
408 for operand in operands {
409 let mine = insisted(operand) == Some(at) && operand.class == class;
410 if !mine || !(operand.role == role || operand.role == Role::EarlyDef) {
411 continue;
412 }
413 named = true;
414 let by = operand.reg.is_virtual().then_some(operand.reg);
415 blocked.push(Blocked { class, at, point, by });
416 }
417 if !named {
418 blocked.push(Blocked { class, at, point, by: None });
419 }
420 }
421 }
422 }
423 }
424 blocked
425}
426
427fn insisted(operand: &Operand) -> Option<PhysReg> {
430 match operand.constraint {
431 Constraint::Fixed(at) => Some(at),
432 _ => operand.reg.phys(),
433 }
434}
435
436fn hints(func: &Func) -> Vec<Option<PhysReg>> {
443 let mut hints = vec![None; func.vregs()];
444 for block in func.blocks() {
445 for inst in func.insts(block) {
446 for operand in &func[func[inst].operands] {
447 let Constraint::Fixed(at) = operand.constraint else { continue };
448 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
449 let Some(number) = number else { continue };
450 if func.class_of(operand.reg) == Some(operand.class) && hints[number].is_none() {
451 hints[number] = Some(at);
452 }
453 }
454 }
455 }
456 hints
457}
458
459fn forced(func: &Func) -> Vec<Reg> {
461 let mut forced = Vec::new();
462 for block in func.blocks() {
463 for inst in func.insts(block) {
464 for operand in &func[func[inst].operands] {
465 if operand.constraint == Constraint::Stack
466 && operand.reg.is_virtual()
467 && !forced.contains(&operand.reg)
468 {
469 forced.push(operand.reg);
470 }
471 }
472 }
473 }
474 forced
475}
476
477fn reuses(func: &Func, order: &Order) -> Vec<Option<Reuse>> {
479 let mut reuses = vec![None; func.vregs()];
480 for block in func.blocks() {
481 for inst in func.insts(block) {
482 let operands = &func[func[inst].operands];
483 for operand in operands {
484 let Constraint::Reuse(other) = operand.constraint else { continue };
485 let number = operand.reg.number().and_then(|number| usize::try_from(number).ok());
486 let Some(number) = number else { continue };
487 let source = operands[usize::from(other)].reg;
488 reuses[number] = Some(Reuse { source, at: order.early(inst) });
489 }
490 }
491 }
492 reuses
493}
494
495fn index(reg: Reg) -> usize {
497 usize::try_from(reg.number().expect("a virtual register")).expect("a register number")
498}
499
500#[cfg(test)]
501mod tests {
502 use rucc_base::Interner;
503 use rucc_mir::{BlockCall, Opcode, Operand};
504 use rucc_target::x86_64::{GPR, R13, R14, R15, RAX, RCX, RDX, REGS, SYSV};
505
506 use super::*;
507
508 fn env() -> Env {
510 let (order, scratch) = SYSV.int_order.split_at(SYSV.int_order.len() - 3);
511 Env::new().with(GPR, order, scratch)
512 }
513
514 fn narrow(count: usize) -> Env {
517 Env::new().with(GPR, &SYSV.int_order[..count], &SYSV.int_order[count..count + 1])
518 }
519
520 fn named(place: Option<Place>) -> String {
522 match place {
523 Some(Place::Reg(reg)) => REGS.name(GPR, reg).expect("a register").to_string(),
524 Some(Place::Slot(slot)) => format!("slot {slot}"),
525 None => "nowhere".to_string(),
526 }
527 }
528
529 fn places(func: &Func, env: &Env) -> Vec<String> {
531 let order = Order::of(func);
532 let live = Live::of(func, &order);
533 let assignment = assign(func, &order, &live, env);
534 (0..func.vregs())
535 .map(|number| {
536 let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
537 named(assignment.place(reg))
538 })
539 .collect()
540 }
541
542 #[test]
543 fn two_values_that_are_never_both_wanted_share_a_register() {
544 let mut names = Interner::new();
545 let mut func = Func::new(names.intern("f"));
546 let opcode = Opcode::new(names.intern("x64.nop"));
547 let block = func.create_block();
548 let first = func.new_vreg(GPR);
549 let second = func.new_vreg(GPR);
550 func.build(block, opcode).def(first, GPR).finish();
551 func.build(block, opcode).uses(first, GPR).finish();
552 func.build(block, opcode).def(second, GPR).finish();
553 func.build(block, opcode).uses(second, GPR).finish();
554
555 assert_eq!(places(&func, &env()), ["rax", "rax"]);
558 }
559
560 #[test]
561 fn two_values_that_are_both_wanted_do_not() {
562 let mut names = Interner::new();
563 let mut func = Func::new(names.intern("f"));
564 let opcode = Opcode::new(names.intern("x64.nop"));
565 let block = func.create_block();
566 let first = func.new_vreg(GPR);
567 let second = func.new_vreg(GPR);
568 func.build(block, opcode).def(first, GPR).finish();
569 func.build(block, opcode).def(second, GPR).finish();
570 func.build(block, opcode).uses(first, GPR).finish();
571 func.build(block, opcode).uses(second, GPR).finish();
572
573 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
574 }
575
576 #[test]
577 fn a_value_written_early_that_nothing_reads_still_holds_its_register() {
578 let mut names = Interner::new();
579 let mut func = Func::new(names.intern("f"));
580 let opcode = Opcode::new(names.intern("x64.nop"));
581 let block = func.create_block();
582 let wanted = func.new_vreg(GPR);
583 let spare = func.new_vreg(GPR);
584 func.build(block, opcode)
587 .def(wanted, GPR)
588 .operand(Operand::write_early(spare, GPR))
589 .finish();
590 func.build(block, opcode).uses(wanted, GPR).finish();
591
592 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
598 }
599
600 #[test]
601 fn the_value_wanted_longest_is_the_one_that_goes_to_the_stack() {
602 let mut names = Interner::new();
603 let mut func = Func::new(names.intern("f"));
604 let opcode = Opcode::new(names.intern("x64.nop"));
605 let block = func.create_block();
606 let long = func.new_vreg(GPR);
607 let short = func.new_vreg(GPR);
608 let third = func.new_vreg(GPR);
609 func.build(block, opcode).def(long, GPR).finish();
610 func.build(block, opcode).def(short, GPR).finish();
611 func.build(block, opcode).def(third, GPR).finish();
612 func.build(block, opcode).uses(short, GPR).finish();
613 func.build(block, opcode).uses(third, GPR).finish();
614 func.build(block, opcode).uses(long, GPR).finish();
615
616 assert_eq!(places(&func, &narrow(2)), ["slot 0", "rcx", "rax"]);
619 }
620
621 #[test]
622 fn a_register_an_instruction_insists_on_goes_to_the_values_that_asked_for_it() {
623 let mut names = Interner::new();
624 let mut func = Func::new(names.intern("f"));
625 let opcode = Opcode::new(names.intern("x64.nop"));
626 let block = func.create_block();
627 let across = func.new_vreg(GPR);
628 let dividend = func.new_vreg(GPR);
629 let quotient = func.new_vreg(GPR);
630 let remainder = func.new_vreg(GPR);
631 func.build(block, opcode).def(across, GPR).finish();
632 func.build(block, opcode).def(dividend, GPR).finish();
633 func.build(block, opcode)
634 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
635 .operand(Operand::write_early(remainder, GPR).with(Constraint::Fixed(RDX)))
636 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
637 .finish();
638 func.build(block, opcode).uses(across, GPR).finish();
639
640 assert_eq!(places(&func, &env()), ["rcx", "rax", "rax", "rdx"]);
645 }
646
647 #[test]
648 fn a_value_wanted_after_the_instruction_that_insists_does_not_get_that_register() {
649 let mut names = Interner::new();
650 let mut func = Func::new(names.intern("f"));
651 let opcode = Opcode::new(names.intern("x64.nop"));
652 let block = func.create_block();
653 let dividend = func.new_vreg(GPR);
654 let quotient = func.new_vreg(GPR);
655 func.build(block, opcode).def(dividend, GPR).finish();
656 func.build(block, opcode)
657 .operand(Operand::write(quotient, GPR).with(Constraint::Fixed(RAX)))
658 .operand(Operand::read(dividend, GPR).with(Constraint::Fixed(RAX)))
659 .finish();
660 func.build(block, opcode).uses(dividend, GPR).finish();
661
662 assert_eq!(places(&func, &env()), ["rcx", "rax"]);
666 }
667
668 #[test]
669 fn a_value_an_instruction_can_only_read_from_memory_is_on_the_stack() {
670 let mut names = Interner::new();
671 let mut func = Func::new(names.intern("f"));
672 let opcode = Opcode::new(names.intern("x64.nop"));
673 let block = func.create_block();
674 let value = func.new_vreg(GPR);
675 func.build(block, opcode).def(value, GPR).finish();
676 func.build(block, opcode)
677 .operand(Operand::read(value, GPR).with(Constraint::Stack))
678 .finish();
679
680 assert_eq!(places(&func, &env()), ["slot 0"]);
681 }
682
683 #[test]
684 fn a_two_address_instruction_writes_the_register_it_read_when_it_can() {
685 let mut names = Interner::new();
686 let mut func = Func::new(names.intern("f"));
687 let opcode = Opcode::new(names.intern("x64.nop"));
688 let block = func.create_block();
689 let left = func.new_vreg(GPR);
690 let right = func.new_vreg(GPR);
691 let sum = func.new_vreg(GPR);
692 func.build(block, opcode).def(left, GPR).finish();
693 func.build(block, opcode).def(right, GPR).finish();
694 func.build(block, opcode)
695 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
696 .uses(left, GPR)
697 .uses(right, GPR)
698 .finish();
699 func.build(block, opcode).uses(right, GPR).finish();
700
701 assert_eq!(places(&func, &env()), ["rax", "rcx", "rax"]);
704 }
705
706 #[test]
707 fn a_two_address_instruction_that_cannot_gets_a_register_nothing_it_reads_is_in() {
708 let mut names = Interner::new();
709 let mut func = Func::new(names.intern("f"));
710 let opcode = Opcode::new(names.intern("x64.nop"));
711 let block = func.create_block();
712 let left = func.new_vreg(GPR);
713 let right = func.new_vreg(GPR);
714 let sum = func.new_vreg(GPR);
715 func.build(block, opcode).def(left, GPR).finish();
716 func.build(block, opcode).def(right, GPR).finish();
717 func.build(block, opcode)
718 .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
719 .uses(left, GPR)
720 .uses(right, GPR)
721 .finish();
722 func.build(block, opcode).uses(left, GPR).finish();
723
724 assert_eq!(places(&func, &env()), ["rax", "rcx", "rdx"]);
728 }
729
730 #[test]
731 fn a_value_live_across_a_whole_loop_holds_its_register_over_all_of_it() {
732 let mut names = Interner::new();
733 let mut func = Func::new(names.intern("f"));
734 let opcode = Opcode::new(names.intern("x64.nop"));
735 let head = func.create_block();
736 let body = func.create_block();
737 let carried = func.new_vreg(GPR);
738 let inside = func.new_vreg(GPR);
739 func.build(head, opcode).def(carried, GPR).finish();
740 *func.succs_mut(head) = vec![BlockCall::to(body)];
741 func.build(body, opcode).def(inside, GPR).finish();
742 func.build(body, opcode).uses(inside, GPR).uses(carried, GPR).finish();
743 *func.succs_mut(body) = vec![BlockCall::to(body)];
744
745 assert_eq!(places(&func, &env()), ["rax", "rcx"]);
748 }
749
750 #[test]
751 fn a_frame_says_what_each_of_its_slots_is_for() {
752 let mut names = Interner::new();
753 let mut func = Func::new(names.intern("f"));
754 let opcode = Opcode::new(names.intern("x64.nop"));
755 let block = func.create_block();
756 let first = func.new_vreg(GPR);
757 let second = func.new_vreg(GPR);
758 func.build(block, opcode).def(first, GPR).finish();
759 func.build(block, opcode).def(second, GPR).finish();
760 func.build(block, opcode).uses(first, GPR).uses(second, GPR).finish();
761
762 let order = Order::of(&func);
763 let live = Live::of(&func, &order);
764 let assignment = assign(&func, &order, &live, &narrow(1));
765 assert_eq!(assignment.spilled(), 1);
766 assert_eq!(assignment.slots(), [GPR]);
767 assert_eq!(assignment.place(Reg::physical(RCX)), None);
770 assert_eq!(env().scratch(GPR), [R13, R14, R15]);
771 }
772}