[dev.regabi] cmd/compile: transform closures during walk

We used to transform directly called closures in a separate pass
before walk, because we couldn't guarantee whether we'd see the
closure call or the closure itself first. As of the last CL, this
ordering is always guaranteed, so we can rewrite calls and the closure
at the same time.

Change-Id: Ia6f4d504c24795e41500108589b53395d301123b
Reviewed-on: https://go-review.googlesource.com/c/go/+/283315
Run-TryBot: Matthew Dempsky <mdempsky@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Trust: Matthew Dempsky <mdempsky@google.com>
Reviewed-by: Keith Randall <khr@golang.org>

3 files changed, 61 insertions, 76 deletions

diff --git a/src/cmd/compile/internal/gc/main.go b/src/cmd/compile/internal/gc/main.go
index 2903d64ff8..9ecdd510b1 100644
--- a/src/cmd/compile/internal/gc/main.go
+++ b/src/cmd/compile/internal/gc/main.go
@@ -22,7 +22,6 @@ import (
 	"cmd/compile/internal/ssagen"
 	"cmd/compile/internal/typecheck"
 	"cmd/compile/internal/types"
-	"cmd/compile/internal/walk"
 	"cmd/internal/dwarf"
 	"cmd/internal/obj"
 	"cmd/internal/objabi"
@@ -269,20 +268,6 @@ func Main(archInit func(*ssagen.ArchInfo)) {
 		ssagen.EnableNoWriteBarrierRecCheck()
 	}
 
-	// Transform closure bodies to properly reference captured variables.
-	// This needs to happen before walk, because closures must be transformed
-	// before walk reaches a call of a closure.
-	base.Timer.Start("fe", "xclosures")
-	for _, n := range typecheck.Target.Decls {
-		if n.Op() == ir.ODCLFUNC {
-			n := n.(*ir.Func)
-			if n.OClosure != nil {
-				ir.CurFunc = n
-				walk.Closure(n)
-			}
-		}
-	}
-
 	// Prepare for SSA compilation.
 	// This must be before peekitabs, because peekitabs
 	// can trigger function compilation.
diff --git a/src/cmd/compile/internal/walk/closure.go b/src/cmd/compile/internal/walk/closure.go
index 7fa63ea9c7..e9b3698080 100644
--- a/src/cmd/compile/internal/walk/closure.go
+++ b/src/cmd/compile/internal/walk/closure.go
@@ -12,50 +12,43 @@ import (
 	"cmd/internal/src"
 )
 
-// Closure is called in a separate phase after escape analysis.
-// It transform closure bodies to properly reference captured variables.
-func Closure(fn *ir.Func) {
-	if len(fn.ClosureVars) == 0 {
-		return
-	}
+// directClosureCall rewrites a direct call of a function literal into
+// a normal function call with closure variables passed as arguments.
+// This avoids allocation of a closure object.
+//
+// For illustration, the following call:
+//
+//	func(a int) {
+//		println(byval)
+//		byref++
+//	}(42)
+//
+// becomes:
+//
+//	func(byval int, &byref *int, a int) {
+//		println(byval)
+//		(*&byref)++
+//	}(byval, &byref, 42)
+func directClosureCall(n *ir.CallExpr) {
+	clo := n.X.(*ir.ClosureExpr)
+	clofn := clo.Func
 
-	if !fn.ClosureCalled() {
-		// The closure is not directly called, so it is going to stay as closure.
-		fn.SetNeedctxt(true)
-		return
+	if ir.IsTrivialClosure(clo) {
+		return // leave for walkClosure to handle
 	}
 
-	lno := base.Pos
-	base.Pos = fn.Pos()
-
-	// If the closure is directly called, we transform it to a plain function call
-	// with variables passed as args. This avoids allocation of a closure object.
-	// Here we do only a part of the transformation. Walk of OCALLFUNC(OCLOSURE)
-	// will complete the transformation later.
-	// For illustration, the following closure:
-	//	func(a int) {
-	//		println(byval)
-	//		byref++
-	//	}(42)
-	// becomes:
-	//	func(byval int, &byref *int, a int) {
-	//		println(byval)
-	//		(*&byref)++
-	//	}(byval, &byref, 42)
-
-	// f is ONAME of the actual function.
-	f := fn.Nname
-
 	// We are going to insert captured variables before input args.
 	var params []*types.Field
 	var decls []*ir.Name
-	for _, v := range fn.ClosureVars {
+	for _, v := range clofn.ClosureVars {
 		if !v.Byval() {
 			// If v of type T is captured by reference,
 			// we introduce function param &v *T
 			// and v remains PAUTOHEAP with &v heapaddr
 			// (accesses will implicitly deref &v).
-			addr := typecheck.NewName(typecheck.Lookup("&" + v.Sym().Name))
+
+			addr := ir.NewNameAt(clofn.Pos(), typecheck.Lookup("&"+v.Sym().Name))
+			addr.Curfn = clofn
 			addr.SetType(types.NewPtr(v.Type()))
 			v.Heapaddr = addr
 			v = addr
@@ -69,32 +62,58 @@ func Closure(fn *ir.Func) {
 		params = append(params, fld)
 	}
 
+	// f is ONAME of the actual function.
+	f := clofn.Nname
+
 	// Prepend params and decls.
-	f.Type().Params().SetFields(append(params, f.Type().Params().FieldSlice()...))
-	fn.Dcl = append(decls, fn.Dcl...)
+	typ := f.Type()
+	typ.Params().SetFields(append(params, typ.Params().FieldSlice()...))
+	clofn.Dcl = append(decls, clofn.Dcl...)
+
+	// Rewrite call.
+	n.X = f
+	n.Args.Prepend(closureArgs(clo)...)
+
+	// Update the call expression's type. We need to do this
+	// because typecheck gave it the result type of the OCLOSURE
+	// node, but we only rewrote the ONAME node's type. Logically,
+	// they're the same, but the stack offsets probably changed.
+	//
+	// TODO(mdempsky): Reuse a single type for both.
+	if typ.NumResults() == 1 {
+		n.SetType(typ.Results().Field(0).Type)
+	} else {
+		n.SetType(typ.Results())
+	}
 
-	base.Pos = lno
+	// Add to Closures for enqueueFunc. It's no longer a proper
+	// closure, but we may have already skipped over it in the
+	// functions list as a non-trivial closure, so this just
+	// ensures it's compiled.
+	ir.CurFunc.Closures = append(ir.CurFunc.Closures, clofn)
 }
 
 func walkClosure(clo *ir.ClosureExpr, init *ir.Nodes) ir.Node {
-	fn := clo.Func
+	clofn := clo.Func
 
 	// If no closure vars, don't bother wrapping.
 	if ir.IsTrivialClosure(clo) {
 		if base.Debug.Closure > 0 {
 			base.WarnfAt(clo.Pos(), "closure converted to global")
 		}
-		return fn.Nname
+		return clofn.Nname
 	}
 
-	ir.CurFunc.Closures = append(ir.CurFunc.Closures, fn)
+	// The closure is not trivial or directly called, so it's going to stay a closure.
 	ir.ClosureDebugRuntimeCheck(clo)
+	clofn.SetNeedctxt(true)
+	ir.CurFunc.Closures = append(ir.CurFunc.Closures, clofn)
 
 	typ := typecheck.ClosureType(clo)
 
 	clos := ir.NewCompLitExpr(base.Pos, ir.OCOMPLIT, ir.TypeNode(typ), nil)
 	clos.SetEsc(clo.Esc())
-	clos.List = append([]ir.Node{ir.NewUnaryExpr(base.Pos, ir.OCFUNC, fn.Nname)}, closureArgs(clo)...)
+	clos.List = append([]ir.Node{ir.NewUnaryExpr(base.Pos, ir.OCFUNC, clofn.Nname)}, closureArgs(clo)...)
 
 	addr := typecheck.NodAddr(clos)
 	addr.SetEsc(clo.Esc())
diff --git a/src/cmd/compile/internal/walk/expr.go b/src/cmd/compile/internal/walk/expr.go
index 508cdd1d06..893a95f403 100644
--- a/src/cmd/compile/internal/walk/expr.go
+++ b/src/cmd/compile/internal/walk/expr.go
@@ -488,27 +488,8 @@ func walkCall(n *ir.CallExpr, init *ir.Nodes) ir.Node {
 		reflectdata.MarkUsedIfaceMethod(n)
 	}
 
-	if n.Op() == ir.OCALLFUNC && n.X.Op() == ir.OCLOSURE && !ir.IsTrivialClosure(n.X.(*ir.ClosureExpr)) {
-		// Transform direct call of a closure to call of a normal function.
-		// transformclosure already did all preparation work.
-		// We leave trivial closures for walkClosure to handle.
-
-		clo := n.X.(*ir.ClosureExpr)
-		ir.CurFunc.Closures = append(ir.CurFunc.Closures, clo.Func)
-
-		// Prepend captured variables to argument list.
-		n.Args.Prepend(closureArgs(clo)...)
-
-		// Replace OCLOSURE with ONAME/PFUNC.
-		n.X = clo.Func.Nname
-
-		// Update type of OCALLFUNC node.
-		// Output arguments had not changed, but their offsets could.
-		if n.X.Type().NumResults() == 1 {
-			n.SetType(n.X.Type().Results().Field(0).Type)
-		} else {
-			n.SetType(n.X.Type().Results())
-		}
+	if n.Op() == ir.OCALLFUNC && n.X.Op() == ir.OCLOSURE {
+		directClosureCall(n)
 	}
 
 	walkCall1(n, init)