From 12894dbcc57288b7a01341ada3adbdf70edcc04d Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Wed, 5 Feb 2025 22:49:39 +0100
Subject: [PATCH] Fixed register allocation on function calls

---
 examples/point/point.q    |  8 ++++----
 src/core/CallSafe.go      | 33 +++++++++++++++++++++++++++++++++
 src/core/CompileCall.go   | 30 +-----------------------------
 src/core/CompileDelete.go | 18 +-----------------
 src/core/CompileNew.go    | 18 +-----------------
 5 files changed, 40 insertions(+), 67 deletions(-)
 create mode 100644 src/core/CallSafe.go

diff --git a/examples/point/point.q b/examples/point/point.q
index bc21184..ccb57c9 100644
--- a/examples/point/point.q
+++ b/examples/point/point.q
@@ -7,15 +7,15 @@ struct Point {
 }
 
 main() {
-	p := construct()
+	p := construct(1, 2)
 	print(p)
 	delete(p)
 }
 
-construct() -> *Point {
+construct(x Int, y Int) -> *Point {
 	p := new(Point)
-	p.x = 1
-	p.y = 2
+	p.x = x
+	p.y = y
 	return p
 }
 
diff --git a/src/core/CallSafe.go b/src/core/CallSafe.go
new file mode 100644
index 0000000..f6756b4
--- /dev/null
+++ b/src/core/CallSafe.go
@@ -0,0 +1,33 @@
+package core
+
+import (
+	"slices"
+
+	"git.akyoto.dev/cli/q/src/asm"
+	"git.akyoto.dev/cli/q/src/cpu"
+)
+
+// CallSafe pushes used registers to the stack, executes the call and restores the original register value.
+func (f *Function) CallSafe(fn *Function, registers []cpu.Register) {
+	for _, register := range f.CPU.Output {
+		f.SaveRegister(register)
+	}
+
+	for _, register := range f.CPU.General {
+		if f.RegisterIsUsed(register) {
+			f.Register(asm.PUSH, register)
+		}
+	}
+
+	f.Call(fn.UniqueName)
+
+	for _, register := range slices.Backward(f.CPU.General) {
+		if f.RegisterIsUsed(register) {
+			f.Register(asm.POP, register)
+		}
+	}
+
+	for _, register := range registers {
+		f.FreeRegister(register)
+	}
+}
diff --git a/src/core/CompileCall.go b/src/core/CompileCall.go
index 2b3f4fa..c88ec87 100644
--- a/src/core/CompileCall.go
+++ b/src/core/CompileCall.go
@@ -3,7 +3,6 @@ package core
 import (
 	"fmt"
 
-	"git.akyoto.dev/cli/q/src/asm"
 	"git.akyoto.dev/cli/q/src/errors"
 	"git.akyoto.dev/cli/q/src/expression"
 	"git.akyoto.dev/cli/q/src/types"
@@ -103,33 +102,6 @@ func (f *Function) CompileCall(root *expression.Expression) (*Function, error) {
 		}
 	}
 
-	for _, register := range f.CPU.Output[:len(fn.Output)] {
-		f.SaveRegister(register)
-	}
-
-	for _, register := range f.CPU.General {
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.PUSH, register)
-		}
-	}
-
-	f.Call(fn.UniqueName)
-
-	for _, register := range registers {
-		if register == f.CPU.Output[0] && root.Parent != nil {
-			continue
-		}
-
-		f.FreeRegister(register)
-	}
-
-	for i := len(f.CPU.General) - 1; i >= 0; i-- {
-		register := f.CPU.General[i]
-
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.POP, register)
-		}
-	}
-
+	f.CallSafe(fn, registers)
 	return fn, nil
 }
diff --git a/src/core/CompileDelete.go b/src/core/CompileDelete.go
index a8f46bd..3d144f8 100644
--- a/src/core/CompileDelete.go
+++ b/src/core/CompileDelete.go
@@ -16,22 +16,6 @@ func (f *Function) CompileDelete(root *expression.Expression) error {
 	f.SaveRegister(f.CPU.Input[1])
 	f.RegisterRegister(asm.MOVE, f.CPU.Input[0], variable.Register)
 	f.RegisterNumber(asm.MOVE, f.CPU.Input[1], int(variable.Type.(*types.Pointer).To.Size()))
-
-	for _, register := range f.CPU.General {
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.PUSH, register)
-		}
-	}
-
-	f.Call("mem.free")
-
-	for i := len(f.CPU.General) - 1; i >= 0; i-- {
-		register := f.CPU.General[i]
-
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.POP, register)
-		}
-	}
-
+	f.CallSafe(f.Functions["mem.free"], f.CPU.Input[:2])
 	return nil
 }
diff --git a/src/core/CompileNew.go b/src/core/CompileNew.go
index cdbf76d..d25818b 100644
--- a/src/core/CompileNew.go
+++ b/src/core/CompileNew.go
@@ -12,22 +12,6 @@ func (f *Function) CompileNew(root *expression.Expression) error {
 	typ := f.Types[structName]
 	f.SaveRegister(f.CPU.Input[0])
 	f.RegisterNumber(asm.MOVE, f.CPU.Input[0], typ.Size())
-
-	for _, register := range f.CPU.General {
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.PUSH, register)
-		}
-	}
-
-	f.Call("mem.alloc")
-
-	for i := len(f.CPU.General) - 1; i >= 0; i-- {
-		register := f.CPU.General[i]
-
-		if f.RegisterIsUsed(register) {
-			f.Register(asm.POP, register)
-		}
-	}
-
+	f.CallSafe(f.Functions["mem.alloc"], f.CPU.Input[:1])
 	return nil
 }