diff --git a/src/build/Result.go b/src/build/Result.go index bfebb2e..158f56f 100644 --- a/src/build/Result.go +++ b/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) + } +} diff --git a/src/build/asm/Assembler.go b/src/build/asm/Assembler.go index e3341a3..9c6f990 100644 --- a/src/build/asm/Assembler.go +++ b/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...) } diff --git a/src/build/compile.go b/src/build/compile.go index 6ffade3..3c1b48a 100644 --- a/src/build/compile.go +++ b/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) -} diff --git a/src/cli/Build.go b/src/cli/Build.go index cd5ac11..c625081 100644 --- a/src/cli/Build.go +++ b/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 {