Improved performance

src/build/Result.go

@@ -8,13 +8,13 @@ import (
 
 // Result contains all the compiled functions in a build.
 type Result struct {
-	Used             []*Function
-	Unused           map[string]*Function
+	Main             *Function
+	Functions        map[string]*Function
 	InstructionCount int
 }
 
 // Finalize generates the final machine code.
-func (r Result) Finalize() ([]byte, []byte) {
+func (r *Result) Finalize() ([]byte, []byte) {
 	// This will be the entry point of the executable.
 	// The only job of the entry function is to call `main` and exit cleanly.
 	// The reason we call `main` instead of using `main` itself is to place
@@ -28,11 +28,38 @@ func (r Result) Finalize() ([]byte, []byte) {
 	final.RegisterNumber(asm.MOVE, x64.SyscallRegisters[1], 0)
 	final.Syscall()
 
-	// Merge all the called functions.
-	for _, f := range r.Used {
-		final.Merge(&f.Assembler)
-	}
+	// 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)
+	})
 
 	code, data := final.Finalize()
 	return code, data
 }
+
+// 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)) {
+	call(caller)
+	traversed[caller] = true
+
+	for _, x := range caller.Assembler.Instructions {
+		if x.Mnemonic != asm.CALL {
+			continue
+		}
+
+		name := x.Data.(*asm.Label).Name
+		callee, exists := r.Functions[name]
+
+		if !exists {
+			continue
+		}
+
+		if traversed[callee] {
+			continue
+		}
+
+		r.EachFunction(callee, traversed, call)
+	}
+}

src/build/asm/Assembler.go

@@ -12,7 +12,7 @@ type Assembler struct {
 }
 
 // Finalize generates the final machine code.
-func (a *Assembler) Finalize() ([]byte, []byte) {
+func (a Assembler) Finalize() ([]byte, []byte) {
 	code := make([]byte, 0, len(a.Instructions)*8)
 	data := make([]byte, 0, 16)
 	labels := map[string]Address{}
@@ -139,7 +139,7 @@ func (a *Assembler) Finalize() ([]byte, []byte) {
 }
 
 // Merge combines the contents of this assembler with another one.
-func (a *Assembler) Merge(b *Assembler) {
+func (a *Assembler) Merge(b Assembler) {
 	a.Instructions = append(a.Instructions, b.Instructions...)
 }
 

src/build/compile.go

@@ -3,21 +3,19 @@ package build
 import (
 	"sync"
 
-	"git.akyoto.dev/cli/q/src/build/asm"
-	fail "git.akyoto.dev/cli/q/src/errors"
+	"git.akyoto.dev/cli/q/src/errors"
 )
 
 // compile waits for the scan to finish and compiles all functions.
-func compile(functions <-chan *Function, errors <-chan error) (Result, error) {
-	result := Result{
-		Unused: map[string]*Function{},
-	}
+func compile(functions <-chan *Function, errs <-chan error) (Result, error) {
+	result := Result{}
+	allFunctions := map[string]*Function{}
 
-	for functions != nil || errors != nil {
+	for functions != nil || errs != nil {
 		select {
-		case err, ok := <-errors:
+		case err, ok := <-errs:
 			if !ok {
-				errors = nil
+				errs = nil
 				continue
 			}
 
@@ -29,28 +27,28 @@ func compile(functions <-chan *Function, errors <-chan error) (Result, error) {
 				continue
 			}
 
-			result.Unused[function.Name] = function
+			allFunctions[function.Name] = function
 		}
 	}
 
-	compileFunctions(result.Unused)
+	compileFunctions(allFunctions)
 
-	for _, function := range result.Unused {
+	for _, function := range allFunctions {
 		if function.Error != nil {
 			return result, function.Error
 		}
+
+		result.InstructionCount += len(function.Assembler.Instructions)
 	}
 
-	main, exists := result.Unused["main"]
+	main, exists := allFunctions["main"]
 
 	if !exists {
-		return result, fail.MissingMainFunction
+		return result, errors.MissingMainFunction
 	}
 
-	result.Used = make([]*Function, 0, len(result.Unused))
-	result.markAlive(main)
-	result.findAliveCode(main)
-
+	result.Main = main
+	result.Functions = allFunctions
 	return result, nil
 }
 
@@ -69,29 +67,3 @@ func compileFunctions(functions map[string]*Function) {
 
 	wg.Wait()
 }
-
-// findAliveCode recursively finds all the calls to external functions and marks them as required.
-func (result *Result) findAliveCode(caller *Function) {
-	for _, x := range caller.Assembler.Instructions {
-		if x.Mnemonic != asm.CALL {
-			continue
-		}
-
-		name := x.Data.(*asm.Label).Name
-		callee, exists := result.Unused[name]
-
-		if !exists {
-			continue
-		}
-
-		result.markAlive(callee)
-		result.findAliveCode(callee)
-	}
-}
-
-// markAlive marks a function as being alive.
-func (result *Result) markAlive(f *Function) {
-	result.Used = append(result.Used, f)
-	result.InstructionCount += len(f.Assembler.Instructions)
-	delete(result.Unused, f.Name)
-}

src/cli/Build.go

@@ -44,9 +44,9 @@ func Build(args []string) int {
 	}
 
 	if config.Verbose {
-		for _, function := range result.Used {
-			function.PrintAsm()
-		}
+		result.EachFunction(result.Main, map[*build.Function]bool{}, func(f *build.Function) {
+			f.PrintAsm()
+		})
 	}
 
 	if dry {