[Refactoring] SR-6051 Expansion of switch statement missing cases

include/swift/IDE/RefactoringKinds.def

@@ -34,6 +34,8 @@ CURSOR_REFACTORING(FillProtocolStub, "Add Missing Protocol Requirements", fillst
 
 CURSOR_REFACTORING(ExpandDefault, "Expand Default", expand.default)
 
+CURSOR_REFACTORING(ExpandSwitchCases, "Expand Switch Cases", expand.switch.cases)
+
 CURSOR_REFACTORING(LocalizeString, "Localize String", localize.string)
 
 CURSOR_REFACTORING(SimplifyNumberLiteral, "Simplify Long Number Literal", simplify.long.number.literal)

include/swift/IDE/Utils.h

@@ -160,6 +160,7 @@ struct ResolvedCursorInfo {
   ValueDecl *ValueD = nullptr;
   TypeDecl *CtorTyRef = nullptr;
   ExtensionDecl *ExtTyRef = nullptr;
+  SourceFile *SF = nullptr;
   ModuleEntity Mod;
   SourceLoc Loc;
   bool IsRef = true;
@@ -174,6 +175,7 @@ struct ResolvedCursorInfo {
   ResolvedCursorInfo(ValueDecl *ValueD,
                      TypeDecl *CtorTyRef,
                      ExtensionDecl *ExtTyRef,
+                     SourceFile *SF,
                      SourceLoc Loc,
                      bool IsRef,
                      Type Ty,
@@ -182,21 +184,26 @@ struct ResolvedCursorInfo {
                         ValueD(ValueD),
                         CtorTyRef(CtorTyRef),
                         ExtTyRef(ExtTyRef),
+                        SF(SF),
                         Loc(Loc),
                         IsRef(IsRef),
                         Ty(Ty),
                         DC(ValueD->getDeclContext()),
                         ContainerType(ContainerType) {}
   ResolvedCursorInfo(ModuleEntity Mod,
+                     SourceFile *SF,
                      SourceLoc Loc) :
                         Kind(CursorInfoKind::ModuleRef),
+                        SF(SF),
                         Mod(Mod),
-                        Loc(Loc) { }
-  ResolvedCursorInfo(Stmt *TrailingStmt) :
+                        Loc(Loc) {}
+  ResolvedCursorInfo(Stmt *TrailingStmt, SourceFile *SF) :
                         Kind(CursorInfoKind::StmtStart),
+                        SF(SF),
                         TrailingStmt(TrailingStmt) {}
-  ResolvedCursorInfo(Expr* TrailingExpr) :
+  ResolvedCursorInfo(Expr* TrailingExpr, SourceFile *SF) :
                         Kind(CursorInfoKind::ExprStart),
+                        SF(SF),
                         TrailingExpr(TrailingExpr) {}
   bool isValid() const { return !isInvalid(); }
   bool isInvalid() const { return Kind == CursorInfoKind::Invalid; }

lib/IDE/Refactoring.cpp

@@ -1800,57 +1800,25 @@ collectAvailableRefactoringsAtCursor(SourceFile *SF, unsigned Line,
   return collectAvailableRefactorings(SF, Tok, Scratch, /*Exclude rename*/false);
 }
 
-bool RefactoringActionExpandDefault::
-isApplicable(ResolvedCursorInfo CursorInfo, DiagnosticEngine &Diag) {
-  auto Exit = [&](bool Applicable) {
-    if (!Applicable)
-      Diag.diagnose(SourceLoc(), diag::invalid_default_location);
-    return Applicable;
-  };
-  if (CursorInfo.Kind != CursorInfoKind::StmtStart)
-    return Exit(false);
-  if (auto *CS = dyn_cast<CaseStmt>(CursorInfo.TrailingStmt)) {
-    return Exit(CS->isDefault());
-  }
-  return Exit(false);
-}
-
-bool RefactoringActionExpandDefault::performChange() {
-  // Try to find the switch statement enclosing the default statement.
-  auto *CS = static_cast<CaseStmt*>(CursorInfo.TrailingStmt);
-  auto IsSwitch = [](ASTNode Node) {
-    return Node.is<Stmt*>() &&
-      Node.get<Stmt*>()->getKind() == StmtKind::Switch;
-  };
-  ContextFinder Finder(*TheFile, CS, IsSwitch);
-  Finder.resolve();
-
-  // If failed to find the switch statement, issue error.
-  if (Finder.getContexts().empty()) {
-    DiagEngine.diagnose(CS->getStartLoc(), diag::no_parent_switch);
-    return true;
-  }
-  auto *SwitchS = static_cast<SwitchStmt*>(Finder.getContexts().back().
-    get<Stmt*>());
-
-  // To find the subject enum decl for this switch statement; if failing,
-  // issue errors.
-  EnumDecl *SubjectED = nullptr;
+static EnumDecl* getEnumDeclFromSwitchStmt(SwitchStmt *SwitchS) {
   if (auto SubjectTy = SwitchS->getSubjectExpr()->getType()) {
-    SubjectED = SubjectTy->getAnyNominal()->getAsEnumOrEnumExtensionContext();
-  }
-  if (!SubjectED) {
-    DiagEngine.diagnose(CS->getStartLoc(), diag::no_subject_enum);
-    return true;
+    return SubjectTy->getAnyNominal()->getAsEnumOrEnumExtensionContext();
   }
+  return nullptr;
+}
 
+static bool performCasesExpansionInSwitchStmt(SwitchStmt *SwitchS,
+                                              DiagnosticEngine &DiagEngine,
+                                              SourceLoc ExpandedStmtLoc,
+                                              EditorConsumerInsertStream &OS
+                                              ) {
   // Assume enum elements are not handled in the switch statement.
+  auto EnumDecl = getEnumDeclFromSwitchStmt(SwitchS);
+  assert(EnumDecl);
   llvm::DenseSet<EnumElementDecl*> UnhandledElements;
-  SubjectED->getAllElements(UnhandledElements);
-  bool FoundDefault = false;
+  EnumDecl->getAllElements(UnhandledElements);
   for (auto Current : SwitchS->getCases()) {
-    if (Current == CS) {
-      FoundDefault = true;
+    if (Current->isDefault()) {
       continue;
     }
     // For each handled enum element, remove it from the bucket.
@@ -1861,27 +1829,119 @@ bool RefactoringActionExpandDefault::performChange() {
     }
   }
 
-  // If we've not seen the default statement inside the switch statement, issue
-  // error.
-  if (!FoundDefault) {
-    DiagEngine.diagnose(CS->getStartLoc(), diag::no_parent_switch);
-    return true;
-  }
-
   // If all enum elements are handled in the switch statement, issue error.
   if (UnhandledElements.empty()) {
-    DiagEngine.diagnose(CS->getStartLoc(), diag::no_remaining_cases);
+    DiagEngine.diagnose(ExpandedStmtLoc, diag::no_remaining_cases);
     return true;
   }
 
-  // Good to go, change the code!
+  printEnumElementsAsCases(UnhandledElements, OS);
+  return false;
+}
+
+// Finds SwitchStmt that contains given CaseStmt.
+static SwitchStmt* findEnclosingSwitchStmt(CaseStmt *CS,
+                                           SourceFile *SF,
+                                           DiagnosticEngine &DiagEngine) {
+  auto IsSwitch = [](ASTNode Node) {
+    return Node.is<Stmt*>() &&
+    Node.get<Stmt*>()->getKind() == StmtKind::Switch;
+  };
+  ContextFinder Finder(*SF, CS, IsSwitch);
+  Finder.resolve();
+
+  // If failed to find the switch statement, issue error.
+  if (Finder.getContexts().empty()) {
+    DiagEngine.diagnose(CS->getStartLoc(), diag::no_parent_switch);
+    return nullptr;
+  }
+  auto *SwitchS = static_cast<SwitchStmt*>(Finder.getContexts().back().
+                                           get<Stmt*>());
+  // Make sure that CaseStmt is included in switch that was found.
+  auto Cases = SwitchS->getCases();
+  auto Default = std::find(Cases.begin(), Cases.end(), CS);
+  if (Default == Cases.end()) {
+    DiagEngine.diagnose(CS->getStartLoc(), diag::no_parent_switch);
+    return nullptr;
+  }
+  return SwitchS;
+}
+
+bool RefactoringActionExpandDefault::
+isApplicable(ResolvedCursorInfo CursorInfo, DiagnosticEngine &Diag) {
+  auto Exit = [&](bool Applicable) {
+    if (!Applicable)
+      Diag.diagnose(SourceLoc(), diag::invalid_default_location);
+    return Applicable;
+  };
+  if (CursorInfo.Kind != CursorInfoKind::StmtStart)
+    return Exit(false);
+  if (auto *CS = dyn_cast<CaseStmt>(CursorInfo.TrailingStmt)) {
+    auto EnclosingSwitchStmt = findEnclosingSwitchStmt(CS,
+                                                       CursorInfo.SF,
+                                                       Diag);
+    if (!EnclosingSwitchStmt)
+      return false;
+    auto EnumD = getEnumDeclFromSwitchStmt(EnclosingSwitchStmt);
+    auto IsApplicable = CS->isDefault() && EnumD != nullptr;
+    return IsApplicable;
+  }
+  return Exit(false);
+}
+
+bool RefactoringActionExpandDefault::performChange() {
+  // If we've not seen the default statement inside the switch statement, issue
+  // error.
+  auto *CS = static_cast<CaseStmt*>(CursorInfo.TrailingStmt);
+  auto *SwitchS = findEnclosingSwitchStmt(CS, TheFile, DiagEngine);
+  assert(SwitchS);
   EditorConsumerInsertStream OS(EditConsumer, SM,
                                 Lexer::getCharSourceRangeFromSourceRange(SM,
                                   CS->getLabelItemsRange()));
-  printEnumElementsAsCases(UnhandledElements, OS);
+  return performCasesExpansionInSwitchStmt(SwitchS,
+                                           DiagEngine,
+                                           CS->getStartLoc(),
+                                           OS);
+}
+
+bool RefactoringActionExpandSwitchCases::
+isApplicable(ResolvedCursorInfo CursorInfo, DiagnosticEngine &DiagEngine) {
+  if (!CursorInfo.TrailingStmt)
+    return false;
+  if (auto *Switch = dyn_cast<SwitchStmt>(CursorInfo.TrailingStmt)) {
+    return getEnumDeclFromSwitchStmt(Switch);
+  }
   return false;
 }
 
+bool RefactoringActionExpandSwitchCases::performChange() {
+  auto *SwitchS = dyn_cast<SwitchStmt>(CursorInfo.TrailingStmt);
+  assert(SwitchS);
+
+  auto InsertRange = CharSourceRange();
+  auto Cases = SwitchS->getCases();
+  auto Default = std::find_if(Cases.begin(), Cases.end(), [](CaseStmt *Stmt) {
+    return Stmt->isDefault();
+  });
+  if (Default != Cases.end()) {
+    auto DefaultRange = (*Default)->getLabelItemsRange();
+    InsertRange = Lexer::getCharSourceRangeFromSourceRange(SM, DefaultRange);
+  } else {
+    auto RBraceLoc = SwitchS->getRBraceLoc();
+    InsertRange = CharSourceRange(SM, RBraceLoc, RBraceLoc);
+  }
+  EditorConsumerInsertStream OS(EditConsumer, SM, InsertRange);
+  if (SM.getLineNumber(SwitchS->getLBraceLoc()) ==
+      SM.getLineNumber(SwitchS->getRBraceLoc())) {
+    OS << "\n";
+  }
+  auto Result = performCasesExpansionInSwitchStmt(SwitchS,
+                                           DiagEngine,
+                                           SwitchS->getStartLoc(),
+                                           OS);
+  return Result;
+}
+
 static Expr *findLocalizeTarget(ResolvedCursorInfo CursorInfo) {
   if (CursorInfo.Kind != CursorInfoKind::ExprStart)
     return nullptr;

lib/IDE/SwiftSourceDocInfo.cpp

@@ -80,15 +80,15 @@ bool CursorInfoResolver::tryResolve(ValueDecl *D, TypeDecl *CtorTyRef,
     return false;
 
   if (Loc == LocToResolve) {
-    CursorInfo = { D, CtorTyRef, ExtTyRef, Loc, IsRef, Ty, ContainerType };
+    CursorInfo = { D, CtorTyRef, ExtTyRef, &SrcFile, Loc, IsRef, Ty, ContainerType };
     return true;
   }
   return false;
 }
 
 bool CursorInfoResolver::tryResolve(ModuleEntity Mod, SourceLoc Loc) {
   if (Loc == LocToResolve) {
-    CursorInfo = { Mod, Loc };
+    CursorInfo = { Mod, &SrcFile, Loc };
     return true;
   }
   return false;
@@ -97,13 +97,13 @@ bool CursorInfoResolver::tryResolve(ModuleEntity Mod, SourceLoc Loc) {
 bool CursorInfoResolver::tryResolve(Stmt *St) {
   if (auto *LST = dyn_cast<LabeledStmt>(St)) {
     if (LST->getStartLoc() == LocToResolve) {
-      CursorInfo = { St };
+      CursorInfo = { St, &SrcFile };
       return true;
     }
   }
   if (auto *CS = dyn_cast<CaseStmt>(St)) {
     if (CS->getStartLoc() == LocToResolve) {
-      CursorInfo = { St };
+      CursorInfo = { St, &SrcFile };
       return true;
     }
   }
@@ -211,7 +211,7 @@ bool CursorInfoResolver::walkToExprPost(Expr *E) {
     return false;
   if (!TrailingExprStack.empty() && TrailingExprStack.back() == E) {
     // We return the outtermost expression in the token info.
-    CursorInfo = { TrailingExprStack.front() };
+    CursorInfo = { TrailingExprStack.front(), &SrcFile };
     return false;
   }
   return true;

test/refactoring/ExpandSwitchCases/Outputs/basic/L10.swift.expected

@@ -0,0 +1,15 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e { 
+case .e1: <#code#>
+case .e2: <#code#>
+case .e3: <#code#>
+case .e4: <#code#>
+}
+}

test/refactoring/ExpandSwitchCases/Outputs/no_space_between_braces/L10.swift.expected

@@ -0,0 +1,15 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e {
+case .e1: <#code#>
+case .e2: <#code#>
+case .e3: <#code#>
+case .e4: <#code#>
+}
+}

test/refactoring/ExpandSwitchCases/Outputs/partially_handled/L11.swift.expected

@@ -0,0 +1,15 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e {
+    case .e1: return 5
+   case .e2: <#code#>
+case .e3: <#code#>
+case .e4: <#code#>
+}
+}

test/refactoring/ExpandSwitchCases/Outputs/with_default/L10.swift.expected

@@ -0,0 +1,16 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e {
+    case .e1: <#code#>
+case .e2: <#code#>
+case .e3: <#code#>
+case .e4: <#code#>
+ return 3
+   }
+}

test/refactoring/ExpandSwitchCases/basic.swift

@@ -0,0 +1,13 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e { }
+}
+// RUN: rm -rf %t.result && mkdir -p %t.result
+// RUN: %refactor -expand-switch-cases -source-filename %s -pos=9:8 >> %t.result/L10.swift
+// RUN: diff -u %S/Outputs/basic/L10.swift.expected %t.result/L10.swift

test/refactoring/ExpandSwitchCases/no_space_between_braces.swift

@@ -0,0 +1,13 @@
+enum E {
+  case e1
+  case e2
+  case e3
+  case e4
+}
+
+func foo(e: E) -> Int {
+  switch e {}
+}
+// RUN: rm -rf %t.result && mkdir -p %t.result
+// RUN: %refactor -expand-switch-cases -source-filename %s -pos=9:8 >> %t.result/L10.swift
+// RUN: diff -u %S/Outputs/no_space_between_braces/L10.swift.expected %t.result/L10.swift