[clangd] Fix fixAll not shown when there is only one unused-include and missing-include diagnostics.

Discovered during the review in https://reviews.llvm.org/D149437#inline-1444851.

Differential Revision: https://reviews.llvm.org/D149822

2 files changed, 84 insertions, 15 deletions

diff --git a/clang-tools-extra/clangd/IncludeCleaner.cpp b/clang-tools-extra/clangd/IncludeCleaner.cpp
index 3b913d851abe..22d0808235f7 100644
--- a/clang-tools-extra/clangd/IncludeCleaner.cpp
+++ b/clang-tools-extra/clangd/IncludeCleaner.cpp
@@ -440,8 +440,9 @@ IncludeCleanerFindings computeIncludeCleanerFindings(ParsedAST &AST) {
   return {std::move(UnusedIncludes), std::move(MissingIncludes)};
 }
 
-Fix removeAllUnusedIncludes(llvm::ArrayRef<Diag> UnusedIncludes) {
-  assert(!UnusedIncludes.empty());
+std::optional<Fix> removeAllUnusedIncludes(llvm::ArrayRef<Diag> UnusedIncludes) {
+  if (UnusedIncludes.empty())
+    return std::nullopt;
 
   Fix RemoveAll;
   RemoveAll.Message = "remove all unused includes";
@@ -465,8 +466,10 @@ Fix removeAllUnusedIncludes(llvm::ArrayRef<Diag> UnusedIncludes) {
   }
   return RemoveAll;
 }
-Fix addAllMissingIncludes(llvm::ArrayRef<Diag> MissingIncludeDiags) {
-  assert(!MissingIncludeDiags.empty());
+std::optional<Fix>
+addAllMissingIncludes(llvm::ArrayRef<Diag> MissingIncludeDiags) {
+  if (MissingIncludeDiags.empty())
+    return std::nullopt;
 
   Fix AddAllMissing;
   AddAllMissing.Message = "add all missing includes";
@@ -516,15 +519,11 @@ std::vector<Diag> generateIncludeCleanerDiagnostic(
     llvm::StringRef Code) {
   std::vector<Diag> UnusedIncludes = generateUnusedIncludeDiagnostics(
       AST.tuPath(), Findings.UnusedIncludes, Code);
-  std::optional<Fix> RemoveAllUnused;;
-  if (UnusedIncludes.size() > 1)
-    RemoveAllUnused = removeAllUnusedIncludes(UnusedIncludes);
+  std::optional<Fix> RemoveAllUnused = removeAllUnusedIncludes(UnusedIncludes);
 
   std::vector<Diag> MissingIncludeDiags = generateMissingIncludeDiagnostics(
       AST, Findings.MissingIncludes, Code);
-  std::optional<Fix> AddAllMissing;
-  if (MissingIncludeDiags.size() > 1)
-    AddAllMissing = addAllMissingIncludes(MissingIncludeDiags);
+  std::optional<Fix> AddAllMissing = addAllMissingIncludes(MissingIncludeDiags);
 
   std::optional<Fix> FixAll;
   if (RemoveAllUnused && AddAllMissing)
@@ -535,11 +534,16 @@ std::vector<Diag> generateIncludeCleanerDiagnostic(
     Out->Fixes.push_back(*F);
   };
   for (auto &Diag : MissingIncludeDiags) {
-    AddBatchFix(AddAllMissing, &Diag);
+    AddBatchFix(MissingIncludeDiags.size() > 1
+                    ? AddAllMissing
+                    : std::nullopt,
+                &Diag);
     AddBatchFix(FixAll, &Diag);
   }
   for (auto &Diag : UnusedIncludes) {
-    AddBatchFix(RemoveAllUnused, &Diag);
+    AddBatchFix(UnusedIncludes.size() > 1 ? RemoveAllUnused
+                                          : std::nullopt,
+                &Diag);
     AddBatchFix(FixAll, &Diag);
   }
 
diff --git a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
index a38c01b43270..c0831c2693c9 100644
--- a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
@@ -197,10 +197,10 @@ $insert_f[[]]$insert_vector[[]]
 
     ns::$bar[[Bar]] bar;
     bar.d();
-    $f[[f]](); 
+    $f[[f]]();
 
     // this should not be diagnosed, because it's ignored in the config
-    buzz(); 
+    buzz();
 
     $foobar[[foobar]]();
 
@@ -244,7 +244,7 @@ $insert_f[[]]$insert_vector[[]]
   TU.AdditionalFiles["foo.h"] = guard(R"cpp(
     #define BAR(x) Foo *x
     #define FOO 1
-    struct Foo{}; 
+    struct Foo{};
   )cpp");
 
   TU.Code = MainFile.code();
@@ -510,6 +510,71 @@ TEST(IncludeCleaner, FirstMatchedProvider) {
   }
 }
 
+TEST(IncludeCleaner, BatchFix) {
+  Config Cfg;
+  Cfg.Diagnostics.MissingIncludes = Config::IncludesPolicy::Strict;
+  Cfg.Diagnostics.UnusedIncludes = Config::IncludesPolicy::Strict;
+  WithContextValue Ctx(Config::Key, std::move(Cfg));
+
+  TestTU TU;
+  TU.Filename = "main.cpp";
+  TU.AdditionalFiles["foo.h"] = guard("class Foo;");
+  TU.AdditionalFiles["bar.h"] = guard("class Bar;");
+  TU.AdditionalFiles["all.h"] = guard(R"cpp(
+    #include "foo.h"
+    #include "bar.h"
+  )cpp");
+
+  TU.Code = R"cpp(
+  #include "all.h"
+
+  Foo* foo;
+  )cpp";
+  auto AST = TU.build();
+  EXPECT_THAT(
+      issueIncludeCleanerDiagnostics(AST, TU.Code),
+      UnorderedElementsAre(withFix({FixMessage("#include \"foo.h\""),
+                                    FixMessage("fix all includes")}),
+                           withFix({FixMessage("remove #include directive"),
+                                    FixMessage("fix all includes")})));
+
+  TU.Code = R"cpp(
+  #include "all.h"
+  #include "bar.h"
+
+  Foo* foo;
+  )cpp";
+  AST = TU.build();
+  EXPECT_THAT(
+      issueIncludeCleanerDiagnostics(AST, TU.Code),
+      UnorderedElementsAre(withFix({FixMessage("#include \"foo.h\""),
+                                    FixMessage("fix all includes")}),
+                           withFix({FixMessage("remove #include directive"),
+                                    FixMessage("remove all unused includes"),
+                                    FixMessage("fix all includes")}),
+                           withFix({FixMessage("remove #include directive"),
+                                   FixMessage("remove all unused includes"),
+                                    FixMessage("fix all includes")})));
+
+  TU.Code = R"cpp(
+  #include "all.h"
+
+  Foo* foo;
+  Bar* bar;
+  )cpp";
+  AST = TU.build();
+  EXPECT_THAT(
+      issueIncludeCleanerDiagnostics(AST, TU.Code),
+      UnorderedElementsAre(withFix({FixMessage("#include \"foo.h\""),
+                                    FixMessage("add all missing includes"),
+                                    FixMessage("fix all includes")}),
+                           withFix({FixMessage("#include \"bar.h\""),
+                                    FixMessage("add all missing includes"),
+                                    FixMessage("fix all includes")}),
+                           withFix({FixMessage("remove #include directive"),
+                                    FixMessage("fix all includes")})));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang