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))