gasm-sdk/asm/link.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

package asm

import (
	"fmt"
	"math"
	"sort"
	"strconv"

	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
)

// Image is an assembled file: the function bodies laid out in source order,
// followed by the file's static data section (GLOBL/DATA).  References to
// file-local static symbols are encoded RIP-relative and resolved within the
// image, so the raw bytes are self-consistent and executable at any base
// address; references to external symbols are recorded as relocations
// (Funcs[i].Relocs, Externals) and left unresolved, the object-file
// emitters turn them into linker relocations.
type Image struct {
	Code      []byte         // concatenated function bodies
	Data      []byte         // static data section
	Funcs     []FuncLayout   // function positions, in source order
	Symbols   map[string]int // static symbol → byte offset within the image
	DataSyms  []DataSymbol   // GLOBL symbols, in layout order
	Externals []string       // referenced but undefined symbols, sorted
	// SourcePath is the assembled file's path, recorded in the DWARF
	// sections in place of a placeholder name.  Empty when the image was
	// not built from a named file.
	SourcePath string
}

// FuncLayout describes one assembled function within an Image.
type FuncLayout struct {
	Name    string
	Pkg     string // explicit package prefix ("" = the current package)
	Static  bool   // the <> marker: file-local, not exported
	Offset  int    // start offset within the image (== offset within Code)
	Size    int
	Args    int            // declared argument/result area (the TEXT size suffix)
	Frame   int            // local frame size (the TEXT $framesize)
	NoSplit bool           // the NOSPLIT flag
	SPWrite bool           // the SPWRITE flag: writes an arbitrary value to SP
	Line    int            // source line of the TEXT directive
	Labels  map[string]int // local labels, function-relative
	Relocs  []Reloc        // static-symbol references, in emission order
	Spadj   []SpadjStep    // stack-adjustment boundaries, ascending by PC
	Lines   []LineEntry    // source-line table: byte offset → source line
}

// SpadjStep is one stack-adjustment boundary: Value is the SP delta from the
// entry state in effect from PC (function-relative) until the next step.
type SpadjStep struct {
	PC    int
	Value int
}

// LineEntry maps a byte offset (function-relative) to a source line number.
type LineEntry struct {
	Offset int
	Line   int
}

// LineAt returns the source line number for the given function-relative byte
// offset, using a binary search on the line table.  Returns 0 if the offset
// is before the first instruction or the table is empty.
func (fl *FuncLayout) LineAt(offset int) int {
	if len(fl.Lines) == 0 {
		return 0
	}
	// Binary search: find the last entry with Offset <= offset.
	lo, hi := 0, len(fl.Lines)-1
	for lo < hi {
		mid := (lo + hi + 1) / 2
		if fl.Lines[mid].Offset <= offset {
			lo = mid
		} else {
			hi = mid - 1
		}
	}
	if fl.Lines[lo].Offset <= offset {
		return fl.Lines[lo].Line
	}
	return 0
}

// RelocKind discriminates the relocation a static-symbol reference needs;
// the encoders record one per SB reference, and the object-file emitters map
// it to their format's relocation type.
type RelocKind int

const (
	RelPCRel32         RelocKind = iota // 32-bit PC-relative (amd64)
	RelCall                             // R_CALL: CALL to a function symbol (amd64)
	RelTLSLE                            // R_TLS_LE: local-exec TLS load, no symbol (amd64 guard)
	RelRISCVPCRELIType                  // R_RISCV_PCREL_ITYPE (AUIPC + I-type pair)
	RelRISCVPCRELSType                  // R_RISCV_PCREL_STYPE (AUIPC + S-type pair)
	RelRISCVJal                         // R_RISCV_JAL (J-type call)
	RelLoong64AddrHi                    // R_LOONG64_ADDR_HI (pcalau12i)
	RelLoong64AddrLo                    // R_LOONG64_ADDR_LO (addi.d/ld/st)
	RelArm64Addr                        // R_ADDRARM64 (ADRP + ADD pair)
	RelArm64Branch                      // R_CALLARM64 (BL instruction)
	RelArm64LDST64                      // R_ARM64_PCREL_LDST64 (ADRP + 64-bit LDR/STR pair)
	RelLoong64Branch                    // R_CALLLOONG64 (BL instruction)
	RelAddr                             // R_ADDR: the absolute address of a symbol held in a DATA field
)

