gopls/internal/golang: support extract variable at top level

This CL causes the refactor.extract.{variable,constant}
logic to support extractions outside of statement context,
inserting a new package-level declaration.
Previously, it returned an error.

Also, use "k" as the basis for choosing names for new constants.

Also, simplify generateAvailableName et al.

Fixes golang/go#70665

Change-Id: I769206b14662e4e70ee04ab6c0040e635ac72820
Reviewed-on: https://go-review.googlesource.com/c/tools/+/633597
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>

gopls/doc/release/v0.17.0.md

@@ -44,6 +44,9 @@ When the selection is a constant expression, gopls now offers "Extract
 constant" instead of "Extract variable", and generates a `const`
 declaration instead of a local variable.
 
+Also, extraction of a constant or variable now works at top-level,
+outside of any function.
+
 ## Pull diagnostics
 
 When initialized with the option `"pullDiagnostics": true`, gopls will advertise support for the

gopls/internal/golang/extract.go

@@ -35,15 +35,16 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 	}
 	constant := info.Types[expr].Value != nil
 
-	// Create new AST node for extracted expression.
+	// Generate name(s) for new declaration.
+	baseName := cond(constant, "k", "x")
 	var lhsNames []string
 	switch expr := expr.(type) {
 	case *ast.CallExpr:
 		tup, ok := info.TypeOf(expr).(*types.Tuple)
 		if !ok {
 			// conversion or single-valued call:
 			// treat it the same as our standard extract variable case.
-			name, _ := generateAvailableName(expr.Pos(), path, pkg, info, "x", 0)
+			name, _ := freshName(info, file, expr.Pos(), baseName, 0)
 			lhsNames = append(lhsNames, name)
 
 		} else {
@@ -52,14 +53,14 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 			for range tup.Len() {
 				// Generate a unique variable for each result.
 				var name string
-				name, idx = generateAvailableName(expr.Pos(), path, pkg, info, "x", idx)
+				name, idx = freshName(info, file, expr.Pos(), baseName, idx)
 				lhsNames = append(lhsNames, name)
 			}
 		}
 
 	default:
 		// TODO: stricter rules for selectorExpr.
-		name, _ := generateAvailableName(expr.Pos(), path, pkg, info, "x", 0)
+		name, _ := freshName(info, file, expr.Pos(), baseName, 0)
 		lhsNames = append(lhsNames, name)
 	}
 
@@ -87,20 +88,38 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 	//   if x := 1; cond {
 	//   } else if y := x1; cond {
 	//   }
