[objcmt] Don't migrate to subscripting syntax if the required methods have not
been declared on NSArray/NSDictionary.

rdar://11581975

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@157951 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/include/clang/AST/NSAPI.h b/include/clang/AST/NSAPI.h
index 3fcd8b4..ec5e542 100644
--- a/include/clang/AST/NSAPI.h
+++ b/include/clang/AST/NSAPI.h
@@ -11,6 +11,7 @@
 #define LLVM_CLANG_AST_NSAPI_H
 
 #include "clang/Basic/IdentifierTable.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 
 namespace clang {
@@ -107,6 +108,30 @@
   llvm::Optional<NSDictionaryMethodKind>
       getNSDictionaryMethodKind(Selector Sel);
 
+  /// \brief Returns selector for "objectForKeyedSubscript:".
+  Selector getObjectForKeyedSubscriptSelector() const {
+    return getOrInitSelector(StringRef("objectForKeyedSubscript"),
+                             objectForKeyedSubscriptSel);
+  }
+
+  /// \brief Returns selector for "objectAtIndexedSubscript:".
+  Selector getObjectAtIndexedSubscriptSelector() const {
+    return getOrInitSelector(StringRef("objectAtIndexedSubscript"),
+                             objectAtIndexedSubscriptSel);
+  }
+
+  /// \brief Returns selector for "setObject:forKeyedSubscript".
+  Selector getSetObjectForKeyedSubscriptSelector() const {
+    StringRef Ids[] = { "setObject", "forKeyedSubscript" };
+    return getOrInitSelector(Ids, setObjectForKeyedSubscriptSel);
+  }
+
+  /// \brief Returns selector for "setObject:atIndexedSubscript".
+  Selector getSetObjectAtIndexedSubscriptSelector() const {
+    StringRef Ids[] = { "setObject", "atIndexedSubscript" };
+    return getOrInitSelector(Ids, setObjectAtIndexedSubscriptSel);
+  }
+
   /// \brief Enumerates the NSNumber methods used to generate literals.
   enum NSNumberLiteralMethodKind {
     NSNumberWithChar,
@@ -159,6 +184,7 @@
   bool isObjCTypedef(QualType T, StringRef name, IdentifierInfo *&II) const;
   bool isObjCEnumerator(const Expr *E,
                         StringRef name, IdentifierInfo *&II) const;
+  Selector getOrInitSelector(ArrayRef<StringRef> Ids, Selector &Sel) const;
 
   ASTContext &Ctx;
 
@@ -176,6 +202,9 @@
   mutable Selector NSNumberClassSelectors[NumNSNumberLiteralMethods];
   mutable Selector NSNumberInstanceSelectors[NumNSNumberLiteralMethods];
 
+  mutable Selector objectForKeyedSubscriptSel, objectAtIndexedSubscriptSel,
+                   setObjectForKeyedSubscriptSel,setObjectAtIndexedSubscriptSel;
+
   mutable IdentifierInfo *BOOLId, *NSIntegerId, *NSUIntegerId;
   mutable IdentifierInfo *NSASCIIStringEncodingId, *NSUTF8StringEncodingId;
 };
diff --git a/lib/AST/NSAPI.cpp b/lib/AST/NSAPI.cpp
index 72251d3..39077d1 100644
--- a/lib/AST/NSAPI.cpp
+++ b/lib/AST/NSAPI.cpp
@@ -399,3 +399,15 @@
 
   return false;
 }
+
+Selector NSAPI::getOrInitSelector(ArrayRef<StringRef> Ids,
+                                  Selector &Sel) const {
+  if (Sel.isNull()) {
+    SmallVector<IdentifierInfo *, 4> Idents;
+    for (ArrayRef<StringRef>::const_iterator
+           I = Ids.begin(), E = Ids.end(); I != E; ++I)
+      Idents.push_back(&Ctx.Idents.get(*I));
+    Sel = Ctx.Selectors.getSelector(Idents.size(), Idents.data());
+  }
+  return Sel;
+}
diff --git a/lib/Edit/RewriteObjCFoundationAPI.cpp b/lib/Edit/RewriteObjCFoundationAPI.cpp
index 9c00d9e..c36793f 100644
--- a/lib/Edit/RewriteObjCFoundationAPI.cpp
+++ b/lib/Edit/RewriteObjCFoundationAPI.cpp
@@ -86,7 +86,8 @@
   }
 }
 