type Reloc struct {
	// Off is the function-relative offset of the field the linker patches
	// and After the address just past the instruction, the base the
	// assembler measures PC-relative displacements from.  Name plus
	// Addend select the target: the symbol plus the byte offset.  An
	// External relocation names a symbol no GLOBL in the file defines;
	// the object-file emitters carry it into the output's relocation
	// table.  Siz is the width of the patched field and is set only for
	// data-field relocations (RelAddr, Off relative to the data symbol),
	// whose width is the DATA line's; code relocations take their width
	// from the architecture's instruction encoding.
	Off      int
	After    int
	Name     string
	Addend   int64
	External bool
	Kind     RelocKind
	Siz      uint8
}

// DataSymbol describes one GLOBL symbol laid out in the data section.
type DataSymbol struct {
	Name   string
	Pkg    string // explicit package prefix ("" = the current package)
	Offset int    // byte offset within Data
	Size   int
	Static bool // the <> marker: file-local, not exported
	Rodata bool // the RODATA flag: read-only data
	Dupok  bool // the DUPOK flag: duplicate-OK
	// Relocs carries the symbol-valued DATA initialisers ("DATA s+0(SB)/8,
	// $other(SB)"): fields of this symbol's data that hold another symbol's
	// address, resolved by the linker.  Off is relative to the symbol's
	// data start.
	Relocs []Reloc
}

// Bytes returns the whole image: code, then data.
func (img *Image) Bytes() []byte {
	out := make([]byte, 0, len(img.Code)+len(img.Data))
	out = append(out, img.Code...)
	return append(out, img.Data...)
}

// AssembleFile assembles every TEXT function of a parsed file and lays out
// its static symbols (GLOBL/DATA) in a data section behind the code.  Each
// reference to a file-local static symbol becomes a RIP-relative load whose
// displacement is resolved against that layout; a reference to a symbol no
// GLOBL defines is recorded as an external relocation (Externals) with its
// displacement left zero, the object-file emitters resolve it at link
// time, while the raw image (Bytes) cannot represent it.
func AssembleFile(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}
	known := make(map[string]bool, len(dataSyms))
	for _, d := range dataSyms {
		known[d.name] = true
	}
	// TEXT symbols are file-level definitions too: a symbol immediate
	// ($fn(SB)) may name one, exactly as a data reference names a GLOBL.
	for _, d := range f.Decls {
		if t, ok := d.(*ast.Text); ok {
			known[t.Name.Name] = true
		}
	}
	link := &linkInfo{symbols: known, allowExternal: true}
	poolSeen := map[string]bool{}

	img := &Image{Symbols: map[string]int{}, SourcePath: f.Path}
	textOff := map[string]int{}
	type asmFunc struct {
		name    string
		patches []sbPatch
	}
	var funcs []asmFunc
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, patches, labels, steps, lines, pool, err := assemble(t, link)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		// The pooled floating-point constants join the declared data as
		// read-only symbols, deduplicated across the file (the toolchain
		// synthesises the same symbols into its rodata).
		for _, entry := range pool {
			if poolSeen[entry.name] {
				continue
			}
			poolSeen[entry.name] = true
			dataSyms = append(dataSyms, dataSym{
				name:   entry.name,
				buf:    entry.data,
				size:   len(entry.data),
				rodata: true,
				dupok:  true,
			})
		}
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		for _, s := range steps {
			fl.Spadj = append(fl.Spadj, SpadjStep{PC: s.pc, Value: s.value})
		}
		textOff[t.Name.Name] = len(img.Code)
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
		funcs = append(funcs, asmFunc{name: t.Name.Name, patches: patches})
	}

	// Lay out the data section behind the code, each symbol 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		if pos := dataStart + len(img.Data); pos != align16(pos) {
			img.Data = append(img.Data, make([]byte, align16(pos)-pos)...)
		}
		img.Symbols[d.name] = dataStart + len(img.Data)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data),
			Size:   len(d.buf),
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
		img.Data = append(img.Data, d.buf...)
	}

	// Resolve the RIP-relative displacements of file-local references now
	// that every address is known, and record every reference (resolved or
	// external) for the object-file emitters.
	externals := map[string]bool{}
	for i, fn := range funcs {
		base := img.Funcs[i].Offset
		code := img.Code[base : base+img.Funcs[i].Size]
		for _, p := range fn.patches {
			reloc := Reloc{Off: p.off, After: p.after, Name: p.name, Addend: p.addend, Kind: p.kind}
			if p.kind == RelTLSLE {
				// The TLS slot has no symbol: the linker fills the offset
				// from the runtime's TLS layout.
				img.Funcs[i].Relocs = append(img.Funcs[i].Relocs, reloc)
				continue
			}
			if imgOff, ok := img.Symbols[p.name]; ok {
				rel := int64(imgOff) + p.addend - int64(base+p.after)
				if rel < -1<<31 || rel >= 1<<31 {
					return nil, fmt.Errorf("%s: displacement to %q out of rel32 range", fn.name, p.name)
				}
				copy(code[p.off:p.off+4], le32(rel))
			} else if imgOff, ok := textOff[p.name]; ok {
				// A CALL to a TEXT function of the same file: resolve the
				// displacement against the function's layout position.
				rel := int64(imgOff) + p.addend - int64(base+p.after)
				if rel < -1<<31 || rel >= 1<<31 {
					return nil, fmt.Errorf("%s: displacement to %q out of rel32 range", fn.name, p.name)
				}
				copy(code[p.off:p.off+4], le32(rel))
			} else {
				reloc.External = true
				externals[p.name] = true
			}
			img.Funcs[i].Relocs = append(img.Funcs[i].Relocs, reloc)
		}
	}
	// The data symbols' symbol-valued DATA fields resolve the same way the
	// code references do: a name the file defines (GLOBL or TEXT) stays an
	// internal reference the emitters resolve, anything else is external.
	// img.DataSyms was laid out in dataSyms order, so the indexes line up.
	for i := range img.DataSyms {
		for _, r := range dataSyms[i].relocs {
			reloc := r
			if _, ok := img.Symbols[reloc.Name]; !ok {
				if _, ok := textOff[reloc.Name]; !ok {
					reloc.External = true
					externals[reloc.Name] = true
				}
			}
			img.DataSyms[i].Relocs = append(img.DataSyms[i].Relocs, reloc)
		}
	}
	for name := range externals {
		img.Externals = append(img.Externals, name)
	}
	sort.Strings(img.Externals)
	return img, nil
}

