internal/refactor: handle qualified names in inlined assignments

Rewrite our call site processing of imports, to simplify, and so that we
can re-use the import lookup logic in the assignStmts strategy for the
purpose of writing out types.

A large TODO is included for a hypothetical refactoring of the inlining
logic that could formalize these types of interactions between call site
analysis and inlining strategy.

Fixes golang/go#65217

Change-Id: Ifd99ea14430deba3a03cdfb936b6edee9e81d0bf
Reviewed-on: https://go-review.googlesource.com/c/tools/+/629435
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

internal/refactor/inline/callee.go

@@ -202,19 +202,19 @@ func AnalyzeCallee(logf func(string, ...any), fset *token.FileSet, pkg *types.Pa
 					objidx, ok := freeObjIndex[obj]
 					if !ok {
 						objidx = len(freeObjIndex)
-						var pkgpath, pkgname string
+						var pkgPath, pkgName string
 						if pn, ok := obj.(*types.PkgName); ok {
-							pkgpath = pn.Imported().Path()
+							pkgPath = pn.Imported().Path()
-							pkgname = pn.Imported().Name()
+							pkgName = pn.Imported().Name()
 						} else if obj.Pkg() != nil {
-							pkgpath = obj.Pkg().Path()
+							pkgPath = obj.Pkg().Path()
-							pkgname = obj.Pkg().Name()
+							pkgName = obj.Pkg().Name()
 						}
 						freeObjs = append(freeObjs, object{
 							Name:     obj.Name(),
 							Kind:     objectKind(obj),
-							PkgName:  pkgname,
+							PkgName:  pkgName,
-							PkgPath:  pkgpath,
+							PkgPath:  pkgPath,
 							ValidPos: obj.Pos().IsValid(),
 						})
 						freeObjIndex[obj] = objidx

internal/refactor/inline/inline.go

@@ -279,7 +279,8 @@ func (st *state) inline() (*Result, error) {
 	// (We can't let imports.Process take care of it as it may
 	// mistake obsolete imports for missing new imports when the
 	// names are similar, as is common during a package migration.)
-	for _, specToDelete := range res.oldImports {
+	for _, oldImport := range res.oldImports {
+		specToDelete := oldImport.spec
 		for _, decl := range f.Decls {
 			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
 				decl.Specs = slicesDeleteFunc(decl.Specs, func(spec ast.Spec) bool {
@@ -375,17 +376,23 @@ func (st *state) inline() (*Result, error) {
 		Content:     newSrc,
 		Literalized: literalized,
 	}, nil
-
 }
 
+// An oldImport is an import that will be deleted from the caller file.
+type oldImport struct {
+	pkgName *types.PkgName
+	spec    *ast.ImportSpec
+}
+
+// A newImport is an import that will be added to the caller file.
 type newImport struct {
 	pkgName string
 	spec    *ast.ImportSpec
 }
 
 type inlineCallResult struct {
-	newImports []newImport       // to add
+	newImports []newImport // to add
-	oldImports []*ast.ImportSpec // to remove
+	oldImports []oldImport // to remove
 
 	// If elideBraces is set, old is an ast.Stmt and new is an ast.BlockStmt to
 	// be spliced in. This allows the inlining analysis to assert that inlining
@@ -412,6 +419,9 @@ type inlineCallResult struct {
 // transformation replacing the call and adding new variable
 // declarations, for example, or replacing a call statement by zero or
 // many statements.)
+// NOTE(rfindley): we've sort-of done this, with the 'elideBraces' flag that
+// allows inlining a statement list. However, due to loss of comments, more
+// sophisticated rewrites are challenging.
 //
 // TODO(adonovan): in earlier drafts, the transformation was expressed
 // by splicing substrings of the two source files because syntax
@@ -421,6 +431,33 @@ type inlineCallResult struct {
 // candidate for evaluating an alternative fully self-contained tree
 // representation, such as any proposed solution to #20744, or even
 // dst or some private fork of go/ast.)
+// TODO(rfindley): see if we can reduce the amount of comment lossiness by
+// using printer.CommentedNode, which has been useful elsewhere.
+//
+// TODO(rfindley): inlineCall is getting very long, and very stateful, making
+// it very hard to read. The following refactoring may improve readability and
+// maintainability:
+//   - Rename 'state' to 'callsite', since that is what it encapsulates.
+//   - Add results of pre-processing analysis into the callsite struct, such as
+//     the effective importMap, new/old imports, arguments, etc. Essentially
+//     anything that resulted from initial analysis of the call site, and which
+//     may be useful to inlining strategies.
+//   - Delegate this call site analysis to a constructor or initializer, such
+//     as 'analyzeCallsite', so that it does not consume bandwidth in the
+//     'inlineCall' logical flow.
+//   - Once analyzeCallsite returns, the callsite is immutable, much in the
+//     same way as the Callee and Caller are immutable.
+//   - Decide on a standard interface for strategies (and substrategies), such
+//     that they may be delegated to a separate method on callsite.
+//
+// In this way, the logical flow of inline call will clearly follow the
+// following structure:
+//  1. Analyze the call site.
+//  2. Try strategies, in order, until one succeeds.
+//  3. Process the results.
+//
+// If any expensive analysis may be avoided by earlier strategies, it can be
+// encapsulated in its own type and passed to subsequent strategies.
 func (st *state) inlineCall() (*inlineCallResult, error) {
 	logf, caller, callee := st.opts.Logf, st.caller, &st.callee.impl
 
@@ -469,39 +506,83 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 		assign1 = func(v *types.Var) bool { return !updatedLocals[v] }
 	}
 
-	// import map, initially populated with caller imports.
+	// import map, initially populated with caller imports, and updated below
+	// with new imports necessary to reference free symbols in the callee.
 	//
-	// For simplicity we ignore existing dot imports, so that a
+	// For simplicity we ignore existing dot imports, so that a qualified
-	// qualified identifier (QI) in the callee is always
+	// identifier (QI) in the callee is always represented by a QI in the caller,
-	// represented by a QI in the caller, allowing us to treat a
+	// allowing us to treat a QI like a selection on a package name.
-	// QI like a selection on a package name.
 	importMap := make(map[string][]string) // maps package path to local name(s)
+	var oldImports []oldImport             // imports referenced only by caller.Call.Fun
+
 	for _, imp := range caller.File.Imports {
-		if pkgname, ok := importedPkgName(caller.Info, imp); ok &&
+		if pkgName, ok := importedPkgName(caller.Info, imp); ok &&
-			pkgname.Name() != "." &&
+			pkgName.Name() != "." &&
-			pkgname.Name() != "_" {
+			pkgName.Name() != "_" {
-			path := pkgname.Imported().Path()
+
-			importMap[path] = append(importMap[path], pkgname.Name())
+			// If the import's sole use is in caller.Call.Fun of the form p.F(...),
+			// where p.F is a qualified identifier, the p import may not be
+			// necessary.
+			//
+			// Only the qualified identifier case matters, as other references to
+			// imported package names in the Call.Fun expression (e.g.
+			// x.after(3*time.Second).f() or time.Second.String()) will remain after
+			// inlining, as arguments.
+			//
+			// If that is the case, proactively check if any of the callee FreeObjs
+			// need this import. Doing so eagerly simplifies the resulting logic.
+			needed := true
+			sel, ok := ast.Unparen(caller.Call.Fun).(*ast.SelectorExpr)
+			if ok && soleUse(caller.Info, pkgName) == sel.X {
+				needed = false // no longer needed by caller
+				// Check to see if any of the inlined free objects need this package.
+				for _, obj := range callee.FreeObjs {
+					if obj.PkgPath == pkgName.Imported().Path() && !obj.Shadow[pkgName.Name()] {
+						needed = true // needed by callee
+						break
+					}
+				}
+			}
+
+			if needed {
+				path := pkgName.Imported().Path()
+				importMap[path] = append(importMap[path], pkgName.Name())
+			} else {
+				oldImports = append(oldImports, oldImport{pkgName: pkgName, spec: imp})
+			}
 		}
 	}
 
-	var oldImports []*ast.ImportSpec // imports referenced only caller.Call.Fun
+	// importName finds an existing import name to use in a particular shadowing
-
+	// context. It is used to determine the set of new imports in
-	// localImportName returns the local name for a given imported package path.
+	// getOrMakeImportName, and is also used for writing out names in inlining
-	var newImports []newImport
+	// strategies below.
-	localImportName := func(obj *object) string {
+	importName := func(pkgPath string, shadow map[string]bool) string {
-		// Does an import exist?
+		for _, name := range importMap[pkgPath] {
-		for _, name := range importMap[obj.PkgPath] {
+			// Check that either the import preexisted, or that it was newly added
-			// Check that either the import preexisted,
+			// (no PkgName) but is not shadowed, either in the callee (shadows) or
-			// or that it was newly added (no PkgName) but is not shadowed,
+			// caller (caller.lookup).
-			// either in the callee (shadows) or caller (caller.lookup).
+			if !shadow[name] {
-			if !obj.Shadow[name] {
 				found := caller.lookup(name)
 				if is[*types.PkgName](found) || found == nil {
 					return name
 				}
 			}
 		}
+		return ""
+	}
+
+	// keep track of new imports that are necessary to reference any free names
+	// in the callee.
+	var newImports []newImport
+
+	// getOrMakeImportName returns the local name for a given imported package path,
+	// adding one if it doesn't exists.
+	getOrMakeImportName := func(pkgPath, pkgName string, shadow map[string]bool) string {
+		// Does an import already exist that works in this shadowing context?
+		if name := importName(pkgPath, shadow); name != "" {
+			return name
+		}
 
 		newlyAdded := func(name string) bool {
 			for _, new := range newImports {
@@ -515,33 +596,17 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 		// shadowedInCaller reports whether a candidate package name
 		// already refers to a declaration in the caller.
 		shadowedInCaller := func(name string) bool {
-			existing := caller.lookup(name)
+			obj := caller.lookup(name)
-
+			if obj == nil {
-			// If the candidate refers to a PkgName p whose sole use is
+				return false
-			// in caller.Call.Fun of the form p.F(...), where p.F is a
+			}
-			// qualified identifier, the p import will be deleted,
+			// If obj will be removed, the name is available.
-			// so it's safe (and better) to recycle the name.
+			for _, old := range oldImports {
-			//
+				if old.pkgName == obj {
-			// Only the qualified identifier case matters, as other
+					return false
-			// references to imported package names in the Call.Fun
-			// expression (e.g. x.after(3*time.Second).f()
-			// or time.Second.String()) will remain after
-			// inlining, as arguments.
-			if pkgName, ok := existing.(*types.PkgName); ok {
-				if sel, ok := ast.Unparen(caller.Call.Fun).(*ast.SelectorExpr); ok {
-					if sole := soleUse(caller.Info, pkgName); sole == sel.X {
-						for _, spec := range caller.File.Imports {
-							pkgName2, ok := importedPkgName(caller.Info, spec)
-							if ok && pkgName2 == pkgName {
-								oldImports = append(oldImports, spec)
-								return false
-							}
-						}
-					}
 				}
 			}
-
+			return true
-			return existing != nil
 		}
 
 		// import added by callee
@@ -555,29 +620,28 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 		// TODO(rfindley): is it worth preserving local package names for callee
 		// imports? Are they likely to be better or worse than the name we choose
 		// here?
-		base := obj.PkgName
+		base := pkgName
 		name := base
-		for n := 0; obj.Shadow[name] || shadowedInCaller(name) || newlyAdded(name) || name == "init"; n++ {
+		for n := 0; shadow[name] || shadowedInCaller(name) || newlyAdded(name) || name == "init"; n++ {
 			name = fmt.Sprintf("%s%d", base, n)
 		}
-
+		logf("adding import %s %q", name, pkgPath)
-		logf("adding import %s %q", name, obj.PkgPath)
 		spec := &ast.ImportSpec{
 			Path: &ast.BasicLit{
 				Kind:  token.STRING,
-				Value: strconv.Quote(obj.PkgPath),
+				Value: strconv.Quote(pkgPath),
 			},
 		}
 		// Use explicit pkgname (out of necessity) when it differs from the declared name,
 		// or (for good style) when it differs from base(pkgpath).
-		if name != obj.PkgName || name != pathpkg.Base(obj.PkgPath) {
+		if name != pkgName || name != pathpkg.Base(pkgPath) {
 			spec.Name = makeIdent(name)
 		}
 		newImports = append(newImports, newImport{
 			pkgName: name,
 			spec:    spec,
 		})
-		importMap[obj.PkgPath] = append(importMap[obj.PkgPath], name)
+		importMap[pkgPath] = append(importMap[pkgPath], name)
 		return name
 	}
 
@@ -607,7 +671,8 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 		var newName ast.Expr
 		if obj.Kind == "pkgname" {
 			// Use locally appropriate import, creating as needed.
-			newName = makeIdent(localImportName(&obj)) // imported package
+			n := getOrMakeImportName(obj.PkgPath, obj.PkgName, obj.Shadow)
+			newName = makeIdent(n) // imported package
 		} else if !obj.ValidPos {
 			// Built-in function, type, or value (e.g. nil, zero):
 			// check not shadowed at caller.
@@ -651,7 +716,7 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 
 			// Form a qualified identifier, pkg.Name.
 			if qualify {
-				pkgName := localImportName(&obj)
+				pkgName := getOrMakeImportName(obj.PkgPath, obj.PkgName, obj.Shadow)
 				newName = &ast.SelectorExpr{
 					X:   makeIdent(pkgName),
 					Sel: makeIdent(obj.Name),
@@ -992,8 +1057,9 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 			(!needBindingDecl || (bindingDecl != nil && len(bindingDecl.names) == 0)) {
 
 			// Reduces to: { var (bindings); lhs... := rhs... }
-			if newStmts, ok := st.assignStmts(stmt, results); ok {
+			if newStmts, ok := st.assignStmts(stmt, results, importName); ok {
 				logf("strategy: reduce assign-context call to { return exprs }")
+
 				clearPositions(calleeDecl.Body)
 
 				block := &ast.BlockStmt{
@@ -2890,6 +2956,15 @@ func declares(stmts []ast.Stmt) map[string]bool {
 	return names
 }
 
+// A importNameFunc is used to query local import names in the caller, in a
+// particular shadowing context.
+//
+// The shadow map contains additional names shadowed in the inlined code, at
+// the position the local import name is to be used. The shadow map only needs
+// to contain newly introduced names in the inlined code; names shadowed at the
+// caller are handled automatically.
+type importNameFunc = func(pkgPath string, shadow map[string]bool) string
+
 // assignStmts rewrites a statement assigning the results of a call into zero
 // or more statements that assign its return operands, or (nil, false) if no
 // such rewrite is possible. The set of bindings created by the result of
@@ -2932,7 +3007,7 @@ func declares(stmts []ast.Stmt) map[string]bool {
 //
 // Note: assignStmts may return (nil, true) if it determines that the rewritten
 // assignment consists only of _ = nil assignments.
-func (st *state) assignStmts(callerStmt *ast.AssignStmt, returnOperands []ast.Expr) ([]ast.Stmt, bool) {
+func (st *state) assignStmts(callerStmt *ast.AssignStmt, returnOperands []ast.Expr, importName importNameFunc) ([]ast.Stmt, bool) {
 	logf, caller, callee := st.opts.Logf, st.caller, &st.callee.impl
 
 	assert(len(callee.Returns) == 1, "unexpected multiple returns")
@@ -3035,18 +3110,23 @@ func (st *state) assignStmts(callerStmt *ast.AssignStmt, returnOperands []ast.Ex
 	// Inlining techniques below will need to write type information in order to
 	// preserve the correct types of LHS identifiers.
 	//
-	// writeType is a simple helper to write out type expressions.
+	// typeExpr is a simple helper to write out type expressions. It currently
+	// handles (possibly qualified) type names.
+	//
 	// TODO(rfindley):
-	//   1. handle qualified type names (potentially adding new imports)
+	//   1. expand this to handle more type expressions.
-	//   2. expand this to handle more type expressions.
+	//   2. refactor to share logic with callee rewriting.
-	//   3. refactor to share logic with callee rewriting.
 	universeAny := types.Universe.Lookup("any")
-	typeExpr := func(typ types.Type, shadows ...map[string]bool) ast.Expr {
+	typeExpr := func(typ types.Type, shadow map[string]bool) ast.Expr {
-		var typeName string
+		var (
+			typeName string
+			obj      *types.TypeName // nil for basic types
+		)
 		switch typ := typ.(type) {
 		case *types.Basic:
 			typeName = typ.Name()
 		case interface{ Obj() *types.TypeName }: // Named, Alias, TypeParam
+			obj = typ.Obj()
 			typeName = typ.Obj().Name()
 		}
 
@@ -3060,15 +3140,20 @@ func (st *state) assignStmts(callerStmt *ast.AssignStmt, returnOperands []ast.Ex
 			return nil
 		}
 
-		for _, shadow := range shadows {
+		if obj == nil || obj.Pkg() == nil || obj.Pkg() == caller.Types { // local type or builtin
 			if shadow[typeName] {
 				logf("cannot write shadowed type name %q", typeName)
 				return nil
 			}
-		}
+			obj, _ := caller.lookup(typeName).(*types.TypeName)
-		obj, _ := caller.lookup(typeName).(*types.TypeName)
+			if obj != nil && types.Identical(obj.Type(), typ) {
-		if obj != nil && types.Identical(obj.Type(), typ) {
+				return ast.NewIdent(typeName)
-			return ast.NewIdent(typeName)
+			}
+		} else if pkgName := importName(obj.Pkg().Path(), shadow); pkgName != "" {
+			return &ast.SelectorExpr{
+				X:   ast.NewIdent(pkgName),
+				Sel: ast.NewIdent(typeName),
+			}
 		}
 		return nil
 	}
@@ -3187,7 +3272,7 @@ func (st *state) assignStmts(callerStmt *ast.AssignStmt, returnOperands []ast.Ex
 		idx := origIdxs[i]
 		if nonTrivial[idx] && defs[idx] != nil {
 			typ := caller.Info.TypeOf(lhs[i])
-			texpr := typeExpr(typ)
+			texpr := typeExpr(typ, nil)
 			if texpr == nil {
 				return nil, false
 			}

internal/refactor/inline/testdata/assignment.txtar

@@ -0,0 +1,139 @@
+Basic tests of inlining a call on the RHS of an assignment.
+
+-- go.mod --
+module testdata
+go 1.12
+
+-- a/a1.go --
+package a
+
+import "testdata/b"
+
+func _() {
+	var y int
+	x, y := b.B1() //@ inline(re"B", b1)
+	_, _ = x, y
+}
+
+-- a/a2.go --
+package a
+
+import "testdata/b"
+
+func _() {
+	var y int
+	x, y := b.B2() //@ inline(re"B", b2)
+	_, _ = x, y
+}
+
+-- a/a3.go --
+package a
+
+import "testdata/b"
+
+func _() {
+	x, y := b.B3() //@ inline(re"B", b3)
+	_, _ = x, y
+}
+
+-- a/a4.go --
+package a
+
+import "testdata/b"
+
+func _() {
+	x, y := b.B4() //@ inline(re"B", b4)
+	_, _ = x, y
+}
+
+-- b/b.go --
+package b
+
+import (
+	"testdata/c"
+)
+
+func B1() (c.C, int) {
+	return 0, 1
+}
+
+func B2() (c.C, int) {
+	return B1()
+}
+
+func B3() (c.C, c.C) {
+	return 0, 1
+}
+
+-- b/b4.go --
+package b
+
+import (
+	c1 "testdata/c"
+	c2 "testdata/c2"
+)
+
+func B4() (c1.C, c2.C) {
+	return 0, 1
+}
+
+-- c/c.go --
+package c
+
+type C int
+
+-- c2/c.go --
+package c
+
+type C int
+
+-- b1 --
+package a
+
+import (
+	"testdata/c"
+)
+
+func _() {
+	var y int
+	x, y := c.C(0), 1 //@ inline(re"B", b1)
+	_, _ = x, y
+}
+-- b2 --
+package a
+
+import (
+	"testdata/b"
+	"testdata/c"
+)
+
+func _() {
+	var y int
+	var x c.C
+	x, y = b.B1() //@ inline(re"B", b2)
+	_, _ = x, y
+}
+-- b3 --
+package a
+
+import (
+	"testdata/c"
+)
+
+func _() {
+	x, y := c.C(0), c.C(1) //@ inline(re"B", b3)
+	_, _ = x, y
+}
+
+-- b4 --
+package a
+
+import (
+	"testdata/c"
+	c0 "testdata/c2"
+)
+
+func _() {
+	x, y := c.C(0), c0.C(1) //@ inline(re"B", b4)
+	_, _ = x, y
+}