-	//
-	// TODO(golang/go#70665): Another bug (or limitation): this
-	// operation fails at top-level:
-	//   package p
-	//   var x = «1 + 2» // error
-	insertBeforeStmt := analysisinternal.StmtToInsertVarBefore(path)
-	if insertBeforeStmt == nil {
-		return nil, nil, fmt.Errorf("cannot find location to insert extraction")
-	}
-	indent, err := calculateIndentation(src, tokFile, insertBeforeStmt)
-	if err != nil {
-		return nil, nil, err
+	var (
+		insertPos   token.Pos
+		indentation string
+		stmtOK      bool // ok to use ":=" instead of var/const decl?
+	)
+	if before := analysisinternal.StmtToInsertVarBefore(path); before != nil {
+		// Within function: compute appropriate statement indentation.
+		indent, err := calculateIndentation(src, tokFile, before)
+		if err != nil {
+			return nil, nil, err
+		}
+		insertPos = before.Pos()
+		indentation = "\n" + indent
+
+		// Currently, we always extract a constant expression
+		// to a const declaration (and logic in CodeAction
+		// assumes that we do so); this is conservative because
+		// it preserves its constant-ness.
+		//
+		// In future, constant expressions used only in
+		// contexts where constant-ness isn't important could
+		// be profitably extracted to a var declaration or :=
+		// statement, especially if the latter is the Init of
+		// an {If,For,Switch}Stmt.
+		stmtOK = !constant
+	} else {
+		// Outside any statement: insert before the current
+		// declaration, without indentation.
+		currentDecl := path[len(path)-2]
+		insertPos = currentDecl.Pos()
+		indentation = "\n"
 	}
-	newLineIndent := "\n" + indent
 
 	// Create statement to declare extracted var/const.
 	//
@@ -121,17 +140,19 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 	//
 	// Conversely, a short var decl stmt is not valid at top level,
 	// so when we fix #70665, we'll need to use a var decl.
-	var declStmt ast.Stmt
-	if constant {
-		// const x = expr
-		declStmt = &ast.DeclStmt{
-			Decl: &ast.GenDecl{
-				Tok: token.CONST,
-				Specs: []ast.Spec{
-					&ast.ValueSpec{
-						Names:  []*ast.Ident{ast.NewIdent(lhsNames[0])}, // there can be only one
-						Values: []ast.Expr{expr},
-					},
+	var newNode ast.Node
+	if !stmtOK {
+		// var/const x1, ..., xn = expr
+		var names []*ast.Ident
+		for _, name := range lhsNames {
+			names = append(names, ast.NewIdent(name))
+		}
+		newNode = &ast.GenDecl{
+			Tok: cond(constant, token.CONST, token.VAR),
+			Specs: []ast.Spec{
+				&ast.ValueSpec{
+					Names:  names,
+					Values: []ast.Expr{expr},
 				},
 			},
 		}
@@ -142,7 +163,7 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 		for _, name := range lhsNames {
 			lhs = append(lhs, ast.NewIdent(name))
 		}
-		declStmt = &ast.AssignStmt{
+		newNode = &ast.AssignStmt{
 			Tok: token.DEFINE,
 			Lhs: lhs,
 			Rhs: []ast.Expr{expr},
@@ -151,17 +172,17 @@ func extractVariable(fset *token.FileSet, start, end token.Pos, src []byte, file
 
 	// Format and indent the declaration.
 	var buf bytes.Buffer
-	if err := format.Node(&buf, fset, declStmt); err != nil {
+	if err := format.Node(&buf, fset, newNode); err != nil {
 		return nil, nil, err
 	}
 	// TODO(adonovan): not sound for `...` string literals containing newlines.
-	assignment := strings.ReplaceAll(buf.String(), "\n", newLineIndent) + newLineIndent
+	assignment := strings.ReplaceAll(buf.String(), "\n", indentation) + indentation
 
 	return fset, &analysis.SuggestedFix{
 		TextEdits: []analysis.TextEdit{
 			{
-				Pos:     insertBeforeStmt.Pos(),
-				End:     insertBeforeStmt.Pos(),
+				Pos:     insertPos,
+				End:     insertPos,
 				NewText: []byte(assignment),
 			},
 			{
@@ -215,40 +236,36 @@ func calculateIndentation(content []byte, tok *token.File, insertBeforeStmt ast.
 	return string(content[lineOffset:stmtOffset]), nil
 }
 
-// generateAvailableName adjusts the new function name until there are no collisions in scope.
-// Possible collisions include other function and variable names. Returns the next index to check for prefix.
-func generateAvailableName(pos token.Pos, path []ast.Node, pkg *types.Package, info *types.Info, prefix string, idx int) (string, int) {
-	scopes := CollectScopes(info, path, pos)
-	scopes = append(scopes, pkg.Scope())
+// freshName returns an identifier based on prefix (perhaps with a
+// numeric suffix) that is not in scope at the specified position
+// within the file. It returns the next numeric suffix to use.
+func freshName(info *types.Info, file *ast.File, pos token.Pos, prefix string, idx int) (string, int) {
+	scope := info.Scopes[file].Innermost(pos)
 	return generateName(idx, prefix, func(name string) bool {
-		for _, scope := range scopes {
-			if scope != nil && scope.Lookup(name) != nil {
-				return true
-			}
-		}
-		return false
+		obj, _ := scope.LookupParent(name, pos)
+		return obj != nil
 	})
 }
 
-// generateNameOutsideOfRange is like generateAvailableName, but ignores names
+// freshNameOutsideRange is like [freshName], but ignores names
 // declared between start and end for the purposes of detecting conflicts.
 //
 // This is used for function extraction, where [start, end) will be extracted
 // to a new scope.
-func generateNameOutsideOfRange(start, end token.Pos, path []ast.Node, pkg *types.Package, info *types.Info, prefix string, idx int) (string, int) {
-	scopes := CollectScopes(info, path, start)
-	scopes = append(scopes, pkg.Scope())
+func freshNameOutsideRange(info *types.Info, file *ast.File, pos, start, end token.Pos, prefix string, idx int) (string, int) {
+	scope := info.Scopes[file].Innermost(pos)
 	return generateName(idx, prefix, func(name string) bool {
-		for _, scope := range scopes {
-			if scope != nil {
-				if obj := scope.Lookup(name); obj != nil {
-					// Only report a collision if the object declaration was outside the
-					// extracted range.
-					if obj.Pos() < start || end <= obj.Pos() {
-						return true
-					}
-				}
+		// Only report a collision if the object declaration
+		// was outside the extracted range.
+		for scope != nil {
+			obj, declScope := scope.LookupParent(name, pos)
+			if obj == nil {
+				return false // undeclared
 			}
+			if !(start <= obj.Pos() && obj.Pos() < end) {
+				return true // declared outside ignored range
+			}
+			scope = declScope.Parent()
 		}
 		return false
 	})
@@ -640,7 +657,7 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
 		// TODO(suzmue): generate a name that does not conflict for "newMethod".
 		funName = "newMethod"
 	} else {
-		funName, _ = generateAvailableName(start, path, pkg, info, "newFunction", 0)
+		funName, _ = freshName(info, file, start, "newFunction", 0)
 	}
 	extractedFunCall := generateFuncCall(hasNonNestedReturn, hasReturnValues, params,
 		append(returns, getNames(retVars)...), funName, sym, receiverName)
@@ -1290,7 +1307,7 @@ func generateReturnInfo(enclosing *ast.FuncType, pkg *types.Package, path []ast.
 	var cond *ast.Ident
 	if !hasNonNestedReturns {
 		// Generate information for the added bool value.
-		name, _ := generateNameOutsideOfRange(start, end, path, pkg, info, "shouldReturn", 0)
+		name, _ := freshNameOutsideRange(info, file, path[0].Pos(), start, end, "shouldReturn", 0)
 		cond = &ast.Ident{Name: name}
 		retVars = append(retVars, &returnVariable{
 			name:    cond,
@@ -1325,7 +1342,7 @@ func generateReturnInfo(enclosing *ast.FuncType, pkg *types.Package, path []ast.
 				} else if n, ok := varNameForType(typ); ok {
 					bestName = n
 				}
-				retName, idx := generateNameOutsideOfRange(start, end, path, pkg, info, bestName, nameIdx[bestName])
+				retName, idx := freshNameOutsideRange(info, file, path[0].Pos(), start, end, bestName, nameIdx[bestName])
 				nameIdx[bestName] = idx
 				z := typesinternal.ZeroExpr(file, pkg, typ)
 				if z == nil {
@@ -1540,3 +1557,11 @@ func getDecls(retVars []*returnVariable) []*ast.Field {
 	}
 	return decls
 }
+
+func cond[T any](cond bool, t, f T) T {
+	if cond {
+		return t
+	} else {
+		return f
+	}
+}

gopls/internal/golang/util.go

@@ -132,6 +132,9 @@ func findFileInDeps(s metadata.Source, mp *metadata.Package, uri protocol.Docume
 
 // CollectScopes returns all scopes in an ast path, ordered as innermost scope
 // first.
+//
+// TODO(adonovan): move this to golang/completion and simplify to use
+// Scopes.Innermost and LookupParent instead.
 func CollectScopes(info *types.Info, path []ast.Node, pos token.Pos) []*types.Scope {
 	// scopes[i], where i<len(path), is the possibly nil Scope of path[i].
 	var scopes []*types.Scope

gopls/internal/test/marker/testdata/codeaction/extract_variable-if.txt

@@ -29,10 +29,10 @@ func variable(y int) {
 
 -- @constant/a.go --
 @@ -4 +4 @@
-+	const x = 1 + 2
++	const k = 1 + 2
 @@ -5 +6 @@
 -	} else if 1 + 2 > 0 { //@ codeaction("1 + 2", "refactor.extract.constant", edit=constant)
-+	} else if x > 0 { //@ codeaction("1 + 2", "refactor.extract.constant", edit=constant)
++	} else if k > 0 { //@ codeaction("1 + 2", "refactor.extract.constant", edit=constant)
 -- @variable/a.go --
 @@ -10 +10 @@
 +	x := y + y

gopls/internal/test/marker/testdata/codeaction/extract_variable-inexact.txt

@@ -17,8 +17,8 @@ func _(ptr *int) {
 -- @spaces/a.go --
 @@ -4 +4,2 @@
 -	var _ = 1 + 2 + 3 //@codeaction("1 + 2 ", "refactor.extract.constant", edit=spaces)
-+	const x = 1 + 2
-+	var _ = x+ 3 //@codeaction("1 + 2 ", "refactor.extract.constant", edit=spaces)
++	const k = 1 + 2
++	var _ = k+ 3 //@codeaction("1 + 2 ", "refactor.extract.constant", edit=spaces)
 -- @funclit/a.go --
 @@ -5 +5,2 @@
 -	var _ = func() {} //@codeaction("func() {}", "refactor.extract.variable", edit=funclit)

gopls/internal/test/marker/testdata/codeaction/extract_variable-toplevel.txt

@@ -0,0 +1,51 @@
+This test checks the behavior of the 'extract variable/constant' code action
+at top level (outside any function). See issue #70665.
+
+-- a.go --
+package a
+
+const length = len("hello") + 2 //@codeaction(`len("hello")`, "refactor.extract.constant", edit=lenhello)
+
+var slice = append([]int{}, 1, 2, 3) //@codeaction("[]int{}", "refactor.extract.variable", edit=sliceliteral)
+
+type SHA256 [32]byte //@codeaction("32", "refactor.extract.constant", edit=arraylen)
+
+func f([2]int) {} //@codeaction("2", "refactor.extract.constant", edit=paramtypearraylen)
+
+-- @lenhello/a.go --
+@@ -3 +3,2 @@
+-const length = len("hello") + 2 //@codeaction(`len("hello")`, "refactor.extract.constant", edit=lenhello)
++const k = len("hello")
++const length = k + 2 //@codeaction(`len("hello")`, "refactor.extract.constant", edit=lenhello)
+-- @sliceliteral/a.go --
+@@ -5 +5,2 @@
+-var slice = append([]int{}, 1, 2, 3) //@codeaction("[]int{}", "refactor.extract.variable", edit=sliceliteral)
++var x = []int{}
++var slice = append(x, 1, 2, 3) //@codeaction("[]int{}", "refactor.extract.variable", edit=sliceliteral)
+-- @arraylen/a.go --
+@@ -7 +7,2 @@
+-type SHA256 [32]byte //@codeaction("32", "refactor.extract.constant", edit=arraylen)
++const k = 32
++type SHA256 [k]byte //@codeaction("32", "refactor.extract.constant", edit=arraylen)
+-- @paramtypearraylen/a.go --
+@@ -9 +9,2 @@
+-func f([2]int) {} //@codeaction("2", "refactor.extract.constant", edit=paramtypearraylen)
++const k = 2
++func f([k]int) {} //@codeaction("2", "refactor.extract.constant", edit=paramtypearraylen)
+-- b/b.go --
+package b
+
+// Check that package- and file-level name collisions are avoided.
+
+import x3 "errors"
+
+var x, x1, x2 any // these names are taken already
+var _ = x3.New("")
+var a, b int
+var c = a + b //@codeaction("a + b", "refactor.extract.variable", edit=fresh)
+
+-- @fresh/b/b.go --
+@@ -10 +10,2 @@
+-var c = a + b //@codeaction("a + b", "refactor.extract.variable", edit=fresh)
++var x4 = a + b
++var c = x4 //@codeaction("a + b", "refactor.extract.variable", edit=fresh)

gopls/internal/test/marker/testdata/codeaction/extract_variable.txt

@@ -15,13 +15,13 @@ func _() {
 -- @basic_lit1/basic_lit.go --
 @@ -4 +4,2 @@
 -	var _ = 1 + 2 //@codeaction("1", "refactor.extract.constant", edit=basic_lit1)
-+	const x = 1
-+	var _ = x + 2 //@codeaction("1", "refactor.extract.constant", edit=basic_lit1)
++	const k = 1
++	var _ = k + 2 //@codeaction("1", "refactor.extract.constant", edit=basic_lit1)
 -- @basic_lit2/basic_lit.go --
 @@ -5 +5,2 @@
 -	var _ = 3 + 4 //@codeaction("3 + 4", "refactor.extract.constant", edit=basic_lit2)
-+	const x = 3 + 4
-+	var _ = x //@codeaction("3 + 4", "refactor.extract.constant", edit=basic_lit2)
++	const k = 3 + 4
++	var _ = k //@codeaction("3 + 4", "refactor.extract.constant", edit=basic_lit2)
 -- func_call.go --
 package extract
 
@@ -54,7 +54,7 @@ func _() {
 		y := ast.CompositeLit{} //@codeaction("ast.CompositeLit{}", "refactor.extract.variable", edit=scope1)
 	}
 	if true {
-		x1 := !false //@codeaction("!false", "refactor.extract.constant", edit=scope2)
+		x := !false //@codeaction("!false", "refactor.extract.constant", edit=scope2)
 	}
 }
 
@@ -65,6 +65,6 @@ func _() {
 +		y := x //@codeaction("ast.CompositeLit{}", "refactor.extract.variable", edit=scope1)
 -- @scope2/scope.go --
 @@ -11 +11,2 @@
--		x1 := !false //@codeaction("!false", "refactor.extract.constant", edit=scope2)
-+		const x = !false
-+		x1 := x //@codeaction("!false", "refactor.extract.constant", edit=scope2)
+-		x := !false //@codeaction("!false", "refactor.extract.constant", edit=scope2)
++		const k = !false
++		x := k //@codeaction("!false", "refactor.extract.constant", edit=scope2)