// AssembleFileRISCV assembles every TEXT function of a parsed RISC-V file
// and lays out its static symbols (GLOBL/DATA) in a data section behind the
// code.  SB references in the code are encoded as AUIPC pairs with zero
// immediates; the object-file emitters record relocations for the linker.
func AssembleFileRISCV(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}
	// The pooled $i64 constants the wide MOV immediate loads refer to join
	// the declared data as read-only symbols, deduplicated across the file
	// (the toolchain synthesises the same symbols into its rodata).
	litSeen := map[string]bool{}

	img := &Image{Symbols: map[string]int{}, SourcePath: f.Path}
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, labels, relocs, lines, spadj, lits, err := assembleRISCV(t)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		for _, lit := range lits {
			if litSeen[lit.Name] {
				continue
			}
			litSeen[lit.Name] = true
			dataSyms = append(dataSyms, dataSym{
				name:   lit.Name,
				buf:    lit.Data,
				size:   len(lit.Data),
				rodata: true,
				dupok:  true,
			})
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
			Spadj:  spadj,
			Relocs: relocs,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
	}

	// Lay out the data section behind the code, 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		pos := dataStart + len(img.Data)
		for pos%16 != 0 {
			img.Data = append(img.Data, 0)
			pos++
		}
		img.Symbols[d.name] = pos
		img.Data = append(img.Data, d.buf...)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data) - len(d.buf), // relative to the data section
			Size:   d.size,
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
	}

	markExternals(img, dataSyms)
	return img, nil
}

