[objcmt] Rewrite a NSDictionary dictionaryWithObjects:forKeys: to a dictionary literal

if we can see the elements of the arrays.

for example:

 NSDictionary *dict = [NSDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"1", @"2", nil] forKeys:[NSArray arrayWithObjects:@"A", @"B", nil]];

-->

 NSDictionary *dict = @{ @"A" : @"1", @"B" : @"2" };

rdar://12428166

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@172679 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/include/clang/AST/NSAPI.h b/include/clang/AST/NSAPI.h
index f9fd1f906b5553c6b08d95be82b628e630685abf..47797d0346b239d7475fa26c5cf1a4fffa85db11 100644 (file)
--- a/include/clang/AST/NSAPI.h
+++ b/include/clang/AST/NSAPI.h
@@ -96,10 +96,11 @@ public:
     NSDict_dictionaryWithObjectsAndKeys,
     NSDict_initWithDictionary,
     NSDict_initWithObjectsAndKeys,
+    NSDict_initWithObjectsForKeys,
     NSDict_objectForKey,
     NSMutableDict_setObjectForKey
   };
-  static const unsigned NumNSDictionaryMethods = 10;
+  static const unsigned NumNSDictionaryMethods = 11;
 
   /// \brief The Objective-C NSDictionary selectors.
   Selector getNSDictionarySelector(NSDictionaryMethodKind MK) const;

