reflect: rewrite ·dispatchLabel for ABIInternal

For golang#38783

src/reflect/asm_amd64.s

@@ -897,19 +897,19 @@ TEXT ·dispatchToMethod(SB),(NOSPLIT|NOFRAME),$0-0
 	CALL ·dispatchLabel(SB)
 // Use TEXT instead of label because assembler does not support call to label.
 TEXT ·dispatchLabel(SB),(NOSPLIT|NOFRAME),$0-0
-	// FIXME: This dispatch code still assumes ABI0.
-	//        Need to rewrite for ABIInternal.
-	//
 	// At this point, the stack layout is:
 	//
-	//     closure_key, ret_addr, recv, arg1, …, argN, ret1, …, retN
+	//     closure_key, ret_addr, possible stack args of caller
 	//
 	// We now look up the closure function pointer, load the context and dispatch
 	// to the implementation.
-	PUSHQ (SP)
-	CALL ·getMethodImpl(SB)
+	POPQ DX
+	SUBQ $0xc8, SP // spill space + return value, see go.dev/s/regabi
+	PUSHQ DX // arg
+	CALL ·getMethodImpl<ABIInternal>(SB)
 	POPQ DX
 	POPQ DX
+	ADDQ $0xc0, SP
 	// Currently, (DX) is always makeFuncStub, but this method of dispatch
 	// might have other clients in the future, so we leave the indirect jump in.
 	JMP (DX)

src/reflect/makefunc.go

@@ -8,7 +8,6 @@ package reflect
 
 import (
 	"internal/abi"
-	"strconv"
 	"sync"
 	"sync/atomic"
 	"unsafe"
@@ -249,12 +248,3 @@ func allocMethodSlot(closure unsafe.Pointer) unsafe.Pointer {
 	methodMap.Store(newKey, closure)
 	return unsafe.Pointer(newKey - dispatchStep)
 }
-
-// getMethodImpl returns the pointer to the closure for the specified key.
-func getMethodImpl(key uintptr) unsafe.Pointer {
-	if x, ok := methodMap.Load(key); !ok {
-		panic("reflect: method " + strconv.Itoa(int(key)) + " not found")
-	} else {
-		return x.(unsafe.Pointer)
-	}
-}

src/reflect/makefunc_amd64.go

@@ -0,0 +1,37 @@
+// Copyright 2021 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build amd64
+
+package reflect
+
+import (
+	"strconv"
+	"unsafe"
+)
+
+// getMethodImpl returns the pointer to the closure for the specified key.
+// It is called from the ASM dispatch code.
+func getMethodImpl(
+	// forward all possible register args,
+	// force key and closure onto the stack.
+	ax, bx, cx, di, si, r8, r9, r10, r11 uintptr,
+	// FIXME: does sync.Map.Load clobber FP registers?
+	// If not, we could omit x0 through x14 here.
+	x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14 float64,
+	key uintptr,
+) (
+	rax, rbx, rcx, rdi, rsi, rr8, rr9, rr10, rr11 uintptr,
+	rx0, rx1, rx2, rx3, rx4, rx5, rx6, rx7, rx8, rx9,
+	rx10, rx11, rx12, rx13, rx14 float64,
+	closure unsafe.Pointer,
+) {
+	if x, ok := methodMap.Load(key); !ok {
+		panic("reflect: method " + strconv.Itoa(int(key)) + " not found")
+	} else {
+		return ax, bx, cx, di, si, r8, r9, r10, r11,
+			x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14,
+			x.(unsafe.Pointer)
+	}
+}