gopls/internal: add code action "extract declarations to new file"

This code action moves selected code sections to a newly created file
within the same package. The created filename is chosen as the first
{function, type, const, var} name encountered. In addition, import
declarations are added or removed as needed.

Fixes golang/go#65707

gopls/internal/golang/codeaction.go

@@ -180,8 +180,16 @@ func fixedByImportFix(fix *imports.ImportFix, diagnostics []protocol.Diagnostic)
 
 // getExtractCodeActions returns any refactor.extract code actions for the selection.
 func getExtractCodeActions(pgf *parsego.File, rng protocol.Range, options *settings.Options) ([]protocol.CodeAction, error) {
+	var actions []protocol.CodeAction
+
+	extractToNewFileActions, err := getExtractToNewFileCodeActions(pgf, rng, options)
+	if err != nil {
+		return nil, err
+	}
+	actions = append(actions, extractToNewFileActions...)
+
 	if rng.Start == rng.End {
-		return nil, nil
+		return actions, nil
 	}
 
 	start, end, err := pgf.RangePos(rng)
@@ -226,7 +234,6 @@ func getExtractCodeActions(pgf *parsego.File, rng protocol.Range, options *setti
 		}
 		commands = append(commands, cmd)
 	}
-	var actions []protocol.CodeAction
 	for i := range commands {
 		actions = append(actions, newCodeAction(commands[i].Title, protocol.RefactorExtract, &commands[i], nil, options))
 	}

gopls/internal/golang/move_to_new_file.go

