Prototype code for the FillInEnumSwitchCases action for Clang's new refactoring engine

FillInEnumSwitchCasesAction.cpp

// clang-refactor add-missing-enum-cases-in-switch -selection=file.cpp:1:2-3:4 file.cpp
// clang-refactor add-missing-enum-cases-in-switch -selection=enum.h:3:4 file.cpp
// clang-refactor add-missing-enum-cases-in-switch -enum-name "MyEnum"
class FillInEnumSwitchCases final: public RefactoringAction {
  StringRef getDescription() const override {
    return "Adds any missing switch cases to a switch over an enum";
  }

  StringRef getCommand() const override {
    return "add-missing-enum-cases-in-switch";
  }

  struct SelectedSwitch {
    const SwitchStmt &S;
    const EnumDecl   *ED;
    const DeclContext *EnclosingContext;
  };
  
  // All of the action rules allowed by this action.
  RefactoringActionRules createActionRules() const override {
    RefactoringActionRules Rules;
    using namespace refactoring_action_rules;
    
    // Single switch selection fills missing cases in just that one switch (the only rule supported by Xcode 9):
    auto FillInSwitchWhenSelected = RefactoringEditorCommand(RefactoringEditorCommand::FillInSwitchWhenSelected,
                                                             "Add Missing Switch Cases");
    Rules.push_back(FillInSwitchWhenSelected.bind(
      apply(fillInSwitch, requiredSelection(mapSelectionToStateOrDiag))));
    
    // Stating continuation requirements upfront enables the engine to dispatch these continuations
    // across process & machine boundaries.
    auto FillInSwitchesInTUCOntinuation = requiredContinuation("fill-in-switches-in-tu", fillInSwitchesInTU);
    using FillInSwitchesInTUType = decltype(FillInSwitchesInTUCOntinuation.get());
    // Stating indexer requirements upfront allows clients to limit the set of supported actions by verifying
    // that the required indexer operations are supported by the client.
    using IndexerInterface = indexer::IndexerCapabilities<indexer::ForEachASTUnitThatCapability,
                                                          indexer::ContainsOccurrenceOfCapability>;
    auto IndexerRequirements = requiredIndexerCapabilities<IndexerInterface>();
                         
    // Project-wide refactoring initiated on a selection of an enum declaration in the editor:
    auto FindAndFixSwitchesWhenEnumSelected = RefactoringEditorCommand(RefactoringEditorCommand::FindAndFixSwitchesWhenEnumSelected, 
                                                                       "Find and Fix Switches With Missing Cases");
    Rules.push_back(FindAndFixSwitchesWhenEnumSelected.bind(
      apply([](const ASTRefactoringOperation &Op, selection::ASTDecl<EnumDecl> Selection,
               IndexerInterface Indexer,
               FillInSwitchesInTUType FillInSwitchesInTU) -> std::unique_ptr<IndexerOperation> {
        EnumDecl &ED = Selection.get();
        return Indexer.foreachASTUnitThat(Indexer.containsOccurrenceOf(indexer::symbolForDecl(ED),
                                          FillInSwitchesInTU));
      },
      requiredSelection(selection::identity<selection::ASTDecl<EnumDecl>>()),
      IndexerRequirements,
      FillInSwitchesInTUCOntinuation
    )));
    
    // Project-wide tool rules:
    auto DispatchForEnumSymbol = [] (indexer::Symbol<EnumDecl> Enum,
                                     IndexerInterface Indexer,
                                     FillInSwitchesInTUType FillInSwitchesInTU) -> std::unique_ptr<IndexerOperation> {
      return Indexer.foreachASTUnitThat(Indexer.containsOccurrenceOf(Enum), FillInSwitchesInTU);
    };
    auto EnumNameOpt = RefactoringOption<std::string>("enum-name");
    auto EnumUSROpt = RefactoringOption<std::string>("enum-usr");
    Rules.push_back(
      apply(DispatchForEnumSymbol,
            requiredOption(EnumNameOpt).map([](std::string Name) {
              // This doesn't actually do anything here,
              // just specifies the mapping that's resolved
              // and verified later by the indexer.
              return indexer::symbolForName<EnumDecl>(Name);
    }), IndexerRequirements, FillInSwitchesInTUCOntinuation));
    Rules.push_back(
      apply(DispatchForEnumSymbol,
            // indexer::symbolForUSRAdaptor is a essentially the function shown above, except that it calls symbolForUSR
            requiredOption(EnumUSROpt).map(indexer::symbolForUSRAdaptor<EnumDecl>()),
            IndexerRequirements, FillInSwitchesInTUCOntinuation));

    return Rules;
  }
  
