gopls/internal/cmd: factor three loops in fix subcommand

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

gopls/internal/cmd/suggested_fix.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
+	"golang.org/x/tools/gopls/internal/util/slices"
 	"golang.org/x/tools/internal/tool"
 )
 
@@ -148,42 +149,22 @@ func (s *suggestedFix) Run(ctx context.Context, args ...string) error {
 			continue
 		}
 
-		// Partially apply CodeAction.Edit, a WorkspaceEdit.
-		// (See also conn.Client.applyWorkspaceEdit(a.Edit)).
-		if !from.HasPosition() {
-			for _, c := range a.Edit.DocumentChanges {
-				if c.TextDocumentEdit != nil {
-					if c.TextDocumentEdit.TextDocument.URI == uri {
-						edits = append(edits, protocol.AsTextEdits(c.TextDocumentEdit.Edits)...)
-					}
-				}
-			}
+		// If the provided span has a position (not just offsets),
+		// and the action has diagnostics, the action must have a
+		// diagnostic with the same range as it.
+		if from.HasPosition() && len(a.Diagnostics) > 0 &&
+			!slices.ContainsFunc(a.Diagnostics, func(diag protocol.Diagnostic) bool {
+				return diag.Range.Start == rng.Start
+			}) {
 			continue
 		}
 
-		// The provided span has a position (not just offsets).
-		// Find the code action that has the same range as it.
-		for _, diag := range a.Diagnostics {
-			if diag.Range.Start == rng.Start {
-				for _, c := range a.Edit.DocumentChanges {
-					if c.TextDocumentEdit != nil {
-						if c.TextDocumentEdit.TextDocument.URI == uri {
-							edits = append(edits, protocol.AsTextEdits(c.TextDocumentEdit.Edits)...)
-						}
-					}
-				}
-				break
-			}
-		}
-
-		// If suggested fix is not a diagnostic, still must collect edits.
-		if len(a.Diagnostics) == 0 {
-			for _, c := range a.Edit.DocumentChanges {
-				if c.TextDocumentEdit != nil {
-					if c.TextDocumentEdit.TextDocument.URI == uri {
-						edits = append(edits, protocol.AsTextEdits(c.TextDocumentEdit.Edits)...)
-					}
-				}
+		// Partially apply CodeAction.Edit, a WorkspaceEdit.
+		// (See also conn.Client.applyWorkspaceEdit(a.Edit)).
+		for _, c := range a.Edit.DocumentChanges {
+			tde := c.TextDocumentEdit
+			if tde != nil && tde.TextDocument.URI == uri {
+				edits = append(edits, protocol.AsTextEdits(tde.Edits)...)
 			}
 		}
 	}

gopls/internal/util/slices/slices.go

@@ -15,6 +15,25 @@ func Contains[S ~[]E, E comparable](slice S, x E) bool {
 	return false
 }
 
+// IndexFunc returns the first index i satisfying f(s[i]),
+// or -1 if none do.
+// TODO(adonovan): use go1.19 slices.IndexFunc.
+func IndexFunc[S ~[]E, E any](s S, f func(E) bool) int {
+	for i := range s {
+		if f(s[i]) {
+			return i
+		}
+	}
+	return -1
+}
+
+// ContainsFunc reports whether at least one
+// element e of s satisfies f(e).
+// TODO(adonovan): use go1.19 slices.ContainsFunc.
+func ContainsFunc[S ~[]E, E any](s S, f func(E) bool) bool {
+	return IndexFunc(s, f) >= 0
+}
+
 // Concat returns a new slice concatenating the passed in slices.
 // TODO(rfindley): use go1.22 slices.Contains.
 func Concat[S ~[]E, E any](slices ...S) S {