feat(riscv): add prologue, epilogue and frame pseudo-register support

Assisted-by: Kimi K3
This commit is contained in:
2026-08-02 00:18:00 +02:00
parent f860bf8ce6
commit a2acc9b5a3
5 changed files with 253 additions and 53 deletions
+38
View File
@@ -220,6 +220,44 @@ func AssembleFile(f *ast.File) (*Image, error) {
return img, nil
}
// AssembleFileRISCV assembles every TEXT function of a parsed RISC-V file.
// It produces an Image with the function bodies laid out in source order.
func AssembleFileRISCV(f *ast.File) (*Image, error) {
img := &Image{Symbols: map[string]int{}}
for _, d := range f.Decls {
t, ok := d.(*ast.Text)
if !ok {
continue
}
code, labels, err := assembleRISCV(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,
}
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...)
}
return img, nil
}
// dataSym is one GLOBL symbol and its DATA initialiser.
type dataSym struct {
name string
+77 -50
View File
@@ -12,10 +12,14 @@ import (
// assembleRISCV assembles a RISC-V TEXT function body into machine code.
// It handles the core RV64I/RV64M instruction set with fixed 32-bit encoding.
func assembleRISCV(t *ast.Text) ([]byte, map[string]int, error) {
fi := riscvComputeFrame(t)
prologue := riscvPrologue(fi)
// First pass: compute label offsets (all RISC-V instructions are 4 bytes).
// Labels are offset by the prologue length.
offsets := map[string]int{}
var instrs []*ast.Instr
pos := 0
pos := len(prologue)
for _, stmt := range t.Body {
switch s := stmt.(type) {
case *ast.Label:
@@ -27,9 +31,9 @@ func assembleRISCV(t *ast.Text) ([]byte, map[string]int, error) {
}
// Second pass: encode instructions.
out := make([]byte, 0, pos)
out := append([]byte(nil), prologue...)
for _, instr := range instrs {
code, err := encodeRISCVInstr(instr, pos, offsets)
code, err := encodeRISCVInstr(instr, pos, offsets, fi)
if err != nil {
return nil, nil, fmt.Errorf("%s: %w", instr.Mnemonic.Text, err)
}
@@ -40,50 +44,91 @@ func assembleRISCV(t *ast.Text) ([]byte, map[string]int, error) {
}
// encodeRISCVInstr encodes a single RISC-V instruction.
func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int) ([]byte, error) {
func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int, fi riscvFrameInfo) ([]byte, error) {
mnem := instr.Mnemonic.Text
ops := instr.Operands
var word uint32
// Handle pseudo-instructions and special cases first.
switch mnem {
case "RET":
// RET = JALR X0, 0(X1)
word = riscvIType(riscvEnc{0x67, 0x0, 0x00}, 0, 1, 0)
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}, nil
case "CALL":
// CALL is a pseudo-instruction; encode as NOP placeholder.
word = riscvIType(riscvEnc{0x13, 0x0, 0x00}, 0, 0, 0)
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}, nil
case "JMP":
// JMP = JAL X0, target
var target string
if len(ops) >= 1 {
target = labelFromOperand(ops[0])
}
targetOff, ok := offsets[target]
if !ok {
return nil, fmt.Errorf("undefined label %q", target)
}
offset := int32(targetOff - pc)
word = riscvJType(0, offset)
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}, nil
case "JAL":
rd := 0
var target string
if len(ops) >= 2 {
rd = regFromOperand(ops[0])
target = labelFromOperand(ops[1])
} else if len(ops) == 1 {
target = labelFromOperand(ops[0])
}
targetOff, ok := offsets[target]
if !ok {
return nil, fmt.Errorf("undefined label %q", target)
}
offset := int32(targetOff - pc)
word = riscvJType(rd, offset)
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}, nil
}
enc, ok := riscvInstrTable[mnem]
if !ok {
return nil, fmt.Errorf("unsupported RISC-V instruction %q", mnem)
}
ops := instr.Operands
var word uint32
switch {
// R-type: rd, rs1, rs2 (3 register operands).
// R-type: Plan 9 order is INSTR src1, src2, dst (destination last).
case len(ops) == 3 && isRTypeInstr(mnem):
rd := regFromOperand(ops[0])
rs1 := regFromOperand(ops[1])
rs2 := regFromOperand(ops[2])
rs1 := regFromOperand(ops[0]) // source 1 (first operand)
rs2 := regFromOperand(ops[1]) // source 2 (second operand)
rd := regFromOperand(ops[2]) // destination (last operand)
if rd < 0 || rs1 < 0 || rs2 < 0 {
return nil, fmt.Errorf("invalid register in %s", mnem)
}
word = riscvRType(enc, rd, rs1, rs2)
// I-type with immediate: rd, rs1, imm.
// I-type with immediate: Plan 9 order is INSTR src, imm, dst.
case len(ops) == 3 && isITypeInstr(mnem):
rd := regFromOperand(ops[0])
rs1 := regFromOperand(ops[1])
imm := immFromOperand(ops[2])
rs1 := regFromOperand(ops[0]) // source register
imm := immFromOperand(ops[1]) // immediate
rd := regFromOperand(ops[2]) // destination
if rd < 0 || rs1 < 0 {
return nil, fmt.Errorf("invalid register in %s", mnem)
}
word = riscvIType(enc, rd, rs1, imm)
// Loads: rd, offset(rs1) — parsed as rd, (rs1)+offset.
// Loads: rd, offset(rs1) — Plan 9 order is LD src, dst.
case len(ops) == 2 && isLoadInstr(mnem):
rd := regFromOperand(ops[0])
rs1, imm := memFromOperand(ops[1])
rd := regFromOperand(ops[1]) // destination (last operand)
rs1, imm := memFromOperandWithFrame(ops[0], fi) // memory source (first operand)
if rd < 0 || rs1 < 0 {
return nil, fmt.Errorf("invalid operand in %s", mnem)
}
word = riscvIType(enc, rd, rs1, imm)
// Stores: rs2, offset(rs1).
// Stores: Plan 9 order is SD src, dst (src=register, dst=memory).
case len(ops) == 2 && isStoreInstr(mnem):
rs2 := regFromOperand(ops[0])
rs1, imm := memFromOperand(ops[1])
rs2 := regFromOperand(ops[0]) // source register (first operand)
rs1, imm := memFromOperandWithFrame(ops[1], fi) // memory dest (last operand)
if rs2 < 0 || rs1 < 0 {
return nil, fmt.Errorf("invalid operand in %s", mnem)
}
@@ -115,35 +160,6 @@ func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int) ([]byte,
}
word = riscvUType(enc, rd, imm)
// JAL: rd, label (or just label for JMP).
case mnem == "JAL" || mnem == "JMP":
rd := 0 // JMP uses rd=0 (X0)
if len(ops) >= 1 && mnem == "JAL" {
rd = regFromOperand(ops[0])
}
var target string
if len(ops) >= 2 {
target = labelFromOperand(ops[len(ops)-1])
} else if len(ops) == 1 && mnem == "JMP" {
target = labelFromOperand(ops[0])
}
targetOff, ok := offsets[target]
if !ok {
return nil, fmt.Errorf("undefined label %q", target)
}
offset := int32(targetOff - pc)
word = riscvJType(rd, offset)
// RET: JALR X0, 0(X1).
case mnem == "RET":
word = riscvIType(riscvEnc{0x67, 0x0, 0x00}, 0, 1, 0)
// CALL: AUIPC + JALR sequence (simplified: just encode as NOP placeholder).
case mnem == "CALL":
// CALL is a pseudo-instruction that expands to AUIPC + JALR.
// For now, encode as a NOP (ADDI X0, X0, 0).
word = riscvIType(riscvEnc{0x13, 0x0, 0x00}, 0, 0, 0)
default:
return nil, fmt.Errorf("cannot encode %s with %d operands", mnem, len(ops))
}
@@ -229,6 +245,17 @@ func memFromOperand(op *ast.Operand) (rs1 int, imm int32) {
return
}
// memFromOperandWithFrame resolves a memory operand, handling FP/SP
// pseudo-registers via the frame mapping.
func memFromOperandWithFrame(op *ast.Operand, fi riscvFrameInfo) (rs1 int, imm int32) {
// Check for a pseudo-register reference (name+offset(FP) or name+offset(SP)).
if op.Addr.Sym != nil && op.Addr.Sym.Pseudo != "" {
return riscvResolvePseudo(op.Addr.Sym, fi)
}
// Plain register+offset memory reference.
return memFromOperand(op)
}
func labelFromOperand(op *ast.Operand) string {
if op.Addr.Sym != nil {
return op.Addr.Sym.Name
+117
View File
@@ -0,0 +1,117 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
package asm
import "sourcedock.dev/petrbalvin/gasm-devkit/ast"
// RISC-V frame mapping: translates Go's FP/SP pseudo-register addressing
// into real RISC-V memory accesses.
//
// In Go's ABI0 (used by assembly functions), arguments are passed on the
// stack. At function entry the return address sits at SP, so the frame
// pointer FP == SP+8 and the first argument is at FP+0 == SP+8.
//
// On RISC-V the hardware registers are:
// SP = X2 (stack pointer)
// FP = S0 = X8 (frame pointer, by convention)
//
// For NOSPLIT $0 functions the prologue is omitted and arguments are read
// directly from SP+8+offset.
// riscvFrameInfo holds the frame parameters computed from a TEXT directive.
type riscvFrameInfo struct {
frameSize int // the $framesize from TEXT
argsSize int // the -argsize from TEXT
noSplit bool // the NOSPLIT flag
}
// riscvComputeFrame extracts frame information from a TEXT directive.
func riscvComputeFrame(t *ast.Text) riscvFrameInfo {
fi := riscvFrameInfo{}
fi.frameSize = frameSize(t)
fi.argsSize = argsSize(t)
for _, f := range t.Flags {
if f == "NOSPLIT" {
fi.noSplit = true
}
}
return fi
}
// riscvPrologue returns the prologue bytes for a RISC-V function.
// For NOSPLIT $0 functions there is no prologue. For functions with a
// frame, we emit: ADDI SP, SP, -framesize; SD S0, (framesize-8)(SP); ...
func riscvPrologue(fi riscvFrameInfo) []byte {
if fi.noSplit && fi.frameSize == 0 {
return nil // no prologue for NOSPLIT $0
}
var out []byte
if fi.frameSize > 0 {
// ADDI SP, SP, -framesize
out = append(out, riscvITypeLE(0x13, 0x0, 2, 2, int32(-fi.frameSize))...)
// Save the frame pointer (S0 = X8) at the top of the new frame.
// SD S0, (framesize-8)(SP)
out = append(out, riscvSTypeLE(0x23, 0x3, 2, 8, int32(fi.frameSize-8))...)
}
return out
}
// riscvEpilogue returns the epilogue bytes for a RISC-V function.
func riscvEpilogue(fi riscvFrameInfo) []byte {
if fi.noSplit && fi.frameSize == 0 {
return nil
}
var out []byte
if fi.frameSize > 0 {
// Restore the frame pointer: LD S0, (framesize-8)(SP)
out = append(out, riscvITypeLE(0x03, 0x3, 8, 2, int32(fi.frameSize-8))...)
// ADDI SP, SP, framesize
out = append(out, riscvITypeLE(0x13, 0x0, 2, 2, int32(fi.frameSize))...)
}
return out
}
// riscvResolvePseudo translates a pseudo-register memory reference into a
// real base register and offset. It handles name+offset(FP) and
// name+offset(SP).
//
// Returns the base register number and the adjusted offset.
func riscvResolvePseudo(sym *ast.Symbol, fi riscvFrameInfo) (base int, off int32) {
if sym == nil {
return -1, 0
}
offset := int32(sym.Offset)
switch sym.Pseudo {
case "FP":
// FP == SP+8 for NOSPLIT $0; arguments are at SP+8+offset.
if fi.noSplit && fi.frameSize == 0 {
return 2, 8 + offset // SP + 8 + argOffset
}
// With a frame, FP points to the saved frame; args are at FP+offset.
return 8, offset // S0 + argOffset
case "SP":
// SP-relative; the offset is from the current SP.
return 2, offset
case "SB":
// Static data reference — needs a relocation (not yet supported).
return -1, offset
default:
return -1, offset
}
}
// riscvITypeLE encodes an I-type instruction and returns little-endian bytes.
func riscvITypeLE(opcode, funct3 uint32, rd, rs1 int, imm int32) []byte {
word := (uint32(imm&0xFFF) << 20) | (uint32(rs1) << 15) |
(funct3 << 12) | (uint32(rd) << 7) | opcode
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}
}
// riscvSTypeLE encodes an S-type instruction and returns little-endian bytes.
func riscvSTypeLE(opcode, funct3 uint32, rs1, rs2 int, imm int32) []byte {
immU := uint32(imm) & 0xFFF
word := ((immU >> 5) << 25) | (uint32(rs2) << 20) | (uint32(rs1) << 15) |
(funct3 << 12) | ((immU & 0x1F) << 7) | opcode
return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}
}