diff --git a/include/clang/Edit/Rewriters.h b/include/clang/Edit/Rewriters.h
index aa7a5b232025602ad16631d649fa2ee92e799070..292878e75695dfb8c2b95293b8b26c5956027d60 100644 (file)
--- a/include/clang/Edit/Rewriters.h
+++ b/include/clang/Edit/Rewriters.h
@@ -13,6 +13,7 @@
 namespace clang {
   class ObjCMessageExpr;
   class NSAPI;
+  class ParentMap;
 
 namespace edit {
   class Commit;
@@ -21,7 +22,8 @@ bool rewriteObjCRedundantCallWithLiteral(const ObjCMessageExpr *Msg,
                                          const NSAPI &NS, Commit &commit);
 
 bool rewriteToObjCLiteralSyntax(const ObjCMessageExpr *Msg,
-                                const NSAPI &NS, Commit &commit);
+                                const NSAPI &NS, Commit &commit,
+                                const ParentMap *PMap);
 
 bool rewriteToObjCSubscriptSyntax(const ObjCMessageExpr *Msg,
                                   const NSAPI &NS, Commit &commit);

diff --git a/lib/ARCMigrate/ObjCMT.cpp b/lib/ARCMigrate/ObjCMT.cpp
index 7fed082759c53b18902e0ee06ba7e525acc6f7ba..57fac0389fc30837e8d83c5dba8787f493e1ae9a 100644 (file)
--- a/lib/ARCMigrate/ObjCMT.cpp
+++ b/lib/ARCMigrate/ObjCMT.cpp
@@ -11,6 +11,7 @@
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/NSAPI.h"
+#include "clang/AST/ParentMap.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Edit/Commit.h"
@@ -120,9 +121,11 @@ bool ObjCMigrateAction::BeginInvocation(CompilerInstance &CI) {
 namespace {
 class ObjCMigrator : public RecursiveASTVisitor<ObjCMigrator> {
   ObjCMigrateASTConsumer &Consumer;
+  ParentMap &PMap;
 
 public:
-  ObjCMigrator(ObjCMigrateASTConsumer &consumer) : Consumer(consumer) { }
+  ObjCMigrator(ObjCMigrateASTConsumer &consumer, ParentMap &PMap)
+    : Consumer(consumer), PMap(PMap) { }
 
   bool shouldVisitTemplateInstantiations() const { return false; }
   bool shouldWalkTypesOfTypeLocs() const { return false; }
@@ -130,7 +133,7 @@ public:
   bool VisitObjCMessageExpr(ObjCMessageExpr *E) {
     if (Consumer.MigrateLiterals) {
       edit::Commit commit(*Consumer.Editor);
-      edit::rewriteToObjCLiteralSyntax(E, *Consumer.NSAPIObj, commit);
+      edit::rewriteToObjCLiteralSyntax(E, *Consumer.NSAPIObj, commit, &PMap);
       Consumer.Editor->commit(commit);
     }
 
@@ -153,6 +156,23 @@ public:
     return WalkUpFromObjCMessageExpr(E);
   }
 };
+
+class BodyMigrator : public RecursiveASTVisitor<BodyMigrator> {
+  ObjCMigrateASTConsumer &Consumer;
+  OwningPtr<ParentMap> PMap;
+
+public:
+  BodyMigrator(ObjCMigrateASTConsumer &consumer) : Consumer(consumer) { }
+
+  bool shouldVisitTemplateInstantiations() const { return false; }
+  bool shouldWalkTypesOfTypeLocs() const { return false; }
+
+  bool TraverseStmt(Stmt *S) {
+    PMap.reset(new ParentMap(S));
+    ObjCMigrator(Consumer, *PMap).TraverseStmt(S);
+    return true;
+  }
+};
 }
 
 void ObjCMigrateASTConsumer::migrateDecl(Decl *D) {
@@ -161,7 +181,7 @@ void ObjCMigrateASTConsumer::migrateDecl(Decl *D) {
   if (isa<ObjCMethodDecl>(D))
     return; // Wait for the ObjC container declaration.
 
-  ObjCMigrator(*this).TraverseDecl(D);
+  BodyMigrator(*this).TraverseDecl(D);
 }
 
 namespace {

diff --git a/lib/AST/NSAPI.cpp b/lib/AST/NSAPI.cpp
index b92b08c9933b078834a211558d9fcb31bae2382a..35cc7d050bdf7a84fea4f2f4ccad6d97da7035e6 100644 (file)
--- a/lib/AST/NSAPI.cpp
+++ b/lib/AST/NSAPI.cpp
@@ -186,6 +186,14 @@ Selector NSAPI::getNSDictionarySelector(
       Sel = Ctx.Selectors.getUnarySelector(
                                      &Ctx.Idents.get("initWithObjectsAndKeys"));
       break;
+    case NSDict_initWithObjectsForKeys: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("initWithObjects"),
+        &Ctx.Idents.get("forKeys")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
     case NSDict_objectForKey:
       Sel = Ctx.Selectors.getUnarySelector(&Ctx.Idents.get("objectForKey"));
       break;

diff --git a/lib/Edit/RewriteObjCFoundationAPI.cpp b/lib/Edit/RewriteObjCFoundationAPI.cpp
index 68d122dfd26f0764453c024e8501aeb6fc18c8f1..312c86fe9bf9885f90558da2ab22eca9130eeb1c 100644 (file)
--- a/lib/Edit/RewriteObjCFoundationAPI.cpp
+++ b/lib/Edit/RewriteObjCFoundationAPI.cpp
@@ -16,6 +16,7 @@
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/ExprObjC.h"
 #include "clang/AST/NSAPI.h"
+#include "clang/AST/ParentMap.h"
 #include "clang/Edit/Commit.h"
 #include "clang/Lex/Lexer.h"
 
@@ -325,7 +326,8 @@ bool edit::rewriteToObjCSubscriptSyntax(const ObjCMessageExpr *Msg,
 //===----------------------------------------------------------------------===//
 
 static bool rewriteToArrayLiteral(const ObjCMessageExpr *Msg,
-                                  const NSAPI &NS, Commit &commit);
+                                  const NSAPI &NS, Commit &commit,
+                                  const ParentMap *PMap);
 static bool rewriteToDictionaryLiteral(const ObjCMessageExpr *Msg,
                                   const NSAPI &NS, Commit &commit);
 static bool rewriteToNumberLiteral(const ObjCMessageExpr *Msg,
@@ -336,13 +338,14 @@ static bool rewriteToStringBoxedExpression(const ObjCMessageExpr *Msg,
                                            const NSAPI &NS, Commit &commit);
 
 bool edit::rewriteToObjCLiteralSyntax(const ObjCMessageExpr *Msg,
-                                      const NSAPI &NS, Commit &commit) {
+                                      const NSAPI &NS, Commit &commit,
+                                      const ParentMap *PMap) {
   IdentifierInfo *II = 0;
   if (!checkForLiteralCreation(Msg, II, NS.getASTContext().getLangOpts()))
     return false;
 
   if (II == NS.getNSClassId(NSAPI::ClassId_NSArray))
-    return rewriteToArrayLiteral(Msg, NS, commit);
+    return rewriteToArrayLiteral(Msg, NS, commit, PMap);
   if (II == NS.getNSClassId(NSAPI::ClassId_NSDictionary))
     return rewriteToDictionaryLiteral(Msg, NS, commit);
   if (II == NS.getNSClassId(NSAPI::ClassId_NSNumber))
@@ -353,6 +356,19 @@ bool edit::rewriteToObjCLiteralSyntax(const ObjCMessageExpr *Msg,
   return false;
 }
 
+/// \brief Returns true if the immediate message arguments of \c Msg should not
+/// be rewritten because it will interfere with the rewrite of the parent
+/// message expression. e.g.
+/// \code
+///   [NSDictionary dictionaryWithObjects:
+///                                 [NSArray arrayWithObjects:@"1", @"2", nil]
+///                         forKeys:[NSArray arrayWithObjects:@"A", @"B", nil]];
+/// \endcode
+/// It will return true for this because we are going to rewrite this directly
+/// to a dictionary literal without any array literals.
+static bool shouldNotRewriteImmediateMessageArgs(const ObjCMessageExpr *Msg,
+                                                 const NSAPI &NS);
+
 //===----------------------------------------------------------------------===//
 // rewriteToArrayLiteral.
 //===----------------------------------------------------------------------===//
@@ -361,7 +377,15 @@ bool edit::rewriteToObjCLiteralSyntax(const ObjCMessageExpr *Msg,
 static void objectifyExpr(const Expr *E, Commit &commit);
 
 static bool rewriteToArrayLiteral(const ObjCMessageExpr *Msg,
-                                  const NSAPI &NS, Commit &commit) {
+                                  const NSAPI &NS, Commit &commit,
+                                  const ParentMap *PMap) {
+  if (PMap) {
+    const ObjCMessageExpr *ParentMsg =
+        dyn_cast_or_null<ObjCMessageExpr>(PMap->getParentIgnoreParenCasts(Msg));
+    if (shouldNotRewriteImmediateMessageArgs(ParentMsg, NS))
+      return false;
+  }
+
   Selector Sel = Msg->getSelector();
   SourceRange MsgRange = Msg->getSourceRange();
 
@@ -411,6 +435,59 @@ static bool rewriteToArrayLiteral(const ObjCMessageExpr *Msg,
 // rewriteToDictionaryLiteral.
 //===----------------------------------------------------------------------===//
 
+/// \brief If \c Msg is an NSArray creation message or literal, this gets the
+/// objects that were used to create it.
+/// \returns true if it is an NSArray and we got objects, or false otherwise.
+static bool getNSArrayObjects(const Expr *E, const NSAPI &NS,
+                              SmallVectorImpl<const Expr *> &Objs) {
+  if (!E)
+    return false;
+
+  E = E->IgnoreParenCasts();
+  if (!E)
+    return false;
+
+  if (const ObjCMessageExpr *Msg = dyn_cast<ObjCMessageExpr>(E)) {
+    IdentifierInfo *Cls = 0;
+    if (!checkForLiteralCreation(Msg, Cls, NS.getASTContext().getLangOpts()))
+      return false;
+
+    if (Cls != NS.getNSClassId(NSAPI::ClassId_NSArray))
+      return false;
+
+    Selector Sel = Msg->getSelector();
+    if (Sel == NS.getNSArraySelector(NSAPI::NSArr_array))
+      return true; // empty array.
+
+    if (Sel == NS.getNSArraySelector(NSAPI::NSArr_arrayWithObject)) {
+      if (Msg->getNumArgs() != 1)
+        return false;
+      Objs.push_back(Msg->getArg(0));
+      return true;
+    }
+
+    if (Sel == NS.getNSArraySelector(NSAPI::NSArr_arrayWithObjects) ||
+        Sel == NS.getNSArraySelector(NSAPI::NSArr_initWithObjects)) {
+      if (Msg->getNumArgs() == 0)
+        return false;
+      const Expr *SentinelExpr = Msg->getArg(Msg->getNumArgs() - 1);
+      if (!NS.getASTContext().isSentinelNullExpr(SentinelExpr))
+        return false;
+
+      for (unsigned i = 0, e = Msg->getNumArgs() - 1; i != e; ++i)
+        Objs.push_back(Msg->getArg(i));
+      return true;
+    }
+
+  } else if (const ObjCArrayLiteral *ArrLit = dyn_cast<ObjCArrayLiteral>(E)) {
+    for (unsigned i = 0, e = ArrLit->getNumElements(); i != e; ++i)
+      Objs.push_back(ArrLit->getElement(i));
+    return true;
+  }
+
+  return false;
+}
+
 static bool rewriteToDictionaryLiteral(const ObjCMessageExpr *Msg,
                                        const NSAPI &NS, Commit &commit) {
   Selector Sel = Msg->getSelector();
@@ -481,6 +558,83 @@ static bool rewriteToDictionaryLiteral(const ObjCMessageExpr *Msg,
     return true;
   }
 
+  if (Sel == NS.getNSDictionarySelector(
+                                  NSAPI::NSDict_dictionaryWithObjectsForKeys) ||
+      Sel == NS.getNSDictionarySelector(NSAPI::NSDict_initWithObjectsForKeys)) {
+    if (Msg->getNumArgs() != 2)
+      return false;
+
+    SmallVector<const Expr *, 8> Vals;
+    if (!getNSArrayObjects(Msg->getArg(0), NS, Vals))
+      return false;
+
+    SmallVector<const Expr *, 8> Keys;
+    if (!getNSArrayObjects(Msg->getArg(1), NS, Keys))
+      return false;
+
+    if (Vals.size() != Keys.size())
+      return false;
+
+    if (Vals.empty()) {
+      commit.replace(MsgRange, "@{}");
+      return true;
+    }
+
+    for (unsigned i = 0, n = Vals.size(); i < n; ++i) {
+      objectifyExpr(Vals[i], commit);
+      objectifyExpr(Keys[i], commit);
+
+      SourceRange ValRange = Vals[i]->getSourceRange();
+      SourceRange KeyRange = Keys[i]->getSourceRange();
+      // Insert value after key.
+      commit.insertAfterToken(KeyRange.getEnd(), ": ");
+      commit.insertFromRange(KeyRange.getEnd(), ValRange, /*afterToken=*/true);
+    }
+    // Range of arguments up until and including the last key.
+    // The first value is cut off, the value will move after the key.
+    SourceRange ArgRange(Keys.front()->getLocStart(),
+                         Keys.back()->getLocEnd());
+    commit.insertWrap("@{", ArgRange, "}");
+    commit.replaceWithInner(MsgRange, ArgRange);
+    return true;
+  }
+
+  return false;
+}
+
+static bool shouldNotRewriteImmediateMessageArgs(const ObjCMessageExpr *Msg,
+                                                 const NSAPI &NS) {
+  if (!Msg)
+    return false;
+
+  IdentifierInfo *II = 0;
+  if (!checkForLiteralCreation(Msg, II, NS.getASTContext().getLangOpts()))
+    return false;
+
+  if (II != NS.getNSClassId(NSAPI::ClassId_NSDictionary))
+    return false;
+
+  Selector Sel = Msg->getSelector();
+  if (Sel == NS.getNSDictionarySelector(
+                                  NSAPI::NSDict_dictionaryWithObjectsForKeys) ||
+      Sel == NS.getNSDictionarySelector(NSAPI::NSDict_initWithObjectsForKeys)) {
+    if (Msg->getNumArgs() != 2)
+      return false;
+
+    SmallVector<const Expr *, 8> Vals;
+    if (!getNSArrayObjects(Msg->getArg(0), NS, Vals))
+      return false;
+
+    SmallVector<const Expr *, 8> Keys;
+    if (!getNSArrayObjects(Msg->getArg(1), NS, Keys))
+      return false;
+
+    if (Vals.size() != Keys.size())
+      return false;
+
+    return true;
+  }
+
   return false;
 }
 

diff --git a/test/ARCMT/objcmt-subscripting-literals-in-arc.m b/test/ARCMT/objcmt-subscripting-literals-in-arc.m
index 4d941626c0d07b9bf9412a4dceb6b30eeeada812..1f56f4a2cf510c1bc63c1e5af29fa3201431798c 100644 (file)
--- a/test/ARCMT/objcmt-subscripting-literals-in-arc.m
+++ b/test/ARCMT/objcmt-subscripting-literals-in-arc.m
@@ -101,6 +101,8 @@ typedef const struct __CFString * CFStringRef;
   dict = [NSDictionary dictionaryWithObjectsAndKeys: @"value1", @"key1", @"value2", @"key2", nil];
   dict = [[NSDictionary alloc] initWithObjectsAndKeys: @"value1", @"key1", @"value2", @"key2", nil];
 
+  dict = [[NSDictionary alloc] initWithObjects:[[NSArray alloc] initWithObjects:@"1", @"2", nil] forKeys:[NSArray arrayWithObjects:@"A", @"B", nil]];
+
   NSNumber *n = [[NSNumber alloc] initWithInt:2];
 }
 @end

diff --git a/test/ARCMT/objcmt-subscripting-literals-in-arc.m.result b/test/ARCMT/objcmt-subscripting-literals-in-arc.m.result
index 6f7a723bc4731ce4e81d2ec41dd4e9c686efbf7b..d974a2564d4302b386b62641efa88c4dc4ece693 100644 (file)
--- a/test/ARCMT/objcmt-subscripting-literals-in-arc.m.result
+++ b/test/ARCMT/objcmt-subscripting-literals-in-arc.m.result
@@ -101,6 +101,8 @@ typedef const struct __CFString * CFStringRef;
   dict = @{@"key1": @"value1", @"key2": @"value2"};
   dict = @{@"key1": @"value1", @"key2": @"value2"};
 
+  dict = @{@"A": @"1", @"B": @"2"};
+
   NSNumber *n = @2;
 }
 @end

diff --git a/test/ARCMT/objcmt-subscripting-literals.m b/test/ARCMT/objcmt-subscripting-literals.m
index 0174fcf060e92f76e62d2b52a5d5b3cb0d16c603..8cef0919bba84f68d4275527dd3626008d6989ef 100644 (file)
--- a/test/ARCMT/objcmt-subscripting-literals.m
+++ b/test/ARCMT/objcmt-subscripting-literals.m
@@ -153,6 +153,10 @@ typedef const struct __CFString * CFStringRef;
   void *hd;
   o = [(NSArray*)hd objectAtIndex:2];
   o = [ivarArr objectAtIndex:2];
+
+  dict = [NSDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"1", [NSArray array], nil] forKeys:[NSArray arrayWithObjects:@"A", [arr objectAtIndex:2], nil]];
+  dict = [NSDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"1", @"2", nil] forKeys:arr];
+  dict = [NSDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"1", @"2", nil] forKeys:@[@"A", @"B"]];
 }
 @end
 

diff --git a/test/ARCMT/objcmt-subscripting-literals.m.result b/test/ARCMT/objcmt-subscripting-literals.m.result
index 9975996524bd3b2055e1315a5c2e4c787e582c35..0ca6dca1fea2a42aba15888272d54200bdc1f7b1 100644 (file)
--- a/test/ARCMT/objcmt-subscripting-literals.m.result
+++ b/test/ARCMT/objcmt-subscripting-literals.m.result
@@ -153,6 +153,10 @@ typedef const struct __CFString * CFStringRef;
   void *hd;
   o = ((NSArray*)hd)[2];
   o = ivarArr[2];
+
+  dict = @{@"A": @"1", arr[2]: @[]};
+  dict = [NSDictionary dictionaryWithObjects:@[@"1", @"2"] forKeys:arr];
+  dict = @{@"A": @"1", @"B": @"2"};
 }
 @end