ObjectiveC migrator: Another use case of enum

declaration which can be migrated to NS_ENUM.

llvm-svn: 186716

3 files changed, 48 insertions, 3 deletions

diff --git a/clang/lib/ARCMigrate/ObjCMT.cpp b/clang/lib/ARCMigrate/ObjCMT.cpp
index 282006b5769..e4ec0c83c50 100644
--- a/clang/lib/ARCMigrate/ObjCMT.cpp
+++ b/clang/lib/ARCMigrate/ObjCMT.cpp
@@ -384,7 +384,7 @@ static bool rewriteToNSEnumDecl(const EnumDecl *EnumDcl,
                                 const TypedefDecl *TypedefDcl,
                                 const NSAPI &NS, edit::Commit &commit,
                                 bool IsNSIntegerType) {
-  std::string ClassString = 
+  std::string ClassString =
     IsNSIntegerType ? "typedef NS_ENUM(NSInteger, " : "typedef NS_OPTIONS(NSUInteger, ";
   ClassString += TypedefDcl->getIdentifier()->getName();
   ClassString += ')';
@@ -399,6 +399,19 @@ static bool rewriteToNSEnumDecl(const EnumDecl *EnumDcl,
   return false;
 }
 
+static bool rewriteToNSEnumDecl(const EnumDecl *EnumDcl,
+                                const TypedefDecl *TypedefDcl,
+                                const NSAPI &NS, edit::Commit &commit) {
+  std::string ClassString = "NS_ENUM(NSInteger, ";
+  ClassString += TypedefDcl->getIdentifier()->getName();
+  ClassString += ')';
+  SourceRange R(EnumDcl->getLocStart(), EnumDcl->getLocStart());
+  commit.replace(R, ClassString);
+  SourceLocation TypedefLoc = TypedefDcl->getLocEnd();
+  commit.remove(SourceRange(TypedefLoc, TypedefLoc));
+  return true;
+}
+
 void ObjCMigrateASTConsumer::migrateProtocolConformance(ASTContext &Ctx,   
                                             const ObjCImplementationDecl *ImpDecl) {
   const ObjCInterfaceDecl *IDecl = ImpDecl->getClassInterface();
@@ -466,8 +479,24 @@ void ObjCMigrateASTConsumer::migrateNSEnumDecl(ASTContext &Ctx,
   QualType qt = TypedefDcl->getTypeSourceInfo()->getType();
   bool IsNSIntegerType = NSAPIObj->isObjCNSIntegerType(qt);
   bool IsNSUIntegerType = !IsNSIntegerType && NSAPIObj->isObjCNSUIntegerType(qt);
-  if (!IsNSIntegerType && !IsNSUIntegerType)
-    return;
+  if (!IsNSIntegerType && !IsNSUIntegerType) {
+    // Also check for typedef enum {...} TD;
+    if (const EnumType *EnumTy = qt->getAs<EnumType>()) {
+      if (EnumTy->getDecl() == EnumDcl) {
+        // NS_ENUM must be available.
+        if (!Ctx.Idents.get("NS_ENUM").hasMacroDefinition())
+          return;
+        edit::Commit commit(*Editor);
+        rewriteToNSEnumDecl(EnumDcl, TypedefDcl, *NSAPIObj, commit);
+        Editor->commit(commit);
+        return;
+      }
+      else
+        return;
+    }
+    else
+      return;
+  }
   
   // NS_ENUM must be available.
   if (IsNSIntegerType && !Ctx.Idents.get("NS_ENUM").hasMacroDefinition())
diff --git a/clang/test/ARCMT/objcmt-ns-macros.m b/clang/test/ARCMT/objcmt-ns-macros.m
index b9069f57923..42d7d2cd8f6 100644
--- a/clang/test/ARCMT/objcmt-ns-macros.m
+++ b/clang/test/ARCMT/objcmt-ns-macros.m
@@ -27,6 +27,14 @@ enum {
 
 typedef NSUInteger UITableViewCellStyle;
 
+typedef enum {
+    UIViewAnimationTransitionNone,
+    UIViewAnimationTransitionFlipFromLeft,
+    UIViewAnimationTransitionFlipFromRight,
+    UIViewAnimationTransitionCurlUp,
+    UIViewAnimationTransitionCurlDown,
+} UIViewAnimationTransition;
+
 enum {
   UNOne,
   UNTwo
diff --git a/clang/test/ARCMT/objcmt-ns-macros.m.result b/clang/test/ARCMT/objcmt-ns-macros.m.result
index fcfe077039c..ee9b60596ff 100644
--- a/clang/test/ARCMT/objcmt-ns-macros.m.result
+++ b/clang/test/ARCMT/objcmt-ns-macros.m.result
@@ -27,6 +27,14 @@ typedef NS_OPTIONS(NSUInteger, UITableViewCellStyle) {
 
 
 
+typedef NS_ENUM(NSInteger, UIViewAnimationTransition) {
+    UIViewAnimationTransitionNone,
+    UIViewAnimationTransitionFlipFromLeft,
+    UIViewAnimationTransitionFlipFromRight,
+    UIViewAnimationTransitionCurlUp,
+    UIViewAnimationTransitionCurlDown,
+} ;
+
 enum {
   UNOne,
   UNTwo