ObjectiveC migrator: Another use case of enum
declaration which can be migrated to NS_ENUM.



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@186716 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/lib/ARCMigrate/ObjCMT.cpp b/lib/ARCMigrate/ObjCMT.cpp
index 282006b..e4ec0c8 100644
--- a/lib/ARCMigrate/ObjCMT.cpp
+++ b/lib/ARCMigrate/ObjCMT.cpp
@@ -384,7 +384,7 @@
                                 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 @@
   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 @@
   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/test/ARCMT/objcmt-ns-macros.m b/test/ARCMT/objcmt-ns-macros.m
index b9069f5..42d7d2c 100644
--- a/test/ARCMT/objcmt-ns-macros.m
+++ b/test/ARCMT/objcmt-ns-macros.m
@@ -27,6 +27,14 @@
 
 typedef NSUInteger UITableViewCellStyle;
 
+typedef enum {
+    UIViewAnimationTransitionNone,
+    UIViewAnimationTransitionFlipFromLeft,
+    UIViewAnimationTransitionFlipFromRight,
+    UIViewAnimationTransitionCurlUp,
+    UIViewAnimationTransitionCurlDown,
+} UIViewAnimationTransition;
+
 enum {
   UNOne,
   UNTwo
diff --git a/test/ARCMT/objcmt-ns-macros.m.result b/test/ARCMT/objcmt-ns-macros.m.result
index fcfe077..ee9b605 100644
--- a/test/ARCMT/objcmt-ns-macros.m.result
+++ b/test/ARCMT/objcmt-ns-macros.m.result
@@ -27,6 +27,14 @@
 
 
 
+typedef NS_ENUM(NSInteger, UIViewAnimationTransition) {
+    UIViewAnimationTransitionNone,
+    UIViewAnimationTransitionFlipFromLeft,
+    UIViewAnimationTransitionFlipFromRight,
+    UIViewAnimationTransitionCurlUp,
+    UIViewAnimationTransitionCurlDown,
+} ;
+
 enum {
   UNOne,
   UNTwo