loov--lensm/internal/goobj/code.go

package goobj

import (
	"fmt"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"loov.dev/lensm/internal/disasm"
	godisasm "loov.dev/lensm/internal/go/src/disasm"
)

var rxRefAbs = regexp.MustCompile(`\s0x[\da-fA-F]+$`)
var rxRefRel = regexp.MustCompile(`\s-?\d+\(PC\)$`)
var rxCallOrJump = regexp.MustCompile(`^(?:CALL|JMP)\s+(.+?)\(SB\)`)

// Disassemble disassembles the specified symbol.
func Disassemble(dis *godisasm.Disasm, sym *Func, opts disasm.Options) (*disasm.Code, error) {
	neededLines := make(map[string]*disasm.LineSet)

	file, _, _ := dis.PCLN().PCToLine(sym.sym.Addr)
	needRefPCs := map[uint64]struct{}{}

	code := &disasm.Code{
		Name: sym.Name(),
		File: file,
		Arch: dis.GOARCH(),
	}
	var instructions []disasm.Inst
	dis.DecodeSyntax(sym.sym.Addr, sym.sym.Addr+uint64(sym.sym.Size), sym.sym.Relocs,
		func(pc, size uint64, file string, line int, text, nativeText, mnemonic string) {
			// TODO: find a better way to calculate the jump target
			var refPC uint64
			var call string
			if match := rxRefAbs.FindString(text); match != "" {
				if target, err := strconv.ParseInt(match[3:], 16, 64); err == nil {
					refPC = uint64(target)
				}
			} else if match := rxRefRel.FindString(text); match != "" {
				// TODO: this calculation seems incorrect
				if target, err := strconv.ParseInt(match[1:len(match)-4], 10, 64); err == nil {
					refPC = uint64(int64(pc) + target*4)
				} else {
					panic(err)
				}
			} else if match := rxCallOrJump.FindStringSubmatch(text); len(match) > 0 {
				call = match[1]
			}

			if refPC != 0 {
				needRefPCs[refPC] = struct{}{}
			}
			instructions = append(instructions, disasm.Inst{
				PC:         pc,
				Text:       text,
				NativeText: nativeText,
				Mnemonic:   mnemonic,
				File:       file,
				Line:       line,
				Call:       call,
				RefPC:      refPC,
			})

			if file != "" && file != "<autogenerated>" {
				lineset, ok := neededLines[file]
				if !ok {
					lineset = &disasm.LineSet{}
					neededLines[file] = lineset
				}
				lineset.Add(line)
			}
		})

	pcToIndex := map[uint64]int{}
	for _, ix := range instructions {
		if _, ok := needRefPCs[ix.PC]; ok {
			// add empty line
			code.Insts = append(code.Insts, disasm.Inst{})
		}
		pcToIndex[ix.PC] = len(code.Insts)
		code.Insts = append(code.Insts, ix)
	}

	type jumpInterval struct {
		index    int
		ix       *disasm.Inst
		min, max uint64
	}

	var jumps []jumpInterval
	for i := range code.Insts {
		ix := &code.Insts[i]
		if ix.RefPC != 0 {
			target, ok := pcToIndex[ix.RefPC]
			if !ok {
				continue
			}
			ix.RefOffset = target - i

			if ix.PC <= ix.RefPC {
				jumps = append(jumps, jumpInterval{
					index: i,
					ix:    ix,
					min:   ix.PC,
					max:   ix.RefPC,
				})
			} else {
				jumps = append(jumps, jumpInterval{
					index: i,
					ix:    ix,
					min:   ix.RefPC,
					max:   ix.PC,
				})
			}
		}
	}

	sort.Slice(jumps, func(i, k int) bool {
		if jumps[i].min == jumps[k].min {
			return jumps[i].max > jumps[k].max
		}
		return jumps[i].min < jumps[k].min
	})

	var stackLayers []uint64
	insertToStack := func(ix *disasm.Inst, max uint64) {
		found := false
		for k, pc := range stackLayers {
			if pc == 0 {
				stackLayers[k] = max
				ix.RefStack = k
				found = true
				break
			}
		}
		if !found {
			code.MaxJump = len(stackLayers)
			ix.RefStack = len(stackLayers)
			stackLayers = append(stackLayers, max)
		}
	}

	for _, jump := range jumps {
		for i, pc := range stackLayers {
			if pc <= jump.min {
				stackLayers[i] = 0
			}
		}
		insertToStack(jump.ix, jump.max)
	}
	for i := range code.Insts {
		ix := &code.Insts[i]
		ix.RefStack = code.MaxJump - ix.RefStack + 1
	}
	code.MaxJump++

	// remove trailing interrupts from funcs
	for len(code.Insts) > 0 &&
		(strings.HasPrefix(code.Insts[len(code.Insts)-1].Text, "INT ") ||
			code.Insts[len(code.Insts)-1].Text == "?") {
		code.Insts = code.Insts[:len(code.Insts)-1]
	}

	// load sources
	code.Source = LoadSources(neededLines, code.File, opts.Context)

	// create a mapping from source code to disassembly
	type fileLine struct {
		file string
		line int
	}

	lineRefs := map[fileLine]*disasm.LineSet{}
	for i, ix := range code.Insts {
		k := fileLine{file: ix.File, line: ix.Line}
		n, ok := lineRefs[k]
		if !ok {
			n = &disasm.LineSet{}
			lineRefs[k] = n
		}
		n.Add(i)
	}
	for i := range code.Source {
		src := &code.Source[i]
		for k := range src.Blocks {
			block := &src.Blocks[k]
			block.Related = make([][]disasm.LineRange, len(block.Lines))
			for line := block.From; line <= block.To; line++ { // todo check: line <= block.To
				if refs, ok := lineRefs[fileLine{file: src.File, line: line}]; ok {
					block.Related[line-block.From] = refs.RangesZero()
				}
			}
		}
	}

	return code, nil
}

var rxEnvVariable = regexp.MustCompile(`\$[a-zA-Z_]+[a-zA-Z0-9_]+\b`)

func replaceEnvironmentVariables(s string) string {
	return rxEnvVariable.ReplaceAllStringFunc(s, func(env string) string {
		replacement := os.Getenv(env[1:])
		if replacement != "" {
			return replacement
		}
		return env
	})
}

// LoadSources loads the specified line sets.
func LoadSources(needed map[string]*disasm.LineSet, symbolFile string, context int) []disasm.Source {
	var sources []disasm.Source
	for file, set := range needed {
		data, err := os.ReadFile(replaceEnvironmentVariables(file))
		if err != nil {
			// TODO: should we create a stub source block instead?
			fmt.Fprintf(os.Stderr, "unable to load source from %q: %v\n", file, err)
			continue
		}
		lines := strings.Split(string(data), "\n")
		source := disasm.Source{
			File: file,
		}
		for _, r := range set.Ranges(context) {
			to := min(r.To-1, len(lines))
			lineBlock := lines[r.From-1 : to]
			for i, v := range lineBlock {
				lineBlock[i] = strings.Replace(v, "\t", "    ", -1)
			}

			source.Blocks = append(source.Blocks, disasm.SourceBlock{
				LineRange: r,
				Lines:     lineBlock,
			})
		}
		sources = append(sources, source)
	}

	// Sort the sources and prioritize the file where the main symbol is located.
	sort.Slice(sources, func(i, j int) bool {
		if sources[i].File == symbolFile {
			return true
		}
		if sources[j].File == symbolFile {
			return false
		}
		return sources[i].File < sources[j].File
	})

	return sources
}