[clangd] Improve PopulateSwitch tweak to work on non-empty switches

Improve the recently-added PopulateSwitch tweak to work on non-empty switches.

Reviewed By: sammccall

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

clang-tools-extra/clangd/refactor/tweaks/PopulateSwitch.cpp

@@ -33,12 +33,15 @@
 #include "AST.h"
 #include "Selection.h"
 #include "refactor/Tweak.h"
+#include "support/Logger.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Stmt.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Core/Replacement.h"
+#include "llvm/ADT/SmallSet.h"
+#include <cassert>
 #include <string>
 
 namespace clang {
@@ -52,18 +55,16 @@ class PopulateSwitch : public Tweak {
   Intent intent() const override { return Refactor; }
 
 private:
-  ASTContext *ASTCtx = nullptr;
   const DeclContext *DeclCtx = nullptr;
   const SwitchStmt *Switch = nullptr;
   const CompoundStmt *Body = nullptr;
+  const EnumType *EnumT = nullptr;
   const EnumDecl *EnumD = nullptr;
 };
 
 REGISTER_TWEAK(PopulateSwitch)
 
 bool PopulateSwitch::prepare(const Selection &Sel) {
-  ASTCtx = &Sel.AST->getASTContext();
-
   const SelectionTree::Node *CA = Sel.ASTSelection.commonAncestor();
   if (!CA)
     return false;
@@ -94,11 +95,6 @@ bool PopulateSwitch::prepare(const Selection &Sel) {
   if (!Body)
     return false;
 
-  // Since we currently always insert all enumerators, don't suggest this tweak
-  // if the body is not empty.
-  if (!Body->body_empty())
-    return false;
-
   const Expr *Cond = Switch->getCond();
   if (!Cond)
     return false;
@@ -106,7 +102,7 @@ bool PopulateSwitch::prepare(const Selection &Sel) {
   // Ignore implicit casts, since enums implicitly cast to integer types.
   Cond = Cond->IgnoreParenImpCasts();
 
-  const EnumType *EnumT = Cond->getType()->getAsAdjusted<EnumType>();
+  EnumT = Cond->getType()->getAsAdjusted<EnumType>();
   if (!EnumT)
     return false;
 
@@ -114,21 +110,65 @@ bool PopulateSwitch::prepare(const Selection &Sel) {
   if (!EnumD)
     return false;
 
-  // If there aren't any enumerators, there's nothing to insert.
-  if (EnumD->enumerator_begin() == EnumD->enumerator_end())
-    return false;
+  // We trigger if there are fewer cases than enum values (and no case covers
+  // multiple values). This guarantees we'll have at least one case to insert.
+  // We don't yet determine what the cases are, as that means evaluating
+  // expressions.
+  auto I = EnumD->enumerator_begin();
+  auto E = EnumD->enumerator_end();
 
-  return true;
+  for (const SwitchCase *CaseList = Switch->getSwitchCaseList();
+       CaseList && I != E; CaseList = CaseList->getNextSwitchCase(), I++) {
+    // Default likely intends to cover cases we'd insert.
+    if (isa<DefaultStmt>(CaseList))
+      return false;
+
+    const CaseStmt *CS = cast<CaseStmt>(CaseList);
+    // Case statement covers multiple values, so just counting doesn't work.
+    if (CS->caseStmtIsGNURange())
+      return false;
+
+    // Case expression is not a constant expression or is value-dependent,
+    // so we may not be able to work out which cases are covered.
+    const ConstantExpr *CE = dyn_cast<ConstantExpr>(CS->getLHS());
+    if (!CE || CE->isValueDependent())
+      return false;
+  }
+
+  // Only suggest tweak if we have more enumerators than cases.
+  return I != E;
 }
 
 Expected<Tweak::Effect> PopulateSwitch::apply(const Selection &Sel) {
-  const SourceManager &SM = ASTCtx->getSourceManager();
+  ASTContext &Ctx = Sel.AST->getASTContext();
+
+  // Get the enum's integer width and signedness, for adjusting case literals.
+  unsigned EnumIntWidth = Ctx.getIntWidth(QualType(EnumT, 0));
+  bool EnumIsSigned = EnumT->isSignedIntegerOrEnumerationType();
+
+  llvm::SmallSet<llvm::APSInt, 32> ExistingEnumerators;
+  for (const SwitchCase *CaseList = Switch->getSwitchCaseList(); CaseList;
+       CaseList = CaseList->getNextSwitchCase()) {
+    const CaseStmt *CS = cast<CaseStmt>(CaseList);
+    assert(!CS->caseStmtIsGNURange());
+    const ConstantExpr *CE = cast<ConstantExpr>(CS->getLHS());
+    assert(!CE->isValueDependent());
+    llvm::APSInt Val = CE->getResultAsAPSInt();
+    Val = Val.extOrTrunc(EnumIntWidth);
+    Val.setIsSigned(EnumIsSigned);
+    ExistingEnumerators.insert(Val);
+  }
+
   SourceLocation Loc = Body->getRBracLoc();
+  ASTContext &DeclASTCtx = DeclCtx->getParentASTContext();
 
   std::string Text;
   for (EnumConstantDecl *Enumerator : EnumD->enumerators()) {
+    if (ExistingEnumerators.contains(Enumerator->getInitVal()))
+      continue;
+
     Text += "case ";
-    Text += getQualification(*ASTCtx, DeclCtx, Loc, EnumD);
+    Text += getQualification(DeclASTCtx, DeclCtx, Loc, EnumD);
     if (EnumD->isScoped()) {
       Text += EnumD->getName();
       Text += "::";
@@ -136,8 +176,11 @@ Expected<Tweak::Effect> PopulateSwitch::apply(const Selection &Sel) {
     Text += Enumerator->getName();
     Text += ":";
   }
+
+  assert(!Text.empty() && "No enumerators to insert!");
   Text += "break;";
 
+  const SourceManager &SM = Ctx.getSourceManager();
   return Effect::mainFileEdit(
       SM, tooling::Replacements(tooling::Replacement(SM, Loc, 0, Text)));
 }

clang-tools-extra/clangd/unittests/TweakTests.cpp

@@ -2829,9 +2829,48 @@ TEST_F(PopulateSwitchTest, Test) {
           "unavailable",
       },
       {
-          // Existing enumerators in switch
+          // All enumerators already in switch (unscoped)
           Function,
-          R""(enum Enum {A}; ^switch ((Enum)0) {case A:break;})"",
+          R""(enum Enum {A,B}; ^switch (A) {case A:break;case B:break;})"",
+          "unavailable",
+      },
+      {
+          // All enumerators already in switch (scoped)
+          Function,
+          R""(
+            enum class Enum {A,B};
+            ^switch (Enum::A) {case Enum::A:break;case Enum::B:break;}
+          )"",
+          "unavailable",
+      },
+      {
+          // Default case in switch
+          Function,
+          R""(
+            enum class Enum {A,B};
+            ^switch (Enum::A) {default:break;}
+          )"",
+          "unavailable",
+      },
+      {
+          // GNU range in switch
+          Function,
+          R""(
+            enum class Enum {A,B};
+            ^switch (Enum::A) {case Enum::A ... Enum::B:break;}
+          )"",
+          "unavailable",
+      },
+      {
+          // Value dependent case expression
+          File,
+          R""(
+            enum class Enum {A,B};
+            template<Enum Value>
+            void function() {
+                ^switch (Enum::A) {case Value:break;}
+            }
+          )"",
           "unavailable",
       },
       {
@@ -2867,9 +2906,53 @@ TEST_F(PopulateSwitchTest, Test) {
       {
           // Scoped enumeration with multiple enumerators
           Function,
-          R""(enum class Enum {A,B}; ^switch (Enum::A) {})"",
-          R""(enum class Enum {A,B}; )""
-          R""(switch (Enum::A) {case Enum::A:case Enum::B:break;})"",
+          R""(
+            enum class Enum {A,B};
+            ^switch (Enum::A) {}
+          )"",
+          R""(
+            enum class Enum {A,B};
+            switch (Enum::A) {case Enum::A:case Enum::B:break;}
+          )"",
+      },
+      {
+          // Only filling in missing enumerators (unscoped)
+          Function,
+          R""(
+            enum Enum {A,B,C};
+            ^switch (A) {case B:break;}
+          )"",
+          R""(
+            enum Enum {A,B,C};
+            switch (A) {case B:break;case A:case C:break;}
+          )"",
+      },
+      {
+          // Only filling in missing enumerators,
+          // even when using integer literals
+          Function,
+          R""(
+            enum Enum {A,B=1,C};
+            ^switch (A) {case 1:break;}
+          )"",
+          R""(
+            enum Enum {A,B=1,C};
+            switch (A) {case 1:break;case A:case C:break;}
+          )"",
+      },
+      {
+          // Only filling in missing enumerators (scoped)
+          Function,
+          R""(
+            enum class Enum {A,B,C};
+            ^switch (Enum::A)
+            {case Enum::B:break;}
+          )"",
+          R""(
+            enum class Enum {A,B,C};
+            switch (Enum::A)
+            {case Enum::B:break;case Enum::A:case Enum::C:break;}
+          )"",
       },
       {
           // Scoped enumerations in namespace