internal/refactor/inline: handle generic functions

This CL is a first step towards support for inlining all forms
of generic functions. Its limitations include:

- No support for methods on generic types.
- Conservative shadowing.
- Unnecessary type conversions (see generic.txtar, a1).
- Conservative parenthesizing (see generic.txtar, file a1a).

For golang/go#68236.

Change-Id: Ib00b89cb61c611e8d1efd0e2f8b5d93032638b83
Reviewed-on: https://go-review.googlesource.com/c/tools/+/666716
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>

Author: jba

internal/refactor/inline/callee.go

@@ -42,6 +42,7 @@ type gobCallee struct {
 	ValidForCallStmt bool                   // function body is "return expr" where expr is f() or <-ch
 	NumResults       int                    // number of results (according to type, not ast.FieldList)
 	Params           []*paramInfo           // information about parameters (incl. receiver)
+	TypeParams       []*paramInfo           // information about type parameters
 	Results          []*paramInfo           // information about result variables
 	Effects          []int                  // order in which parameters are evaluated (see calleefx)
 	HasDefer         bool                   // uses defer
@@ -113,17 +114,6 @@ func AnalyzeCallee(logf func(string, ...any), fset *token.FileSet, pkg *types.Pa
 		return nil, fmt.Errorf("cannot inline function %s as it has no body", name)
 	}
 
-	// TODO(adonovan): support inlining of instantiated generic
-	// functions by replacing each occurrence of a type parameter
-	// T by its instantiating type argument (e.g. int). We'll need
-	// to wrap the instantiating type in parens when it's not an
-	// ident or qualified ident to prevent "if x == struct{}"
-	// parsing ambiguity, or "T(x)" where T = "*int" or "func()"
-	// from misparsing.
-	if funcHasTypeParams(decl) {
-		return nil, fmt.Errorf("cannot inline generic function %s: type parameters are not yet supported", name)
-	}
-
 	// Record the location of all free references in the FuncDecl.
 	// (Parameters are not free by this definition.)
 	var (
@@ -347,6 +337,7 @@ func AnalyzeCallee(logf func(string, ...any), fset *token.FileSet, pkg *types.Pa
 	}
 
 	params, results, effects, falcon := analyzeParams(logf, fset, info, decl)
+	tparams := analyzeTypeParams(logf, fset, info, decl)
 	return &Callee{gobCallee{
 		Content:          content,
 		PkgPath:          pkg.Path(),
@@ -357,6 +348,7 @@ func AnalyzeCallee(logf func(string, ...any), fset *token.FileSet, pkg *types.Pa
 		ValidForCallStmt: validForCallStmt,
 		NumResults:       sig.Results().Len(),
 		Params:           params,
+		TypeParams:       tparams,
 		Results:          results,
 		Effects:          effects,
 		HasDefer:         hasDefer,
@@ -404,20 +396,15 @@ type refInfo struct {
 	IsSelectionOperand bool
 }
 
-// analyzeParams computes information about parameters of function fn,
+// analyzeParams computes information about parameters of the function declared by decl,
 // including a simple "address taken" escape analysis.
 //
 // It returns two new arrays, one of the receiver and parameters, and
-// the other of the result variables of function fn.
+// the other of the result variables of the function.
 //
 // The input must be well-typed.
 func analyzeParams(logf func(string, ...any), fset *token.FileSet, info *types.Info, decl *ast.FuncDecl) (params, results []*paramInfo, effects []int, _ falconResult) {
-	fnobj, ok := info.Defs[decl.Name]
-	if !ok {
-		panic(fmt.Sprintf("%s: no func object for %q",
-			fset.PositionFor(decl.Name.Pos(), false), decl.Name)) // ill-typed?
-	}
-	sig := fnobj.Type().(*types.Signature)
+	sig := signature(fset, info, decl)
 
 	paramInfos := make(map[*types.Var]*paramInfo)
 	{
@@ -504,6 +491,52 @@ func analyzeParams(logf func(string, ...any), fset *token.FileSet, info *types.I
 	return params, results, effects, falcon
 }
 
+// analyzeTypeParams computes information about the type parameters of the function declared by decl.
+func analyzeTypeParams(_ logger, fset *token.FileSet, info *types.Info, decl *ast.FuncDecl) []*paramInfo {
+	sig := signature(fset, info, decl)
+	paramInfos := make(map[*types.TypeName]*paramInfo)
+	var params []*paramInfo
+	collect := func(tpl *types.TypeParamList) {
+		for i := range tpl.Len() {
+			typeName := tpl.At(i).Obj()
+			info := &paramInfo{Name: typeName.Name()}
+			params = append(params, info)
+			paramInfos[typeName] = info
+		}
+	}
+	collect(sig.RecvTypeParams())
+	collect(sig.TypeParams())
+
+	// Find references.
+	// We don't care about most of the properties that matter for parameter references:
+	// a type is immutable, cannot have its address taken, and does not undergo conversions.
+	// TODO(jba): can we nevertheless combine this with the traversal in analyzeParams?
+	var stack []ast.Node
+	stack = append(stack, decl.Type) // for scope of function itself
+	astutil.PreorderStack(decl.Body, stack, func(n ast.Node, stack []ast.Node) bool {
+		if id, ok := n.(*ast.Ident); ok {
+			if v, ok := info.Uses[id].(*types.TypeName); ok {
+				if pinfo, ok := paramInfos[v]; ok {
+					ref := refInfo{Offset: int(n.Pos() - decl.Pos())}
+					pinfo.Refs = append(pinfo.Refs, ref)
+					pinfo.Shadow = pinfo.Shadow.add(info, nil, pinfo.Name, stack)
+				}
+			}
+		}
+		return true
+	})
+	return params
+}
+
+func signature(fset *token.FileSet, info *types.Info, decl *ast.FuncDecl) *types.Signature {
+	fnobj, ok := info.Defs[decl.Name]
+	if !ok {
+		panic(fmt.Sprintf("%s: no func object for %q",
+			fset.PositionFor(decl.Name.Pos(), false), decl.Name)) // ill-typed?
+	}
+	return fnobj.Type().(*types.Signature)
+}
+
 // -- callee helpers --
 
 // analyzeAssignment looks at the the given stack, and analyzes certain

internal/refactor/inline/inline.go

@@ -839,6 +839,14 @@ func (st *state) inlineCall() (*inlineCallResult, error) {
 		}
 	}
 
+	typeArgs := st.typeArguments(caller.Call)
+	if len(typeArgs) != len(callee.TypeParams) {
+		return nil, fmt.Errorf("cannot inline: type parameter inference is not yet supported")
+	}
+	if err := substituteTypeParams(logf, callee.TypeParams, typeArgs, params, replaceCalleeID); err != nil {
+		return nil, err
+	}
+
 	// Log effective arguments.
 	for i, arg := range args {
 		logf("arg #%d: %s pure=%t effects=%t duplicable=%t free=%v type=%v",
@@ -1378,6 +1386,35 @@ type argument struct {
 	desugaredRecv bool            // is *recv or &recv, where operator was elided
 }
 
+// typeArguments returns the type arguments of the call.
+// It only collects the arguments that are explicitly provided; it does
+// not attempt type inference.
+func (st *state) typeArguments(call *ast.CallExpr) []*argument {
+	var exprs []ast.Expr
+	switch d := ast.Unparen(call.Fun).(type) {
+	case *ast.IndexExpr:
+		exprs = []ast.Expr{d.Index}
+	case *ast.IndexListExpr:
+		exprs = d.Indices
+	default:
+		// No type  arguments
+		return nil
+	}
+	var args []*argument
+	for _, e := range exprs {
+		arg := &argument{expr: e, freevars: freeVars(st.caller.Info, e)}
+		// Wrap the instantiating type in parens when it's not an
+		// ident or qualified ident to prevent "if x == struct{}"
+		// parsing ambiguity, or "T(x)" where T = "*int" or "func()"
+		// from misparsing.
+		if _, ok := arg.expr.(*ast.Ident); !ok {
+			arg.expr = &ast.ParenExpr{X: arg.expr}
+		}
+		args = append(args, arg)
+	}
+	return args
+}
+
 // arguments returns the effective arguments of the call.
 //
 // If the receiver argument and parameter have
@@ -1413,6 +1450,9 @@ func (st *state) arguments(caller *Caller, calleeDecl *ast.FuncDecl, assign1 fun
 
 	callArgs := caller.Call.Args
 	if calleeDecl.Recv != nil {
+		if len(st.callee.impl.TypeParams) > 0 {
+			return nil, fmt.Errorf("cannot inline: generic methods not yet supported")
+		}
 		sel := ast.Unparen(caller.Call.Fun).(*ast.SelectorExpr)
 		seln := caller.Info.Selections[sel]
 		var recvArg ast.Expr
@@ -1536,9 +1576,52 @@ type parameter struct {
 // A replacer replaces an identifier at the given offset in the callee.
 // The replacement tree must not belong to the caller; use cloneNode as needed.
 // If unpackVariadic is set, the replacement is a composite resulting from
-// variadic elimination, and may be unpackeded into variadic calls.
+// variadic elimination, and may be unpacked into variadic calls.
 type replacer = func(offset int, repl ast.Expr, unpackVariadic bool)
 
+// substituteTypeParams replaces type parameters in the callee with the corresponding type arguments
+// from the call.
+func substituteTypeParams(logf logger, typeParams []*paramInfo, typeArgs []*argument, params []*parameter, replace replacer) error {
+	assert(len(typeParams) == len(typeArgs), "mismatched number of type params/args")
+	for i, paramInfo := range typeParams {
+		arg := typeArgs[i]
+		// Perform a simplified, conservative shadow analysis: fail if there is any shadowing.
+		for free := range arg.freevars {
+			if paramInfo.Shadow[free] != 0 {
+				return fmt.Errorf("cannot inline: type argument #%d (type parameter %s) is shadowed", i, paramInfo.Name)
+			}
+		}
+		logf("replacing type param %s with %s", paramInfo.Name, debugFormatNode(token.NewFileSet(), arg.expr))
+		for _, ref := range paramInfo.Refs {
+			replace(ref.Offset, internalastutil.CloneNode(arg.expr), false)
+		}
+		// Also replace parameter field types.
+		// TODO(jba): find a way to do this that is not so slow and clumsy.
+		// Ideally, we'd walk each p.fieldType once, replacing all type params together.
+		for _, p := range params {
+			if id, ok := p.fieldType.(*ast.Ident); ok && id.Name == paramInfo.Name {
+				p.fieldType = arg.expr
+			} else {
+				for _, id := range identsNamed(p.fieldType, paramInfo.Name) {
+					replaceNode(p.fieldType, id, arg.expr)
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func identsNamed(n ast.Node, name string) []*ast.Ident {
+	var ids []*ast.Ident
+	ast.Inspect(n, func(n ast.Node) bool {
+		if id, ok := n.(*ast.Ident); ok && id.Name == name {
+			ids = append(ids, id)
+		}
+		return true
+	})
+	return ids
+}
+
 // substitute implements parameter elimination by substitution.
 //
 // It considers each parameter and its corresponding argument in turn
@@ -1664,7 +1747,6 @@ next:
 		// parameter is also removed by substitution.
 
 		sg[arg] = nil // Absent shadowing, the arg is substitutable.
-
 		for free := range arg.freevars {
 			switch s := param.info.Shadow[free]; {
 			case s < 0:

internal/refactor/inline/inline_test.go

@@ -308,14 +308,22 @@ func doInlineNote(logf func(string, ...any), pkg *packages.Package, file *ast.Fi
 	if want, ok := want.([]byte); ok {
 		got = append(bytes.TrimSpace(got), '\n')
 		want = append(bytes.TrimSpace(want), '\n')
-		if diff := diff.Unified("want", "got", string(want), string(got)); diff != "" {
-			return fmt.Errorf("Inline returned wrong output:\n%s\nWant:\n%s\nDiff:\n%s",
-				got, want, diff)
+		// If the "want" file begins "...", it need only be a substring of the "got" result,
+		// rather than an exact match.
+		if rest, ok := bytes.CutPrefix(want, []byte("...\n")); ok {
+			want = rest
+			if !bytes.Contains(got, want) {
+				return fmt.Errorf("Inline returned wrong output:\n%s\nWant substring:\n%s", got, want)
+			}
+		} else {
+			if diff := diff.Unified("want", "got", string(want), string(got)); diff != "" {
+				return fmt.Errorf("Inline returned wrong output:\n%s\nWant:\n%s\nDiff:\n%s",
+					got, want, diff)
+			}
 		}
 		return nil
 	}
 	return fmt.Errorf("Inline succeeded unexpectedly: want error matching %q, got <<%s>>", want, got)
-
 }
 
 // findFuncByPosition returns the FuncDecl at the specified (package-agnostic) position.
@@ -364,16 +372,16 @@ type testcase struct {
 func TestErrors(t *testing.T) {
 	runTests(t, []testcase{
 		{
-			"Generic functions are not yet supported.",
+			"Inference of type parameters is not yet supported.",
 			`func f[T any](x T) T { return x }`,
 			`var _ = f(0)`,
-			`error: type parameters are not yet supported`,
+			`error: type parameter inference is not yet supported`,
 		},
 		{
 			"Methods on generic types are not yet supported.",
 			`type G[T any] struct{}; func (G[T]) f(x T) T { return x }`,
 			`var _ = G[int]{}.f(0)`,
-			`error: type parameters are not yet supported`,
+			`error: generic methods not yet supported`,
 		},
 	})
 }
@@ -434,6 +442,13 @@ func TestBasics(t *testing.T) {
 	}
 }`,
 		},
+		{
+			"Explicit type parameters.",
+			`func f[T any](x T) T { return x }`,
+			`var _ = f[int](0)`,
+			// TODO(jba): remove the unnecessary conversion.
+			`var _ = int(0)`,
+		},
 	})
 }
 
@@ -602,7 +617,6 @@ func f1(i int) int { return i + 1 }`,
 				`func _() { print(F(f1), F(f1)) }`,
 			},
 		})
-
 	})
 }
 
@@ -1832,6 +1846,14 @@ func runTests(t *testing.T, tests []testcase) {
 						if fun.Name == funcName {
 							call = n
 						}
+					case *ast.IndexExpr:
+						if id, ok := fun.X.(*ast.Ident); ok && id.Name == funcName {
+							call = n
+						}
+					case *ast.IndexListExpr:
+						if id, ok := fun.X.(*ast.Ident); ok && id.Name == funcName {
+							call = n
+						}
 					}
 				}
 				return call == nil

internal/refactor/inline/testdata/err-basic.txtar

@@ -11,15 +11,6 @@ it doesn't even reach the Indent function.
 module testdata
 go 1.12
 
--- a/generic.go --
-package a
-
-func _() {
-	f[int]() //@ inline(re"f", re"type parameters are not yet supported")
-}
-
-func f[T any]() {}
-
 -- a/nobody.go --
 package a
 

internal/refactor/inline/testdata/err-shadow-builtin.txtar

@@ -14,6 +14,12 @@ func _() {
 }
 
 func f() *int { return nil }
+-- a/nil-type-param.go --
+package a
+
+func _[nil any]() {
+	_ = f() //@ inline(re"f", re"nil.*shadowed.*by.*typename.*line 3")
+}
 
 -- a/nil-typename.go --
 package a