From 1204591cdca2fc43cf27b53239f3eef1c349abc2 Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Tue, 9 Jul 2024 12:43:29 +0200
Subject: [PATCH] Improved code generation

---
 src/build/core/CompileCall.go          | 6 ++++++
 src/build/core/ExpressionToRegister.go | 2 ++
 src/build/core/Function.go             | 7 -------
 src/build/core/SaveRegister.go         | 4 ++++
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/build/core/CompileCall.go b/src/build/core/CompileCall.go
index a70c4b6..ade18b5 100644
--- a/src/build/core/CompileCall.go
+++ b/src/build/core/CompileCall.go
@@ -35,6 +35,12 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 		return err
 	}
 
+	f.SaveRegister(f.cpu.Output[0])
+
+	for _, register := range registers {
+		f.SaveRegister(register)
+	}
+
 	// Push
 	for _, register := range f.cpu.General {
 		if f.cpu.IsUsed(register) {
diff --git a/src/build/core/ExpressionToRegister.go b/src/build/core/ExpressionToRegister.go
index b532757..580e96f 100644
--- a/src/build/core/ExpressionToRegister.go
+++ b/src/build/core/ExpressionToRegister.go
@@ -34,6 +34,8 @@ func (f *Function) ExpressionToRegister(node *expression.Expression, register cp
 
 	if f.UsesRegister(right, register) {
 		register = f.cpu.MustFindFree(f.cpu.General)
+	} else {
+		f.SaveRegister(register)
 	}
 
 	f.cpu.Reserve(register)
diff --git a/src/build/core/Function.go b/src/build/core/Function.go
index 4eff587..a9d90b8 100644
--- a/src/build/core/Function.go
+++ b/src/build/core/Function.go
@@ -47,13 +47,6 @@ func NewFunction(name string, file *fs.File, body token.List) *Function {
 func (f *Function) Compile() {
 	defer close(f.finished)
 	f.assembler.Label(asm.LABEL, f.Name)
-
-	for _, register := range f.cpu.Input {
-		if f.cpu.IsUsed(register) {
-			f.SaveRegister(register)
-		}
-	}
-
 	f.err = f.CompileTokens(f.Body)
 	f.assembler.Return()
 }
diff --git a/src/build/core/SaveRegister.go b/src/build/core/SaveRegister.go
index 95ec385..44c5129 100644
--- a/src/build/core/SaveRegister.go
+++ b/src/build/core/SaveRegister.go
@@ -10,6 +10,10 @@ import (
 
 // SaveRegister attempts to move a variable occupying this register to another register.
 func (f *Function) SaveRegister(register cpu.Register) {
+	if !f.cpu.IsUsed(register) {
+		return
+	}
+
 	var variable *Variable
 
 	for _, v := range f.variables {