[objcmt] Allow migrating to subscripting syntax for other classes

(apart from NSDictionary/NSArray) that implement objectForKey:/objectAtIndex/etc.
and the subscripting methods as well.

Part of rdar://11734969

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

3 files changed, 109 insertions, 15 deletions

diff --git a/lib/Edit/RewriteObjCFoundationAPI.cpp b/lib/Edit/RewriteObjCFoundationAPI.cpp
index 0e7b877994..36704f66dc 100644
--- a/lib/Edit/RewriteObjCFoundationAPI.cpp
+++ b/lib/Edit/RewriteObjCFoundationAPI.cpp
@@ -94,6 +94,15 @@ bool edit::rewriteObjCRedundantCallWithLiteral(const ObjCMessageExpr *Msg,
 // rewriteToObjCSubscriptSyntax.
 //===----------------------------------------------------------------------===//
 
+static bool canRewriteToSubscriptSyntax(const ObjCInterfaceDecl *IFace,
+                                        Selector subscriptSel) {
+  if (const ObjCMethodDecl *MD = IFace->lookupInstanceMethod(subscriptSel)) {
+    if (!MD->isUnavailable())
+      return true;
+  }
+  return false;
+}
+
 static bool subscriptOperatorNeedsParens(const Expr *FullExpr);
 
 static void maybePutParensOnReceiver(const Expr *Receiver, Commit &commit) {
@@ -129,7 +138,8 @@ static bool rewriteToArraySubscriptGet(const ObjCInterfaceDecl *IFace,
                                        const ObjCMessageExpr *Msg,
                                        const NSAPI &NS,
                                        Commit &commit) {
-  if (!IFace->lookupInstanceMethod(NS.getObjectAtIndexedSubscriptSelector()))
+  if (!canRewriteToSubscriptSyntax(IFace,
+                                   NS.getObjectAtIndexedSubscriptSelector()))
     return false;
   return rewriteToSubscriptGetCommon(Msg, commit);
 }
@@ -138,7 +148,8 @@ static bool rewriteToDictionarySubscriptGet(const ObjCInterfaceDecl *IFace,
                                             const ObjCMessageExpr *Msg,
                                             const NSAPI &NS,
                                             Commit &commit) {
-  if (!IFace->lookupInstanceMethod(NS.getObjectForKeyedSubscriptSelector()))
+  if (!canRewriteToSubscriptSyntax(IFace,
+                                  NS.getObjectForKeyedSubscriptSelector()))
     return false;
   return rewriteToSubscriptGetCommon(Msg, commit);
 }
@@ -147,13 +158,15 @@ static bool rewriteToArraySubscriptSet(const ObjCInterfaceDecl *IFace,
                                        const ObjCMessageExpr *Msg,
                                        const NSAPI &NS,
                                        Commit &commit) {
+  if (!canRewriteToSubscriptSyntax(IFace,
+                                   NS.getSetObjectAtIndexedSubscriptSelector()))
+    return false;
+
   if (Msg->getNumArgs() != 2)
     return false;
   const Expr *Rec = Msg->getInstanceReceiver();
   if (!Rec)
     return false;
-  if (!IFace->lookupInstanceMethod(NS.getSetObjectAtIndexedSubscriptSelector()))
-    return false;
 
   SourceRange MsgRange = Msg->getSourceRange();
   SourceRange RecRange = Rec->getSourceRange();
@@ -179,13 +192,15 @@ static bool rewriteToDictionarySubscriptSet(const ObjCInterfaceDecl *IFace,
                                             const ObjCMessageExpr *Msg,
                                             const NSAPI &NS,
                                             Commit &commit) {
+  if (!canRewriteToSubscriptSyntax(IFace,
+                                   NS.getSetObjectForKeyedSubscriptSelector()))
+    return false;
+
   if (Msg->getNumArgs() != 2)
     return false;
   const Expr *Rec = Msg->getInstanceReceiver();
   if (!Rec)
     return false;
-  if (!IFace->lookupInstanceMethod(NS.getSetObjectForKeyedSubscriptSelector()))
-    return false;
 
   SourceRange MsgRange = Msg->getSourceRange();
   SourceRange RecRange = Rec->getSourceRange();
@@ -220,26 +235,21 @@ bool edit::rewriteToObjCSubscriptSyntax(const ObjCMessageExpr *Msg,
                                           const_cast<ObjCMethodDecl *>(Method));
   if (!IFace)
     return false;
-  IdentifierInfo *II = IFace->getIdentifier();
   Selector Sel = Msg->getSelector();
 
-  if (II == NS.getNSClassId(NSAPI::ClassId_NSArray) &&
-      Sel == NS.getNSArraySelector(NSAPI::NSArr_objectAtIndex))
+  if (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))
+  if (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))
+  if (Sel == NS.getNSArraySelector(NSAPI::NSMutableArr_replaceObjectAtIndex))
     return rewriteToArraySubscriptSet(IFace, Msg, NS, commit);
 
-  if (II == NS.getNSClassId(NSAPI::ClassId_NSMutableDictionary) &&
-      Sel == NS.getNSDictionarySelector(NSAPI::NSMutableDict_setObjectForKey))
+  if (Sel == NS.getNSDictionarySelector(NSAPI::NSMutableDict_setObjectForKey))
     return rewriteToDictionarySubscriptSet(IFace, Msg, NS, commit);
 
   return false;
diff --git a/test/ARCMT/objcmt-subscripting-literals.m b/test/ARCMT/objcmt-subscripting-literals.m
index e9ed788275..ea33bbae55 100644
--- a/test/ARCMT/objcmt-subscripting-literals.m
+++ b/test/ARCMT/objcmt-subscripting-literals.m
@@ -169,3 +169,45 @@ void test1(NSString *str) {
   arr = [NSArray arrayWithObjects: globStr, str, nil];
   arr = [NSArray arrayWithObject:globStr];
 }
+
+@interface Custom : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface Custom (Extended)
+- (id)objectAtIndexedSubscript:(unsigned)idx;
+@end
+
+@interface MutableCustom : Custom
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface MutableCustom (Extended)
+- (void)setObject:(id)obj atIndexedSubscript:(unsigned)idx;
+@end
+
+@interface CustomUnavail : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface CustomUnavail (Extended)
+- (id)objectAtIndexedSubscript:(unsigned)idx __attribute__((unavailable));
+@end
+
+@interface MutableCustomUnavail : CustomUnavail
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface MutableCustomUnavail (Extended)
+- (void)setObject:(id)obj atIndexedSubscript:(unsigned)idx __attribute__((unavailable));
+@end
+
+void test2() {
+  MutableCustom *mutc;
+  id o = [mutc objectAtIndex:4];
+  [mutc replaceObjectAtIndex:2 withObject:@"val"];
+
+  MutableCustomUnavail *mutcunaval;
+  o = [mutcunaval objectAtIndex:4];
+  [mutcunaval replaceObjectAtIndex:2 withObject:@"val"];
+}
diff --git a/test/ARCMT/objcmt-subscripting-literals.m.result b/test/ARCMT/objcmt-subscripting-literals.m.result
index 4e49c5e344..8a9e6f0b92 100644
--- a/test/ARCMT/objcmt-subscripting-literals.m.result
+++ b/test/ARCMT/objcmt-subscripting-literals.m.result
@@ -169,3 +169,45 @@ void test1(NSString *str) {
   arr = @[(id)globStr, str];
   arr = @[(id)globStr];
 }
+
+@interface Custom : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface Custom (Extended)
+- (id)objectAtIndexedSubscript:(unsigned)idx;
+@end
+
+@interface MutableCustom : Custom
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface MutableCustom (Extended)
+- (void)setObject:(id)obj atIndexedSubscript:(unsigned)idx;
+@end
+
+@interface CustomUnavail : NSObject
+- (id)objectAtIndex:(unsigned long)index;
+@end
+
+@interface CustomUnavail (Extended)
+- (id)objectAtIndexedSubscript:(unsigned)idx __attribute__((unavailable));
+@end
+
+@interface MutableCustomUnavail : CustomUnavail
+- (void)replaceObjectAtIndex:(unsigned long)index withObject:(id)anObject;
+@end
+
+@interface MutableCustomUnavail (Extended)
+- (void)setObject:(id)obj atIndexedSubscript:(unsigned)idx __attribute__((unavailable));
+@end
+
+void test2() {
+  MutableCustom *mutc;
+  id o = mutc[4];
+  mutc[2] = @"val";
+
+  MutableCustomUnavail *mutcunaval;
+  o = [mutcunaval objectAtIndex:4];
+  [mutcunaval replaceObjectAtIndex:2 withObject:@"val"];
+}