cmd/compile: compute ptrBytes during CalcSize instead of on demand

Compute ptrBytes while computing the size of a type.
Requires an extra field on the type, but means that we don't
have potentially exponential behavior in the PtrDataSize computation.

For #65540.

Change-Id: Ia23c72bbd996730baddd32d9ed46cfc00c3472ee
Reviewed-on: https://go-review.googlesource.com/c/go/+/571543
Reviewed-by: Keith Randall <khr@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Matthew Dempsky <mdempsky@google.com>

3 files changed, 64 insertions, 67 deletions

diff --git a/src/cmd/compile/internal/types/size.go b/src/cmd/compile/internal/types/size.go
index a8fff0a84f..1ae8b6988a 100644
--- a/src/cmd/compile/internal/types/size.go
+++ b/src/cmd/compile/internal/types/size.go
@@ -202,7 +202,8 @@ func isAtomicStdPkg(p *Pkg) bool {
 	return p.Prefix == "sync/atomic" || p.Prefix == "runtime/internal/atomic"
 }
 
-// CalcSize calculates and stores the size, alignment, and eq/hash algorithm for t.
+// CalcSize calculates and stores the size, alignment, eq/hash algorithm,
+// and ptrBytes for t.
 // If CalcSizeDisabled is set, and the size/alignment
 // have not already been calculated, it calls Fatal.
 // This is used to prevent data races in the back end.