// AssembleFileLOONG64 assembles every TEXT function of a parsed loong64 file
// and lays out its static symbols (GLOBL/DATA) in a data section behind the
// code.  SB references in the code are encoded as pcalau12i pairs with zero
// immediates; the object-file emitters record R_LOONG64_ADDR_HI/LO
// relocations for the linker.
func AssembleFileLOONG64(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}

	img := &Image{Symbols: map[string]int{}, SourcePath: f.Path}
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, labels, relocs, lines, spadj, err := assembleLOONG64(t)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
			Spadj:  spadj,
			Relocs: relocs,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
	}

	// Lay out the data section behind the code, 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		pos := dataStart + len(img.Data)
		for pos%16 != 0 {
			img.Data = append(img.Data, 0)
			pos++
		}
		img.Symbols[d.name] = pos
		img.Data = append(img.Data, d.buf...)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data) - len(d.buf), // relative to the data section
			Size:   d.size,
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
	}

	markExternals(img, dataSyms)
	return img, nil
}

// markExternals identifies relocations that reference symbols not defined in
// the file (neither a GLOBL/DATA symbol nor a TEXT function) and records them
// as external.  The non-amd64 architectures emit relocations for every SB
// reference; this post-processing step distinguishes file-local from external.
func markExternals(img *Image, dataSyms []dataSym) {
	known := make(map[string]bool, len(dataSyms)+len(img.Funcs))
	for _, d := range dataSyms {
		known[d.name] = true
	}
	for _, fn := range img.Funcs {
		known[fn.Name] = true
	}
	externals := map[string]bool{}
	for i := range img.Funcs {
		for j := range img.Funcs[i].Relocs {
			r := &img.Funcs[i].Relocs[j]
			if !known[r.Name] {
				r.External = true
				externals[r.Name] = true
			}
		}
	}
	// The declared data symbols carry the file's own relocations (the
	// symbol-valued DATA fields); the layouts appended img.DataSyms in
	// dataSyms order, so the indexes line up.  The trailing entries (the
	// pooled arm64 literals) have no source relocations.
	for i := range img.DataSyms {
		if i >= len(dataSyms) {
			break
		}
		for _, r := range dataSyms[i].relocs {
			reloc := r
			if !known[reloc.Name] {
				reloc.External = true
				externals[reloc.Name] = true
			}
			img.DataSyms[i].Relocs = append(img.DataSyms[i].Relocs, reloc)
		}
	}
	for name := range externals {
		img.Externals = append(img.Externals, name)
	}
	sort.Strings(img.Externals)
}

// dataSym is one GLOBL symbol and its DATA initialiser.
type dataSym struct {
	name   string
	pkg    string
	buf    []byte
	size   int
	static bool
	rodata bool
	dupok  bool
	// relocs are the symbol-valued DATA fields, in declaration order; Off
	// is relative to the symbol's data start.
	relocs []Reloc
}