@@ -0,0 +1,361 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package golang
+
+// This file defines the code action "extract to a new file".
+
+// todo: rename file to extract_to_new_file.go after code review
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"go/ast"
+	"go/format"
+	"go/token"
+	"go/types"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"golang.org/x/tools/gopls/internal/cache"
+	"golang.org/x/tools/gopls/internal/cache/parsego"
+	"golang.org/x/tools/gopls/internal/file"
+	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/protocol/command"
+	"golang.org/x/tools/gopls/internal/settings"
+)
+
+func getExtractToNewFileCodeActions(pgf *parsego.File, rng protocol.Range, _ *settings.Options) ([]protocol.CodeAction, error) {
+	ok := canExtractToNewFile(pgf, rng)
+	if !ok {
+		return nil, nil
+	}
+	cmd, err := command.NewExtractToNewFileCommand(
+		"Extract declarations to new file",
+		command.ExtractToNewFileArgs{URI: pgf.URI, Range: rng},
+	)
+	if err != nil {
+		return nil, err
+	}
+	return []protocol.CodeAction{{
+		Title:   "Extract declarations to new file",
+		Kind:    protocol.RefactorExtract,
+		Command: &cmd,
+	}}, nil
+}
+
+// canExtractToNewFile reports whether the code in the given range can be extracted to a new file.
+func canExtractToNewFile(pgf *parsego.File, rng protocol.Range) bool {
+	_, err := extractToNewFileInternal(nil, nil, pgf, rng, true)
+	if err != nil {
+		return false
+	} else {
+		return true
+	}
+}
+
+// ExtractToNewFile moves selected declarations into a new file.
+func ExtractToNewFile(
+	ctx context.Context,
+	snapshot *cache.Snapshot,
+	fh file.Handle,
+	rng protocol.Range,
+) (*protocol.WorkspaceEdit, error) {
+	pkg, pgf, err := NarrowestPackageForFile(ctx, snapshot, fh.URI())
+	if err != nil {
+		return nil, err
+	}
+	return extractToNewFileInternal(fh, pkg, pgf, rng, false)
+}
+
+// findImportEdits finds imports specs that needs to be added to the new file
+// or deleted from the old file if the range is extracted to a new file.
+func findImportEdits(pkg *cache.Package, pgf *parsego.File, start, end token.Pos) (adds []*ast.ImportSpec, deletes []*ast.ImportSpec) {
+	var (
+		foundInSelection    = make(map[*types.PkgName]bool)
+		foundInNonSelection = make(map[*types.PkgName]bool)
+	)
+	for ident, use := range pkg.GetTypesInfo().Uses {
+		if pkgName, ok := use.(*types.PkgName); ok {
+			if contain(start, end, ident.Pos(), ident.End()) {
+				foundInSelection[pkgName] = true
+			} else {
+				foundInNonSelection[pkgName] = true
+			}
+		}
+	}
+	type NamePath struct {
+		Name string
+		Path string
+	}
+
+	imports := make(map[NamePath]*ast.ImportSpec)
+	for _, v := range pgf.File.Imports {
+		path := strings.Trim(v.Path.Value, `"`)
+		if v.Name != nil {
+			imports[NamePath{v.Name.Name, path}] = v
+		} else {
+			imports[NamePath{"", path}] = v
+		}
+	}
+
+	for pkgName := range foundInSelection {
+		importSpec := imports[NamePath{pkgName.Name(), pkgName.Imported().Path()}]
+		if importSpec == nil {
+			importSpec = imports[NamePath{"", pkgName.Imported().Path()}]
+		}
+
+		adds = append(adds, importSpec)
+		if !foundInNonSelection[pkgName] {
+			deletes = append(deletes, importSpec)
+		}
+	}
+
+	return adds, deletes
+}
+
+// extractToNewFileInternal moves selected declarations into a new file.
+func extractToNewFileInternal(
+	fh file.Handle,
+	pkg *cache.Package,
+	pgf *parsego.File,
+	rng protocol.Range,
+	dry bool,
+) (*protocol.WorkspaceEdit, error) {
+	errorPrefix := "moveToANewFileInternal"
+
+	start, end, err := pgf.RangePos(rng)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	start, end, filename, err := findRangeAndFilename(pgf, start, end)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	if dry {
+		return nil, nil
+	}
+
+	end = skipWhiteSpaces(pgf, end)
+
+	replaceRange, err := pgf.PosRange(start, end)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	adds, deletes := findImportEdits(pkg, pgf, start, end)
+
+	var importDeletes []protocol.TextEdit
+	parenthesisFreeImports := findParenthesisFreeImports(pgf)
+	for _, importSpec := range deletes {
+		if decl := parenthesisFreeImports[importSpec]; decl != nil {
+			importDeletes = append(importDeletes, removeNode(pgf, decl))
+		} else {
+			importDeletes = append(importDeletes, removeNode(pgf, importSpec))
+		}
+	}
+
+	importAdds := ""
+	if len(adds) > 0 {
+		importAdds += "import ("
+		for _, importSpec := range adds {
+			if importSpec.Name != nil {
+				importAdds += importSpec.Name.Name + " " + importSpec.Path.Value + "\n"
+			} else {
+				importAdds += importSpec.Path.Value + "\n"
+			}
+		}
+		importAdds += ")"
+	}
+
+	createFileURI, err := resolveCreateFileURI(pgf.URI.Dir().Path(), filename)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	creatFileText, err := format.Source([]byte(
+		"package " + pgf.File.Name.Name + "\n" +
+			importAdds + "\n" +
+			string(pgf.Src[start-pgf.File.FileStart:end-pgf.File.FileStart]),
+	))
+	if err != nil {
+		return nil, err
+	}
+
+	return &protocol.WorkspaceEdit{
+		DocumentChanges: []protocol.DocumentChanges{
+			// original file edits
+			protocol.TextEditsToDocumentChanges(fh.URI(), fh.Version(), append(
+				importDeletes,
+				protocol.TextEdit{
+					Range:   replaceRange,
+					NewText: "",
+				},
+			))[0],
+			{
+				CreateFile: &protocol.CreateFile{
+					Kind: "create",
+					URI:  createFileURI,
+				},
+			},
+			// created file edits
+			protocol.TextEditsToDocumentChanges(createFileURI, 0, []protocol.TextEdit{
+				{
+					Range:   protocol.Range{},
+					NewText: string(creatFileText),
+				},
+			})[0],
+		},
+	}, nil
+}
+
+// resolveCreateFileURI checks that basename.go does not exists in dir, otherwise
+// select basename.{1,2,3,4,5}.go as filename.
+func resolveCreateFileURI(dir string, basename string) (protocol.DocumentURI, error) {
+	basename = strings.ToLower(basename)
+	newPath := filepath.Join(dir, basename+".go")
+	for count := 1; ; count++ {
+		if _, err := os.Stat(newPath); errors.Is(err, os.ErrNotExist) {
+			break
+		}
+		if count >= 5 {
+			return "", fmt.Errorf("resolveNewFileURI: exceeded retry limit")
+		}
+		filename := fmt.Sprintf("%s.%d.go", basename, count)
+		newPath = filepath.Join(dir, filename)
+	}
+	return protocol.URIFromPath(newPath), nil
+}
+
+// findRangeAndFilename checks the selection is valid and extends range as needed and returns adjusted
+// range and selected filename.
+func findRangeAndFilename(pgf *parsego.File, start, end token.Pos) (token.Pos, token.Pos, string, error) {
+	if intersect(start, end, pgf.File.Package, pgf.File.Name.End()) {
+		return 0, 0, "", errors.New("selection cannot intersect a package declaration")
+	}
+	firstName := ""
+	for _, node := range pgf.File.Decls {
+		if intersect(start, end, node.Pos(), node.End()) {
+			if v, ok := node.(*ast.GenDecl); ok && v.Tok == token.IMPORT {
+				return 0, 0, "", errors.New("selection cannot intersect an import declaration")
+			}
+			if _, ok := node.(*ast.BadDecl); ok {
+				return 0, 0, "", errors.New("selection cannot intersect a bad declaration")
+			}
+			// should work when only selecting keyword "func" or function name
+			if v, ok := node.(*ast.FuncDecl); ok && contain(v.Pos(), v.Name.End(), start, end) {
+				start, end = v.Pos(), v.End()
+			}
+			// should work when only selecting keyword "type", "var", "const"
+			if v, ok := node.(*ast.GenDecl); ok && (v.Tok == token.TYPE && contain(v.Pos(), v.Pos()+4, start, end) ||
+				v.Tok == token.CONST && contain(v.Pos(), v.Pos()+5, start, end) ||
+				v.Tok == token.VAR && contain(v.Pos(), v.Pos()+3, start, end)) {
+				start, end = v.Pos(), v.End()
+			}
+			if !contain(start, end, node.Pos(), node.End()) {
+				return 0, 0, "", errors.New("selection cannot partially intersect a node")
+			} else {
+				if firstName == "" {
+					firstName = getNodeName(node)
+				}
+				// extends selection to docs comments
+				if c := getCommentGroup(node); c != nil {
+					if c.Pos() < start {
+						start = c.Pos()
+					}
+				}
+			}
+		}
+	}
+	for _, node := range pgf.File.Comments {
+		if intersect(start, end, node.Pos(), node.End()) {
+			if !contain(start, end, node.Pos(), node.End()) {
+				return 0, 0, "", errors.New("selection cannot partially intersect a comment")
+			}
+		}
+	}
+	if firstName == "" {
+		return 0, 0, "", errors.New("nothing selected")
+	}
+	return start, end, firstName, nil
+}
+
+func skipWhiteSpaces(pgf *parsego.File, pos token.Pos) token.Pos {
+	i := pos
+	for ; i-pgf.File.FileStart < token.Pos(len(pgf.Src)); i++ {
+		c := pgf.Src[i-pgf.File.FileStart]
+		if c == ' ' || c == '\t' || c == '\n' {
+			continue
+		} else {
+			break
+		}
+	}
+	return i
+}
+
+func getCommentGroup(node ast.Node) *ast.CommentGroup {
+	switch n := node.(type) {
+	case *ast.GenDecl:
+		return n.Doc
+	case *ast.FuncDecl:
+		return n.Doc
+	}
+	return nil
+}
+
+func findParenthesisFreeImports(pgf *parsego.File) map[*ast.ImportSpec]*ast.GenDecl {
+	decls := make(map[*ast.ImportSpec]*ast.GenDecl)
+	for _, decl := range pgf.File.Decls {
+		if g, ok := decl.(*ast.GenDecl); ok {
+			if !g.Lparen.IsValid() && len(g.Specs) > 0 {
+				if v, ok := g.Specs[0].(*ast.ImportSpec); ok {
+					decls[v] = g
+				}
+			}
+		}
+	}
+	return decls
+}
+
+// removeNode returns a TextEdit that removes the node
+func removeNode(pgf *parsego.File, node ast.Node) protocol.TextEdit {
+	rng, _ := pgf.PosRange(node.Pos(), node.End())
+	return protocol.TextEdit{Range: rng, NewText: ""}
+}
+
+// getNodeName returns the first func name or variable name
+func getNodeName(node ast.Node) string {
+	switch n := node.(type) {
+	case *ast.FuncDecl:
+		return n.Name.Name
+	case *ast.GenDecl:
+		if len(n.Specs) == 0 {
+			return ""
+		}
+		switch m := n.Specs[0].(type) {
+		case *ast.TypeSpec:
+			return m.Name.Name
+		case *ast.ValueSpec:
+			if len(m.Names) == 0 {
+				return ""
+			}
+			return m.Names[0].Name
+		}
+	}
+	return ""
+}
+
+// intersect checks if [a, b) and [c, d) intersect, assuming a <= b and c <= d
+func intersect(a, b, c, d token.Pos) bool {
+	return !(b <= c || d <= a)
+}
+
+// contain checks if [a, b) contains [c, d), assuming a <= b and c <= d
+func contain(a, b, c, d token.Pos) bool {
+	return a <= c && d <= b
+}