feat(asm): expand riscv64 memory offsets beyond the 12-bit immediate

Assisted-by: GLM 5.3 Flash

asm/riscv_assemble.go

@@ -475,6 +475,11 @@ func riscvInstrSize(instr *ast.Instr, fi riscvFrameInfo, tlsSyms map[string]bool
 			return 4
 		}
 	}
+	// Plain loads and stores whose offset leaves the signed 12-bit span
+	// expand to the X31 materialisation plus the access word.
+	if n, ok := riscvMemInstrSize(mnem, ops, fi); ok {
+		return n
+	}
 	// I-type arithmetic with a large immediate expands to several instructions.
 	if (mnem == "ADDI" || mnem == "ANDI" || mnem == "ORI" || mnem == "XORI") && len(ops) >= 1 && isImmOperand(ops[0]) {
 		imm := immFromOperand(ops[0])
@@ -1822,11 +1827,11 @@ func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int, fi riscv
 		if err != nil {
 			return nil, err
 		}
-		rs1, imm := memFromOperandWithFrame(ops[0], fi)
-		if rs1 < 0 {
-			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
+		rs1, imm, err := riscvAccessMem(mnem, ops[0], fi)
+		if err != nil {
+			return nil, err
 		}
-		word = riscvIType(enc, rd, rs1, imm)
+		return riscvFrameMemOp(enc, false, rd, rs1, imm), nil
 
 	// FP stores: INSTR freg, addr (Plan 9: source first).
 	case len(ops) == 2 && isFPStoreInstr(mnem):
@@ -1834,11 +1839,11 @@ func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int, fi riscv
 		if err != nil {
 			return nil, err
 		}
-		rs1, imm := memFromOperandWithFrame(ops[1], fi)
-		if rs1 < 0 {
-			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
+		rs1, imm, err := riscvAccessMem(mnem, ops[1], fi)
+		if err != nil {
+			return nil, err
 		}
-		word = riscvSType(enc, rs1, rs2, imm)
+		return riscvFrameMemOp(enc, true, rs2, rs1, imm), nil
 
 	// LR (load-reserved): INSTR (addr), dst.  The toolchain reads the
 	// operands positionally, so the base register comes from the first
@@ -1920,11 +1925,11 @@ func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int, fi riscv
 		if err != nil {
 			return nil, err
 		}
-		rs1, imm := memFromOperandWithFrame(ops[0], fi) // memory source (first operand)
-		if rs1 < 0 {
-			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
+		rs1, imm, err := riscvAccessMem(mnem, ops[0], fi) // memory source (first operand)
+		if err != nil {
+			return nil, err
 		}
-		word = riscvIType(enc, rd, rs1, imm)
+		return riscvFrameMemOp(enc, false, rd, rs1, imm), nil
 
 	// Stores: Plan 9 order is SD src, dst (src=register, dst=memory).
 	case len(ops) == 2 && isStoreInstr(mnem):
@@ -1932,14 +1937,11 @@ func encodeRISCVInstr(instr *ast.Instr, pc int, offsets map[string]int, fi riscv
 		if err != nil {
 			return nil, err
 		}
-		rs1, imm := memFromOperandWithFrame(ops[1], fi) // memory dest (last operand)
-		if rs1 < 0 {
-			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
+		rs1, imm, err := riscvAccessMem(mnem, ops[1], fi) // memory dest (last operand)
+		if err != nil {
+			return nil, err
 		}
-		if off := int64(ops[1].Addr.Offset); ops[1].Addr.Sym == nil && (off < math.MinInt32 || off > math.MaxInt32) {
-			return nil, fmt.Errorf("%s: constant %d too large", mnem, off)
-		}
-		word = riscvSType(enc, rs1, rs2, imm)
+		return riscvFrameMemOp(enc, true, rs2, rs1, imm), nil
 
 	// Branches: rs1, rs2, label.  BGT/BLE/BGTU/BLEU are the swapped-spelling
 	// forms of BLT/BGE/BLTU/BGEU (bgt rs1, rs2 is blt rs2, rs1).
@@ -2211,6 +2213,9 @@ func encodeRISCVMov(instr *ast.Instr, fi riscvFrameInfo, relocs *[]Reloc, lits *
 		if err := riscvWantMemBase(mnem, "rs1", src); err != nil {
 			return nil, err
 		}
+		if err := riscvWantMemOffset(mnem, src, fi); err != nil {
+			return nil, err
+		}
 		rs1, off := memFromOperandWithFrame(src, fi)
 		if rs1 < 0 {
 			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
@@ -2240,6 +2245,9 @@ func encodeRISCVMov(instr *ast.Instr, fi riscvFrameInfo, relocs *[]Reloc, lits *
 		if err := riscvWantMemBase(mnem, "rs1", dst); err != nil {
 			return nil, err
 		}
+		if err := riscvWantMemOffset(mnem, dst, fi); err != nil {
+			return nil, err
+		}
 		rs1, off := memFromOperandWithFrame(dst, fi)
 		if rs1 < 0 {
 			return nil, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
@@ -2421,6 +2429,75 @@ func riscvFrameMemSize(op *ast.Operand, fi riscvFrameInfo) int {
 	return len(riscvAddressInX31WithBase(off, rs1)) + 4
 }
 
+// riscvMemInstrSize returns the encoded size of a plain load or store
+// (integer and FP widths, one register end and one memory end) for the layout
+// pass: 4 bytes when the offset fits the signed 12-bit span, otherwise the
+// X31 materialisation plus the access, exactly what riscvFrameMemOp emits
+// for them.  A constant beyond the signed 32-bit span encodes never: the 4
+// bytes guessed here are never emitted, because the encode pass rejects the
+// instruction with the toolchain's "constant too large" diagnostic first.
+func riscvMemInstrSize(mnem string, ops []*ast.Operand, fi riscvFrameInfo) (int, bool) {
+	var mem *ast.Operand
+	switch {
+	case len(ops) == 2 && (isLoadInstr(mnem) || isFPLoadInstr(mnem)):
+		mem = ops[0]
+	case len(ops) == 2 && (isStoreInstr(mnem) || isFPStoreInstr(mnem)):
+		mem = ops[1]
+	default:
+		return 0, false
+	}
+	if err := riscvWantMemOffset(mnem, mem, fi); err != nil {
+		return 4, true
+	}
+	rs1, off := memFromOperandWithFrame(mem, fi)
+	if rs1 < 0 || fits12(off) {
+		return 4, true
+	}
+	return len(riscvAddressInX31WithBase(off, rs1)) + 4, true
+}
+
+// riscvWantMemOffset rejects a memory operand whose byte offset leaves the
+// signed 32-bit span, the point where the toolchain's Split32BitImmediate
+// stops and reports "constant %d too large".  The frame pseudo-registers
+// resolve against the frame first, so their offsets are checked resolved,
+// exactly as the toolchain's stackOffset feeds Split32BitImmediate the
+// adjusted address; an SB reference rides the relocation paths and is never
+// this check's subject.
+func riscvWantMemOffset(mnem string, op *ast.Operand, fi riscvFrameInfo) error {
+	var off int64
+	switch {
+	case op.Addr.Sym != nil && op.Addr.Sym.Pseudo == "FP":
+		off = op.Addr.Offset + int64(fi.autosize) + 8
+	case op.Addr.Sym != nil && op.Addr.Sym.Pseudo == "SP":
+		off = op.Addr.Offset + int64(fi.autosize)
+	case op.Addr.Sym != nil:
+		return nil
+	default:
+		off = op.Addr.Offset
+	}
+	if off < math.MinInt32 || off > math.MaxInt32 {
+		return fmt.Errorf("%s: constant %d too large", mnem, off)
+	}
+	return nil
+}
+
+// riscvAccessMem resolves the memory end of a plain load or store: the base
+// register and the byte offset, range-checked.  The access itself is emitted
+// by riscvFrameMemOp, whose hi/lo split matches the toolchain's
+// instructionsForLoad and instructionsForStore: the high part materialises in
+// the assembler's temporary register (X31) and the access reads or writes
+// through it.
+func riscvAccessMem(mnem string, op *ast.Operand, fi riscvFrameInfo) (int, int32, error) {
+	if err := riscvWantMemOffset(mnem, op, fi); err != nil {
+		return -1, 0, err
+	}
+	rs1, off := memFromOperandWithFrame(op, fi)
+	if rs1 < 0 {
+		return -1, 0, fmt.Errorf("%s: expected integer register in rs1 position", mnem)
+	}
+	return rs1, off, nil
+}
+
 // encodeRISCVLoadImm encodes loading an immediate into a register (MOV $imm,
 // rd), matching the toolchain's instructionsForMOVConst.  For 12-bit
 // immediates it emits ADDI $imm, ZERO, rd (compressed to C.LI when it fits
@@ -2946,6 +3023,15 @@ func word16(w uint16) []byte {
 // form), or JALR offset(rs1) (memory → rd=X1).
 func encodeRISCVJALR(instr *ast.Instr, fi riscvFrameInfo) ([]byte, error) {
 	ops := instr.Operands
+	// The toolchain's I-type validation bounds the JALR displacement to the
+	// signed 12-bit span and rejects the rest; a wider one would truncate
+	// silently into a jump somewhere else entirely.
+	checkImm := func(imm int32) error {
+		if imm < -2048 || imm > 2047 {
+			return fmt.Errorf("JALR: signed immediate %d must be in range [-2048, 2047] (12 bits)", imm)
+		}
+		return nil
+	}
 	// JALR rd, offset(rs1): the memory operand's base is the jump-target
 	// register, not the destination.
 	if len(ops) == 2 && isMemOperand(ops[1]) {
@@ -2954,6 +3040,9 @@ func encodeRISCVJALR(instr *ast.Instr, fi riscvFrameInfo) ([]byte, error) {
 		if rd < 0 || rs1 < 0 {
 			return nil, fmt.Errorf("JALR: invalid register operand")
 		}
+		if err := checkImm(imm); err != nil {
+			return nil, err
+		}
 		return wordLE(riscvIType(riscvEnc{0x67, 0x0, 0x00}, rd, rs1, imm)), nil
 	}
 	if len(ops) == 2 {
@@ -2969,6 +3058,9 @@ func encodeRISCVJALR(instr *ast.Instr, fi riscvFrameInfo) ([]byte, error) {
 		if rs1 < 0 {
 			return nil, fmt.Errorf("JALR: invalid memory operand")
 		}
+		if err := checkImm(imm); err != nil {
+			return nil, err
+		}
 		return wordLE(riscvIType(riscvEnc{0x67, 0x0, 0x00}, 1, rs1, imm)), nil
 	}
 	return nil, fmt.Errorf("JALR expects 1 or 2 operands, got %d", len(ops))

asm/riscv_memoffset_test.go

@@ -0,0 +1,128 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: BSD-3-Clause
+
+package asm
+
+import (
+	"strings"
+	"testing"
+)
+
+// TestRISCVBigMemOffset_Differential proves the memory-offset expansion
+// against the oracle: a load or store whose offset leaves the signed 12-bit
+// span materialises the high part in the assembler's temporary register
+// (C.LUI, or LUI beyond its six-bit immediate) followed by C.ADD of the base,
+// and accesses through it, exactly as the toolchain's instructionsForLoad and
+// instructionsForStore synthesise.  The battery spans both sides of the
+// 12-bit edge, both int32 boundaries, the compressed and uncompressed LUI
+// halves, every width family and the MOV spellings, under a prologue too so
+// the size accounting keeps the layout honest.
+func TestRISCVBigMemOffset_Differential(t *testing.T) {
+	src := `#include "textflag.h"
+
+TEXT ·bigmem(SB), NOSPLIT, $16
+	LD 2047(X6), X5			// single word: the edge that still fits
+	LD 2048(X6), X5			// hi 1, lo -2048
+	LD 4096(X6), X5			// hi 1, lo 0
+	SD X5, 8192(X6)			// store side of the same split
+	FLD 4096(X6), F5		// FP widths share the expansion
+	FSD F5, 8192(X6)
+	MOVB 4097(X6), X8		// the MOV widths route through the same load
+	MOVHU 4098(X6), X9
+	MOV X7, 16384(X6)		// and the same store
+	LD 1048576(X6), X10		// hi 256: LUI, not C.LUI
+	LD -4097(X6), X11		// hi -1, lo -1
+	SD X11, -8192(X6)
+	LD 2147483647(X6), X12	// int32 upper boundary
+	LD -2147483648(X6), X13	// int32 lower boundary
+	RET
+`
+	path := writeRISCVSrc(t, "bigmem_riscv64.s", src)
+	assertRISCVDifferential(t, path, src, "bigmem")
+}
+
+// TestRISCVBigMemOffsetFrame_Differential proves the expansion under the
+// stack-split guard: a function that opens a frame and calls out gets the
+// guard, the prologue and the epilogue ahead of the body, so every body
+// offset now depends on the expansion's size accounting having kept the
+// layout in step with the emitted bytes.
+func TestRISCVBigMemOffsetFrame_Differential(t *testing.T) {
+	src := `#include "textflag.h"
+
+TEXT ·bigframe(SB), $16
+	CALL extcal(SB)
+	LD 4096(X6), X5
+	SD X5, 4096(X6)
+	FLD 8192(X6), F5
+	FSD F5, 8192(X6)
+	RET
+`
+	path := writeRISCVSrc(t, "bigframe_riscv64.s", src)
+	assertRISCVDifferential(t, path, src, "bigframe")
+}
+
+// TestRISCVBigMemOffsetRejections pins the honest rejections the expansion
+// cannot carry: a constant beyond the signed 32-bit span is the toolchain's
+// "constant too large" (its Split32BitImmediate has no wider split), and a
+// JALR displacement beyond the signed 12-bit span is its I-type range check,
+// because the jump would land somewhere else entirely.
+func TestRISCVBigMemOffsetRejections(t *testing.T) {
+	cases := []struct {
+		name string
+		src  string
+		want string
+	}{
+		{
+			name: "store beyond int32",
+			src:  "\tSD X5, 4294967295(X6)\n",
+			want: "SD: constant 4294967295 too large",
+		},
+		{
+			name: "load below int32",
+			src:  "\tLD -2147483649(X6), X5\n",
+			want: "LD: constant -2147483649 too large",
+		},
+		{
+			name: "MOV load beyond int32",
+			src:  "\tMOV 4294967296(X6), X5\n",
+			want: "MOV: constant 4294967296 too large",
+		},
+		{
+			name: "MOV store beyond int32",
+			src:  "\tMOV X5, 4294967296(X6)\n",
+			want: "MOV: constant 4294967296 too large",
+		},
+		{
+			name: "FP store beyond int32",
+			src:  "\tFSD F5, 4294967296(X6)\n",
+			want: "FSD: constant 4294967296 too large",
+		},
+		{
+			name: "FP load beyond int32",
+			src:  "\tFLD 4294967296(X6), F5\n",
+			want: "FLD: constant 4294967296 too large",
+		},
+		{
+			name: "JALR displacement above 12 bits",
+			src:  "\tJALR 4096(X7)\n",
+			want: "JALR: signed immediate 4096 must be in range [-2048, 2047] (12 bits)",
+		},
+		{
+			name: "JALR displacement below 12 bits",
+			src:  "\tJALR X5, -2049(X7)\n",
+			want: "JALR: signed immediate -2049 must be in range [-2048, 2047] (12 bits)",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			fn := firstTextRISCV(t, "#include \"textflag.h\"\n\nTEXT ·r(SB), NOSPLIT, $0\n"+tc.src+"\tRET\n")
+			_, _, _, _, _, _, err := assembleRISCV(fn, nil)
+			if err == nil {
+				t.Fatalf("source assembled, want rejection %q", tc.want)
+			}
+			if !strings.Contains(err.Error(), tc.want) {
+				t.Errorf("error %q does not carry %q", err.Error(), tc.want)
+			}
+		})
+	}
+}