From 724794b4aab53c6b363a9349c329b8b6cde530c7 Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Thu, 18 Jul 2024 10:08:38 +0200
Subject: [PATCH] Separated compiler into its own package

---
 src/build/Build.go                     |   6 +-
 src/build/compiler/Compile.go          |  75 +++++++++++++++++
 src/build/{core => compiler}/Result.go |  20 ++---
 src/build/core/Compile.go              | 107 ++++++++++++++-----------
 src/build/core/CompileAllFunctions.go  |  19 -----
 src/build/core/CompileCall.go          |   6 +-
 src/build/core/CompileDefinition.go    |   8 +-
 src/build/core/Function.go             |  67 +---------------
 src/build/core/Instructions.go         |  20 ++---
 src/build/core/Scope.go                |   2 +-
 src/build/core/{state.go => State.go}  |  14 ++--
 src/build/core/TokenToRegister.go      |   2 +-
 src/build/core/Variable.go             |   4 +-
 tests/examples_test.go                 |  16 ++++
 14 files changed, 199 insertions(+), 167 deletions(-)
 create mode 100644 src/build/compiler/Compile.go
 rename src/build/{core => compiler}/Result.go (79%)
 delete mode 100644 src/build/core/CompileAllFunctions.go
 rename src/build/core/{state.go => State.go} (85%)

diff --git a/src/build/Build.go b/src/build/Build.go
index bf90d2c..b292f4b 100644
--- a/src/build/Build.go
+++ b/src/build/Build.go
@@ -4,7 +4,7 @@ import (
 	"path/filepath"
 	"strings"
 
-	"git.akyoto.dev/cli/q/src/build/core"
+	"git.akyoto.dev/cli/q/src/build/compiler"
 	"git.akyoto.dev/cli/q/src/build/scanner"
 )
 
@@ -21,9 +21,9 @@ func New(files ...string) *Build {
 }
 
 // Run parses the input files and generates an executable file.
