Moved register state to scopes

src/build/asm/Finalize.go

@@ -169,7 +169,7 @@ func (a Assembler) Finalize() ([]byte, []byte) {
 			code = x64.Syscall(code)
 
 		default:
-			panic("Unknown mnemonic: " + x.Mnemonic.String())
+			panic("unknown mnemonic: " + x.Mnemonic.String())
 		}
 	}
 

src/build/ast/Count.go

@@ -0,0 +1,36 @@
+package ast
+
+import "git.akyoto.dev/cli/q/src/build/token"
+
+// Count counts how often the given token appears in the AST.
+func Count(body AST, kind token.Kind, name string) int {
+	count := 0
+
+	for _, node := range body {
+		switch node := node.(type) {
+		case *Assign:
+			count += node.Value.Count(kind, name)
+
+		case *Call:
+			count += node.Expression.Count(kind, name)
+
+		case *Define:
+			count += node.Value.Count(kind, name)
+
+		case *Return:
+			count += node.Value.Count(kind, name)
+
+		case *If:
+			count += node.Condition.Count(kind, name)
+			count += Count(node.Body, kind, name)
+
+		case *Loop:
+			count += Count(node.Body, kind, name)
+
+		default:
+			panic("unknown AST type")
+		}
+	}
+
+	return count
+}

src/build/core/Compare.go

@@ -34,13 +34,13 @@ func (f *Function) Compare(comparison *expression.Expression) error {
 		return f.Execute(comparison.Token, f.cpu.Output[0], right)
 	}
 
-	tmp := f.cpu.MustFindFree(f.cpu.General)
+	tmp := f.Scope().MustFindFree(f.cpu.General)
 	err := f.ExpressionToRegister(left, tmp)
 
 	if err != nil {
 		return err
 	}
 
-	defer f.cpu.Free(tmp)
+	defer f.Scope().Free(tmp)
 	return f.Execute(comparison.Token, tmp, right)
 }

src/build/core/CompileCall.go

