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

[tdeo]

Summary

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

[sammccall]

Very nice, thank you!
A few nits about comments and asserts.

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

113–129

Somewhere we should have a high-level comment about what we're doing here:

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.

120

Code is obvious here, say why instead?

"Default likely intends to cover cases we'd insert."?

(Sometimes this is still interesting in that case, so we could consider downgrading from "fix" to "refactoring" or so in this case. But it makes sense to be conservative for now)

124

nit: just cast<CaseStmt> and don't check for null (it'll automatically assert). There are no other options.

128

, so our counting argument doesn't work.

132

"We may not be able to work out which cases are covered"?

153

this can be an assert (and previous line a cast) - apply() never gets called unless prepare() returns true

157

similarly here

180

this is also an assert, right? I don't think we can get here.

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

2838

is there an interesting difference between the single/multiple case that wouldn't be covered by just the multiple case?
(and below)

2838

can we add a case with default?
(And optionally template/gnu range, though those are pretty rare)

[tdeo]

• Improve comments
• In apply(), assert on the conditions we established in prepare() instead of proceeding.
• Remove redundant tests and add new ones for value-dependent and GNU range cases.

[tdeo]

Thanks for the review! Please land this if there are no more issues.

[sammccall]

Thank you!

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

 //     }
 //   }
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 namespace clangd {
 namespace {
 class PopulateSwitch : public Tweak {
   const char *id() const override;
   bool prepare(const Selection &Sel) override;
   Expected<Effect> apply(const Selection &Sel) override;
   std::string title() const override { return "Populate switch"; }
   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;
 
   const Stmt *CAStmt = CA->ASTNode.get<Stmt>();
   if (!CAStmt)
     return false;
 
bool PopulateSwitch::prepare(const Selection &Sel) {
   Switch = dyn_cast<SwitchStmt>(CAStmt);
   if (!Switch)
     return false;
 
   Body = dyn_cast<CompoundStmt>(Switch->getBody());
   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;
 
   // 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;
 
   EnumD = EnumT->getDecl();
   if (!EnumD)
     return false;
 
-  // If there aren't any enumerators, there's nothing to insert.
-  if (EnumD->enumerator_begin() == EnumD->enumerator_end())
+  // 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();
+
+  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;
+  }
 
-  return true;
+  // 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 += "::";
     }
     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)));
 }
 } // namespace
 } // namespace clangd
 } // namespace clang

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

TEST_F(PopulateSwitchTest, Test) {
   Case Cases[]{
       {
           // No enumerators
           Function,
           R""(enum Enum {}; ^switch ((Enum)0) {})"",
           "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",
       },
       {
           // Body not CompoundStmt
           Function,
           R""(enum Enum {A}; ^switch (A);)"",
           "unavailable",
       },
Case Cases[]{
           // Scoped enumeration
           Function,
           R""(enum class Enum {A}; ^switch (Enum::A) {})"",
           R""(enum class Enum {A}; switch (Enum::A) {case Enum::A:break;})"",
       },
       {
           // 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
           File,
           R""(
             namespace ns { enum class Enum {A}; }
             void function() { ^switch (ns::Enum::A) {} }
           )"",