From cb6b3a4cd0535702b2cf683a49e82730be3cb2c3 Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Mon, 3 Feb 2025 13:17:48 +0100
Subject: [PATCH] Fixed missing register move in if statements

---
 lib/time/time.q              | 22 +++++++++-------------
 src/core/CompileAssign.go    | 27 ++++++++++++++++++++++++++-
 src/core/CompileIf.go        |  4 ++++
 tests/programs/branch-save.q | 23 +++++++++++++++++++++++
 tests/programs_test.go       |  1 +
 5 files changed, 63 insertions(+), 14 deletions(-)
 create mode 100644 tests/programs/branch-save.q

diff --git a/lib/time/time.q b/lib/time/time.q
index 0d2d749..9890e67 100644
--- a/lib/time/time.q
+++ b/lib/time/time.q
@@ -4,18 +4,14 @@ import sys
 sleep(nanoseconds Int) {
 	seconds := 0
 
-	loop {
-		if nanoseconds >= 1000000000 {
-			nanoseconds = nanoseconds - 1000000000
-			seconds += 1
-		} else {
-			timespec := mem.alloc(16)
-			store(timespec, 8, seconds)
-			offset := timespec + 8
-			store(offset, 8, nanoseconds)
-			sys.nanosleep(timespec)
-			mem.free(timespec, 16)
-			return
-		}
+	if nanoseconds >= 1000000000 {
+		seconds, nanoseconds = nanoseconds / 1000000000
 	}
+
+	timespec := mem.alloc(16)
+	store(timespec, 8, seconds)
+	offset := timespec + 8
+	store(offset, 8, nanoseconds)
+	sys.nanosleep(timespec)
+	mem.free(timespec, 16)
 }
\ No newline at end of file
diff --git a/src/core/CompileAssign.go b/src/core/CompileAssign.go
index d048b80..8949cbc 100644
--- a/src/core/CompileAssign.go
+++ b/src/core/CompileAssign.go
@@ -1,8 +1,10 @@
 package core
 
 import (
+	"git.akyoto.dev/cli/q/src/asm"
 	"git.akyoto.dev/cli/q/src/ast"
 	"git.akyoto.dev/cli/q/src/errors"
+	"git.akyoto.dev/cli/q/src/expression"
 	"git.akyoto.dev/cli/q/src/token"
 )
 
@@ -32,5 +34,28 @@ func (f *Function) CompileAssign(node *ast.Assign) error {
 		return f.CompileAssignDivision(node)
 	}
 
-	return errors.New(errors.NotImplemented, f.File, left.Token.Position)
+	if !ast.IsFunctionCall(right) {
+		return errors.New(errors.NotImplemented, f.File, node.Expression.Token.Position)
+	}
+
+	count := 0
+	_, err := f.CompileCall(right)
+
+	if err != nil {
+		return err
+	}
+
+	return left.EachLeaf(func(leaf *expression.Expression) error {
+		name := leaf.Token.Text(f.File.Bytes)
+		variable := f.VariableByName(name)
+
+		if variable == nil {
+			return errors.New(&errors.UnknownIdentifier{Name: name}, f.File, left.Token.Position)
+		}
+
+		f.RegisterRegister(asm.MOVE, variable.Register, f.CPU.Output[count])
+		f.UseVariable(variable)
+		count++
+		return nil
+	})
 }
diff --git a/src/core/CompileIf.go b/src/core/CompileIf.go
index e7ce92c..029d7e1 100644
--- a/src/core/CompileIf.go
+++ b/src/core/CompileIf.go
@@ -9,6 +9,10 @@ import (
 
 // CompileIf compiles a branch instruction.
 func (f *Function) CompileIf(branch *ast.If) error {
+	for _, register := range f.CPU.Input {
+		f.SaveRegister(register)
+	}
+
 	f.count.branch++
 
 	var (
diff --git a/tests/programs/branch-save.q b/tests/programs/branch-save.q
new file mode 100644
index 0000000..75ad7bf
--- /dev/null
+++ b/tests/programs/branch-save.q
@@ -0,0 +1,23 @@
+main() {
+	a, b := f(5)
+	assert a == 0
+	assert b == 5
+
+	a, b = f(15)
+	assert a == 1
+	assert b == 5
+
+	a, b = f(25)
+	assert a == 2
+	assert b == 5
+}
+
+f(b Int) -> (Int, Int) {
+	a := 0
+
+	if b >= 10 {
+		a, b = b / 10
+	}
+
+	return a, b
+}
\ No newline at end of file
diff --git a/tests/programs_test.go b/tests/programs_test.go
index 23fc44a..8a6bbb3 100644
--- a/tests/programs_test.go
+++ b/tests/programs_test.go
@@ -54,6 +54,7 @@ var programs = []struct {
 	{"branch-and", "", "", 0},
 	{"branch-or", "", "", 0},
 	{"branch-both", "", "", 0},
+	{"branch-save", "", "", 0},
 	{"jump-near", "", "", 0},
 	{"switch", "", "", 0},
 	{"loop", "", "", 0},