Merge pull request #58611 from ahoppen/pr/fix-switch-refactoring

[Refactoring] Fix a crash that occurred when converting if/else of enum with raw value to switch statement

Author: ahoppen

lib/IDE/Refactoring.cpp

@@ -2514,14 +2514,24 @@ bool RefactoringActionConvertToSwitchStmt::performChange() {
     SourceManager &SM;
 
     bool isFunctionNameAllowed(BinaryExpr *E) {
-      auto FunctionBody = dyn_cast<DotSyntaxCallExpr>(E->getFn())->getFn();
-      auto FunctionDeclaration = dyn_cast<DeclRefExpr>(FunctionBody)->getDecl();
-      const auto FunctionName = dyn_cast<FuncDecl>(FunctionDeclaration)
-          ->getBaseIdentifier().str();
-      return FunctionName == "~="
-      || FunctionName == "=="
-      || FunctionName == "__derived_enum_equals"
-      || FunctionName == "__derived_struct_equals";
+      Expr *Fn = E->getFn();
+      if (auto DotSyntaxCall = dyn_cast_or_null<DotSyntaxCallExpr>(Fn)) {
+        Fn = DotSyntaxCall->getFn();
+      }
+      DeclRefExpr *DeclRef = dyn_cast_or_null<DeclRefExpr>(Fn);
+      if (!DeclRef) {
+        return false;
+      }
+      auto FunctionDeclaration = dyn_cast_or_null<FuncDecl>(DeclRef->getDecl());
+      if (!FunctionDeclaration) {
+        return false;
+      }
+      auto &ASTCtx = FunctionDeclaration->getASTContext();
+      const auto FunctionName = FunctionDeclaration->getBaseIdentifier();
+      return FunctionName == ASTCtx.Id_MatchOperator ||
+             FunctionName == ASTCtx.Id_EqualsOperator ||
+             FunctionName == ASTCtx.Id_derived_enum_equals ||
+             FunctionName == ASTCtx.Id_derived_struct_equals;
     }
 
     void appendPattern(Expr *LHS, Expr *RHS) {

test/refactoring/ConvertToSwitchStmt/Outputs/enum_with_raw_value/L8-3.swift.expected

@@ -0,0 +1,21 @@
+enum EnumWithUnderlyingValue: String {
+    case north, south, east, west
+}
+
+
+
+func foo(test: EnumWithUnderlyingValue) {
+  switch test {
+case .north:
+print("north")
+case .south:
+print("south")
+case .east:
+print("east")
+default:
+break
+}
+}
+
+
+

test/refactoring/ConvertToSwitchStmt/basic.swift

@@ -131,7 +131,7 @@ func checkEmptyBody(e: E) {
   }
 }
 
-// RUN: rm -rf %t.result && mkdir -p %t.result
+// RUN: %empty-directory(%t.result)
 
 // RUN: %refactor -convert-to-switch-stmt -source-filename %s -pos=9:3 -end-pos=16:4 > %t.result/L9-3.swift
 // RUN: %target-swift-frontend-typecheck %t.result/L9-3.swift

test/refactoring/ConvertToSwitchStmt/enum_with_raw_value.swift

@@ -0,0 +1,22 @@
+enum EnumWithUnderlyingValue: String {
+    case north, south, east, west
+}
+
+
+
+func foo(test: EnumWithUnderlyingValue) {
+  if test == .north {
+    print("north")
+  } else if test == .south {
+    print("south")
+  } else if test == .east {
+    print("east")
+  }
+}
+
+
+// RUN: %empty-directory(%t.result)
+
+// RUN: %refactor -convert-to-switch-stmt -source-filename %s -pos=8:3 -end-pos=14:4 > %t.result/L8-3.swift
+// RUN: %target-swift-frontend-typecheck %t.result/L8-3.swift
+// RUN: diff -u %S/Outputs/enum_with_raw_value/L8-3.swift.expected %t.result/L8-3.swift