[dev.regabi] cmd/compile: reorganize escape analysis somewhat

To do closure conversion during escape analysis, we need to walk the
AST in order. So this CL makes a few changes:

1. Function literals are walked where they appear in their enclosing
function, rather than as independent functions.

2. Walking "range" and "switch" statements is reordered to visit the
X/Tag expression up front, before the body.

3. Most assignments are refactored to use a new assignList helper,
which handles 1:1, 2:1, and N:N assignments. N:1 function call
assignments are still handled directly by the OAS2FUNC case.

4. A latent missed-optimization in escape.addr is fixed: the
ONAMEOFFSET case was failing to update k with the result of calling
e.addr(n.Name_). In partice, this probably wasn't an issue because
ONAMEOFFSET is likely only used for PEXTERN variables (which are
treated as heap memory anyway) or code generated by walk (which has
already gone through escape analysis).

5. Finally, don't replace k with discardHole at the end of
escape.addr. This is already handled at the start of escape.expr, and
we'll want to be able to access the hole's location after escape.expr
returns.

Passes toolstash -cmp.

Change-Id: I2325234346b12b10056a360c489692bab8fdbd93
Reviewed-on: https://go-review.googlesource.com/c/go/+/281003
Trust: Matthew Dempsky <mdempsky@google.com>
Run-TryBot: Matthew Dempsky <mdempsky@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Cuong Manh Le <cuong.manhle.vn@gmail.com>

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

@@ -201,10 +201,12 @@ func Batch(fns []*ir.Func, recursive bool) {
 
 	// Construct data-flow graph from syntax trees.
 	for _, fn := range fns {
-		b.with(fn).initFunc()
+		b.initFunc(fn)
 	}
 	for _, fn := range fns {
-		b.with(fn).walkFunc()
+		if !fn.IsHiddenClosure() {
+			b.walkFunc(fn)
+		}
 	}
 
 	b.walkAll()
@@ -219,8 +221,8 @@ func (b *batch) with(fn *ir.Func) *escape {
 	}
 }
 
