From 22c6479a543f9f6d2257db47c607daac7f8c2a31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dominik=20Inf=C3=BChr?= <dominik.infuehr@gmail.com>
Date: Tue, 22 Oct 2024 07:08:27 +0200
Subject: [PATCH] regalloc: Simplify initializeRegisters a bit

---
 dora-parser/src/parser.rs | 23 +++++++++--------------
 pkgs/boots/regalloc.dora  | 38 ++++++++++++++++++--------------------
 2 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/dora-parser/src/parser.rs b/dora-parser/src/parser.rs
index 3f0d55425..5bec9325e 100644
--- a/dora-parser/src/parser.rs
+++ b/dora-parser/src/parser.rs
@@ -1171,7 +1171,7 @@ impl Parser {
         self.start_node();
         self.builder.start_node();
 
-        if self.is_pair(IDENTIFIER, EQ) {
+        if self.is2(IDENTIFIER, EQ) {
             let name = self.expect_identifier().expect("identifier expected");
             self.assert(EQ);
             let ty = self.parse_type();
@@ -1508,21 +1508,21 @@ impl Parser {
                 span: self.finish_node(),
                 expr,
             }))
-        } else if self.is(INT_LITERAL) || self.is_pair(SUB, INT_LITERAL) {
+        } else if self.is(INT_LITERAL) || self.is2(SUB, INT_LITERAL) {
             let expr = self.parse_lit_int_minus();
             Arc::new(PatternAlt::LitInt(PatternLit {
                 id: self.new_node_id(),
                 span: self.finish_node(),
                 expr,
             }))
-        } else if self.is(FLOAT_LITERAL) || self.is_pair(SUB, FLOAT_LITERAL) {
+        } else if self.is(FLOAT_LITERAL) || self.is2(SUB, FLOAT_LITERAL) {
             let expr = self.parse_lit_float_minus();
             Arc::new(PatternAlt::LitFloat(PatternLit {
                 id: self.new_node_id(),
                 span: self.finish_node(),
                 expr,
             }))
-        } else if self.is_pair(MUT_KW, IDENTIFIER) {
+        } else if self.is2(MUT_KW, IDENTIFIER) {
             self.assert(MUT_KW);
             let name = self.expect_identifier().expect("identifier expected");
             Arc::new(PatternAlt::Ident(PatternIdent {
@@ -2321,23 +2321,18 @@ impl Parser {
         self.current() == kind
     }
 
-    fn is_pair(&self, fst: TokenKind, snd: TokenKind) -> bool {
+    fn is2(&self, fst: TokenKind, snd: TokenKind) -> bool {
         if !self.is(fst) {
             return false;
         }
 
         let mut idx = 1;
 
-        loop {
-            let curr = self.nth(idx);
-
-            if curr.is_trivia() {
-                idx += 1;
-                continue;
-            }
-
-            return curr == snd;
+        while self.nth(idx).is_trivia() {
+            idx += 1;
         }
+
+        self.nth(idx) == snd
     }
 
     fn is_set(&self, set: TokenSet) -> bool {
diff --git a/pkgs/boots/regalloc.dora b/pkgs/boots/regalloc.dora
index fae2c2c31..e7f4592b1 100644
--- a/pkgs/boots/regalloc.dora
+++ b/pkgs/boots/regalloc.dora
@@ -86,44 +86,43 @@ impl SimpleRegisterAllocator {
         self.registers = Some(RegisterTracker[Register]::new(self.codegen.allocatableRegisters()));
         self.float_registers = Some(RegisterTracker[FloatRegister]::new(self.codegen.allocatableFloatRegisters()));
 
+        let predecessorCount = block.predecessors.size();
+
         if block.isLoopHeader() {
             // For now everything is spilled at the beginning of a loop.
             for inst in block.getLiveIn() {
                 self.spill(inst);
             }
+        } else if predecessorCount == 0 {
+            // Do nothing.
         } else {
             let candidates = HashMap[Inst, Int64]::new();
-            let all_gp = HashMap[Inst, Register]::new();
-            let all_fp = HashMap[Inst, FloatRegister]::new();
-            let predecessorCount = block.predecessors.size();
+            let mut idx = 0;
 
             for predecessorEdge in block.predecessors {
                 let predecessor = predecessorEdge.source;
                 let trackers = predecessor.getTrackerOut();
 
-                self.processLiveRegisters[Register](all_gp, trackers.gp, block, candidates);
-                self.processLiveRegisters[FloatRegister](all_fp, trackers.fp, block, candidates);
-            }
-
-            for (inst, reg) in all_gp {
-                assert(self.registers().allocateFixedRegister(reg, inst));
-            }
-
-            for (inst, reg) in all_fp {
-                assert(self.float_registers().allocateFixedRegister(reg, inst));
+                self.processLiveRegisters[Register](block, self.registers(), idx, trackers.gp, candidates);
+                self.processLiveRegisters[FloatRegister](block, self.float_registers(),idx, trackers.fp, candidates);
+                idx = idx + 1;
             }
 
             for inst in block.getLiveIn() {
-                if all_gp.contains(inst) || all_fp.contains(inst) {
-                    continue;
-                }
+                let inRegister = if inst.getValueType().isAnyFloat() {
+                    self.float_registers().locations.contains(inst)
+                } else {
+                    self.registers().locations.contains(inst)
+                };
 
-                self.spill(inst);
+                if !inRegister {
+                    self.spill(inst);
+                }
             }
         }
     }
 
-    fn processLiveRegisters[T: RegisterType](all: HashMap[Inst, T], tracker: RegisterTracker[T], block: Block, candidates: HashMap[Inst, Int64]) {
+    fn processLiveRegisters[T: RegisterType](block: Block, registers: RegisterTracker[T], idx: Int, tracker: RegisterTracker[T], candidates: HashMap[Inst, Int64]) {
         let predecessorCount = block.predecessors.size();
 
         for (reg, inst) in tracker.data {
@@ -138,9 +137,8 @@ impl SimpleRegisterAllocator {
                 1
             };
             if current == predecessorCount {
-                assert(all.insert(inst, reg).isNone());
+                assert(registers.allocateFixedRegister(reg, inst));
             }
-            candidates.insert(inst, current);
         }
     }