gopls/internal: add code action "move to a 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

Author: golopot

gopls/internal/cache/snapshot.go

@@ -1658,6 +1658,22 @@ func inVendor(uri protocol.DocumentURI) bool {
 	return found && strings.Contains(after, "/")
 }
 
+// Sandbox clones the receiver, applying given file modifications as overlays.
+func (s *Snapshot) Sandbox(ctx, bgCtx context.Context, modifications []file.Modification) *Snapshot {
+	updatedFiles := make(map[protocol.DocumentURI]file.Handle)
+	for _, m := range modifications {
+		updatedFiles[m.URI] = &overlay{
+			uri:     m.URI,
+			content: m.Text,
+		}
+	}
+	cloned, _ := s.clone(ctx, bgCtx, StateChange{
+		Modifications: modifications,
+		Files:         updatedFiles,
+	}, func() {})
+	return cloned
+}
+
 // clone copies state from the receiver into a new Snapshot, applying the given
 // state changes.
 //

gopls/internal/golang/codeaction.go

@@ -86,6 +86,11 @@ func CodeActions(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle,
 				return nil, err
 			}
 			actions = append(actions, extractions...)
+			moves, err := getMoveToNewFileCodeAction(pgf, rng, snapshot.Options())
+			if err != nil {
+				return nil, err
+			}
+			actions = append(actions, moves)
 		}
 	}
 

gopls/internal/golang/format.go

@@ -92,6 +92,38 @@ func Format(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle) ([]pr
 	return computeTextEdits(ctx, pgf, formatted)
 }
 
