cmd/compile: optimize escape graph construction and walking

This CL implements several optimizations for the escape analysis flow
graph:

1. Instead of recognizing heapLoc specially within Escape.outlives,
set heapLoc.escapes = true and recognize any location with escapes
set. This allows us to skip adding edges from the heap to escaped
variables in two cases:

1a. In newLoc, if the location is for a variable or allocation too
large to fit on the stack.

1b. During walkOne, if we discover that an object's address flows
somewhere that naturally outlives it.

2. When recording edges in Escape.flow, if x escapes and we're adding
an edge like "x = &y", we can simply mark that y escapes too.

3. During walkOne, if we reach a location that's marked as escaping,
we can skip visiting it again: we've either already walked from it, or
it's in queue to be walked from again.

On average, reduces the number of visited locations by 15%. Reduces
time spent in escape analysis for particularly hairy packages like
runtime and gc by about 8%. Reduces escape.go's TODO count by 22%.

Passes toolstash-check.

Change-Id: Iaf86a29d76044e4b4c8ab581b916ef5bb5df4437
Reviewed-on: https://go-review.googlesource.com/c/go/+/196811
Reviewed-by: Cherry Zhang <[email protected]>

src/cmd/compile/internal/gc/escape.go

@@ -140,6 +140,7 @@ func escapeFuncs(fns []*Node, recursive bool) {
 	}
 
 	var e Escape
+	e.heapLoc.escapes = true
 
 	// Construct data-flow graph from syntax trees.
 	for _, fn := range fns {
@@ -984,9 +985,6 @@ func (e *Escape) dcl(n *Node) EscHole {
 // its address to k, and returns a hole that flows values to it. It's
 // intended for use with most expressions that allocate storage.
 func (e *Escape) spill(k EscHole, n *Node) EscHole {
-	// TODO(mdempsky): Optimize. E.g., if k is the heap or blank,
-	// then we already know whether n leaks, and we can return a
-	// more optimized hole.
 	loc := e.newLoc(n, true)
 	e.flow(k.addr(n, "spill"), loc)
 	return loc.asHole()
@@ -1037,9 +1035,8 @@ func (e *Escape) newLoc(n *Node, transient bool) *EscLocation {
 		}
 		n.SetOpt(loc)
 
-		// TODO(mdempsky): Perhaps set n.Esc and then just return &HeapLoc?
 		if mustHeapAlloc(n) && !loc.isName(PPARAM) && !loc.isName(PPARAMOUT) {
-			e.flow(e.heapHole().addr(nil, ""), loc)
+			loc.escapes = true
 		}
 	}
 	return loc
@@ -1059,10 +1056,13 @@ func (e *Escape) flow(k EscHole, src *EscLocation) {
 	if dst == &e.blankLoc {
 		return
 	}
-	if dst == src && k.derefs >= 0 {
+	if dst == src && k.derefs >= 0 { // dst = dst, dst = *dst, ...
+		return
+	}
+	if dst.escapes && k.derefs < 0 { // dst = &src
+		src.escapes = true
 		return
 	}
-	// TODO(mdempsky): More optimizations?
 
 	// TODO(mdempsky): Deduplicate edges?
 	dst.edges = append(dst.edges, EscEdge{src: src, derefs: k.derefs})
@@ -1076,6 +1076,11 @@ func (e *Escape) discardHole() EscHole { return e.blankLoc.asHole() }
 func (e *Escape) walkAll() {
 	// We use a work queue to keep track of locations that we need
 	// to visit, and repeatedly walk until we reach a fixed point.
+	//
+	// We walk once from each location (including the heap), and
+	// then re-enqueue each location on its transition from
+	// transient->!transient and !escapes->escapes, which can each
+	// happen at most once. So we take Θ(len(e.allLocs)) walks.
 
 	var todo []*EscLocation // LIFO queue
 	enqueue := func(loc *EscLocation) {
@@ -1085,10 +1090,10 @@ func (e *Escape) walkAll() {
 		}
 	}
 
-	enqueue(&e.heapLoc)
 	for _, loc := range e.allLocs {
 		enqueue(loc)
 	}
+	enqueue(&e.heapLoc)
 
 	var walkgen uint32
 	for len(todo) > 0 {
@@ -1138,22 +1143,6 @@ func (e *Escape) walkOne(root *EscLocation, walkgen uint32, enqueue func(*EscLoc
 		}
 
 		if e.outlives(root, l) {
-			// If l's address flows somewhere that
-			// outlives it, then l needs to be heap
-			// allocated.
-			if addressOf && !l.escapes {
-				l.escapes = true
-
-				// If l is heap allocated, then any
-				// values stored into it flow to the
-				// heap too.
-				// TODO(mdempsky): Better way to handle this?
-				if root != &e.heapLoc {
-					e.flow(e.heapHole(), l)
-					enqueue(&e.heapLoc)
-				}
-			}
-
 			// l's value flows to root. If l is a function
 			// parameter and root is the heap or a
 			// corresponding result parameter, then record
@@ -1162,9 +1151,21 @@ func (e *Escape) walkOne(root *EscLocation, walkgen uint32, enqueue func(*EscLoc
 			if l.isName(PPARAM) {
 				l.leakTo(root, base)
 			}
+
+			// If l's address flows somewhere that
+			// outlives it, then l needs to be heap
+			// allocated.
+			if addressOf && !l.escapes {
+				l.escapes = true
+				enqueue(l)
+				continue
+			}
 		}
 
 		for _, edge := range l.edges {
+			if edge.src.escapes {
+				continue
+			}
 			derefs := base + edge.derefs
 			if edge.src.walkgen != walkgen || edge.src.derefs > derefs {
 				edge.src.walkgen = walkgen
@@ -1179,7 +1180,7 @@ func (e *Escape) walkOne(root *EscLocation, walkgen uint32, enqueue func(*EscLoc
 // other's lifetime if stack allocated.
 func (e *Escape) outlives(l, other *EscLocation) bool {
 	// The heap outlives everything.
-	if l == &e.heapLoc {
+	if l.escapes {
 		return true
 	}