@@ -247,6 +248,7 @@ func CalcSize(t *Type) {
 	t.width = -2
 	t.align = 0  // 0 means use t.Width, below
 	t.alg = AMEM // default
+	// default t.ptrBytes is 0.
 	if t.Noalg() {
 		t.setAlg(ANOALG)
 	}
@@ -314,10 +316,12 @@ func CalcSize(t *Type) {
 		w = int64(PtrSize)
 		t.intRegs = 1
 		CheckSize(t.Elem())
+		t.ptrBytes = int64(PtrSize) // See PtrDataSize
 
 	case TUNSAFEPTR:
 		w = int64(PtrSize)
 		t.intRegs = 1
+		t.ptrBytes = int64(PtrSize)
 
 	case TINTER: // implemented as 2 pointers
 		w = 2 * int64(PtrSize)
@@ -329,10 +333,12 @@ func CalcSize(t *Type) {
 		} else {
 			t.setAlg(AINTER)
 		}
+		t.ptrBytes = int64(2 * PtrSize)
 
 	case TCHAN: // implemented as pointer
 		w = int64(PtrSize)
 		t.intRegs = 1
+		t.ptrBytes = int64(PtrSize)
 
 		CheckSize(t.Elem())
 
@@ -360,6 +366,7 @@ func CalcSize(t *Type) {
 		CheckSize(t.Elem())
 		CheckSize(t.Key())
 		t.setAlg(ANOEQ)
+		t.ptrBytes = int64(PtrSize)
 
 	case TFORW: // should have been filled in
 		base.Fatalf("invalid recursive type %v", t)
@@ -375,6 +382,7 @@ func CalcSize(t *Type) {
 		t.align = uint8(PtrSize)
 		t.intRegs = 2
 		t.setAlg(ASTRING)
+		t.ptrBytes = int64(PtrSize)
 
 	case TARRAY:
 		if t.Elem() == nil {
@@ -420,6 +428,12 @@ func CalcSize(t *Type) {
 				t.setAlg(ASPECIAL)
 			}
 		}
+		if t.NumElem() > 0 {
+			x := PtrDataSize(t.Elem())
+			if x > 0 {
+				t.ptrBytes = t.Elem().width*(t.NumElem()-1) + x
+			}
+		}
 
 	case TSLICE:
 		if t.Elem() == nil {
@@ -430,6 +444,9 @@ func CalcSize(t *Type) {
 		t.align = uint8(PtrSize)
 		t.intRegs = 3
 		t.setAlg(ANOEQ)
+		if !t.Elem().NotInHeap() {
+			t.ptrBytes = int64(PtrSize)
+		}
 
 	case TSTRUCT:
 		if t.IsFuncArgStruct() {
@@ -446,6 +463,7 @@ func CalcSize(t *Type) {
 		w = int64(PtrSize) // width of func type is pointer
 		t.intRegs = 1
 		t.setAlg(ANOEQ)
+		t.ptrBytes = int64(PtrSize)
 
 	// function is 3 cated structures;
 	// compute their widths as side-effect.
@@ -490,8 +508,6 @@ func CalcStructSize(t *Type) {
 		switch {
 		case sym.Name == "align64" && isAtomicStdPkg(sym.Pkg):
 			maxAlign = 8
-		case sym.Pkg.Path == "runtime/internal/sys" && sym.Name == "nih":
-			t.SetNotInHeap(true)
 		}
 	}
 
@@ -560,6 +576,14 @@ func CalcStructSize(t *Type) {
 			t.setAlg(a)
 		}
 	}
+	// Compute ptrBytes.
+	for i := len(fields) - 1; i >= 0; i-- {
+		f := fields[i]
+		if size := PtrDataSize(f.Type); size > 0 {
+			t.ptrBytes = f.Offset + size
+			break
+		}
+	}
 }
 
 func (t *Type) widthCalculated() bool {
@@ -630,67 +654,13 @@ func ResumeCheckSize() {
 // PtrDataSize is only defined for actual Go types. It's an error to
 // use it on compiler-internal types (e.g., TSSA, TRESULTS).
 func PtrDataSize(t *Type) int64 {
-	switch t.Kind() {
-	case TBOOL, TINT8, TUINT8, TINT16, TUINT16, TINT32,
-		TUINT32, TINT64, TUINT64, TINT, TUINT,
-		TUINTPTR, TCOMPLEX64, TCOMPLEX128, TFLOAT32, TFLOAT64:
-		return 0
-
-	case TPTR:
-		if t.Elem().NotInHeap() {
-			return 0
-		}
-		return int64(PtrSize)
-
-	case TUNSAFEPTR, TFUNC, TCHAN, TMAP:
-		return int64(PtrSize)
-
-	case TSTRING:
-		// struct { byte *str; intgo len; }
-		return int64(PtrSize)
-
-	case TINTER:
-		// struct { Itab *tab;	void *data; } or
-		// struct { Type *type; void *data; }
-		// Note: see comment in typebits.Set
-		return 2 * int64(PtrSize)
-
-	case TSLICE:
-		if t.Elem().NotInHeap() {
-			return 0
-		}
-		// struct { byte *array; uintgo len; uintgo cap; }
-		return int64(PtrSize)
-
-	case TARRAY:
-		if t.NumElem() == 0 {
-			return 0
-		}
-		// t.NumElem() > 0
-		size := PtrDataSize(t.Elem())
-		if size == 0 {
-			return 0
-		}
-		return (t.NumElem()-1)*t.Elem().Size() + size
-
-	case TSTRUCT:
-		// Find the last field that has pointers, if any.
-		fs := t.Fields()
-		for i := len(fs) - 1; i >= 0; i-- {
-			if size := PtrDataSize(fs[i].Type); size > 0 {
-				return fs[i].Offset + size
-			}
-		}
-		return 0
-
-	case TSSA:
-		if t != TypeInt128 {
-			base.Fatalf("PtrDataSize: unexpected ssa type %v", t)
-		}
-		return 0
-
-	default:
-		base.Fatalf("PtrDataSize: unexpected type, %v", t)
-		return 0
+	CalcSize(t)
+	x := t.ptrBytes
+	if t.Kind() == TPTR && t.Elem().NotInHeap() {
+		// Note: this is done here instead of when we're setting
+		// the ptrBytes field, because at that time (in NewPtr, usually)
+		// the NotInHeap bit of the element type might not be set yet.
+		x = 0
 	}
+	return x
 }
diff --git a/src/cmd/compile/internal/types/sizeof_test.go b/src/cmd/compile/internal/types/sizeof_test.go
index 8a6f24124a..7e3e7769d7 100644
--- a/src/cmd/compile/internal/types/sizeof_test.go
+++ b/src/cmd/compile/internal/types/sizeof_test.go
@@ -21,7 +21,7 @@ func TestSizeof(t *testing.T) {
 		_64bit uintptr     // size on 64bit platforms
 	}{
 		{Sym{}, 32, 64},
-		{Type{}, 56, 96},
+		{Type{}, 64, 104},
 		{Map{}, 12, 24},
 		{Forward{}, 20, 32},
 		{Func{}, 32, 56},
diff --git a/src/cmd/compile/internal/types/type.go b/src/cmd/compile/internal/types/type.go
index 0ea20ae262..f372af32b5 100644
--- a/src/cmd/compile/internal/types/type.go
+++ b/src/cmd/compile/internal/types/type.go
@@ -203,6 +203,11 @@ type Type struct {
 	flags bitset8
 	alg   AlgKind // valid if Align > 0
 
+	// size of prefix of object that contains all pointers. valid if Align > 0.
+	// Note that for pointers, this is always PtrSize even if the element type
+	// is NotInHeap. See size.go:PtrDataSize for details.
+	ptrBytes int64
+
 	// For defined (named) generic types, a pointer to the list of type params
 	// (in order) of this type that need to be instantiated. For instantiated
 	// generic types, this is the targs used to instantiate them. These targs
@@ -549,6 +554,9 @@ func NewArray(elem *Type, bound int64) *Type {
 	if elem.HasShape() {
 		t.SetHasShape(true)
 	}
+	if elem.NotInHeap() {
+		t.SetNotInHeap(true)
+	}
 	return t
 }
 
@@ -663,6 +671,9 @@ func NewPtr(elem *Type) *Type {
 		t.SetNoalg(true)
 		t.alg = ANOALG
 	}
+	// Note: we can't check elem.NotInHeap here because it might
+	// not be set yet. See size.go:PtrDataSize.
+	t.ptrBytes = int64(PtrSize)
 	return t
 }
 
@@ -1634,10 +1645,18 @@ func init() {
 func NewNamed(obj Object) *Type {
 	t := newType(TFORW)
 	t.obj = obj
-	if obj.Sym().Pkg == ShapePkg {
+	sym := obj.Sym()
+	if sym.Pkg == ShapePkg {
 		t.SetIsShape(true)
 		t.SetHasShape(true)
 	}
+	if sym.Pkg.Path == "runtime/internal/sys" && sym.Name == "nih" {
+		// Recognize the special not-in-heap type. Any type including
+		// this type will also be not-in-heap.
+		// This logic is duplicated in go/types and
+		// cmd/compile/internal/types2.
+		t.SetNotInHeap(true)
+	}
 	return t
 }
 
@@ -1664,6 +1683,7 @@ func (t *Type) SetUnderlying(underlying *Type) {
 	t.width = underlying.width
 	t.align = underlying.align
 	t.alg = underlying.alg
+	t.ptrBytes = underlying.ptrBytes
 	t.intRegs = underlying.intRegs
 	t.floatRegs = underlying.floatRegs
 	t.underlying = underlying.underlying
@@ -1772,6 +1792,13 @@ func NewStruct(fields []*Field) *Type {
 	if fieldsHasShape(fields) {
 		t.SetHasShape(true)
 	}
+	for _, f := range fields {
+		if f.Type.NotInHeap() {
+			t.SetNotInHeap(true)
+			break
+		}
+	}
+
 	return t
 }