cmd/compile/internal, cmd/internal: implement jump table on riscv64

src/cmd/compile/internal/riscv64/ssa.go

@@ -1042,6 +1042,46 @@ func ssaGenBlock(s *ssagen.State, b, next *ssa.Block) {
 			p.From.Reg = b.Controls[0].Reg()
 		}
 
+	case ssa.BlockRISCV64JUMPTABLE:
+		// Jump table:
+		// TMP = base + index*8 (SH3ADD if Zba else SLLI+ADD).
+		// Load slot into TMP, then indirect JMP through TMP.
+		var p *obj.Prog
+		if buildcfg.GORISCV64 >= 22 {
+			p = s.Prog(riscv.ASH3ADD)
+			p.From.Type = obj.TYPE_REG
+			p.From.Reg = b.Controls[1].Reg()
+			p.Reg = b.Controls[0].Reg()
+			p.To.Type = obj.TYPE_REG
+			p.To.Reg = riscv.REG_TMP
+		} else {
+			p = s.Prog(riscv.ASLLI)
+			p.From.Type = obj.TYPE_CONST
+			p.From.Offset = 3
+			p.Reg = b.Controls[0].Reg()
+			p.To.Type = obj.TYPE_REG
+			p.To.Reg = riscv.REG_TMP
+
+			p = s.Prog(riscv.AADD)
+			p.From.Type = obj.TYPE_REG
+			p.From.Reg = riscv.REG_TMP
+			p.Reg = b.Controls[1].Reg()
+			p.To.Type = obj.TYPE_REG
+			p.To.Reg = riscv.REG_TMP
+		}
+
+		p = s.Prog(riscv.AMOV)
+		p.From.Type = obj.TYPE_MEM
+		p.From.Reg = riscv.REG_TMP
+		p.To.Type = obj.TYPE_REG
+		p.To.Reg = riscv.REG_TMP
+
+		p = s.Prog(obj.AJMP)
+		p.To.Type = obj.TYPE_MEM
+		p.To.Reg = riscv.REG_TMP
+		// Save jump tables for later resolution of the target blocks.
+		s.JumpTables = append(s.JumpTables, b)
+
 	default:
 		b.Fatalf("Unhandled block: %s", b.LongString())
 	}

src/cmd/compile/internal/ssa/_gen/RISCV64.rules

@@ -39,6 +39,8 @@
 (Select1 (Sub64borrow x y c)) =>
 	(OR (SLTU <typ.UInt64> x s:(SUB <typ.UInt64> x y)) (SLTU <typ.UInt64> s (SUB <typ.UInt64> s c)))
 
+(JumpTable idx) => (JUMPTABLE {makeJumpTableSym(b)} idx (MOVaddr <typ.Uintptr> {makeJumpTableSym(b)} (SB)))
+
 // (x + y) / 2 => (x / 2) + (y / 2) + (x & y & 1)
 (Avg64u <t> x y) => (ADD (ADD <t> (SRLI <t> [1] x) (SRLI <t> [1] y)) (ANDI <t> [1] (AND <t> x y)))
 

src/cmd/compile/internal/ssa/_gen/RISCV64Ops.go

@@ -550,6 +550,11 @@ func init() {
 		{name: "BGEZ", controls: 1},
 		{name: "BLTZ", controls: 1},
 		{name: "BGTZ", controls: 1},
+		// JUMPTABLE implements jump tables.
+		// Aux is the symbol (an *obj.LSym) for the jump table.
+		// control[0] is the index into the jump table.
+		// control[1] is the address of the jump table (the address of the symbol stored in Aux).
+		{name: "JUMPTABLE", controls: 2, aux: "Sym"},
 	}
 
 	archs = append(archs, arch{

src/cmd/internal/obj/riscv/obj.go

@@ -5063,6 +5063,13 @@ func assemble(ctxt *obj.Link, cursym *obj.LSym, newprog obj.ProgAlloc) {
 	}
 
 	obj.MarkUnsafePoints(ctxt, cursym.Func().Text, newprog, isUnsafePoint, nil)
+
+	// generate jump table entries.
+	for _, jt := range cursym.Func().JumpTables {
+		for i, p := range jt.Targets {
+			jt.Sym.WriteAddr(ctxt, int64(i)*8, 8, cursym, p.Pc)
+		}
+	}
 }
 
 func isUnsafePoint(p *obj.Prog) bool {

src/cmd/internal/sys/arch.go

@@ -239,6 +239,7 @@ var ArchRISCV64 = &Arch{
 	MinLC:          2,
 	Alignment:      8, // riscv unaligned loads work, but are really slow (trap + simulated by OS)
 	CanMergeLoads:  false,
+	CanJumpTable:   true,
 	HasLR:          true,
 	FixedFrameSize: 8, // LR
 }

test/codegen/switch.go

@@ -26,6 +26,8 @@ func square(x int) int {
 	// amd64:`JMP \(.*\)\(.*\)$`
 	// arm64:`MOVD \(R.*\)\(R.*<<3\)` `JMP \(R.*\)$`
 	// loong64: `ALSLV` `MOVV` `JMP`
+	// riscv64/rva20u64:`SLLI`,`ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
+	// riscv64/rva22u64,riscv64/rva23u64:`SH3ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
 	switch x {
 	case 1:
 		return 1
@@ -53,6 +55,8 @@ func length(x string) int {
 	// amd64:`JMP \(.*\)\(.*\)$`
 	// arm64:`MOVD \(R.*\)\(R.*<<3\)` `JMP \(R.*\)$`
 	// loong64:`ALSLV` `MOVV` `JMP`
+	// riscv64/rva20u64:`SLLI`,`ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
+	// riscv64/rva22u64,riscv64/rva23u64:`SH3ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
 	switch x {
 	case "a":
 		return 1
@@ -106,6 +110,8 @@ func mimetype(ext string) string {
 func typeSwitch(x any) int {
 	// amd64:`JMP \(.*\)\(.*\)$`
 	// arm64:`MOVD \(R.*\)\(R.*<<3\)` `JMP \(R.*\)$`
+	// riscv64/rva20u64:`SLLI`,`ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
+	// riscv64/rva22u64,riscv64/rva23u64:`SH3ADD`,`MOV\s\(X31\), X31`,`JALR\sX0, \(X31\)$`
 	switch x.(type) {
 	case int:
 		return 0