  // This function actually creates source replacements.
  static RefactoringResult fillInSwitch(const ASTRefactoringOperation &Op, SelectedSwitch Switch) {
    AtomicChange Change(Op.getSourceManager(), Switch.Switch->getLocation());
    // Dispatch to a lower-level API that creates the fix-its that are also reused by the -Wswitch warning.
    edit::fillInMissingSwitchEnumCases(
        Op.getASTContext(), Switch.Switch, Switch.ED, Switch.EnclosingContext,
        FixitToAtomicChangeAdaptor(Change));
    return std::move(Change);
  }

  // Editor selection support.
  static DiagOr<SelectedSwitch>
  mapSelectionToSwitchOrDiag(const ASTRefactoringOperation &Op, selection::Code<SwitchStmt> Selection) {
    const SwitchStmt &Switch = Selection.get();
    const DeclContext *EnclosingContext = Selection.getEnclosingDeclContext();

    // Ensure that the type is an enum.
    const Expr *Cond = Switch.getCond()->IgnoreImpCasts();
    const EnumDecl *ED = nullptr;
    if (const auto *ET = Cond->getType()->getAs<EnumType>())
      ED = ET->getDecl();
    else {
      // Enum literals are 'int' in C.
      if (const auto *DRE = dyn_cast<DeclRefExpr>(Cond)) {
        if (const auto *EC = dyn_cast<EnumConstantDecl>(DRE->getDecl()))
          ED = dyn_cast<EnumDecl>(EC->getDeclContext());
      }
    }

    if (!ED || !ED->isCompleteDefinition())
       // "The switch doesn't operate on a complete enum"
       return Op.Diag(diag::refactoring_fill_switch_not_enum);
    if (Switch.isAllEnumCasesCovered())
       return Op.Diag(diag::refactoring_fill_switch_all_cases_covered);
    return SelectedSwitch { Switch, ED, EnclosingContext };
  }
  
  // Project-wide refactoring support.
  static RefactoringResult
  fillInSwitchesInTU(const ASTRefactoringOperation &Op, indexer::Symbol<EnumDecl> Symbol) {
    using namespace ast_matchers;
    // Find all switches that operate on the target enum.
    // NB: I didn't test this matcher, this is just sample code. The matching
    // should probably be done using the enum's USR.
    // FIXME: This doesn't account for C EnumConstant case.
    auto Matches = match(
          switchStmt(hasCondition(hasType(enumDecl(isSymbol(Symbol).bind("ed")))).bind("switch"), Op.getASTContext());
    if (Matches.empty())
      return None;
    RefactoringResult Result;
    for (const auto &Match: Matches) {
      const EnumDecl *ED = Match.get<EnumDecl>("ed");
      const SwitchStmt *S = Match.get<SwitchStmt>("switch");
      const DeclContext *EnclosingContext = ... ;
      Result.push_back(fillInSwitch(Op, { *S, ED, EnclosingContext }));
    }
    return std::move(Result);
  }
};
      
// Declared automatically with macros using RefactoringActionRegistry.def, but we have to manually implement it.
std::unique<RefactoringAction> createFillInEnumSwitchCasesAction() {
  return llvm::make_unique<FillInEnumSwitchCases>();
}
      
// Snippet in RefactoringActionRegistry.def:
REFACTORING_ACTION(FillInEnumSwitchCases) 
EDITOR_COMMAND(FillInSwitchWhenSelected) 
EDITOR_COMMAND(FindAndFixSwitchesWhenEnumSelected)