From f9d72fe4905039564ebb21736e3a494a272c8daf Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Tue, 16 Jul 2024 20:22:28 +0200
Subject: [PATCH] Fixed variable lifetime in loops

---
 src/build/asm/Instructions.go       |  8 ++++++--
 src/build/ast/Count.go              |  8 +++++++-
 src/build/ast/Parse.go              |  4 ++++
 src/build/core/CompileCall.go       |  4 ++++
 src/build/core/CompileDefinition.go | 23 ++++++++++++++++-------
 src/build/core/CompileLoop.go       |  5 +++--
 src/build/core/Scope.go             | 25 +++++++++++++++++++------
 src/build/core/Variable.go          |  1 +
 tests/programs/loop-lifetime.q      | 17 +++++++++++++++++
 tests/programs/loop.q               | 11 +++++++++++
 tests/programs_test.go              |  2 ++
 11 files changed, 90 insertions(+), 18 deletions(-)
 create mode 100644 tests/programs/loop-lifetime.q
 create mode 100644 tests/programs/loop.q

diff --git a/src/build/asm/Instructions.go b/src/build/asm/Instructions.go
index 4c51e87..ced8cf5 100644
--- a/src/build/asm/Instructions.go
+++ b/src/build/asm/Instructions.go
@@ -22,8 +22,12 @@ func (a *Assembler) Call(name string) {
 
 // Return returns back to the caller.
 func (a *Assembler) Return() {
-	if len(a.Instructions) > 0 && a.Instructions[len(a.Instructions)-1].Mnemonic == RETURN {
-		return
+	if len(a.Instructions) > 0 {
+		lastMnemonic := a.Instructions[len(a.Instructions)-1].Mnemonic
+
+		if lastMnemonic == RETURN || lastMnemonic == JUMP {
+			return
+		}
 	}
 
 	a.Instructions = append(a.Instructions, Instruction{Mnemonic: RETURN})
diff --git a/src/build/ast/Count.go b/src/build/ast/Count.go
index b14482d..d15a29e 100644
--- a/src/build/ast/Count.go
+++ b/src/build/ast/Count.go
@@ -9,6 +9,10 @@ func Count(body AST, kind token.Kind, name string) int {
 	for _, node := range body {
 		switch node := node.(type) {
 		case *Assign:
+			if node.Name.Kind == kind && node.Name.Text() == name {
+				count++
+			}
+
 			count += node.Value.Count(kind, name)
 
 		case *Call:
@@ -18,7 +22,9 @@ func Count(body AST, kind token.Kind, name string) int {
 			count += node.Value.Count(kind, name)
 
 		case *Return:
-			count += node.Value.Count(kind, name)
+			if node.Value != nil {
+				count += node.Value.Count(kind, name)
+			}
 
 		case *If:
 			count += node.Condition.Count(kind, name)
diff --git a/src/build/ast/Parse.go b/src/build/ast/Parse.go
index eaaad01..eb55aec 100644
--- a/src/build/ast/Parse.go
+++ b/src/build/ast/Parse.go
@@ -30,6 +30,10 @@ func toASTNode(tokens token.List) (Node, error) {
 		word := tokens[0].Text()
 
 		if word == keyword.Return {
+			if len(tokens) == 1 {
+				return &Return{}, nil
+			}
+
 			value := expression.Parse(tokens[1:])
 			return &Return{Value: value}, nil
 		}
diff --git a/src/build/core/CompileCall.go b/src/build/core/CompileCall.go
index f7b9dfb..9bb4f1f 100644
--- a/src/build/core/CompileCall.go
+++ b/src/build/core/CompileCall.go
@@ -51,6 +51,10 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 		f.Call(funcName)
 	}
 
+	for _, register := range registers {
+		f.Scope().Free(register)
+	}
+
 	// Pop
 	for i := len(f.cpu.General) - 1; i >= 0; i-- {
 		register := f.cpu.General[i]
diff --git a/src/build/core/CompileDefinition.go b/src/build/core/CompileDefinition.go
index 6a481a9..e591efa 100644
--- a/src/build/core/CompileDefinition.go
+++ b/src/build/core/CompileDefinition.go
@@ -44,26 +44,35 @@ func (f *Function) AddVariable(variable *Variable) {
 		f.Comment("%s = %s (%s, %d uses)", variable.Name, variable.Value, variable.Register, variable.Alive)
 	}
 
-	f.Scope().variables[variable.Name] = variable
-	f.Scope().Reserve(variable.Register)
-	f.Scope().Use(variable.Register)
+	scope := f.Scope()
+	variable.Scope = scope
+
+	scope.variables[variable.Name] = variable
+	scope.Reserve(variable.Register)
+	scope.Use(variable.Register)
 }
 
 func (f *Function) useVariable(variable *Variable) {
-	for _, scope := range f.scopes {
+	for i, scope := range f.scopes {
+		if scope.inLoop && variable.Scope != scope {
+			continue
+		}
+
 		local := scope.variables[variable.Name]
 
-		if local != nil {
-			local.Alive--
+		if local == nil {
+			continue
 		}
 
+		local.Alive--
+
 		if local.Alive < 0 {
 			panic("incorrect number of variable use calls")
 		}
 
 		if local.Alive == 0 {
 			if config.Comments {
-				f.Comment("%s died (%s)", local.Name, local.Register)
+				f.Comment("%s died (%s) in scope %d", local.Name, local.Register, i)
 			}
 
 			scope.Free(local.Register)
diff --git a/src/build/core/CompileLoop.go b/src/build/core/CompileLoop.go
index 6a798e5..c8e8258 100644
--- a/src/build/core/CompileLoop.go
+++ b/src/build/core/CompileLoop.go
@@ -12,9 +12,10 @@ func (f *Function) CompileLoop(loop *ast.Loop) error {
 	f.count.loop++
 	label := fmt.Sprintf("%s_loop_%d", f.Name, f.count.loop)
 	f.AddLabel(label)
-	f.pushScope(loop.Body)
+	scope := f.pushScope(loop.Body)
+	scope.inLoop = true
 	err := f.CompileAST(loop.Body)
-	f.popScope()
 	f.Jump(asm.JUMP, label)
+	f.popScope()
 	return err
 }
diff --git a/src/build/core/Scope.go b/src/build/core/Scope.go
index a88187b..4cfcaeb 100644
--- a/src/build/core/Scope.go
+++ b/src/build/core/Scope.go
@@ -2,6 +2,7 @@ package core
 
 import (
 	"git.akyoto.dev/cli/q/src/build/ast"
+	"git.akyoto.dev/cli/q/src/build/config"
 	"git.akyoto.dev/cli/q/src/build/cpu"
 	"git.akyoto.dev/cli/q/src/build/token"
 )
@@ -10,6 +11,7 @@ import (
 type Scope struct {
 	cpu.State
 	variables map[string]*Variable
+	inLoop    bool
 }
 
 // Scope returns the current scope.
@@ -18,13 +20,14 @@ func (s *state) Scope() *Scope {
 }
 
 // pushScope pushes a new scope to the top of the stack.
-func (s *state) pushScope(body ast.AST) {
+func (f *Function) pushScope(body ast.AST) *Scope {
 	scope := &Scope{}
 
-	if len(s.scopes) > 0 {
-		lastScope := s.scopes[len(s.scopes)-1]
+	if len(f.scopes) > 0 {
+		lastScope := f.scopes[len(f.scopes)-1]
 		scope.State = lastScope.State
 		scope.variables = make(map[string]*Variable, len(lastScope.variables))
+		scope.inLoop = lastScope.inLoop
 
 		for k, v := range lastScope.variables {
 			count := ast.Count(body, token.Identifier, v.Name)
@@ -44,10 +47,20 @@ func (s *state) pushScope(body ast.AST) {
 		scope.variables = map[string]*Variable{}
 	}
 
-	s.scopes = append(s.scopes, scope)
+	f.scopes = append(f.scopes, scope)
+
+	if config.Comments {
+		f.Comment("scope %d start", len(f.scopes))
+	}
+
+	return scope
 }
 
 // popScope removes the scope at the top of the stack.
-func (s *state) popScope() {
-	s.scopes = s.scopes[:len(s.scopes)-1]
+func (f *Function) popScope() {
+	if config.Comments {
+		f.Comment("scope %d end", len(f.scopes))
+	}
+
+	f.scopes = f.scopes[:len(f.scopes)-1]
 }
diff --git a/src/build/core/Variable.go b/src/build/core/Variable.go
index b7dd6ad..4f31806 100644
--- a/src/build/core/Variable.go
+++ b/src/build/core/Variable.go
@@ -11,6 +11,7 @@ type Variable struct {
 	Name     string
 	Register cpu.Register
 	Alive    int
+	Scope    *Scope
 }
 
 // Variable returns the variable with the given name or `nil` if it doesn't exist.
diff --git a/tests/programs/loop-lifetime.q b/tests/programs/loop-lifetime.q
new file mode 100644
index 0000000..1262b7d
--- /dev/null
+++ b/tests/programs/loop-lifetime.q
@@ -0,0 +1,17 @@
+main() {
+	n := 10
+	x := 1
+
+	loop {
+		if n == 0 {
+			return
+		}
+
+		f(x)
+		n -= 1
+	}
+}
+
+f(x) {
+	return x
+}
\ No newline at end of file
diff --git a/tests/programs/loop.q b/tests/programs/loop.q
new file mode 100644
index 0000000..c6456dc
--- /dev/null
+++ b/tests/programs/loop.q
@@ -0,0 +1,11 @@
+main() {
+	n := 10
+
+	loop {
+		if n == 0 {
+			return
+		}
+
+		n -= 1
+	}
+}
\ No newline at end of file
diff --git a/tests/programs_test.go b/tests/programs_test.go
index 93dbd50..aa3a6d5 100644
--- a/tests/programs_test.go
+++ b/tests/programs_test.go
@@ -31,6 +31,8 @@ var programs = []struct {
 	{"branch-or", "", 0},
 	{"branch-both", "", 0},
 	{"jump-near", "", 0},
+	{"loop", "", 0},
+	{"loop-lifetime", "", 0},
 }
 
 func TestPrograms(t *testing.T) {