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