+func formatImportsBytes(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle) ([]byte, error) {
+	_, done := event.Start(ctx, "golang.formatImportsBytes")
+	defer done()
+
+	errorPrefix := "formatImportsBytes"
+
+	text, err := fh.Content()
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	var out []byte
+	if err := snapshot.RunProcessEnvFunc(ctx, func(ctx context.Context, opts *imports.Options) error {
+		fixes, err := imports.FixImports(ctx, fh.URI().Path(), text, opts)
+		if err != nil {
+			return fmt.Errorf("%s: %w", errorPrefix, err)
+		}
+		out, err = imports.ApplyFixes(fixes, fh.URI().Path(), text, opts, parsego.Full)
+		if err != nil {
+			return fmt.Errorf("%s: %w", errorPrefix, err)
+		}
+		return nil
+	}); err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+	out, err = format.Source(out)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+	return out, nil
+}
+
 func formatSource(ctx context.Context, fh file.Handle) ([]byte, error) {
 	_, done := event.Start(ctx, "golang.formatSource")
 	defer done()

gopls/internal/golang/move_to_new_file.go

@@ -0,0 +1,325 @@
+// 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 "move to a new file".
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"go/ast"
+	"go/token"
+	"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 getMoveToNewFileCodeAction(pgf *parsego.File, rng protocol.Range, _ *settings.Options) (protocol.CodeAction, error) {
+	ok := canMoveToANewFile(pgf, rng)
+	if !ok {
+		return protocol.CodeAction{}, nil
+	}
+	cmd, err := command.NewMoveToANewFileCommand("m", command.MoveToANewFileArgs{URI: pgf.URI, Range: rng})
+	if err != nil {
+		return protocol.CodeAction{}, err
+	}
+	return protocol.CodeAction{
+		Title:   "Move to a new file",
+		Kind:    protocol.RefactorExtract,
+		Command: &cmd,
+	}, nil
+}
+
+// canMoveToANewFile reports whether the code in the given range can be move to a new file.
+func canMoveToANewFile(pgf *parsego.File, rng protocol.Range) bool {
+	_, err := moveToANewFileInternal(context.Background(), nil, nil, pgf, rng, true)
+	if err != nil {
+		return false
+	} else {
+		return true
+	}
+}
+
+// MoveToANewFile moves selected declarations into a new file.
+func MoveToANewFile(
+	ctx context.Context,
+	snapshot *cache.Snapshot,
+	fh file.Handle,
+	rng protocol.Range,
+) (*protocol.WorkspaceEdit, error) {
+	pgf, err := snapshot.ParseGo(ctx, fh, parsego.Full)
+	if err != nil {
+		return nil, fmt.Errorf("MoveToANewFile: %w", err)
+	}
+	return moveToANewFileInternal(ctx, snapshot, fh, pgf, rng, false)
+}
+
+// moveToANewFileInternal moves selected declarations into a new file.
+func moveToANewFileInternal(
+	ctx context.Context,
+	snapshot *cache.Snapshot,
+	fh file.Handle,
+	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
+	}
+
+	start, end = adjustRangeForEmptyLines(pgf, start, end)
+
+	createFileURI, err := resolveCreateFileURI(pgf.URI.Dir().Path(), filename)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	replaceRange, err := pgf.PosRange(start, end)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	start -= pgf.File.FileStart
+	end -= pgf.File.FileStart
+
+	modifiedText := make([]byte, 0)
+	modifiedText = append(modifiedText, pgf.Src[:start]...)
+	modifiedText = append(modifiedText, pgf.Src[end:]...)
+
+	packageName := pgf.File.Name.Name
+	createFileTextWithoutImports := []byte("package " + packageName + "\n" + string(pgf.Src[start:end]))
+
+	modifications := []file.Modification{
+		{
+			URI:    fh.URI(),
+			Action: file.Change,
+			Text:   modifiedText,
+		},
+		{
+			URI:    createFileURI,
+			Action: file.Change,
+			Text:   createFileTextWithoutImports,
+		},
+	}
+
+	// apply edits into a cloned snapshot and calculate import edits
+	snapshot = snapshot.Sandbox(ctx, ctx, modifications)
+	newFh, err := snapshot.ReadFile(ctx, fh.URI())
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	newPgf, err := snapshot.ParseGo(ctx, newFh, parsego.Full)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	importEdits, _, err := allImportsFixes(ctx, snapshot, newPgf)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	createFh, err := snapshot.ReadFile(ctx, createFileURI)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	createFileText, err := formatImportsBytes(ctx, snapshot, createFh)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", errorPrefix, err)
+	}
+
+	return &protocol.WorkspaceEdit{
+		DocumentChanges: []protocol.DocumentChanges{
+			{
+				// original file edits
+				TextDocumentEdit: &protocol.TextDocumentEdit{
+					TextDocument: protocol.OptionalVersionedTextDocumentIdentifier{
+						TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+							URI: fh.URI(),
+						},
+						Version: fh.Version(),
+					},
+					Edits: protocol.AsAnnotatedTextEdits(
+						append(
+							importEdits,
+							protocol.TextEdit{
+								Range:   replaceRange,
+								NewText: "",
+							}),
+					),
+				},
+			},
+			{
+				CreateFile: &protocol.CreateFile{
+					Kind: "create",
+					URI:  createFileURI,
+				},
+			},
+			{
+				// created file edits
+				TextDocumentEdit: &protocol.TextDocumentEdit{
+					TextDocument: protocol.OptionalVersionedTextDocumentIdentifier{
+						TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+							URI: createFileURI,
+						},
+						Version: 0,
+					},
+					Edits: protocol.AsAnnotatedTextEdits([]protocol.TextEdit{
+						{
+							Range:   protocol.Range{},
+							NewText: string(createFileText),
+						},
+					})},
+			},
+		},
+	}, 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 adjustRangeForEmptyLines(pgf *parsego.File, start, end token.Pos) (token.Pos, token.Pos) {
+	i := int(end)
+	for ; i-int(pgf.File.FileStart) < len(pgf.Src); i++ {
+		c := pgf.Src[i-int(pgf.File.FileStart)]
+		if c == ' ' || c == '\t' || c == '\n' {
+			continue
+		} else {
+			break
+		}
+	}
+	return start, token.Pos(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
+}
+
+// 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
+}