// collectData gathers the file's static symbols (GLOBL) and their initial
// contents (DATA) into byte buffers.  Two passes: the Plan 9 convention puts
// every DATA line before its symbol's GLOBL, so the symbols are registered
// before the initialisers are applied.
func collectData(f *ast.File) ([]dataSym, error) {
	index := map[string]int{}
	var syms []dataSym
	for _, d := range f.Decls {
		gd, ok := d.(*ast.Globl)
		if !ok {
			continue
		}
		if gd.Name == nil || gd.Name.Pseudo != "SB" {
			continue
		}
		name := gd.Name.Name
		if _, dup := index[name]; dup {
			return nil, fmt.Errorf("duplicate GLOBL %q", name)
		}
		size := 0
		if gd.Size != nil && gd.Size.Imm.HasVal {
			size = int(gd.Size.Imm.Val)
		}
		index[name] = len(syms)
		ds := dataSym{
			name:   name,
			pkg:    gd.Name.Pkg,
			buf:    make([]byte, size),
			size:   size,
			static: gd.Name.Static,
		}
		for _, f := range gd.Flags {
			switch f {
			case "RODATA":
				ds.rodata = true
			case "DUPOK":
				ds.dupok = true
			default:
				// Legacy numeric flag constants (runtime/textflag.h):
				// DUPOK is 2, RODATA is 8; combinations arrive as one
				// number (e.g. 10 = RODATA|DUPOK).
				if n, err := strconv.Atoi(f); err == nil {
					if n&2 != 0 {
						ds.dupok = true
					}
					if n&8 != 0 {
						ds.rodata = true
					}
				}
			}
		}
		syms = append(syms, ds)
	}
	for _, d := range f.Decls {
		dd, ok := d.(*ast.Data)
		if !ok {
			continue
		}
		if dd.Name == nil || dd.Name.Pseudo != "SB" {
			continue
		}
		i, ok := index[dd.Name.Name]
		if !ok {
			return nil, fmt.Errorf("DATA %q: no matching GLOBL", dd.Name.Name)
		}
		if dd.Value == nil {
			return nil, fmt.Errorf("DATA %q: missing value", dd.Name.Name)
		}
		w := dd.Width
		off := dd.Name.Offset
		buf := syms[i].buf
		if off < 0 || off+int64(w) > int64(len(buf)) {
			return nil, fmt.Errorf("DATA %q+%d/%d exceeds GLOBL size %d", dd.Name.Name, off, w, len(buf))
		}
		// A symbol value ("DATA s+0(SB)/8, $other(SB)", the rt0 spelling)
		// leaves the field zero and records a relocation against the named
		// symbol: the linker patches the absolute address at this data
		// offset.  The toolchain emits the same shape, an R_ADDR of the
		// DATA width with the value's offset as the addend, on every
		// architecture.
		if sym := dd.Value.Imm.Sym; !dd.Value.Imm.HasVal && sym != nil {
			syms[i].relocs = append(syms[i].relocs, Reloc{
				Off:    int(off),
				Name:   sym.Name,
				Addend: sym.Offset,
				Kind:   RelAddr,
				Siz:    uint8(w),
			})
			continue
		}
		// A string or rune value ("DATA s+0(SB)/20, $"text"") writes its
		// bytes into the field and leaves the rest zero, the toolchain's
		// WriteString: the declared width must hold every byte, and any
		// width is legal.
		if s := dd.Value.Imm.Str; s != "" && !dd.Value.Imm.HasVal {
			text, err := strconv.Unquote(s)
			if err != nil {
				return nil, fmt.Errorf("DATA %q: invalid string value %s", dd.Name.Name, s)
			}
			if len(text) > w {
				return nil, fmt.Errorf("DATA %q: string of %d bytes does not fit width %d", dd.Name.Name, len(text), w)
			}
			copy(buf[off:], text)
			continue
		}
		// A floating-point value stores its IEEE-754 bits: /4 the float32
		// rounding of the parsed double, /8 the full 64 bits, the
		// toolchain's WriteFloat32 and WriteFloat64.
		if f := dd.Value.Imm.Float; f != "" && !dd.Value.Imm.HasVal {
			num, err := strconv.ParseFloat(f, 64)
			if err != nil {
				return nil, fmt.Errorf("DATA %q: invalid floating-point value %q", dd.Name.Name, f)
			}
			if dd.Value.Imm.Neg {
				num = -num
			}
			var v uint64
			switch w {
			case 4:
				v = uint64(math.Float32bits(float32(num)))
			case 8:
				v = math.Float64bits(num)
			default:
				return nil, fmt.Errorf("DATA %q: invalid width %d for a float (want 4 or 8)", dd.Name.Name, w)
			}
			for j := range w {
				buf[off+int64(j)] = byte(v >> (8 * j))
			}
			continue
		}
		if !dd.Value.Imm.HasVal {
			return nil, fmt.Errorf("DATA %q: value must be an integer immediate or a symbol address", dd.Name.Name)
		}
		switch w {
		case 1, 2, 4, 8:
		default:
			return nil, fmt.Errorf("DATA %q: invalid width %d (want 1, 2, 4 or 8)", dd.Name.Name, w)
		}
		v := dd.Value.Imm.Val
		if dd.Value.Imm.Neg {
			v = -v
		}
		for j := range w {
			buf[off+int64(j)] = byte(v >> (8 * j))
		}
	}
	return syms, nil
}

// align16 rounds n up to the next multiple of 16.
func align16(n int) int {
	return (n + 15) &^ 15
}

// frameSize returns the local frame size declared on the TEXT directive.
func frameSize(t *ast.Text) int {
	if t.Frame != nil && t.Frame.Imm.HasVal {
		return int(t.Frame.Imm.Val)
	}
	return 0
}

// argsSize returns the argument/result area declared on the TEXT directive.
func argsSize(t *ast.Text) int {
	if t.Args != nil && t.Args.Imm.HasVal {
		return int(t.Args.Imm.Val)
	}
	return 0
}