-func (e *escape) initFunc() {
+func (b *batch) initFunc(fn *ir.Func) {
-	fn := e.curfn
+	e := b.with(fn)
 	if fn.Esc() != escFuncUnknown {
 		base.Fatalf("unexpected node: %v", fn)
 	}
@@ -237,8 +239,8 @@ func (e *escape) initFunc() {
 	}
 }
 
-func (e *escape) walkFunc() {
+func (b *batch) walkFunc(fn *ir.Func) {
-	fn := e.curfn
+	e := b.with(fn)
 	fn.SetEsc(escFuncStarted)
 
 	// Identify labels that mark the head of an unstructured loop.
@@ -366,54 +368,52 @@ func (e *escape) stmt(n ir.Node) {
 	case ir.ORANGE:
 		// for Key, Value = range X { Body }
 		n := n.(*ir.RangeStmt)
+
+		// X is evaluated outside the loop.
+		tmp := e.newLoc(nil, false)
+		e.expr(tmp.asHole(), n.X)
+
 		e.loopDepth++
-		e.addr(n.Key)
+		ks := e.addrs([]ir.Node{n.Key, n.Value})
-		k := e.addr(n.Value)
+		if n.X.Type().IsArray() {
+			e.flow(ks[1].note(n, "range"), tmp)
+		} else {
+			e.flow(ks[1].deref(n, "range-deref"), tmp)
+		}
+
 		e.block(n.Body)
 		e.loopDepth--
 
-		// X is evaluated outside the loop.
-		if n.X.Type().IsArray() {
-			k = k.note(n, "range")
-		} else {
-			k = k.deref(n, "range-deref")
-		}
-		e.expr(e.later(k), n.X)
-
 	case ir.OSWITCH:
 		n := n.(*ir.SwitchStmt)
-		typesw := n.Tag != nil && n.Tag.Op() == ir.OTYPESW
 
+		if guard, ok := n.Tag.(*ir.TypeSwitchGuard); ok {
 			var ks []hole
-		for _, cas := range n.Cases { // cases
+			if guard.Tag != nil {
-			if typesw && n.Tag.(*ir.TypeSwitchGuard).Tag != nil {
+				for _, cas := range n.Cases {
 					cv := cas.Var
 					k := e.dcl(cv) // type switch variables have no ODCL.
 					if cv.Type().HasPointers() {
 						ks = append(ks, k.dotType(cv.Type(), cas, "switch case"))
 					}
 				}
-
-			e.discards(cas.List)
-			e.block(cas.Body)
 			}
-
-		if typesw {
 			e.expr(e.teeHole(ks...), n.Tag.(*ir.TypeSwitchGuard).X)
 		} else {
 			e.discard(n.Tag)
 		}
 
+		for _, cas := range n.Cases {
+			e.discards(cas.List)
+			e.block(cas.Body)
+		}
+
 	case ir.OSELECT:
 		n := n.(*ir.SelectStmt)
 		for _, cas := range n.Cases {
 			e.stmt(cas.Comm)
 			e.block(cas.Body)
 		}
-	case ir.OSELRECV2:
-		n := n.(*ir.AssignListStmt)
-		e.assign(n.Lhs[0], n.Rhs[0], "selrecv", n)
-		e.assign(n.Lhs[1], nil, "selrecv", n)
 	case ir.ORECV:
 		// TODO(mdempsky): Consider e.discard(n.Left).
 		n := n.(*ir.UnaryExpr)
@@ -425,28 +425,24 @@ func (e *escape) stmt(n ir.Node) {
 
 	case ir.OAS:
 		n := n.(*ir.AssignStmt)
-		e.assign(n.X, n.Y, "assign", n)
+		e.assignList([]ir.Node{n.X}, []ir.Node{n.Y}, "assign", n)
 	case ir.OASOP:
 		n := n.(*ir.AssignOpStmt)
-		e.assign(n.X, n.Y, "assign", n)
+		// TODO(mdempsky): Worry about OLSH/ORSH?
+		e.assignList([]ir.Node{n.X}, []ir.Node{n.Y}, "assign", n)
 	case ir.OAS2:
 		n := n.(*ir.AssignListStmt)
-		for i, nl := range n.Lhs {
+		e.assignList(n.Lhs, n.Rhs, "assign-pair", n)
-			e.assign(nl, n.Rhs[i], "assign-pair", n)
-		}
 
 	case ir.OAS2DOTTYPE: // v, ok = x.(type)
 		n := n.(*ir.AssignListStmt)
-		e.assign(n.Lhs[0], n.Rhs[0], "assign-pair-dot-type", n)
+		e.assignList(n.Lhs, n.Rhs, "assign-pair-dot-type", n)
-		e.assign(n.Lhs[1], nil, "assign-pair-dot-type", n)
 	case ir.OAS2MAPR: // v, ok = m[k]
 		n := n.(*ir.AssignListStmt)
-		e.assign(n.Lhs[0], n.Rhs[0], "assign-pair-mapr", n)
+		e.assignList(n.Lhs, n.Rhs, "assign-pair-mapr", n)
-		e.assign(n.Lhs[1], nil, "assign-pair-mapr", n)
+	case ir.OAS2RECV, ir.OSELRECV2: // v, ok = <-ch
-	case ir.OAS2RECV: // v, ok = <-ch
 		n := n.(*ir.AssignListStmt)
-		e.assign(n.Lhs[0], n.Rhs[0], "assign-pair-receive", n)
+		e.assignList(n.Lhs, n.Rhs, "assign-pair-receive", n)
-		e.assign(n.Lhs[1], nil, "assign-pair-receive", n)
 
 	case ir.OAS2FUNC:
 		n := n.(*ir.AssignListStmt)
@@ -455,9 +451,11 @@ func (e *escape) stmt(n ir.Node) {
 	case ir.ORETURN:
 		n := n.(*ir.ReturnStmt)
 		results := e.curfn.Type().Results().FieldSlice()
-		for i, v := range n.Results {
+		dsts := make([]ir.Node, len(results))
-			e.assign(ir.AsNode(results[i].Nname), v, "return", n)
+		for i, res := range results {
+			dsts[i] = res.Nname.(*ir.Name)
 		}
+		e.assignList(dsts, n.Results, "return", n)
 	case ir.OCALLFUNC, ir.OCALLMETH, ir.OCALLINTER, ir.OCLOSE, ir.OCOPY, ir.ODELETE, ir.OPANIC, ir.OPRINT, ir.OPRINTN, ir.ORECOVER:
 		e.call(nil, n, nil)
 	case ir.OGO, ir.ODEFER:
@@ -694,6 +692,10 @@ func (e *escape) exprSkipInit(k hole, n ir.Node) {
 	case ir.OCLOSURE:
 		n := n.(*ir.ClosureExpr)
 
+		if fn := n.Func; fn.IsHiddenClosure() {
+			e.walkFunc(fn)
+		}
+
 		// Link addresses of captured variables to closure.
 		k = e.spill(k, n)
 		for _, v := range n.Func.ClosureVars {
@@ -795,7 +797,7 @@ func (e *escape) addr(n ir.Node) hole {
 		k = e.oldLoc(n).asHole()
 	case ir.ONAMEOFFSET:
 		n := n.(*ir.NameOffsetExpr)
-		e.addr(n.Name_)
+		k = e.addr(n.Name_)
 	case ir.ODOT:
 		n := n.(*ir.SelectorExpr)
 		k = e.addr(n.X)
@@ -815,10 +817,6 @@ func (e *escape) addr(n ir.Node) hole {
 		e.assignHeap(n.Index, "key of map put", n)
 	}
 
-	if !n.Type().HasPointers() {
-		k = e.discardHole()
-	}
-
 	return k
 }
 
@@ -830,6 +828,16 @@ func (e *escape) addrs(l ir.Nodes) []hole {
 	return ks
 }
 
+func (e *escape) assignList(dsts, srcs []ir.Node, why string, where ir.Node) {
+	for i, dst := range dsts {
+		var src ir.Node
+		if i < len(srcs) {
+			src = srcs[i]
+		}
+		e.assign(dst, src, why, where)
+	}
+}
+
 // assign evaluates the assignment dst = src.
 func (e *escape) assign(dst, src ir.Node, why string, where ir.Node) {
 	// Filter out some no-op assignments for escape analysis.