reflect: refactor funcLayout tests

This change refactors the existing funcLayout tests and sets them up to
support the new register ABI by explicitly setting the register counts
to zero. This allows the test to pass if GOEXPERIMENT=regabiargs is set.

A follow-up change will add tests for a non-zero register count.

For #40724.

Change-Id: Ibbe061b4ed4fd70566eb38b9e6182dca32b81127
Reviewed-on: https://go-review.googlesource.com/c/go/+/307869
Trust: Michael Knyszek <mknyszek@google.com>
Run-TryBot: Michael Knyszek <mknyszek@google.com>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

3 files changed, 158 insertions, 188 deletions

diff --git a/src/reflect/abi_test.go b/src/reflect/abi_test.go
index 62f6bd2e3e..224143c9bf 100644
--- a/src/reflect/abi_test.go
+++ b/src/reflect/abi_test.go
@@ -22,23 +22,7 @@ func TestMethodValueCallABI(t *testing.T) {
 	// Enable register-based reflect.Call and ensure we don't
 	// use potentially incorrect cached versions by clearing
 	// the cache before we start and after we're done.
-	var oldRegs struct {
-		ints, floats int
-		floatSize    uintptr
-	}
-	oldRegs.ints = *reflect.IntArgRegs
-	oldRegs.floats = *reflect.FloatArgRegs
-	oldRegs.floatSize = *reflect.FloatRegSize
-	*reflect.IntArgRegs = abi.IntArgRegs
-	*reflect.FloatArgRegs = abi.FloatArgRegs
-	*reflect.FloatRegSize = uintptr(abi.EffectiveFloatRegSize)
-	reflect.ClearLayoutCache()
-	defer func() {
-		*reflect.IntArgRegs = oldRegs.ints
-		*reflect.FloatArgRegs = oldRegs.floats
-		*reflect.FloatRegSize = oldRegs.floatSize
-		reflect.ClearLayoutCache()
-	}()
+	defer reflect.SetArgRegs(reflect.SetArgRegs(abi.IntArgRegs, abi.FloatArgRegs, abi.EffectiveFloatRegSize))
 
 	// This test is simple. Calling a method value involves
 	// pretty much just plumbing whatever arguments in whichever
@@ -129,23 +113,7 @@ func TestReflectCallABI(t *testing.T) {
 	// Enable register-based reflect.Call and ensure we don't
 	// use potentially incorrect cached versions by clearing
 	// the cache before we start and after we're done.
-	var oldRegs struct {
-		ints, floats int
-		floatSize    uintptr
-	}
-	oldRegs.ints = *reflect.IntArgRegs
-	oldRegs.floats = *reflect.FloatArgRegs
-	oldRegs.floatSize = *reflect.FloatRegSize
-	*reflect.IntArgRegs = abi.IntArgRegs
-	*reflect.FloatArgRegs = abi.FloatArgRegs
-	*reflect.FloatRegSize = uintptr(abi.EffectiveFloatRegSize)
-	reflect.ClearLayoutCache()
-	defer func() {
-		*reflect.IntArgRegs = oldRegs.ints
-		*reflect.FloatArgRegs = oldRegs.floats
-		*reflect.FloatRegSize = oldRegs.floatSize
-		reflect.ClearLayoutCache()
-	}()
+	defer reflect.SetArgRegs(reflect.SetArgRegs(abi.IntArgRegs, abi.FloatArgRegs, abi.EffectiveFloatRegSize))
 
 	// Execute the functions defined below which all have the
 	// same form and perform the same function: pass all arguments
@@ -182,23 +150,7 @@ func TestReflectMakeFuncCallABI(t *testing.T) {
 	// Enable register-based reflect.MakeFunc and ensure we don't
 	// use potentially incorrect cached versions by clearing
 	// the cache before we start and after we're done.
-	var oldRegs struct {
-		ints, floats int
-		floatSize    uintptr
-	}
-	oldRegs.ints = *reflect.IntArgRegs
-	oldRegs.floats = *reflect.FloatArgRegs
-	oldRegs.floatSize = *reflect.FloatRegSize
-	*reflect.IntArgRegs = abi.IntArgRegs
-	*reflect.FloatArgRegs = abi.FloatArgRegs
-	*reflect.FloatRegSize = uintptr(abi.EffectiveFloatRegSize)
-	reflect.ClearLayoutCache()
-	defer func() {
-		*reflect.IntArgRegs = oldRegs.ints
-		*reflect.FloatArgRegs = oldRegs.floats
-		*reflect.FloatRegSize = oldRegs.floatSize
-		reflect.ClearLayoutCache()
-	}()
+	defer reflect.SetArgRegs(reflect.SetArgRegs(abi.IntArgRegs, abi.FloatArgRegs, abi.EffectiveFloatRegSize))
 
 	// Execute the functions defined below which all have the
 	// same form and perform the same function: pass all arguments
diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index e4b74f72d9..241f6b0b5a 100644
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -6396,144 +6396,135 @@ func clobber() {
 	runtime.GC()
 }
 
-type funcLayoutTest struct {
-	rcvr, t                  Type
-	size, argsize, retOffset uintptr
-	stack                    []byte // pointer bitmap: 1 is pointer, 0 is scalar
-	gc                       []byte
-}
-
-var funcLayoutTests []funcLayoutTest
-
-func init() {
-	var argAlign uintptr = PtrSize
-	roundup := func(x uintptr, a uintptr) uintptr {
-		return (x + a - 1) / a * a
-	}
-
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func(a, b string) string { return "" }).Type(),
-			6 * PtrSize,
-			4 * PtrSize,
-			4 * PtrSize,
-			[]byte{1, 0, 1, 0, 1},
-			[]byte{1, 0, 1, 0, 1},
-		})
-
+func TestFuncLayout(t *testing.T) {
+	align := func(x uintptr) uintptr {
+		return (x + PtrSize - 1) &^ (PtrSize - 1)
+	}
 	var r []byte
 	if PtrSize == 4 {
 		r = []byte{0, 0, 0, 1}
 	} else {
 		r = []byte{0, 0, 1}
 	}
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func(a, b, c uint32, p *byte, d uint16) {}).Type(),
-			roundup(roundup(3*4, PtrSize)+PtrSize+2, argAlign),
-			roundup(3*4, PtrSize) + PtrSize + 2,
-			roundup(roundup(3*4, PtrSize)+PtrSize+2, argAlign),
-			r,
-			r,
-		})
-
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func(a map[int]int, b uintptr, c interface{}) {}).Type(),
-			4 * PtrSize,
-			4 * PtrSize,
-			4 * PtrSize,
-			[]byte{1, 0, 1, 1},
-			[]byte{1, 0, 1, 1},
-		})
 
 	type S struct {
 		a, b uintptr
 		c, d *byte
 	}
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func(a S) {}).Type(),
-			4 * PtrSize,
-			4 * PtrSize,
-			4 * PtrSize,
-			[]byte{0, 0, 1, 1},
-			[]byte{0, 0, 1, 1},
-		})
-
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			ValueOf((*byte)(nil)).Type(),
-			ValueOf(func(a uintptr, b *int) {}).Type(),
-			roundup(3*PtrSize, argAlign),
-			3 * PtrSize,
-			roundup(3*PtrSize, argAlign),
-			[]byte{1, 0, 1},
-			[]byte{1, 0, 1},
-		})
 
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func(a uintptr) {}).Type(),
-			roundup(PtrSize, argAlign),
-			PtrSize,
-			roundup(PtrSize, argAlign),
-			[]byte{},
-			[]byte{},
-		})
-
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			nil,
-			ValueOf(func() uintptr { return 0 }).Type(),
-			PtrSize,
-			0,
-			0,
-			[]byte{},
-			[]byte{},
-		})
-
-	funcLayoutTests = append(funcLayoutTests,
-		funcLayoutTest{
-			ValueOf(uintptr(0)).Type(),
-			ValueOf(func(a uintptr) {}).Type(),
-			2 * PtrSize,
-			2 * PtrSize,
-			2 * PtrSize,
-			[]byte{1},
-			[]byte{1},
+	type test struct {
+		rcvr, typ                  Type
+		size, argsize, retOffset   uintptr
+		stack, gc, inRegs, outRegs []byte // pointer bitmap: 1 is pointer, 0 is scalar
+		intRegs, floatRegs         int
+		floatRegSize               uintptr
+	}
+	tests := []test{
+		{
+			typ:       ValueOf(func(a, b string) string { return "" }).Type(),
+			size:      6 * PtrSize,
+			argsize:   4 * PtrSize,
+			retOffset: 4 * PtrSize,
+			stack:     []byte{1, 0, 1, 0, 1},
+			gc:        []byte{1, 0, 1, 0, 1},
+		},
+		{
+			typ:       ValueOf(func(a, b, c uint32, p *byte, d uint16) {}).Type(),
+			size:      align(align(3*4) + PtrSize + 2),
+			argsize:   align(3*4) + PtrSize + 2,
+			retOffset: align(align(3*4) + PtrSize + 2),
+			stack:     r,
+			gc:        r,
+		},
+		{
+			typ:       ValueOf(func(a map[int]int, b uintptr, c interface{}) {}).Type(),
+			size:      4 * PtrSize,
+			argsize:   4 * PtrSize,
+			retOffset: 4 * PtrSize,
+			stack:     []byte{1, 0, 1, 1},
+			gc:        []byte{1, 0, 1, 1},
+		},
+		{
+			typ:       ValueOf(func(a S) {}).Type(),
+			size:      4 * PtrSize,
+			argsize:   4 * PtrSize,
+			retOffset: 4 * PtrSize,
+			stack:     []byte{0, 0, 1, 1},
+			gc:        []byte{0, 0, 1, 1},
+		},
+		{
+			rcvr:      ValueOf((*byte)(nil)).Type(),
+			typ:       ValueOf(func(a uintptr, b *int) {}).Type(),
+			size:      3 * PtrSize,
+			argsize:   3 * PtrSize,
+			retOffset: 3 * PtrSize,
+			stack:     []byte{1, 0, 1},
+			gc:        []byte{1, 0, 1},
+		},
+		{
+			typ:       ValueOf(func(a uintptr) {}).Type(),
+			size:      PtrSize,
+			argsize:   PtrSize,
+			retOffset: PtrSize,
+			stack:     []byte{},
+			gc:        []byte{},
+		},
+		{
+			typ:       ValueOf(func() uintptr { return 0 }).Type(),
+			size:      PtrSize,
+			argsize:   0,
+			retOffset: 0,
+			stack:     []byte{},
+			gc:        []byte{},
+		},
+		{
+			rcvr:      ValueOf(uintptr(0)).Type(),
+			typ:       ValueOf(func(a uintptr) {}).Type(),
+			size:      2 * PtrSize,
+			argsize:   2 * PtrSize,
+			retOffset: 2 * PtrSize,
+			stack:     []byte{1},
+			gc:        []byte{1},
 			// Note: this one is tricky, as the receiver is not a pointer. But we
 			// pass the receiver by reference to the autogenerated pointer-receiver
 			// version of the function.
-		})
-}
-
-func TestFuncLayout(t *testing.T) {
-	for _, lt := range funcLayoutTests {
-		typ, argsize, retOffset, stack, gc, ptrs := FuncLayout(lt.t, lt.rcvr)
-		if typ.Size() != lt.size {
-			t.Errorf("funcLayout(%v, %v).size=%d, want %d", lt.t, lt.rcvr, typ.Size(), lt.size)
-		}
-		if argsize != lt.argsize {
-			t.Errorf("funcLayout(%v, %v).argsize=%d, want %d", lt.t, lt.rcvr, argsize, lt.argsize)
-		}
-		if retOffset != lt.retOffset {
-			t.Errorf("funcLayout(%v, %v).retOffset=%d, want %d", lt.t, lt.rcvr, retOffset, lt.retOffset)
-		}
-		if !bytes.Equal(stack, lt.stack) {
-			t.Errorf("funcLayout(%v, %v).stack=%v, want %v", lt.t, lt.rcvr, stack, lt.stack)
-		}
-		if !bytes.Equal(gc, lt.gc) {
-			t.Errorf("funcLayout(%v, %v).gc=%v, want %v", lt.t, lt.rcvr, gc, lt.gc)
-		}
-		if ptrs && len(stack) == 0 || !ptrs && len(stack) > 0 {
-			t.Errorf("funcLayout(%v, %v) pointers flag=%v, want %v", lt.t, lt.rcvr, ptrs, !ptrs)
+		},
+		// TODO(mknyszek): Add tests for non-zero register count.
+	}
+	for _, lt := range tests {
+		name := lt.typ.String()
+		if lt.rcvr != nil {
+			name = lt.rcvr.String() + "." + name
 		}
+		t.Run(name, func(t *testing.T) {
+			defer SetArgRegs(SetArgRegs(lt.intRegs, lt.floatRegs, lt.floatRegSize))
+
+			typ, argsize, retOffset, stack, gc, inRegs, outRegs, ptrs := FuncLayout(lt.typ, lt.rcvr)
+			if typ.Size() != lt.size {
+				t.Errorf("funcLayout(%v, %v).size=%d, want %d", lt.typ, lt.rcvr, typ.Size(), lt.size)
+			}
+			if argsize != lt.argsize {
+				t.Errorf("funcLayout(%v, %v).argsize=%d, want %d", lt.typ, lt.rcvr, argsize, lt.argsize)
+			}
+			if retOffset != lt.retOffset {
+				t.Errorf("funcLayout(%v, %v).retOffset=%d, want %d", lt.typ, lt.rcvr, retOffset, lt.retOffset)
+			}
+			if !bytes.Equal(stack, lt.stack) {
+				t.Errorf("funcLayout(%v, %v).stack=%v, want %v", lt.typ, lt.rcvr, stack, lt.stack)
+			}
+			if !bytes.Equal(gc, lt.gc) {
+				t.Errorf("funcLayout(%v, %v).gc=%v, want %v", lt.typ, lt.rcvr, gc, lt.gc)
+			}
+			if !bytes.Equal(inRegs, lt.inRegs) {
+				t.Errorf("funcLayout(%v, %v).inRegs=%v, want %v", lt.typ, lt.rcvr, inRegs, lt.inRegs)
+			}
+			if !bytes.Equal(outRegs, lt.outRegs) {
+				t.Errorf("funcLayout(%v, %v).outRegs=%v, want %v", lt.typ, lt.rcvr, outRegs, lt.outRegs)
+			}
+			if ptrs && len(stack) == 0 || !ptrs && len(stack) > 0 {
+				t.Errorf("funcLayout(%v, %v) pointers flag=%v, want %v", lt.typ, lt.rcvr, ptrs, !ptrs)
+			}
+		})
 	}
 }
 
diff --git a/src/reflect/export_test.go b/src/reflect/export_test.go
index 3a5ed5af3c..b6830a9802 100644
--- a/src/reflect/export_test.go
+++ b/src/reflect/export_test.go
@@ -20,33 +20,49 @@ func IsRO(v Value) bool {
 	return v.flag&flagStickyRO != 0
 }
 
-var (
-	IntArgRegs   = &intArgRegs
-	FloatArgRegs = &floatArgRegs
-	FloatRegSize = &floatRegSize
-)
-
 var CallGC = &callGC
 
 const PtrSize = ptrSize
 
-func FuncLayout(t Type, rcvr Type) (frametype Type, argSize, retOffset uintptr, stack []byte, gc []byte, ptrs bool) {
+// FuncLayout calls funcLayout and returns a subset of the results for testing.
+//
+// Bitmaps like stack, gc, inReg, and outReg are expanded such that each bit
+// takes up one byte, so that writing out test cases is a little clearer.
+// If ptrs is false, gc will be nil.
+func FuncLayout(t Type, rcvr Type) (frametype Type, argSize, retOffset uintptr, stack, gc, inReg, outReg []byte, ptrs bool) {
 	var ft *rtype
-	var abi abiDesc
+	var abid abiDesc
 	if rcvr != nil {
-		ft, _, abi = funcLayout((*funcType)(unsafe.Pointer(t.(*rtype))), rcvr.(*rtype))
+		ft, _, abid = funcLayout((*funcType)(unsafe.Pointer(t.(*rtype))), rcvr.(*rtype))
 	} else {
-		ft, _, abi = funcLayout((*funcType)(unsafe.Pointer(t.(*rtype))), nil)
+		ft, _, abid = funcLayout((*funcType)(unsafe.Pointer(t.(*rtype))), nil)
 	}
-	argSize = abi.stackCallArgsSize
-	retOffset = abi.retOffset
+	// Extract size information.
+	argSize = abid.stackCallArgsSize
+	retOffset = abid.retOffset
 	frametype = ft
-	for i := uint32(0); i < abi.stackPtrs.n; i++ {
-		stack = append(stack, abi.stackPtrs.data[i/8]>>(i%8)&1)
+
+	// Expand stack pointer bitmap into byte-map.
+	for i := uint32(0); i < abid.stackPtrs.n; i++ {
+		stack = append(stack, abid.stackPtrs.data[i/8]>>(i%8)&1)
+	}
+
+	// Expand register pointer bitmaps into byte-maps.
+	bool2byte := func(b bool) byte {
+		if b {
+			return 1
+		}
+		return 0
+	}
+	for i := 0; i < intArgRegs; i++ {
+		inReg = append(inReg, bool2byte(abid.inRegPtrs.Get(i)))
+		outReg = append(outReg, bool2byte(abid.outRegPtrs.Get(i)))
 	}
 	if ft.kind&kindGCProg != 0 {
 		panic("can't handle gc programs")
 	}
+
+	// Expand frame type's GC bitmap into byte-map.
 	ptrs = ft.ptrdata != 0
 	if ptrs {
 		nptrs := ft.ptrdata / ptrSize
@@ -132,6 +148,17 @@ type Buffer struct {
 	buf []byte
 }
 
-func ClearLayoutCache() {
+func clearLayoutCache() {
 	layoutCache = sync.Map{}
 }
+
+func SetArgRegs(ints, floats int, floatSize uintptr) (oldInts, oldFloats int, oldFloatSize uintptr) {
+	oldInts = intArgRegs
+	oldFloats = floatArgRegs
+	oldFloatSize = floatRegSize
+	intArgRegs = ints
+	floatArgRegs = floats
+	floatRegSize = floatSize
+	clearLayoutCache()
+	return
+}