From 008f097186437b3ad3dda5083f5d3a9782c13a6d Mon Sep 17 00:00:00 2001
From: Eduard Urbach <eduard@urbach.dev>
Date: Mon, 31 Mar 2025 14:36:42 +0200
Subject: [PATCH] Simplified CPU state

---
 src/core/CompileAssign.go       | 13 +++++++++----
 src/cpu/State.go                | 24 +++++-------------------
 src/cpu/State_test.go           | 16 +++++-----------
 src/register/Machine.go         |  2 +-
 src/register/NewRegister.go     |  2 +-
 src/register/Return.go          |  4 +---
 src/register/postInstruction.go |  2 +-
 7 files changed, 23 insertions(+), 40 deletions(-)

diff --git a/src/core/CompileAssign.go b/src/core/CompileAssign.go
index 1a2f2e6..af628ef 100644
--- a/src/core/CompileAssign.go
+++ b/src/core/CompileAssign.go
@@ -37,12 +37,17 @@ func (f *Function) CompileAssign(node *ast.Assign) error {
 
 	switch leftValue := leftValue.(type) {
 	case *eval.Register:
-		// TODO: Reservation needs to be canceled on defer
-		f.CurrentScope().Reserve(leftValue.Register)
+		if !f.RegisterIsUsed(leftValue.Register) {
+			f.UseRegister(leftValue.Register)
+			defer f.FreeRegister(leftValue.Register)
+		}
+
 		f.Execute(operation, leftValue.Register, right)
 	case *eval.Memory:
-		// TODO: Reservation needs to be canceled on defer
-		f.CurrentScope().Reserve(leftValue.Memory.Base)
+		if !f.RegisterIsUsed(leftValue.Memory.Base) {
+			f.UseRegister(leftValue.Memory.Base)
+			defer f.FreeRegister(leftValue.Memory.Base)
+		}
 
 		if operation.Kind == token.Assign {
 			rightValue, err := f.Evaluate(right)
diff --git a/src/cpu/State.go b/src/cpu/State.go
index fb8659b..481b585 100644
--- a/src/cpu/State.go
+++ b/src/cpu/State.go
@@ -1,41 +1,27 @@
 package cpu
 
 // State contains information about which registers are currently in use.
-type State struct {
-	Reserved uint64
-	Used     uint64
-}
+type State uint64
 
 // Free will reset the reserved and used status which means the register can be allocated again.
 func (s *State) Free(reg Register) {
-	s.Reserved &= ^(1 << reg)
-	s.Used &= ^(1 << reg)
-}
-
-// IsReserved returns true if the register was marked for future use.
-func (s *State) IsReserved(reg Register) bool {
-	return s.Reserved&(1<<reg) != 0
+	*s &= ^(1 << reg)
 }
 
 // IsUsed returns true if the register is currently in use and holds a value.
 func (s *State) IsUsed(reg Register) bool {
-	return s.Used&(1<<reg) != 0
-}
-
-// Reserve reserves a register for future use.
-func (s *State) Reserve(reg Register) {
-	s.Reserved |= (1 << reg)
+	return *s&(1<<reg) != 0
 }
 
 // Use marks a register to be currently in use which means it must be preserved across function calls.
 func (s *State) Use(reg Register) {
-	s.Used |= (1 << reg)
+	*s |= (1 << reg)
 }
 
 // FindFree tries to find a free register in the given slice of registers.
 func (s *State) FindFree(registers []Register) (Register, bool) {
 	for _, reg := range registers {
-		if !s.IsReserved(reg) && !s.IsUsed(reg) {
+		if !s.IsUsed(reg) {
 			return reg, true
 		}
 	}
diff --git a/src/cpu/State_test.go b/src/cpu/State_test.go
index 6fdbf0b..97311be 100644
--- a/src/cpu/State_test.go
+++ b/src/cpu/State_test.go
@@ -8,23 +8,17 @@ import (
 )
 
 func TestRegisterState(t *testing.T) {
-	s := cpu.State{}
-	assert.False(t, s.IsReserved(0))
-	assert.False(t, s.IsUsed(0))
-	s.Reserve(0)
-	assert.True(t, s.IsReserved(0))
+	var s cpu.State
 	assert.False(t, s.IsUsed(0))
 	s.Use(0)
-	assert.True(t, s.IsReserved(0))
 	assert.True(t, s.IsUsed(0))
 	s.Free(0)
-	assert.False(t, s.IsReserved(0))
 	assert.False(t, s.IsUsed(0))
 }
 
 func TestFindFree(t *testing.T) {
-	s := cpu.State{}
-	s.Reserve(0)
+	var s cpu.State
+	s.Use(0)
 	s.Use(1)
 
 	reg, found := s.FindFree([]cpu.Register{0, 1, 2, 3})
@@ -36,8 +30,8 @@ func TestFindFree(t *testing.T) {
 }
 
 func TestMustFindFree(t *testing.T) {
-	s := cpu.State{}
-	s.Reserve(0)
+	var s cpu.State
+	s.Use(0)
 	s.Use(1)
 
 	reg := s.MustFindFree([]cpu.Register{0, 1, 2, 3})
diff --git a/src/register/Machine.go b/src/register/Machine.go
index 04d398a..8306339 100644
--- a/src/register/Machine.go
+++ b/src/register/Machine.go
@@ -11,5 +11,5 @@ type Machine struct {
 	scope.Stack
 	Assembler       asm.Assembler
 	CPU             *cpu.CPU
-	RegisterHistory []uint64
+	RegisterHistory []cpu.State
 }
diff --git a/src/register/NewRegister.go b/src/register/NewRegister.go
index b515b67..d94b9d1 100644
--- a/src/register/NewRegister.go
+++ b/src/register/NewRegister.go
@@ -5,6 +5,6 @@ import "git.urbach.dev/cli/q/src/cpu"
 // NewRegister reserves a register.
 func (f *Machine) NewRegister() cpu.Register {
 	register := f.CurrentScope().MustFindFree(f.CPU.General)
-	f.CurrentScope().Reserve(register)
+	f.CurrentScope().Use(register)
 	return register
 }
diff --git a/src/register/Return.go b/src/register/Return.go
index 032b74c..592d259 100644
--- a/src/register/Return.go
+++ b/src/register/Return.go
@@ -6,8 +6,6 @@ func (f *Machine) Return() {
 	}
 
 	f.Assembler.Return()
-	scope := f.CurrentScope()
-	scope.Reserved = 0
-	scope.Used = 0
+	f.CurrentScope().State = 0
 	f.postInstruction()
 }
diff --git a/src/register/postInstruction.go b/src/register/postInstruction.go
index 18520b5..d36dc6e 100644
--- a/src/register/postInstruction.go
+++ b/src/register/postInstruction.go
@@ -7,5 +7,5 @@ func (f *Machine) postInstruction() {
 		return
 	}
 
-	f.RegisterHistory = append(f.RegisterHistory, f.CurrentScope().Used)
+	f.RegisterHistory = append(f.RegisterHistory, f.CurrentScope().State)
 }