gasm-sdk/asm/riscv_assemble.go

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

package asm

import (
	"fmt"

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

// 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 := len(prologue)
	for _, stmt := range t.Body {
		switch s := stmt.(type) {
		case *ast.Label:
			offsets[s.Name.Text] = pos
		case *ast.Instr:
			instrs = append(instrs, s)
			pos += 4 // all RISC-V instructions are 4 bytes
		}
	}

	// Second pass: encode instructions.
	out := append([]byte(nil), prologue...)
	for _, instr := range instrs {
		code, err := encodeRISCVInstr(instr, pos, offsets, fi)
		if err != nil {
			return nil, nil, fmt.Errorf("%s: %w", instr.Mnemonic.Text, err)
		}
		out = append(out, code...)
		pos -= 4 // track remaining
	}
	return out, offsets, nil
}

// encodeRISCVInstr encodes a single RISC-V instruction.
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)
	}

	switch {
	// R-type: Plan 9 order is INSTR src1, src2, dst (destination last).
	case len(ops) == 3 && isRTypeInstr(mnem):
		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: Plan 9 order is INSTR src, imm, dst.
	case len(ops) == 3 && isITypeInstr(mnem):
		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) — Plan 9 order is LD src, dst.
	case len(ops) == 2 && isLoadInstr(mnem):
		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: Plan 9 order is SD src, dst (src=register, dst=memory).
	case len(ops) == 2 && isStoreInstr(mnem):
		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)
		}
		word = riscvSType(enc, rs1, rs2, imm)

	// Branches: rs1, rs2, label.
	case len(ops) == 3 && isBranchInstr(mnem):
		rs1 := regFromOperand(ops[0])
		rs2 := regFromOperand(ops[1])
		target := labelFromOperand(ops[2])
		targetOff, ok := offsets[target]
		if !ok {
			return nil, fmt.Errorf("undefined label %q", target)
		}
		rs1Off := 0 // placeholder
		_ = rs1Off
		offset := int32(targetOff - pc)
		if rs1 < 0 || rs2 < 0 {
			return nil, fmt.Errorf("invalid register in %s", mnem)
		}
		word = riscvBType(enc, rs1, rs2, offset)

	// U-type: rd, imm.
	case len(ops) == 2 && isUTypeInstr(mnem):
		rd := regFromOperand(ops[0])
		imm := immFromOperand(ops[1])
		if rd < 0 {
			return nil, fmt.Errorf("invalid register in %s", mnem)
		}
		word = riscvUType(enc, rd, imm)

	default:
		return nil, fmt.Errorf("cannot encode %s with %d operands", mnem, len(ops))
	}

	// Emit as little-endian 32-bit word.
	return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}, nil
}

// Instruction type classifiers.
func isRTypeInstr(m string) bool {
	switch m {
	case "ADD", "SUB", "SLL", "SLT", "SLTU", "XOR", "SRL", "SRA", "OR", "AND",
		"ADDW", "SUBW", "SLLW", "SRLW", "SRAW",
		"MUL", "MULH", "MULHSU", "MULHU", "DIV", "DIVU", "REM", "REMU",
		"MULW", "DIVW", "DIVUW", "REMW", "REMUW":
		return true
	}
	return false
}

func isITypeInstr(m string) bool {
	switch m {
	case "ADDI", "ADDIW", "SLTI", "SLTIU", "XORI", "ORI", "ANDI", "JALR":
		return true
	}
	return false
}

func isLoadInstr(m string) bool {
	switch m {
	case "LB", "LH", "LW", "LD", "LBU", "LHU", "LWU":
		return true
	}
	return false
}

func isStoreInstr(m string) bool {
	switch m {
	case "SB", "SH", "SW", "SD":
		return true
	}
	return false
}

func isBranchInstr(m string) bool {
	switch m {
	case "BEQ", "BNE", "BLT", "BGE", "BLTU", "BGEU":
		return true
	}
	return false
}

func isUTypeInstr(m string) bool {
	return m == "LUI" || m == "AUIPC"
}

// Operand helpers.
func regFromOperand(op *ast.Operand) int {
	// Register is in Addr.Base (from (base) syntax) or Addr.Sym.Name (bare ident).
	if op.Addr.Base != "" {
		return riscvRegNum(op.Addr.Base)
	}
	if op.Addr.Sym != nil && op.Addr.Sym.Name != "" {
		return riscvRegNum(op.Addr.Sym.Name)
	}
	return -1
}

func immFromOperand(op *ast.Operand) int32 {
	if op.Imm.HasVal {
		v := op.Imm.Val
		if op.Imm.Neg {
			v = -v
		}
		return int32(v)
	}
	return 0
}

func memFromOperand(op *ast.Operand) (rs1 int, imm int32) {
	rs1 = riscvRegNum(op.Addr.Base)
	imm = int32(op.Addr.Offset)
	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
	}
	return op.Raw
}