-static bool rewriteToSubscriptGet(const ObjCMessageExpr *Msg, Commit &commit) {
+static bool rewriteToSubscriptGetCommon(const ObjCMessageExpr *Msg,
+                                        Commit &commit) {
   if (Msg->getNumArgs() != 1)
     return false;
   const Expr *Rec = Msg->getInstanceReceiver();
@@ -107,13 +108,35 @@
   return true;
 }
 
-static bool rewriteToArraySubscriptSet(const ObjCMessageExpr *Msg,
+static bool rewriteToArraySubscriptGet(const ObjCInterfaceDecl *IFace,
+                                       const ObjCMessageExpr *Msg,
+                                       const NSAPI &NS,
+                                       Commit &commit) {
+  if (!IFace->getInstanceMethod(NS.getObjectAtIndexedSubscriptSelector()))
+    return false;
+  return rewriteToSubscriptGetCommon(Msg, commit);
+}
+
+static bool rewriteToDictionarySubscriptGet(const ObjCInterfaceDecl *IFace,
+                                            const ObjCMessageExpr *Msg,
+                                            const NSAPI &NS,
+                                            Commit &commit) {
+  if (!IFace->getInstanceMethod(NS.getObjectForKeyedSubscriptSelector()))
+    return false;
+  return rewriteToSubscriptGetCommon(Msg, commit);
+}
+
+static bool rewriteToArraySubscriptSet(const ObjCInterfaceDecl *IFace,
+                                       const ObjCMessageExpr *Msg,
+                                       const NSAPI &NS,
                                        Commit &commit) {
   if (Msg->getNumArgs() != 2)
     return false;
   const Expr *Rec = Msg->getInstanceReceiver();
   if (!Rec)
     return false;
+  if (!IFace->getInstanceMethod(NS.getSetObjectAtIndexedSubscriptSelector()))
+    return false;
 
   SourceRange MsgRange = Msg->getSourceRange();
   SourceRange RecRange = Rec->getSourceRange();
@@ -135,13 +158,17 @@
   return true;
 }
 
-static bool rewriteToDictionarySubscriptSet(const ObjCMessageExpr *Msg,
+static bool rewriteToDictionarySubscriptSet(const ObjCInterfaceDecl *IFace,
+                                            const ObjCMessageExpr *Msg,
+                                            const NSAPI &NS,
                                             Commit &commit) {
   if (Msg->getNumArgs() != 2)
     return false;
   const Expr *Rec = Msg->getInstanceReceiver();
   if (!Rec)
     return false;
+  if (!IFace->getInstanceMethod(NS.getSetObjectForKeyedSubscriptSelector()))
+    return false;
 
   SourceRange MsgRange = Msg->getSourceRange();
   SourceRange RecRange = Rec->getSourceRange();
@@ -163,7 +190,7 @@
 }
 
 bool edit::rewriteToObjCSubscriptSyntax(const ObjCMessageExpr *Msg,
-                                           const NSAPI &NS, Commit &commit) {
+                                        const NSAPI &NS, Commit &commit) {
   if (!Msg || Msg->isImplicit() ||
       Msg->getReceiverKind() != ObjCMessageExpr::Instance)
     return false;
@@ -179,22 +206,24 @@
   IdentifierInfo *II = IFace->getIdentifier();
   Selector Sel = Msg->getSelector();
 
-  if ((II == NS.getNSClassId(NSAPI::ClassId_NSArray) &&
-       Sel == NS.getNSArraySelector(NSAPI::NSArr_objectAtIndex)) ||
-      (II == NS.getNSClassId(NSAPI::ClassId_NSDictionary) &&
-       Sel == NS.getNSDictionarySelector(NSAPI::NSDict_objectForKey)))
-    return rewriteToSubscriptGet(Msg, commit);
+  if (II == NS.getNSClassId(NSAPI::ClassId_NSArray) &&
+      Sel == NS.getNSArraySelector(NSAPI::NSArr_objectAtIndex))
+    return rewriteToArraySubscriptGet(IFace, Msg, NS, commit);
+
+  if (II == NS.getNSClassId(NSAPI::ClassId_NSDictionary) &&
+      Sel == NS.getNSDictionarySelector(NSAPI::NSDict_objectForKey))
+    return rewriteToDictionarySubscriptGet(IFace, Msg, NS, commit);
 
   if (Msg->getNumArgs() != 2)
     return false;
 
   if (II == NS.getNSClassId(NSAPI::ClassId_NSMutableArray) &&
       Sel == NS.getNSArraySelector(NSAPI::NSMutableArr_replaceObjectAtIndex))
-    return rewriteToArraySubscriptSet(Msg, commit);
+    return rewriteToArraySubscriptSet(IFace, Msg, NS, commit);
 
   if (II == NS.getNSClassId(NSAPI::ClassId_NSMutableDictionary) &&
       Sel == NS.getNSDictionarySelector(NSAPI::NSMutableDict_setObjectForKey))
-    return rewriteToDictionarySubscriptSet(Msg, commit);
+    return rewriteToDictionarySubscriptSet(IFace, Msg, NS, commit);
 
   return false;
 }
diff --git a/test/ARCMT/objcmt-subscripting-unavailable.m b/test/ARCMT/objcmt-subscripting-unavailable.m
new file mode 100644
index 0000000..d72c362
--- /dev/null
+++ b/test/ARCMT/objcmt-subscripting-unavailable.m
@@ -0,0 +1,79 @@
+// RUN: rm -rf %t
+// RUN: %clang_cc1 -objcmt-migrate-literals -objcmt-migrate-subscripting -mt-migrate-directory %t %s -x objective-c -triple x86_64-apple-darwin11 
+// RUN: c-arcmt-test -mt-migrate-directory %t | arcmt-test -verify-transformed-files %s.result
+// RUN: %clang_cc1 -triple x86_64-apple-darwin10 -fsyntax-only -x objective-c %s.result
+
+typedef signed char BOOL;
+#define nil ((void*) 0)
+
+@interface NSObject
++ (id)alloc;
+@end
+
+@interface NSArray : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface NSArray (NSArrayCreation)
++ (id)array;
++ (id)arrayWithObject:(id)anObject;
++ (id)arrayWithObjects:(const id [])objects count:(unsigned long)cnt;
++ (id)arrayWithObjects:(id)firstObj, ...;
++ (id)arrayWithArray:(NSArray *)array;
+
+- (id)initWithObjects:(const id [])objects count:(unsigned long)cnt;
+- (id)initWithObjects:(id)firstObj, ...;
+- (id)initWithArray:(NSArray *)array;
+@end
+
+@interface NSMutableArray : NSArray
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface NSDictionary : NSObject
+@end
+
+@interface NSDictionary (NSDictionaryCreation)
++ (id)dictionary;
++ (id)dictionaryWithObject:(id)object forKey:(id)key;
++ (id)dictionaryWithObjects:(const id [])objects forKeys:(const id [])keys count:(unsigned long)cnt;
++ (id)dictionaryWithObjectsAndKeys:(id)firstObject, ...;
++ (id)dictionaryWithDictionary:(NSDictionary *)dict;
++ (id)dictionaryWithObjects:(NSArray *)objects forKeys:(NSArray *)keys;
+
+- (id)initWithObjects:(const id [])objects forKeys:(const id [])keys count:(unsigned long)cnt;
+- (id)initWithObjectsAndKeys:(id)firstObject, ...;
+- (id)initWithDictionary:(NSDictionary *)otherDictionary;
+- (id)initWithObjects:(NSArray *)objects forKeys:(NSArray *)keys;
+
+- (id)objectForKey:(id)aKey;
+@end
+
+@interface NSMutableDictionary : NSDictionary
+- (void)setObject:(id)anObject forKey:(id)aKey;
+@end
+
+@interface I
+@end
+@implementation I
+-(void) foo {
+  id str;
+  NSArray *arr;
+  NSDictionary *dict;
+
+  arr = [NSArray array];
+  arr = [NSArray arrayWithObject:str];
+  arr = [NSArray arrayWithObjects:str, str, nil];
+  dict = [NSDictionary dictionary];
+  dict = [NSDictionary dictionaryWithObject:arr forKey:str];
+
+  id o = [arr objectAtIndex:2];
+  o = [dict objectForKey:@"key"];
+  NSMutableArray *marr = 0;
+  NSMutableDictionary *mdict = 0;
+  [marr replaceObjectAtIndex:2 withObject:@"val"];
+  [mdict setObject:@"value" forKey:@"key"];
+  [marr replaceObjectAtIndex:2 withObject:[arr objectAtIndex:4]];
+  [mdict setObject:[dict objectForKey:@"key2"] forKey:@"key"];
+}
+@end
diff --git a/test/ARCMT/objcmt-subscripting-unavailable.m.result b/test/ARCMT/objcmt-subscripting-unavailable.m.result
new file mode 100644
index 0000000..bd74d55
--- /dev/null
+++ b/test/ARCMT/objcmt-subscripting-unavailable.m.result
@@ -0,0 +1,79 @@
+// RUN: rm -rf %t
+// RUN: %clang_cc1 -objcmt-migrate-literals -objcmt-migrate-subscripting -mt-migrate-directory %t %s -x objective-c -triple x86_64-apple-darwin11 
+// RUN: c-arcmt-test -mt-migrate-directory %t | arcmt-test -verify-transformed-files %s.result
+// RUN: %clang_cc1 -triple x86_64-apple-darwin10 -fsyntax-only -x objective-c %s.result
+
+typedef signed char BOOL;
+#define nil ((void*) 0)
+
+@interface NSObject
++ (id)alloc;
+@end
+
+@interface NSArray : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface NSArray (NSArrayCreation)
++ (id)array;
++ (id)arrayWithObject:(id)anObject;
++ (id)arrayWithObjects:(const id [])objects count:(unsigned long)cnt;
++ (id)arrayWithObjects:(id)firstObj, ...;
++ (id)arrayWithArray:(NSArray *)array;
+
+- (id)initWithObjects:(const id [])objects count:(unsigned long)cnt;
+- (id)initWithObjects:(id)firstObj, ...;
+- (id)initWithArray:(NSArray *)array;
+@end
+
+@interface NSMutableArray : NSArray
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface NSDictionary : NSObject
+@end
+
+@interface NSDictionary (NSDictionaryCreation)
++ (id)dictionary;
++ (id)dictionaryWithObject:(id)object forKey:(id)key;
++ (id)dictionaryWithObjects:(const id [])objects forKeys:(const id [])keys count:(unsigned long)cnt;
++ (id)dictionaryWithObjectsAndKeys:(id)firstObject, ...;
++ (id)dictionaryWithDictionary:(NSDictionary *)dict;
++ (id)dictionaryWithObjects:(NSArray *)objects forKeys:(NSArray *)keys;
+
+- (id)initWithObjects:(const id [])objects forKeys:(const id [])keys count:(unsigned long)cnt;
+- (id)initWithObjectsAndKeys:(id)firstObject, ...;
+- (id)initWithDictionary:(NSDictionary *)otherDictionary;
+- (id)initWithObjects:(NSArray *)objects forKeys:(NSArray *)keys;
+
+- (id)objectForKey:(id)aKey;
+@end
+
+@interface NSMutableDictionary : NSDictionary
+- (void)setObject:(id)anObject forKey:(id)aKey;
+@end
+
+@interface I
+@end
+@implementation I
+-(void) foo {
+  id str;
+  NSArray *arr;
+  NSDictionary *dict;
+
+  arr = @[];
+  arr = @[str];
+  arr = @[str, str];
+  dict = @{};
+  dict = @{str: arr};
+
+  id o = [arr objectAtIndex:2];
+  o = [dict objectForKey:@"key"];
+  NSMutableArray *marr = 0;
+  NSMutableDictionary *mdict = 0;
+  [marr replaceObjectAtIndex:2 withObject:@"val"];
+  [mdict setObject:@"value" forKey:@"key"];
+  [marr replaceObjectAtIndex:2 withObject:[arr objectAtIndex:4]];
+  [mdict setObject:[dict objectForKey:@"key2"] forKey:@"key"];
+}
+@end