Fixed variable lifetime in loops

src/build/asm/Instructions.go

@@ -22,9 +22,13 @@ 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 {
+	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})
 }

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:
+			if node.Value != nil {
 				count += node.Value.Count(kind, name)
+			}
 
 		case *If:
 			count += node.Condition.Count(kind, name)

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
 		}

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]

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
+	scope := f.Scope()
-	f.Scope().Reserve(variable.Register)
+	variable.Scope = scope
-	f.Scope().Use(variable.Register)
+
+	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 {
+		if local == nil {
-			local.Alive--
+			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)

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
 }

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 {
+	if len(f.scopes) > 0 {
-		lastScope := s.scopes[len(s.scopes)-1]
+		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() {
+func (f *Function) popScope() {
-	s.scopes = s.scopes[:len(s.scopes)-1]
+	if config.Comments {
+		f.Comment("scope %d end", len(f.scopes))
+	}
+
+	f.scopes = f.scopes[:len(f.scopes)-1]
 }

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.

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
+}

tests/programs/loop.q

@@ -0,0 +1,11 @@
+main() {
+	n := 10
+
+	loop {
+		if n == 0 {
+			return
+		}
+
+		n -= 1
+	}
+}

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) {