@@ -39,7 +39,7 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 
 	// Push
 	for _, register := range f.cpu.General {
-		if f.cpu.IsUsed(register) {
+		if f.Scope().IsUsed(register) {
 			f.Register(asm.PUSH, register)
 		}
 	}
@@ -55,7 +55,7 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 	for i := len(f.cpu.General) - 1; i >= 0; i-- {
 		register := f.cpu.General[i]
 
-		if f.cpu.IsUsed(register) {
+		if f.Scope().IsUsed(register) {
 			f.Register(asm.POP, register)
 		}
 	}

src/build/core/CompileDefinition.go

@@ -16,7 +16,7 @@ func (f *Function) CompileDefinition(node *ast.Define) error {
 		return errors.New(&errors.VariableAlreadyExists{Name: name}, f.File, node.Name.Position)
 	}
 
-	uses := CountIdentifier(f.Body, name) - 1
+	uses := token.Count(f.Body, token.Identifier, name) - 1
 
 	if uses == 0 {
 		return errors.New(&errors.UnusedVariable{Name: name}, f.File, node.Name.Position)
@@ -44,24 +44,30 @@ func (f *Function) AddVariable(variable *Variable) {
 		f.Comment("%s = %s (%s, %d uses)", variable.Name, variable.Value, variable.Register, variable.Alive)
 	}
 
-	f.Scope()[variable.Name] = variable
-	f.cpu.Reserve(variable.Register)
-	f.cpu.Use(variable.Register)
+	f.Scope().variables[variable.Name] = variable
+	f.Scope().Reserve(variable.Register)
+	f.Scope().Use(variable.Register)
 }
 
 func (f *Function) useVariable(variable *Variable) {
-	variable.Alive--
+	for _, scope := range f.scopes {
+		local := scope.variables[variable.Name]
 
-	if variable.Alive < 0 {
-		panic("incorrect number of variable use calls")
-	}
-
-	if variable.Alive == 0 {
-		if config.Comments {
-			f.Comment("%s died (%s)", variable.Name, variable.Register)
+		if local != nil {
+			local.Alive--
 		}
 
-		f.cpu.Free(variable.Register)
+		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)
+			}
+
+			scope.Free(local.Register)
+		}
 	}
 }
 
@@ -78,13 +84,13 @@ func (f *Function) identifierExists(name string) bool {
 }
 
 func (f *Function) storeVariableInRegister(name string, value *expression.Expression, uses int) error {
-	reg, exists := f.cpu.FindFree(f.cpu.General)
+	reg, exists := f.Scope().FindFree(f.cpu.General)
 
 	if !exists {
 		panic("no free registers")
 	}
 
-	f.cpu.Reserve(reg)
+	f.Scope().Reserve(reg)
 	err := f.ExpressionToRegister(value, reg)
 
 	f.AddVariable(&Variable{
@@ -95,15 +101,3 @@ func (f *Function) storeVariableInRegister(name string, value *expression.Expres
 
 	return err
 }
-
-func CountIdentifier(tokens token.List, name string) int {
-	count := 0
-
-	for _, t := range tokens {
-		if t.Kind == token.Identifier && t.Text() == name {
-			count++
-		}
-	}
-
-	return count
-}

src/build/core/CompileIf.go

@@ -18,7 +18,7 @@ func (f *Function) CompileIf(branch *ast.If) error {
 	}
 
 	f.AddLabel(success)
-	f.pushScope()
+	f.pushScope(branch.Body)
 	err = f.CompileAST(branch.Body)
 	f.popScope()
 	f.AddLabel(fail)

src/build/core/CompileLoop.go

@@ -12,7 +12,7 @@ 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()
+	f.pushScope(loop.Body)
 	err := f.CompileAST(loop.Body)
 	f.popScope()
 	f.Jump(asm.JUMP, label)

src/build/core/Execute.go

@@ -23,8 +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.MustFindFree(f.cpu.General)
-	defer f.cpu.Free(tmp)
+	tmp := f.Scope().MustFindFree(f.cpu.General)
+	defer f.Scope().Free(tmp)
 
 	err := f.ExpressionToRegister(value, tmp)
 

src/build/core/ExpressionToRegister.go

@@ -34,14 +34,14 @@ func (f *Function) ExpressionToRegister(node *expression.Expression, register cp
 	final := register
 
 	if f.UsesRegister(right, register) {
-		register = f.cpu.MustFindFree(f.cpu.General)
+		register = f.Scope().MustFindFree(f.cpu.General)
 
 		if config.Comments {
 			f.Comment("temporary register %s", register)
 		}
 	}
 
-	f.cpu.Reserve(register)
+	f.Scope().Reserve(register)
 	err := f.ExpressionToRegister(left, register)
 
 	if err != nil {
@@ -52,7 +52,7 @@ func (f *Function) ExpressionToRegister(node *expression.Expression, register cp
 
 	if register != final {
 		f.RegisterRegister(asm.MOVE, final, register)
-		f.cpu.Free(register)
+		f.Scope().Free(register)
 	}
 
 	return err

src/build/core/Function.go

@@ -36,7 +36,9 @@ func NewFunction(name string, file *fs.File, body token.List) *Function {
 				Syscall: x64.SyscallRegisters,
 				Output:  x64.ReturnValueRegisters,
 			},
-			scopes:   []Scope{{}},
+			scopes: []*Scope{
+				{variables: map[string]*Variable{}},
+			},
 			finished: make(chan struct{}),
 		},
 	}
@@ -52,14 +54,14 @@ func (f *Function) Compile() {
 
 // CompileTokens compiles a token list.
 func (f *Function) CompileTokens(tokens token.List) error {
-	tree, err := ast.Parse(tokens)
+	body, err := ast.Parse(tokens)
 
 	if err != nil {
 		err.(*errors.Error).File = f.File
 		return err
 	}
 
-	return f.CompileAST(tree)
+	return f.CompileAST(body)
 }
 
 // CompileAST compiles an abstract syntax tree.
@@ -97,7 +99,7 @@ func (f *Function) CompileASTNode(node ast.Node) error {
 		return f.CompileLoop(node)
 
 	default:
-		panic("Unknown AST type")
+		panic("unknown AST type")
 	}
 }
 

src/build/core/Instructions.go

@@ -15,7 +15,7 @@ func (f *Function) AddLabel(label string) {
 
 func (f *Function) Call(label string) {
 	f.assembler.Call(label)
-	f.cpu.Use(f.cpu.Output[0])
+	f.Scope().Use(f.cpu.Output[0])
 	f.postInstruction()
 }
 
@@ -33,35 +33,35 @@ func (f *Function) Register(mnemonic asm.Mnemonic, a cpu.Register) {
 	f.assembler.Register(mnemonic, a)
 
 	if mnemonic == asm.POP {
-		f.cpu.Use(a)
+		f.Scope().Use(a)
 	}
 
 	f.postInstruction()
 }
 
 func (f *Function) RegisterNumber(mnemonic asm.Mnemonic, a cpu.Register, b int) {
-	if f.cpu.IsUsed(a) && isDestructive(mnemonic) {
+	if f.Scope().IsUsed(a) && isDestructive(mnemonic) {
 		f.SaveRegister(a)
 	}
 
 	f.assembler.RegisterNumber(mnemonic, a, b)
 
 	if mnemonic == asm.MOVE {
-		f.cpu.Use(a)
+		f.Scope().Use(a)
 	}
 
 	f.postInstruction()
 }
 
 func (f *Function) RegisterLabel(mnemonic asm.Mnemonic, register cpu.Register, label string) {
-	if f.cpu.IsUsed(register) && isDestructive(mnemonic) {
+	if f.Scope().IsUsed(register) && isDestructive(mnemonic) {
 		f.SaveRegister(register)
 	}
 
 	f.assembler.RegisterLabel(mnemonic, register, label)
 
 	if mnemonic == asm.MOVE {
-		f.cpu.Use(register)
+		f.Scope().Use(register)
 	}
 
 	f.postInstruction()
@@ -72,14 +72,14 @@ func (f *Function) RegisterRegister(mnemonic asm.Mnemonic, a cpu.Register, b cpu
 		return
 	}
 
-	if f.cpu.IsUsed(a) && isDestructive(mnemonic) {
+	if f.Scope().IsUsed(a) && isDestructive(mnemonic) {
 		f.SaveRegister(a)
 	}
 
 	f.assembler.RegisterRegister(mnemonic, a, b)
 
 	if mnemonic == asm.MOVE {
-		f.cpu.Use(a)
+		f.Scope().Use(a)
 	}
 
 	f.postInstruction()
@@ -92,7 +92,7 @@ func (f *Function) Return() {
 
 func (f *Function) Syscall() {
 	f.assembler.Syscall()
-	f.cpu.Use(f.cpu.Output[0])
+	f.Scope().Use(f.cpu.Output[0])
 	f.postInstruction()
 }
 
@@ -101,7 +101,7 @@ func (f *Function) postInstruction() {
 		return
 	}
 
-	f.registerHistory = append(f.registerHistory, f.cpu.Used)
+	f.registerHistory = append(f.registerHistory, f.Scope().Used)
 }
 
 func isDestructive(mnemonic asm.Mnemonic) bool {

src/build/core/SaveRegister.go

@@ -8,7 +8,7 @@ import (
 
 // SaveRegister attempts to move a variable occupying this register to another register.
 func (f *Function) SaveRegister(register cpu.Register) {
-	if !f.cpu.IsUsed(register) {
+	if !f.Scope().IsUsed(register) {
 		return
 	}
 
@@ -24,8 +24,8 @@ func (f *Function) SaveRegister(register cpu.Register) {
 		return
 	}
 
-	newRegister := f.cpu.MustFindFree(f.cpu.General)
-	f.cpu.Reserve(newRegister)
+	newRegister := f.Scope().MustFindFree(f.cpu.General)
+	f.Scope().Reserve(newRegister)
 
 	if config.Comments {
 		f.Comment("save %s to %s", register, newRegister)

src/build/core/Scope.go

@@ -1,28 +1,44 @@
 package core
 
-// Scope represents a map of variables.
-type Scope map[string]*Variable
+import (
+	"git.akyoto.dev/cli/q/src/build/ast"
+	"git.akyoto.dev/cli/q/src/build/cpu"
+	"git.akyoto.dev/cli/q/src/build/token"
+)
+
+// Scope represents an independent code block.
+type Scope struct {
+	cpu.State
+	variables map[string]*Variable
+}
 
 // Scope returns the current scope.
-func (s *state) Scope() Scope {
+func (s *state) Scope() *Scope {
 	return s.scopes[len(s.scopes)-1]
 }
 
 // pushScope pushes a new scope to the top of the stack.
-func (s *state) pushScope() {
-	lastScope := s.scopes[len(s.scopes)-1]
-	newScope := make(Scope, len(lastScope))
+func (s *state) pushScope(body ast.AST) {
+	scope := &Scope{}
 
-	for k, v := range lastScope {
-		newScope[k] = &Variable{
-			Value:    v.Value,
-			Name:     v.Name,
-			Register: v.Register,
-			Alive:    v.Alive,
+	if len(s.scopes) > 0 {
+		lastScope := s.scopes[len(s.scopes)-1]
+		scope.State = lastScope.State
+		scope.variables = make(map[string]*Variable, len(lastScope.variables))
+
+		for k, v := range lastScope.variables {
+			scope.variables[k] = &Variable{
+				Value:    v.Value,
+				Name:     v.Name,
+				Register: v.Register,
+				Alive:    ast.Count(body, token.Identifier, v.Name),
+			}
 		}
+	} else {
+		scope.variables = map[string]*Variable{}
 	}
 
-	s.scopes = append(s.scopes, newScope)
+	s.scopes = append(s.scopes, scope)
 }
 
 // popScope removes the scope at the top of the stack.

src/build/core/Variable.go

@@ -15,12 +15,12 @@ type Variable struct {
 
 // Variable returns the variable with the given name or `nil` if it doesn't exist.
 func (s *state) Variable(name string) *Variable {
-	return s.Scope()[name]
+	return s.Scope().variables[name]
 }
 
 // VariableInRegister returns the variable that occupies the given register or `nil` if none occupy the register.
 func (s *state) VariableInRegister(register cpu.Register) *Variable {
-	for _, v := range s.Scope() {
+	for _, v := range s.Scope().variables {
 		if v.Register == register {
 			return v
 		}

src/build/core/state.go

@@ -12,7 +12,7 @@ import (
 // state is the data structure we embed in each function to preserve compilation state.
 type state struct {
 	err             error
-	scopes          []Scope
+	scopes          []*Scope
 	functions       map[string]*Function
 	registerHistory []uint64
 	finished        chan struct{}

src/build/cpu/CPU.go

@@ -2,59 +2,9 @@ package cpu
 
 // CPU represents the processor.
 type CPU struct {
-	All      []Register
-	General  []Register
-	Syscall  []Register
-	Input    []Register
-	Output   []Register
-	Reserved uint64
-	Used     uint64
-}
-
-// Free will reset the reserved and used status which means the register can be allocated again.
-func (c *CPU) Free(reg Register) {
-	c.Used &= ^(1 << reg)
-	c.Reserved &= ^(1 << reg)
-}
-
-// IsReserved returns true if the register was marked for future use.
-func (c *CPU) IsReserved(reg Register) bool {
-	return c.Reserved&(1<<reg) != 0
-}
-
-// IsUsed returns true if the register is currently in use and holds a value.
-func (c *CPU) IsUsed(reg Register) bool {
-	return c.Used&(1<<reg) != 0
-}
-
-// Reserve reserves a register for future use.
-func (c *CPU) Reserve(reg Register) {
-	c.Reserved |= (1 << reg)
-}
-
-// Use marks a register to be currently in use which means it must be preserved across function calls.
-func (c *CPU) Use(reg Register) {
-	c.Used |= (1 << reg)
-}
-
-// FindFree tries to find a free register in the given slice of registers.
-func (c *CPU) FindFree(registers []Register) (Register, bool) {
-	for _, reg := range registers {
-		if !c.IsReserved(reg) {
-			return reg, true
-		}
-	}
-
-	return 0, false
-}
-
-// MustFindFree tries to find a free register and panics if it could not be found.
-func (c *CPU) MustFindFree(registers []Register) Register {
-	register, exists := c.FindFree(registers)
-
-	if !exists {
-		panic("no free registers")
-	}
-
-	return register
+	All     []Register
+	General []Register
+	Syscall []Register
+	Input   []Register
+	Output  []Register
 }

src/build/cpu/State.go

@@ -0,0 +1,55 @@
+package cpu
+
+// State contains information about which registers are currently in use.
+type State struct {
+	Reserved uint64
+	Used     uint64
+}
+
+// Free will reset the reserved and used status which means the register can be allocated again.
+func (s *State) Free(reg Register) {
+	s.Reserved &= ^(1 << reg)
+	s.Used &= ^(1 << reg)
+}
+
+// IsReserved returns true if the register was marked for future use.
+func (s *State) IsReserved(reg Register) bool {
+	return s.Reserved&(1<<reg) != 0
+}
+
+// IsUsed returns true if the register is currently in use and holds a value.
+func (s *State) IsUsed(reg Register) bool {
+	return s.Used&(1<<reg) != 0
+}
+
+// Reserve reserves a register for future use.
+func (s *State) Reserve(reg Register) {
+	s.Reserved |= (1 << reg)
+}
+
+// Use marks a register to be currently in use which means it must be preserved across function calls.
+func (s *State) Use(reg Register) {
+	s.Used |= (1 << reg)
+}
+
+// FindFree tries to find a free register in the given slice of registers.
+func (s *State) FindFree(registers []Register) (Register, bool) {
+	for _, reg := range registers {
+		if !s.IsReserved(reg) {
+			return reg, true
+		}
+	}
+
+	return 0, false
+}
+
+// MustFindFree tries to find a free register and panics if it could not be found.
+func (s *State) MustFindFree(registers []Register) Register {
+	register, exists := s.FindFree(registers)
+
+	if !exists {
+		panic("no free registers")
+	}
+
+	return register
+}

src/build/expression/Expression.go

@@ -32,6 +32,21 @@ func (expr *Expression) AddChild(child *Expression) {
 	child.Parent = expr
 }
 
+// Count counts how often the given token appears in the expression.
+func (expr *Expression) Count(kind token.Kind, name string) int {
+	count := 0
+
+	expr.EachLeaf(func(leaf *Expression) error {
+		if leaf.Token.Kind == kind && leaf.Token.Text() == name {
+			count++
+		}
+
+		return nil
+	})
+
+	return count
+}
+
 // Reset resets all values to the default.
 func (expr *Expression) Reset() {
 	for _, child := range expr.Children {

src/build/expression/Expression_test.go

@@ -9,7 +9,7 @@ import (
 	"git.akyoto.dev/go/assert"
 )
 
-func TestExpressionParse(t *testing.T) {
+func TestParse(t *testing.T) {
 	tests := []struct {
 		Name       string
 		Expression string
@@ -96,6 +96,17 @@ func TestExpressionParse(t *testing.T) {
 	}
 }
 
+func TestCount(t *testing.T) {
+	src := []byte("(a+b-c*d)+(a*b-c+d)")
+	tokens := token.Tokenize(src)
+	expr := expression.Parse(tokens)
+	assert.Equal(t, expr.Count(token.Identifier, "a"), 2)
+	assert.Equal(t, expr.Count(token.Identifier, "b"), 2)
+	assert.Equal(t, expr.Count(token.Identifier, "c"), 2)
+	assert.Equal(t, expr.Count(token.Identifier, "d"), 2)
+	assert.Equal(t, expr.Count(token.Identifier, "e"), 0)
+}
+
 func TestEachLeaf(t *testing.T) {
 	src := []byte("(1+2-3*4)+(5*6-7+8)")
 	tokens := token.Tokenize(src)

src/build/scanner/Scan.go

@@ -248,7 +248,7 @@ func scanFile(path string, functions chan<- *core.Function) error {
 
 			name := tokens[0].Text()
 			register := x64.CallRegisters[count]
-			uses := core.CountIdentifier(function.Body, name)
+			uses := token.Count(function.Body, token.Identifier, name)
 
 			if uses == 0 {
 				return errors.New(&errors.UnusedVariable{Name: name}, file, tokens[0].Position)

src/build/token/Count.go

@@ -0,0 +1,14 @@
+package token
+
+// Count counts how often the given token appears in the token list.
+func Count(tokens List, kind Kind, name string) int {
+	count := 0
+
+	for _, t := range tokens {
+		if t.Kind == Identifier && t.Text() == name {
+			count++
+		}
+	}
+
+	return count
+}

src/build/token/Token_test.go

@@ -95,6 +95,14 @@ func TestArray(t *testing.T) {
 	})
 }
 
+func TestCount(t *testing.T) {
+	tokens := token.Tokenize([]byte(`a b b c c c`))
+	assert.Equal(t, token.Count(tokens, token.Identifier, "a"), 1)
+	assert.Equal(t, token.Count(tokens, token.Identifier, "b"), 2)
+	assert.Equal(t, token.Count(tokens, token.Identifier, "c"), 3)
+	assert.Equal(t, token.Count(tokens, token.Identifier, "d"), 0)
+}
+
 func TestNewline(t *testing.T) {
 	tokens := token.Tokenize([]byte("\n\n"))
 	assert.DeepEqual(t, tokens, token.List{