cmd/compile: use topo-sort in prove to correctly learn facts while walking once

Fixes #68857

Change-Id: Ideb359cc6f1550afb4c79f02d25a00d0ae5e5c50
Reviewed-on: https://go-review.googlesource.com/c/go/+/714920
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Reviewed-by: Keith Randall <khr@golang.org>
Auto-Submit: Jorropo <jorropo.pgm@gmail.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@google.com>

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

@@ -122,6 +122,11 @@ func genAllocators() {
 			typ:  "[]ID",
 			base: "LimitSlice",
 		},
+		{
+			name: "UintSlice",
+			typ:  "[]uint",
+			base: "LimitSlice",
+		},
 	}
 
 	w := new(bytes.Buffer)

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

@@ -331,3 +331,29 @@ func (c *Cache) freeIDSlice(s []ID) {
 	}
 	c.freeLimitSlice(*(*[]limit)(unsafe.Pointer(&b)))
 }
+func (c *Cache) allocUintSlice(n int) []uint {
+	var base limit
+	var derived uint
+	if unsafe.Sizeof(base)%unsafe.Sizeof(derived) != 0 {
+		panic("bad")
+	}
+	scale := unsafe.Sizeof(base) / unsafe.Sizeof(derived)
+	b := c.allocLimitSlice(int((uintptr(n) + scale - 1) / scale))
+	s := unsafeheader.Slice{
+		Data: unsafe.Pointer(&b[0]),
+		Len:  n,
+		Cap:  cap(b) * int(scale),
+	}
+	return *(*[]uint)(unsafe.Pointer(&s))
+}
+func (c *Cache) freeUintSlice(s []uint) {
+	var base limit
+	var derived uint
+	scale := unsafe.Sizeof(base) / unsafe.Sizeof(derived)
+	b := unsafeheader.Slice{
+		Data: unsafe.Pointer(&s[0]),
+		Len:  int((uintptr(len(s)) + scale - 1) / scale),
+		Cap:  int((uintptr(cap(s)) + scale - 1) / scale),
+	}
+	c.freeLimitSlice(*(*[]limit)(unsafe.Pointer(&b)))
+}

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

@@ -7,9 +7,11 @@ package ssa
 import (
 	"cmd/compile/internal/types"
 	"cmd/internal/src"
+	"cmp"
 	"fmt"
 	"math"
 	"math/bits"
+	"slices"
 )
 
 type branch int
@@ -412,6 +414,9 @@ type factsTable struct {
 	// more than one len(s) for a slice. We could keep a list if necessary.
 	lens map[ID]*Value
 	caps map[ID]*Value
+
+	// reusedTopoSortScoresTable recycle allocations for topo-sort
+	reusedTopoSortScoresTable []uint
 }
 
 // checkpointBound is an invalid value used for checkpointing
@@ -1238,6 +1243,9 @@ func (ft *factsTable) cleanup(f *Func) {
 	}
 	f.Cache.freeLimitSlice(ft.limits)
 	f.Cache.freeBoolSlice(ft.recurseCheck)
+	if cap(ft.reusedTopoSortScoresTable) > 0 {
+		f.Cache.freeUintSlice(ft.reusedTopoSortScoresTable)
+	}
 }
 
 // addSlicesOfSameLen finds the slices that are in the same block and whose Op
@@ -2478,23 +2486,13 @@ func checkForChunkedIndexBounds(ft *factsTable, b *Block, index, bound *Value, i
 }
 
 func addLocalFacts(ft *factsTable, b *Block) {
-	// Propagate constant ranges among values in this block.
-	// We do this before the second loop so that we have the
-	// most up-to-date constant bounds for isNonNegative calls.
-	for {
-		changed := false
-		for _, v := range b.Values {
-			changed = ft.flowLimit(v) || changed
-		}
-		if !changed {
-			break
-		}
-	}
+	ft.topoSortValuesInBlock(b)
 
-	// Add facts about individual operations.
 	for _, v := range b.Values {
-		// FIXME(go.dev/issue/68857): this loop only set up limits properly when b.Values is in topological order.
-		// flowLimit can also depend on limits given by this loop which right now is not handled.
+		// Propagate constant ranges before relative relations to get
+		// the most up-to-date constant bounds for isNonNegative calls.
+		ft.flowLimit(v)
+
 		switch v.Op {
 		case OpAdd64, OpAdd32, OpAdd16, OpAdd8:
 			x := ft.limits[v.Args[0].ID]
@@ -2973,3 +2971,57 @@ func isCleanExt(v *Value) bool {
 	}
 	return false
 }
+
+func getDependencyScore(scores []uint, v *Value) (score uint) {
+	if score = scores[v.ID]; score != 0 {
+		return score
+	}
+	defer func() {
+		scores[v.ID] = score
+	}()
+	if v.Op == OpPhi {
+		return 1
+	}
+	score = 2 // NIT(@Jorropo): always order phis first to make GOSSAFUNC pretty.
+	for _, a := range v.Args {
+		if a.Block != v.Block {
+			continue
+		}
+		score = max(score, getDependencyScore(scores, a)+1)
+	}
+	return score
+}
+
+// topoSortValuesInBlock ensure ranging over b.Values visit values before they are being used.
+// It does not consider dependencies with other blocks; thus Phi nodes are considered to not have any dependecies.
+// The result is always determistic and does not depend on the previous slice ordering.
+func (ft *factsTable) topoSortValuesInBlock(b *Block) {
+	f := b.Func
+	want := f.NumValues()
+
+	scores := ft.reusedTopoSortScoresTable
+	if len(scores) < want {
+		if want <= cap(scores) {
+			scores = scores[:want]
+		} else {
+			if cap(scores) > 0 {
+				f.Cache.freeUintSlice(scores)
+			}
+			scores = f.Cache.allocUintSlice(want)
+			ft.reusedTopoSortScoresTable = scores
+		}
+	}
+
+	for _, v := range b.Values {
+		scores[v.ID] = 0 // sentinel
+	}
+
+	slices.SortFunc(b.Values, func(a, b *Value) int {
+		dependencyScoreA := getDependencyScore(scores, a)
+		dependencyScoreB := getDependencyScore(scores, b)
+		if dependencyScoreA != dependencyScoreB {
+			return cmp.Compare(dependencyScoreA, dependencyScoreB)
+		}
+		return cmp.Compare(a.ID, b.ID)
+	})
+}