From c53418bbdf328b5ba4f00bb9bd9f5936feaf5fdc Mon Sep 17 00:00:00 2001
From: Eduard Urbach <eduard@urbach.dev>
Date: Wed, 27 Aug 2025 17:10:51 +0200
Subject: [PATCH] Fixed incorrect register allocation for binary operations

---
 src/asm/Assembler_test.go       |  4 ++--
 src/asm/compilerX86.go          |  4 ++--
 src/codegen/findFreeRegister.go | 23 +++++++++++++++++++++--
 tests/register-restrictions.q   | 12 ++++++++++++
 tests/tests_test.go             |  1 +
 5 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 tests/register-restrictions.q

diff --git a/src/asm/Assembler_test.go b/src/asm/Assembler_test.go
index e3eea24..276e760 100644
--- a/src/asm/Assembler_test.go
+++ b/src/asm/Assembler_test.go
@@ -34,11 +34,11 @@ func TestAssembler(t *testing.T) {
 	b.Append(&asm.Add{Destination: 0, Source: 1, Operand: 2})
 	b.Append(&asm.Subtract{Destination: 0, Source: 1, Operand: 2})
 	b.Append(&asm.Multiply{Destination: 0, Source: 1, Operand: 2})
-	b.Append(&asm.Divide{Destination: 0, Source: 1, Operand: 2})
+	b.Append(&asm.Divide{Destination: 0, Source: 1, Operand: 3})
 	b.Append(&asm.Add{Destination: 1, Source: 1, Operand: 2})
 	b.Append(&asm.Subtract{Destination: 1, Source: 1, Operand: 2})
 	b.Append(&asm.Multiply{Destination: 1, Source: 1, Operand: 2})
-	b.Append(&asm.Divide{Destination: 1, Source: 1, Operand: 2})
+	b.Append(&asm.Divide{Destination: 1, Source: 1, Operand: 3})
 	b.Append(&asm.Return{})
 
 	c := &asm.Assembler{}
diff --git a/src/asm/compilerX86.go b/src/asm/compilerX86.go
index 3035940..90c1c84 100644
--- a/src/asm/compilerX86.go
+++ b/src/asm/compilerX86.go
@@ -87,8 +87,8 @@ func (c *compilerX86) Compile(instr Instruction) {
 			c.code = x86.CompareRegisterNumber(c.code, instr.Destination, instr.Number)
 		}
 	case *Divide:
-		if instr.Operand == x86.R0 {
-			panic("divide operand register cannot be R0")
+		if instr.Operand == x86.R2 {
+			panic("divisor register cannot be R2")
 		}
 
 		if instr.Source != x86.R0 {
diff --git a/src/codegen/findFreeRegister.go b/src/codegen/findFreeRegister.go
index 7d4a4bc..c9e1d2e 100644
--- a/src/codegen/findFreeRegister.go
+++ b/src/codegen/findFreeRegister.go
@@ -11,6 +11,18 @@ import (
 // findFreeRegister finds a free register for the given value.
 func (f *Function) findFreeRegister(step *Step) cpu.Register {
 	usedRegisters := 0
+	binaryOp, isBinaryOp := step.Value.(*ssa.BinaryOp)
+
+	if isBinaryOp {
+		switch binaryOp.Op {
+		case token.Add, token.Mul, token.Sub:
+			right := f.ValueToStep[binaryOp.Right]
+
+			if right.Register != -1 {
+				usedRegisters |= (1 << right.Register)
+			}
+		}
+	}
 
 	for _, current := range f.Steps {
 		// These checks need to happen regardless of whether the value is alive after execution.
@@ -19,9 +31,16 @@ func (f *Function) findFreeRegister(step *Step) cpu.Register {
 
 		if isBinaryOp {
 			switch binaryOp.Op {
+			case token.Add, token.Mul, token.Sub:
+				if current.Register != -1 && binaryOp.Right == step.Value {
+					usedRegisters |= (1 << current.Register)
+				}
+
 			case token.Div, token.Mod:
-				for _, reg := range f.CPU.DivisorRestricted {
-					usedRegisters |= (1 << reg)
+				if binaryOp.Right == step.Value {
+					for _, reg := range f.CPU.DivisorRestricted {
+						usedRegisters |= (1 << reg)
+					}
 				}
 
 			case token.Shl, token.Shr:
diff --git a/tests/register-restrictions.q b/tests/register-restrictions.q
new file mode 100644
index 0000000..ed5371d
--- /dev/null
+++ b/tests/register-restrictions.q
@@ -0,0 +1,12 @@
+main() {
+	a := 0
+	b := 1
+
+	if true {
+		a += 1
+		b += 1
+	}
+
+	assert a == 1
+	assert b == 2
+}
\ No newline at end of file
diff --git a/tests/tests_test.go b/tests/tests_test.go
index e6c3112..e44177e 100644
--- a/tests/tests_test.go
+++ b/tests/tests_test.go
@@ -45,6 +45,7 @@ var tests = []run{
 	{"phi", "", "", 0},
 	{"phi-simple", "", "", 0},
 	{"phi-advanced", "", "", 0},
+	{"register-restrictions", "", "", 0},
 	{"hello", "", "Hello\nHello\nHello\n", 0},
 	{"escape", "", "a\tb\nc\td\n", 0},
 	{"script", "", "Hello\n", 0},
-- 
2.51.2