internal/lsp/source: update SuggestedFixFunc to accept source.Snapshot

As part of the work for implementing method-stub code generation,
this CL updates the function signature of source/command.go's SuggestedFixFunc
so that a command can operate on the entire source.Snapshot to analyze
and change multiple packages and their files.

Updates golang/go#37537

Change-Id: I8430b2189ce7d91d37ab991f87ba368245293e56
Reviewed-on: https://go-review.googlesource.com/c/tools/+/279412
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Trust: Rebecca Stambler <rstambler@golang.org>
Trust: Robert Findley <rfindley@google.com>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>

internal/lsp/source/fix.go

@@ -19,13 +19,16 @@ import (
 	errors "golang.org/x/xerrors"
 )
 
-// SuggestedFixFunc is a function used to get the suggested fixes for a given
-// gopls command, some of which are provided by go/analysis.Analyzers. Some of
-// the analyzers in internal/lsp/analysis are not efficient enough to include
-// suggested fixes with their diagnostics, so we have to compute them
-// separately. Such analyzers should provide a function with a signature of
-// SuggestedFixFunc.
-type SuggestedFixFunc func(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, pkg *types.Package, info *types.Info) (*analysis.SuggestedFix, error)
+type (
+	// SuggestedFixFunc is a function used to get the suggested fixes for a given
+	// gopls command, some of which are provided by go/analysis.Analyzers. Some of
+	// the analyzers in internal/lsp/analysis are not efficient enough to include
+	// suggested fixes with their diagnostics, so we have to compute them
+	// separately. Such analyzers should provide a function with a signature of
+	// SuggestedFixFunc.
+	SuggestedFixFunc  func(ctx context.Context, snapshot Snapshot, fh VersionedFileHandle, pRng protocol.Range) (*analysis.SuggestedFix, error)
+	singleFileFixFunc func(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, pkg *types.Package, info *types.Info) (*analysis.SuggestedFix, error)
+)
 
 const (
 	FillStruct      = "fill_struct"
@@ -37,11 +40,22 @@ const (
 
 // suggestedFixes maps a suggested fix command id to its handler.
 var suggestedFixes = map[string]SuggestedFixFunc{
-	FillStruct:      fillstruct.SuggestedFix,
-	UndeclaredName:  undeclaredname.SuggestedFix,
-	ExtractVariable: extractVariable,
-	ExtractFunction: extractFunction,
-	ExtractMethod:   extractMethod,
+	FillStruct:      singleFile(fillstruct.SuggestedFix),
+	UndeclaredName:  singleFile(undeclaredname.SuggestedFix),
+	ExtractVariable: singleFile(extractVariable),
+	ExtractFunction: singleFile(extractFunction),
+	ExtractMethod:   singleFile(extractMethod),
+}
+
+// singleFile calls analyzers that expect inputs for a single file
+func singleFile(sf singleFileFixFunc) SuggestedFixFunc {
+	return func(ctx context.Context, snapshot Snapshot, fh VersionedFileHandle, pRng protocol.Range) (*analysis.SuggestedFix, error) {
+		fset, rng, src, file, pkg, info, err := getAllSuggestedFixInputs(ctx, snapshot, fh, pRng)
+		if err != nil {
+			return nil, err
+		}
+		return sf(fset, rng, src, file, pkg, info)
+	}
 }
 
 func SuggestedFixFromCommand(cmd protocol.Command, kind protocol.CodeActionKind) SuggestedFix {
@@ -59,55 +73,66 @@ func ApplyFix(ctx context.Context, fix string, snapshot Snapshot, fh VersionedFi
 	if !ok {
 		return nil, fmt.Errorf("no suggested fix function for %s", fix)
 	}
-	fset, rng, src, file, m, pkg, info, err := getAllSuggestedFixInputs(ctx, snapshot, fh, pRng)
-	if err != nil {
-		return nil, err
-	}
-	suggestion, err := handler(fset, rng, src, file, pkg, info)
+	suggestion, err := handler(ctx, snapshot, fh, pRng)
 	if err != nil {
 		return nil, err
 	}
 	if suggestion == nil {
 		return nil, nil
 	}
-
-	var edits []protocol.TextEdit
+	fset := snapshot.FileSet()
+	editsPerFile := map[span.URI]*protocol.TextDocumentEdit{}
 	for _, edit := range suggestion.TextEdits {
-		rng := span.NewRange(fset, edit.Pos, edit.End)
-		spn, err := rng.Span()
+		spn, err := span.NewRange(fset, edit.Pos, edit.End).Span()
 		if err != nil {
 			return nil, err
 		}
-		clRng, err := m.Range(spn)
+		fh, err := snapshot.GetVersionedFile(ctx, spn.URI())
+		if err != nil {
+			return nil, err
+		}
+		te, ok := editsPerFile[spn.URI()]
+		if !ok {
+			te = &protocol.TextDocumentEdit{
+				TextDocument: protocol.OptionalVersionedTextDocumentIdentifier{
+					Version: fh.Version(),
+					TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+						URI: protocol.URIFromSpanURI(fh.URI()),
+					},
+				},
+			}
+			editsPerFile[spn.URI()] = te
+		}
+		_, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
 		if err != nil {
 			return nil, err
 		}
-		edits = append(edits, protocol.TextEdit{
-			Range:   clRng,
+		rng, err := pgf.Mapper.Range(spn)
+		if err != nil {
+			return nil, err
+		}
+		te.Edits = append(te.Edits, protocol.TextEdit{
+			Range:   rng,
 			NewText: string(edit.NewText),
 		})
 	}
-	return []protocol.TextDocumentEdit{{
-		TextDocument: protocol.OptionalVersionedTextDocumentIdentifier{
-			Version: fh.Version(),
-			TextDocumentIdentifier: protocol.TextDocumentIdentifier{
-				URI: protocol.URIFromSpanURI(fh.URI()),
-			},
-		},
-		Edits: edits,
-	}}, nil
+	var edits []protocol.TextDocumentEdit
+	for _, edit := range editsPerFile {
+		edits = append(edits, *edit)
+	}
+	return edits, nil
 }
 
 // getAllSuggestedFixInputs is a helper function to collect all possible needed
 // inputs for an AppliesFunc or SuggestedFixFunc.
-func getAllSuggestedFixInputs(ctx context.Context, snapshot Snapshot, fh FileHandle, pRng protocol.Range) (*token.FileSet, span.Range, []byte, *ast.File, *protocol.ColumnMapper, *types.Package, *types.Info, error) {
+func getAllSuggestedFixInputs(ctx context.Context, snapshot Snapshot, fh FileHandle, pRng protocol.Range) (*token.FileSet, span.Range, []byte, *ast.File, *types.Package, *types.Info, error) {
 	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
 	if err != nil {
-		return nil, span.Range{}, nil, nil, nil, nil, nil, errors.Errorf("getting file for Identifier: %w", err)
+		return nil, span.Range{}, nil, nil, nil, nil, errors.Errorf("getting file for Identifier: %w", err)
 	}
 	rng, err := pgf.Mapper.RangeToSpanRange(pRng)
 	if err != nil {
-		return nil, span.Range{}, nil, nil, nil, nil, nil, err
+		return nil, span.Range{}, nil, nil, nil, nil, err
 	}
-	return snapshot.FileSet(), rng, pgf.Src, pgf.File, pgf.Mapper, pkg.GetTypes(), pkg.GetTypesInfo(), nil
+	return snapshot.FileSet(), rng, pgf.Src, pgf.File, pkg.GetTypes(), pkg.GetTypesInfo(), nil
 }