-func (build *Build) Run() (core.Result, error) {
+func (build *Build) Run() (compiler.Result, error) {
 	functions, errors := scanner.Scan(build.Files)
-	return core.Compile(functions, errors)
+	return compiler.Compile(functions, errors)
 }
 
 // Executable returns the path to the executable.
diff --git a/src/build/compiler/Compile.go b/src/build/compiler/Compile.go
new file mode 100644
index 0000000..565fe28
--- /dev/null
+++ b/src/build/compiler/Compile.go
@@ -0,0 +1,75 @@
+package compiler
+
+import (
+	"sync"
+
+	"git.akyoto.dev/cli/q/src/build/core"
+	"git.akyoto.dev/cli/q/src/build/errors"
+)
+
+// Compile waits for the scan to finish and compiles all functions.
+func Compile(functions <-chan *core.Function, errs <-chan error) (Result, error) {
+	result := Result{}
+	all := map[string]*core.Function{}
+
+	for functions != nil || errs != nil {
+		select {
+		case err, ok := <-errs:
+			if !ok {
+				errs = nil
+				continue
+			}
+
+			return result, err
+
+		case function, ok := <-functions:
+			if !ok {
+				functions = nil
+				continue
+			}
+
+			function.Functions = all
+			all[function.Name] = function
+		}
+	}
+
+	// Start parallel compilation
+	CompileFunctions(all)
+
+	// Report errors if any occurred
+	for _, function := range all {
+		if function.Err != nil {
+			return result, function.Err
+		}
+
+		result.InstructionCount += len(function.Assembler.Instructions)
+		result.DataCount += len(function.Assembler.Data)
+	}
+
+	// Check for existence of `main`
+	main, exists := all["main"]
+
+	if !exists {
+		return result, errors.MissingMainFunction
+	}
+
+	result.Main = main
+	result.Functions = all
+	return result, nil
+}
+
+// CompileFunctions starts a goroutine for each function compilation and waits for completion.
+func CompileFunctions(functions map[string]*core.Function) {
+	wg := sync.WaitGroup{}
+
+	for _, function := range functions {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+			function.Compile()
+		}()
+	}
+
+	wg.Wait()
+}
diff --git a/src/build/core/Result.go b/src/build/compiler/Result.go
similarity index 79%
rename from src/build/core/Result.go
rename to src/build/compiler/Result.go
index 671a90a..2762d93 100644
--- a/src/build/core/Result.go
+++ b/src/build/compiler/Result.go
@@ -1,4 +1,4 @@
-package core
+package compiler
 
 import (
 	"bufio"
@@ -6,15 +6,17 @@ import (
 
 	"git.akyoto.dev/cli/q/src/build/arch/x64"
 	"git.akyoto.dev/cli/q/src/build/asm"
+	"git.akyoto.dev/cli/q/src/build/core"
 	"git.akyoto.dev/cli/q/src/build/elf"
 	"git.akyoto.dev/cli/q/src/build/os/linux"
 )
 
 // Result contains all the compiled functions in a build.
 type Result struct {
-	Main             *Function
-	Functions        map[string]*Function
+	Main             *core.Function
+	Functions        map[string]*core.Function
 	InstructionCount int
+	DataCount        int
 }
 
 // finalize generates the final machine code.
@@ -25,7 +27,7 @@ func (r *Result) finalize() ([]byte, []byte) {
 	// a return address on the stack, which allows return statements in `main`.
 	final := asm.Assembler{
 		Instructions: make([]asm.Instruction, 0, r.InstructionCount+4),
-		Data:         map[string][]byte{},
+		Data:         make(map[string][]byte, r.DataCount),
 	}
 
 	final.Call("main")
@@ -35,8 +37,8 @@ func (r *Result) finalize() ([]byte, []byte) {
 
 	// This will place the main function immediately after the entry point
 	// and also add everything the main function calls recursively.
-	r.eachFunction(r.Main, map[*Function]bool{}, func(f *Function) {
-		final.Merge(f.assembler)
+	r.eachFunction(r.Main, map[*core.Function]bool{}, func(f *core.Function) {
+		final.Merge(f.Assembler)
 	})
 
 	code, data := final.Finalize()
@@ -45,11 +47,11 @@ func (r *Result) finalize() ([]byte, []byte) {
 
 // eachFunction recursively finds all the calls to external functions.
 // It avoids calling the same function twice with the help of a hashmap.
-func (r *Result) eachFunction(caller *Function, traversed map[*Function]bool, call func(*Function)) {
+func (r *Result) eachFunction(caller *core.Function, traversed map[*core.Function]bool, call func(*core.Function)) {
 	call(caller)
 	traversed[caller] = true
 
-	for _, x := range caller.assembler.Instructions {
+	for _, x := range caller.Assembler.Instructions {
 		if x.Mnemonic != asm.CALL {
 			continue
 		}
@@ -71,7 +73,7 @@ func (r *Result) eachFunction(caller *Function, traversed map[*Function]bool, ca
 
 // PrintInstructions prints out the generated instructions.
 func (r *Result) PrintInstructions() {
-	r.eachFunction(r.Main, map[*Function]bool{}, func(f *Function) {
+	r.eachFunction(r.Main, map[*core.Function]bool{}, func(f *core.Function) {
 		f.PrintInstructions()
 	})
 }
diff --git a/src/build/core/Compile.go b/src/build/core/Compile.go
index edd8712..c581921 100644
--- a/src/build/core/Compile.go
+++ b/src/build/core/Compile.go
@@ -1,55 +1,66 @@
 package core
 
 import (
+	"git.akyoto.dev/cli/q/src/build/ast"
 	"git.akyoto.dev/cli/q/src/build/errors"
+	"git.akyoto.dev/cli/q/src/build/token"
 )
 
-// Compile waits for the scan to finish and compiles all functions.
-func Compile(functions <-chan *Function, errs <-chan error) (Result, error) {
-	result := Result{}
-	allFunctions := map[string]*Function{}
-
-	for functions != nil || errs != nil {
-		select {
-		case err, ok := <-errs:
-			if !ok {
-				errs = nil
-				continue
-			}
-
-			return result, err
-
-		case function, ok := <-functions:
-			if !ok {
-				functions = nil
-				continue
-			}
-
-			function.functions = allFunctions
-			allFunctions[function.Name] = function
-		}
-	}
-
-	// Start parallel compilation
-	CompileAllFunctions(allFunctions)
-
-	// Report errors if any occurred
-	for _, function := range allFunctions {
-		if function.err != nil {
-			return result, function.err
-		}
-
-		result.InstructionCount += len(function.assembler.Instructions)
-	}
-
-	// Check for existence of `main`
-	main, exists := allFunctions["main"]
-
-	if !exists {
-		return result, errors.MissingMainFunction
-	}
-
-	result.Main = main
-	result.Functions = allFunctions
-	return result, nil
+// Compile turns a function into machine code.
+func (f *Function) Compile() {
+	defer close(f.finished)
+	f.AddLabel(f.Name)
+	f.Err = f.CompileTokens(f.Body)
+	f.Return()
+}
+
+// CompileTokens compiles a token list.
+func (f *Function) CompileTokens(tokens token.List) error {
+	body, err := ast.Parse(tokens)
+
+	if err != nil {
+		err.(*errors.Error).File = f.File
+		return err
+	}
+
+	return f.CompileAST(body)
+}
+
+// CompileAST compiles an abstract syntax tree.
+func (f *Function) CompileAST(tree ast.AST) error {
+	for _, node := range tree {
+		err := f.CompileASTNode(node)
+
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// CompileASTNode compiles a node in the AST.
+func (f *Function) CompileASTNode(node ast.Node) error {
+	switch node := node.(type) {
+	case *ast.Assign:
+		return f.CompileAssign(node)
+
+	case *ast.Call:
+		return f.CompileCall(node.Expression)
+
+	case *ast.Define:
+		return f.CompileDefinition(node)
+
+	case *ast.Return:
+		return f.CompileReturn(node)
+
+	case *ast.If:
+		return f.CompileIf(node)
+
+	case *ast.Loop:
+		return f.CompileLoop(node)
+
+	default:
+		panic("unknown AST type")
+	}
 }
diff --git a/src/build/core/CompileAllFunctions.go b/src/build/core/CompileAllFunctions.go
deleted file mode 100644
index 66d3ef9..0000000
--- a/src/build/core/CompileAllFunctions.go
+++ /dev/null
@@ -1,19 +0,0 @@
-package core
-
-import "sync"
-
-// CompileAllFunctions starts a goroutine for each function compilation and waits for completion.
-func CompileAllFunctions(functions map[string]*Function) {
-	wg := sync.WaitGroup{}
-
-	for _, function := range functions {
-		wg.Add(1)
-
-		go func() {
-			defer wg.Done()
-			function.Compile()
-		}()
-	}
-
-	wg.Wait()
-}
diff --git a/src/build/core/CompileCall.go b/src/build/core/CompileCall.go
index 9bb4f1f..ae04215 100644
--- a/src/build/core/CompileCall.go
+++ b/src/build/core/CompileCall.go
@@ -15,7 +15,7 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 	isSyscall := funcName == "syscall"
 
 	if !isSyscall {
-		_, exists := f.functions[funcName]
+		_, exists := f.Functions[funcName]
 
 		if !exists {
 			return errors.New(&errors.UnknownFunction{Name: funcName}, f.File, root.Children[0].Token.Position)
@@ -52,6 +52,10 @@ func (f *Function) CompileCall(root *expression.Expression) error {
 	}
 
 	for _, register := range registers {
+		if register == f.cpu.Output[0] && root.Parent != nil {
+			continue
+		}
+
 		f.Scope().Free(register)
 	}
 
diff --git a/src/build/core/CompileDefinition.go b/src/build/core/CompileDefinition.go
index a731841..52bafe3 100644
--- a/src/build/core/CompileDefinition.go
+++ b/src/build/core/CompileDefinition.go
@@ -63,6 +63,10 @@ func (f *Function) useVariable(variable *Variable) {
 			continue
 		}
 
+		if config.Comments {
+			f.Comment("%s (%s) used in scope %d", local.Name, local.Register, i)
+		}
+
 		local.Alive--
 
 		if local.Alive < 0 {
@@ -71,7 +75,7 @@ func (f *Function) useVariable(variable *Variable) {
 
 		if local.Alive == 0 {
 			if config.Comments {
-				f.Comment("%s died (%s) in scope %d", local.Name, local.Register, i)
+				f.Comment("%s (%s) died in scope %d", local.Name, local.Register, i)
 			}
 
 			scope.Free(local.Register)
@@ -87,7 +91,7 @@ func (f *Function) identifierExists(name string) bool {
 		return true
 	}
 
-	_, exists := f.functions[name]
+	_, exists := f.Functions[name]
 	return exists
 }
 
diff --git a/src/build/core/Function.go b/src/build/core/Function.go
index eef5bcd..f154692 100644
--- a/src/build/core/Function.go
+++ b/src/build/core/Function.go
@@ -3,9 +3,7 @@ package core
 import (
 	"git.akyoto.dev/cli/q/src/build/arch/x64"
 	"git.akyoto.dev/cli/q/src/build/asm"
-	"git.akyoto.dev/cli/q/src/build/ast"
 	"git.akyoto.dev/cli/q/src/build/cpu"
-	"git.akyoto.dev/cli/q/src/build/errors"
 	"git.akyoto.dev/cli/q/src/build/fs"
 	"git.akyoto.dev/cli/q/src/build/token"
 )
@@ -15,7 +13,7 @@ type Function struct {
 	Name string
 	File *fs.File
 	Body token.List
-	state
+	State
 }
 
 // NewFunction creates a new function.
@@ -24,8 +22,8 @@ func NewFunction(name string, file *fs.File, body token.List) *Function {
 		Name: name,
 		File: file,
 		Body: body,
-		state: state{
-			assembler: asm.Assembler{
+		State: State{
+			Assembler: asm.Assembler{
 				Instructions: make([]asm.Instruction, 0, 32),
 			},
 			cpu: cpu.CPU{
@@ -43,65 +41,6 @@ func NewFunction(name string, file *fs.File, body token.List) *Function {
 	}
 }
 
-// Compile turns a function into machine code.
-func (f *Function) Compile() {
-	defer close(f.finished)
-	f.AddLabel(f.Name)
-	f.err = f.CompileTokens(f.Body)
-	f.Return()
-}
-
-// CompileTokens compiles a token list.
-func (f *Function) CompileTokens(tokens token.List) error {
-	body, err := ast.Parse(tokens)
-
-	if err != nil {
-		err.(*errors.Error).File = f.File
-		return err
-	}
-
-	return f.CompileAST(body)
-}
-
-// CompileAST compiles an abstract syntax tree.
-func (f *Function) CompileAST(tree ast.AST) error {
-	for _, node := range tree {
-		err := f.CompileASTNode(node)
-
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-// CompileASTNode compiles a node in the AST.
-func (f *Function) CompileASTNode(node ast.Node) error {
-	switch node := node.(type) {
-	case *ast.Assign:
-		return f.CompileAssign(node)
-
-	case *ast.Call:
-		return f.CompileCall(node.Expression)
-
-	case *ast.Define:
-		return f.CompileDefinition(node)
-
-	case *ast.Return:
-		return f.CompileReturn(node)
-
-	case *ast.If:
-		return f.CompileIf(node)
-
-	case *ast.Loop:
-		return f.CompileLoop(node)
-
-	default:
-		panic("unknown AST type")
-	}
-}
-
 // String returns the function name.
 func (f *Function) String() string {
 	return f.Name
diff --git a/src/build/core/Instructions.go b/src/build/core/Instructions.go
index 7a35222..9bdf8fa 100644
--- a/src/build/core/Instructions.go
+++ b/src/build/core/Instructions.go
@@ -9,28 +9,28 @@ import (
 )
 
 func (f *Function) AddLabel(label string) {
-	f.assembler.Label(asm.LABEL, label)
+	f.Assembler.Label(asm.LABEL, label)
 	f.postInstruction()
 }
 
 func (f *Function) Call(label string) {
-	f.assembler.Call(label)
+	f.Assembler.Call(label)
 	f.Scope().Use(f.cpu.Output[0])
 	f.postInstruction()
 }
 
 func (f *Function) Comment(format string, args ...any) {
-	f.assembler.Comment(fmt.Sprintf(format, args...))
+	f.Assembler.Comment(fmt.Sprintf(format, args...))
 	f.postInstruction()
 }
 
 func (f *Function) Jump(mnemonic asm.Mnemonic, label string) {
-	f.assembler.Label(mnemonic, label)
+	f.Assembler.Label(mnemonic, label)
 	f.postInstruction()
 }
 
 func (f *Function) Register(mnemonic asm.Mnemonic, a cpu.Register) {
-	f.assembler.Register(mnemonic, a)
+	f.Assembler.Register(mnemonic, a)
 
 	if mnemonic == asm.POP {
 		f.Scope().Use(a)
@@ -44,7 +44,7 @@ func (f *Function) RegisterNumber(mnemonic asm.Mnemonic, a cpu.Register, b int)
 		f.SaveRegister(a)
 	}
 
-	f.assembler.RegisterNumber(mnemonic, a, b)
+	f.Assembler.RegisterNumber(mnemonic, a, b)
 
 	if mnemonic == asm.MOVE {
 		f.Scope().Use(a)
@@ -58,7 +58,7 @@ func (f *Function) RegisterLabel(mnemonic asm.Mnemonic, register cpu.Register, l
 		f.SaveRegister(register)
 	}
 
-	f.assembler.RegisterLabel(mnemonic, register, label)
+	f.Assembler.RegisterLabel(mnemonic, register, label)
 
 	if mnemonic == asm.MOVE {
 		f.Scope().Use(register)
@@ -76,7 +76,7 @@ func (f *Function) RegisterRegister(mnemonic asm.Mnemonic, a cpu.Register, b cpu
 		f.SaveRegister(a)
 	}
 
-	f.assembler.RegisterRegister(mnemonic, a, b)
+	f.Assembler.RegisterRegister(mnemonic, a, b)
 
 	if mnemonic == asm.MOVE {
 		f.Scope().Use(a)
@@ -86,12 +86,12 @@ func (f *Function) RegisterRegister(mnemonic asm.Mnemonic, a cpu.Register, b cpu
 }
 
 func (f *Function) Return() {
-	f.assembler.Return()
+	f.Assembler.Return()
 	f.postInstruction()
 }
 
 func (f *Function) Syscall() {
-	f.assembler.Syscall()
+	f.Assembler.Syscall()
 	f.Scope().Use(f.cpu.Output[0])
 	f.postInstruction()
 }
diff --git a/src/build/core/Scope.go b/src/build/core/Scope.go
index 7aeb163..655ab27 100644
--- a/src/build/core/Scope.go
+++ b/src/build/core/Scope.go
@@ -15,7 +15,7 @@ type Scope struct {
 }
 
 // Scope returns the current scope.
-func (s *state) Scope() *Scope {
+func (s *State) Scope() *Scope {
 	return s.scopes[len(s.scopes)-1]
 }
 
diff --git a/src/build/core/state.go b/src/build/core/State.go
similarity index 85%
rename from src/build/core/state.go
rename to src/build/core/State.go
index b345c69..a75c58c 100644
--- a/src/build/core/state.go
+++ b/src/build/core/State.go
@@ -9,14 +9,14 @@ import (
 	"git.akyoto.dev/go/color/ansi"
 )
 
-// state is the data structure we embed in each function to preserve compilation state.
-type state struct {
-	err             error
+// State is the data structure we embed in each function to preserve compilation State.
+type State struct {
+	Assembler       asm.Assembler
+	Functions       map[string]*Function
+	Err             error
 	scopes          []*Scope
-	functions       map[string]*Function
 	registerHistory []uint64
 	finished        chan struct{}
-	assembler       asm.Assembler
 	cpu             cpu.CPU
 	count           counter
 }
@@ -30,10 +30,10 @@ type counter struct {
 }
 
 // PrintInstructions shows the assembly instructions.
-func (s *state) PrintInstructions() {
+func (s *State) PrintInstructions() {
 	ansi.Dim.Println("╭──────────────────────────────────────────────────────────────────────────────╮")
 
-	for i, x := range s.assembler.Instructions {
+	for i, x := range s.Assembler.Instructions {
 		ansi.Dim.Print("│ ")
 
 		switch x.Mnemonic {
diff --git a/src/build/core/TokenToRegister.go b/src/build/core/TokenToRegister.go
index 7f81b17..85e8bf9 100644
--- a/src/build/core/TokenToRegister.go
+++ b/src/build/core/TokenToRegister.go
@@ -41,7 +41,7 @@ func (f *Function) TokenToRegister(t token.Token, register cpu.Register) error {
 		f.count.data++
 		label := fmt.Sprintf("%s_data_%d", f.Name, f.count.data)
 		value := t.Bytes[1 : len(t.Bytes)-1]
-		f.assembler.SetData(label, value)
+		f.Assembler.SetData(label, value)
 		f.RegisterLabel(asm.MOVE, register, label)
 		return nil
 
diff --git a/src/build/core/Variable.go b/src/build/core/Variable.go
index cca96e8..2045aef 100644
--- a/src/build/core/Variable.go
+++ b/src/build/core/Variable.go
@@ -13,12 +13,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 {
+func (s *State) Variable(name string) *Variable {
 	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 {
+func (s *State) VariableInRegister(register cpu.Register) *Variable {
 	for _, v := range s.Scope().variables {
 		if v.Register == register {
 			return v
diff --git a/tests/examples_test.go b/tests/examples_test.go
index d9f19f4..8af877d 100644
--- a/tests/examples_test.go
+++ b/tests/examples_test.go
@@ -3,6 +3,9 @@ package tests_test
 import (
 	"path/filepath"
 	"testing"
+
+	"git.akyoto.dev/cli/q/src/build"
+	"git.akyoto.dev/go/assert"
 )
 
 var examples = []struct {
@@ -22,3 +25,16 @@ func TestExamples(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkExamples(b *testing.B) {
+	for _, test := range examples {
+		b.Run(test.Name, func(b *testing.B) {
+			compiler := build.New(filepath.Join("..", "examples", test.Name))
+
+			for i := 0; i < b.N; i++ {
+				_, err := compiler.Run()
+				assert.Nil(b, err)
+			}
+		})
+	}
+}