cmd/compile: prevent IsNewObject from taking quadratic time

As part of IsNewObject, we need to go from the SelectN[0] use of
a call to the SelectN[1] use of a call. The current code does this
by just looking through the block. If the block is very large,
this ends up taking quadratic time.

Instead, prepopulate a map from call -> SelectN[1] user of that call.
That lets us find the SelectN[1] user in constant time.

Fixes #57657

Change-Id: Ie2e0b660af5c080314f4f17ba2838510a1147f9e
Reviewed-on: https://go-review.googlesource.com/c/go/+/461080
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>
Run-TryBot: Keith Randall <khr@golang.org>
Reviewed-by: Cuong Manh Le <cuong.manhle.vn@gmail.com>

1 files changed, 24 insertions, 14 deletions

diff --git a/src/cmd/compile/internal/ssa/writebarrier.go b/src/cmd/compile/internal/ssa/writebarrier.go
index 3b2f781cbe..1676a9347c 100644
--- a/src/cmd/compile/internal/ssa/writebarrier.go
+++ b/src/cmd/compile/internal/ssa/writebarrier.go
@@ -27,7 +27,7 @@ type ZeroRegion struct {
 // needwb reports whether we need write barrier for store op v.
 // v must be Store/Move/Zero.
 // zeroes provides known zero information (keyed by ID of memory-type values).
-func needwb(v *Value, zeroes map[ID]ZeroRegion) bool {
+func needwb(v *Value, zeroes map[ID]ZeroRegion, select1 []*Value) bool {
 	t, ok := v.Aux.(*types.Type)
 	if !ok {
 		v.Fatalf("store aux is not a type: %s", v.LongString())
@@ -39,7 +39,7 @@ func needwb(v *Value, zeroes map[ID]ZeroRegion) bool {
 		return false // write on stack doesn't need write barrier
 	}
 	if v.Op == OpMove && IsReadOnlyGlobalAddr(v.Args[1]) {
-		if mem, ok := IsNewObject(v.Args[0]); ok && mem == v.MemoryArg() {
+		if mem, ok := IsNewObject(v.Args[0], select1); ok && mem == v.MemoryArg() {
 			// Copying data from readonly memory into a fresh object doesn't need a write barrier.
 			return false
 		}
@@ -99,7 +99,22 @@ func writebarrier(f *Func) {
 	var sset *sparseSet
 	var storeNumber []int32
 
-	zeroes := f.computeZeroMap()
+	// Compute map from a value to the SelectN [1] value that uses it.
+	select1 := f.Cache.allocValueSlice(f.NumValues())
+	defer func() { f.Cache.freeValueSlice(select1) }()
+	for _, b := range f.Blocks {
+		for _, v := range b.Values {
+			if v.Op != OpSelectN {
+				continue
+			}
+			if v.AuxInt != 1 {
+				continue
+			}
+			select1[v.Args[0].ID] = v
+		}
+	}
+
+	zeroes := f.computeZeroMap(select1)
 	for _, b := range f.Blocks { // range loop is safe since the blocks we added contain no stores to expand
 		// first, identify all the stores that need to insert a write barrier.
 		// mark them with WB ops temporarily. record presence of WB ops.
@@ -107,7 +122,7 @@ func writebarrier(f *Func) {
 		for _, v := range b.Values {
 			switch v.Op {
 			case OpStore, OpMove, OpZero:
-				if needwb(v, zeroes) {
+				if needwb(v, zeroes, select1) {
 					switch v.Op {
 					case OpStore:
 						v.Op = OpStoreWB
@@ -376,7 +391,8 @@ func writebarrier(f *Func) {
 
 // computeZeroMap returns a map from an ID of a memory value to
 // a set of locations that are known to be zeroed at that memory value.
-func (f *Func) computeZeroMap() map[ID]ZeroRegion {
+func (f *Func) computeZeroMap(select1 []*Value) map[ID]ZeroRegion {
+
 	ptrSize := f.Config.PtrSize
 	// Keep track of which parts of memory are known to be zero.
 	// This helps with removing write barriers for various initialization patterns.
@@ -386,7 +402,7 @@ func (f *Func) computeZeroMap() map[ID]ZeroRegion {
 	// Find new objects.
 	for _, b := range f.Blocks {
 		for _, v := range b.Values {
-			if mem, ok := IsNewObject(v); ok {
+			if mem, ok := IsNewObject(v, select1); ok {
 				// While compiling package runtime itself, we might see user
 				// calls to newobject, which will have result type
 				// unsafe.Pointer instead. We can't easily infer how large the
@@ -584,20 +600,14 @@ func IsReadOnlyGlobalAddr(v *Value) bool {
 
 // IsNewObject reports whether v is a pointer to a freshly allocated & zeroed object,
 // if so, also returns the memory state mem at which v is zero.
-func IsNewObject(v *Value) (mem *Value, ok bool) {
+func IsNewObject(v *Value, select1 []*Value) (mem *Value, ok bool) {
 	f := v.Block.Func
 	c := f.Config
 	if f.ABIDefault == f.ABI1 && len(c.intParamRegs) >= 1 {
 		if v.Op != OpSelectN || v.AuxInt != 0 {
 			return nil, false
 		}
-		// Find the memory
-		for _, w := range v.Block.Values {
-			if w.Op == OpSelectN && w.AuxInt == 1 && w.Args[0] == v.Args[0] {
-				mem = w
-				break
-			}
-		}
+		mem = select1[v.Args[0].ID]
 		if mem == nil {
 			return nil, false
 		}