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"
	"sort"

	"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
}

// 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
	Labels map[string]int // local labels, function-relative
	Relocs []Reloc        // static-symbol references, in emission order
}

// Reloc is one static-symbol reference within a function body: the disp32
// field at Off (function-relative) must reach the symbol plus Addend,
// measured from After, the address just past the instruction.  An External
// relocation names a symbol no GLOBL in the file defines; the object-file
// emitters carry it into the output's relocation table.
type Reloc struct {
	Off      int
	After    int
	Name     string
	Addend   int64
	External bool
}

// 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
}

// 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
	}
	link := &linkInfo{symbols: known, allowExternal: true}

	img := &Image{Symbols: 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, err := assemble(t, link)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		img.Funcs = append(img.Funcs, FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Labels: labels,
		})
		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,
		})
		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}
			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 {
				reloc.External = true
				externals[p.name] = true
			}
			img.Funcs[i].Relocs = append(img.Funcs[i].Relocs, reloc)
		}
	}
	for name := range externals {
		img.Externals = append(img.Externals, name)
	}
	sort.Strings(img.Externals)
	return img, nil
}

// dataSym is one GLOBL symbol and its DATA initialiser.
type dataSym struct {
	name   string
	pkg    string
	buf    []byte
	static bool
}

// collectData gathers the file's static symbols (GLOBL) and their initial
// contents (DATA) into byte buffers, in declaration order.
func collectData(f *ast.File) ([]dataSym, error) {
	index := map[string]int{}
	var syms []dataSym
	for _, d := range f.Decls {
		switch dd := d.(type) {
		case *ast.Globl:
			if dd.Name == nil || dd.Name.Pseudo != "SB" {
				continue
			}
			name := dd.Name.Name
			if _, dup := index[name]; dup {
				return nil, fmt.Errorf("duplicate GLOBL %q", name)
			}
			size := 0
			if dd.Size != nil && dd.Size.Imm.HasVal {
				size = int(dd.Size.Imm.Val)
			}
			index[name] = len(syms)
			syms = append(syms, dataSym{
				name:   name,
				pkg:    dd.Name.Pkg,
				buf:    make([]byte, size),
				static: dd.Name.Static,
			})

		case *ast.Data:
			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 || !dd.Value.Imm.HasVal {
				return nil, fmt.Errorf("DATA %q: value must be an integer immediate", dd.Name.Name)
			}
			w := dd.Width
			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)
			}
			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))
			}
			v := dd.Value.Imm.Val
			if dd.Value.Imm.Neg {
				v = -v
			}
			for j := 0; j < w; j++ {
				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
}