Improved code generation

examples/hello/hello.q

@@ -1,5 +1,9 @@
 main() {
-	syscall(60, f(1) + f(2) + f(3))
+	exit(f(1) + f(2) + f(3))
+}
+
+exit(code) {
+	syscall(60, code)
 }
 
 f(x) {

src/build/core/CompileCall.go

@@ -25,18 +25,21 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 		return err
 	}
 
+	// Push
 	for _, register := range f.cpu.General {
 		if !f.cpu.IsFree(register) {
 			f.assembler.Register(asm.PUSH, register)
 		}
 	}
 
+	// Call
 	if isSyscall {
 		f.assembler.Syscall()
 	} else {
 		f.assembler.Call(funcName)
 	}
 
+	// Pop
 	for i := len(f.cpu.General) - 1; i >= 0; i-- {
 		register := f.cpu.General[i]
 

src/build/core/Execute.go

@@ -23,7 +23,8 @@ func (f *Function) Execute(operation token.Token, register cpu.Register, value *
 		return f.ExecuteRegisterRegister(operation, register, f.cpu.Output[0])
 	}
 
-	tmp := f.cpu.MustUseFree(f.cpu.General)
+	tmp := f.cpu.MustFindFree(f.cpu.General)
+	f.cpu.Use(tmp)
 	defer f.cpu.Free(tmp)
 
 	err := f.ExpressionToRegister(value, tmp)

src/build/core/ExpressionToRegister.go

@@ -8,19 +8,28 @@ import (
 )
 
 // ExpressionToRegister puts the result of an expression into the specified register.
-func (f *Function) ExpressionToRegister(root *expression.Expression, register cpu.Register) error {
+func (f *Function) ExpressionToRegister(node *expression.Expression, register cpu.Register) error {
-	if root.IsLeaf() {
+	if node.IsLeaf() {
-		return f.TokenToRegister(root.Token, register)
+		return f.TokenToRegister(node.Token, register)
 	}
 
-	if ast.IsFunctionCall(root) {
+	if ast.IsFunctionCall(node) {
-		err := f.CompileCall(root)
+		err := f.CompileCall(node)
+
+		if register != f.cpu.Output[0] {
 			f.assembler.RegisterRegister(asm.MOVE, register, f.cpu.Output[0])
+		}
+
 		return err
 	}
 
-	left := root.Children[0]
+	left := node.Children[0]
-	right := root.Children[1]
+	right := node.Children[1]
+	final := register
+
+	if OverwritesRegister(right, register) {
+		register = f.cpu.MustFindFree(f.cpu.General)
+	}
 
 	err := f.ExpressionToRegister(left, register)
 
@@ -28,6 +37,13 @@ func (f *Function) ExpressionToRegister(root *expression.Expression, register cp
 		return err
 	}
 
-	f.SaveRegister(register)
+	f.cpu.Use(register)
-	return f.Execute(root.Token, register, right)
+	err = f.Execute(node.Token, register, right)
+
+	if register != final {
+		f.assembler.RegisterRegister(asm.MOVE, final, register)
+	}
+
+	f.cpu.Free(register)
+	return err
 }

src/build/core/OverwritesRegister.go

@@ -0,0 +1,11 @@
+package core
+
+import (
+	"git.akyoto.dev/cli/q/src/build/cpu"
+	"git.akyoto.dev/cli/q/src/build/expression"
+)
+
+// OverwritesRegister returns true if evaluating the expression would overwrite the given register.
+func OverwritesRegister(expr *expression.Expression, register cpu.Register) bool {
+	return !expr.IsLeaf()
+}

src/build/core/SaveRegister.go

@@ -27,7 +27,8 @@ func (f *Function) SaveRegister(register cpu.Register) {
 		return
 	}
 
-	newRegister := f.cpu.MustUseFree(f.cpu.General)
+	newRegister := f.cpu.MustFindFree(f.cpu.General)
+	f.cpu.Use(newRegister)
 
 	if config.Comments {
 		f.assembler.Comment(fmt.Sprintf("save %s to %s", register, newRegister))

src/build/cpu/CPU.go

@@ -31,13 +31,12 @@ func (c *CPU) FindFree(registers []Register) (Register, bool) {
 	return 0, false
 }
 
-func (c *CPU) MustUseFree(registers []Register) Register {
+func (c *CPU) MustFindFree(registers []Register) Register {
 	register, exists := c.FindFree(registers)
 
 	if !exists {
 		panic("no free registers")
 	}
 
-	c.Use(register)
 	return register
 }

tests/programs/exit.q

@@ -0,0 +1,11 @@
+main() {
+	exit(f(1) + f(2) + f(3))
+}
+
+exit(code) {
+	syscall(60, code)
+}
+
+f(x) {
+	return x + 1
+}

tests/programs_test.go

@@ -18,6 +18,15 @@ var programs = []struct {
 	{"empty.q", "", 0},
 	{"square-sum.q", "", 25},
 	{"multi-calls.q", "", 9},
+	{"exit.q", "", 9},
+}
+
+func TestPrograms(t *testing.T) {
+	for _, test := range programs {
+		t.Run(test.Name, func(t *testing.T) {
+			run(t, filepath.Join("programs", test.Name), test.ExpectedOutput, test.ExpectedExitCode)
+		})
+	}
 }
 
 func BenchmarkPrograms(b *testing.B) {
@@ -33,14 +42,6 @@ func BenchmarkPrograms(b *testing.B) {
 	}
 }
 
-func TestPrograms(t *testing.T) {
-	for _, test := range programs {
-		t.Run(test.Name, func(t *testing.T) {
-			run(t, filepath.Join("programs", test.Name), test.ExpectedOutput, test.ExpectedExitCode)
-		})
-	}
-}
-
 // run builds and runs the file to check if the output matches the expected output.
 func run(t *testing.T, name string, expectedOutput string, expectedExitCode int) {
 	b := build.New(name)