modern objective-c translation: writing container

subscripting. // rdar://11203853

llvm-svn: 154441

2 files changed, 71 insertions, 11 deletions

diff --git a/clang/lib/Rewrite/RewriteModernObjC.cpp b/clang/lib/Rewrite/RewriteModernObjC.cpp
index 162f7a52894..4be7eb806f4 100644
--- a/clang/lib/Rewrite/RewriteModernObjC.cpp
+++ b/clang/lib/Rewrite/RewriteModernObjC.cpp
@@ -1325,7 +1325,8 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitSetter(PseudoObjectExpr *Pseud
 
   // Because the rewriter doesn't allow us to rewrite rewritten code,
   // we need to suppress rewriting the sub-statements.
-  Expr *Base, *RHS;
+  Expr *Base;
+  SmallVector<Expr*, 2> Args;
   {
     DisableReplaceStmtScope S(*this);
 
@@ -1336,11 +1337,15 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitSetter(PseudoObjectExpr *Pseud
       Base = cast<OpaqueValueExpr>(Base)->getSourceExpr();
       Base = cast<Expr>(RewriteFunctionBodyOrGlobalInitializer(Base));
     }
-
-    // Rebuild the RHS.
-    RHS = cast<BinaryOperator>(PseudoOp->getSyntacticForm())->getRHS();
-    RHS = cast<OpaqueValueExpr>(RHS)->getSourceExpr();
-    RHS = cast<Expr>(RewriteFunctionBodyOrGlobalInitializer(RHS));
+  
+    unsigned numArgs = OldMsg->getNumArgs();
+    for (unsigned i = 0; i < numArgs; i++) {
+      Expr *Arg = OldMsg->getArg(i);
+      if (isa<OpaqueValueExpr>(Arg))
+        Arg = cast<OpaqueValueExpr>(Arg)->getSourceExpr();
+      Arg = cast<Expr>(RewriteFunctionBodyOrGlobalInitializer(Arg));
+      Args.push_back(Arg);
+    }
   }
 
   // TODO: avoid this copy.
@@ -1357,7 +1362,7 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitSetter(PseudoObjectExpr *Pseud
                                      OldMsg->getSelector(),
                                      SelLocs,
                                      OldMsg->getMethodDecl(),
-                                     RHS,
+                                     Args,
                                      OldMsg->getRightLoc(),
                                      OldMsg->isImplicit());
     break;
@@ -1370,7 +1375,7 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitSetter(PseudoObjectExpr *Pseud
                                      OldMsg->getSelector(),
                                      SelLocs,
                                      OldMsg->getMethodDecl(),
-                                     RHS,
+                                     Args,
                                      OldMsg->getRightLoc(),
                                      OldMsg->isImplicit());
     break;
@@ -1386,7 +1391,7 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitSetter(PseudoObjectExpr *Pseud
                                      OldMsg->getSelector(),
                                      SelLocs,
                                      OldMsg->getMethodDecl(),
-                                     RHS,
+                                     Args,
                                      OldMsg->getRightLoc(),
                                      OldMsg->isImplicit());
     break;
@@ -1408,20 +1413,27 @@ Stmt *RewriteModernObjC::RewritePropertyOrImplicitGetter(PseudoObjectExpr *Pseud
   // Because the rewriter doesn't allow us to rewrite rewritten code,
   // we need to suppress rewriting the sub-statements.
   Expr *Base = 0;
+  SmallVector<Expr*, 1> Args;
   {
     DisableReplaceStmtScope S(*this);
-
     // Rebuild the base expression if we have one.
     if (OldMsg->getReceiverKind() == ObjCMessageExpr::Instance) {
       Base = OldMsg->getInstanceReceiver();
       Base = cast<OpaqueValueExpr>(Base)->getSourceExpr();
       Base = cast<Expr>(RewriteFunctionBodyOrGlobalInitializer(Base));
     }
+    unsigned numArgs = OldMsg->getNumArgs();
+    for (unsigned i = 0; i < numArgs; i++) {
+      Expr *Arg = OldMsg->getArg(i);
+      if (isa<OpaqueValueExpr>(Arg))
+        Arg = cast<OpaqueValueExpr>(Arg)->getSourceExpr();
+      Arg = cast<Expr>(RewriteFunctionBodyOrGlobalInitializer(Arg));
+      Args.push_back(Arg);
+    }
   }
 
   // Intentionally empty.
   SmallVector<SourceLocation, 1> SelLocs;
-  SmallVector<Expr*, 1> Args;
 
   ObjCMessageExpr *NewMsg = 0;
   switch (OldMsg->getReceiverKind()) {
diff --git a/clang/test/Rewriter/objc-modern-container-subscript.mm b/clang/test/Rewriter/objc-modern-container-subscript.mm
new file mode 100644
index 00000000000..d6bb9c2eb31
--- /dev/null
+++ b/clang/test/Rewriter/objc-modern-container-subscript.mm
@@ -0,0 +1,48 @@
+// RUN: %clang_cc1 -x objective-c++ -fblocks -fms-extensions -rewrite-objc %s -o %t-rw.cpp
+// RUN: %clang_cc1 -fsyntax-only -fblocks -Wno-address-of-temporary -D"Class=void*" -D"id=void*" -D"SEL=void*" -D"__declspec(X)=" %t-rw.cpp
+// rdar://11203853
+
+void *sel_registerName(const char *);
+
+typedef unsigned int size_t;
+@protocol P @end
+
+@interface NSMutableArray
+#if __has_feature(objc_subscripting)
+- (id)objectAtIndexedSubscript:(size_t)index;
+- (void)setObject:(id)object atIndexedSubscript:(size_t)index;
+#endif
+@end
+
+#if __has_feature(objc_subscripting)
+@interface XNSMutableArray
+- (id)objectAtIndexedSubscript:(size_t)index;
+- (void)setObject:(id)object atIndexedSubscript:(size_t)index;
+#endif
+@end
+
+@interface NSMutableDictionary
+- (id)objectForKeyedSubscript:(id)key;
+- (void)setObject:(id)object forKeyedSubscript:(id)key;
+@end
+
+@class NSString;
+
+int main() {
+  NSMutableArray<P> * array;
+  id oldObject = array[10];
+
+  array[10] = oldObject;
+
+  id unknown_array;
+  oldObject = unknown_array[1];
+
+  unknown_array[1] = oldObject;
+
+  NSMutableDictionary *dictionary;
+  NSString *key;
+  id newObject;
+  oldObject = dictionary[key];
+  dictionary[key] = newObject;  // replace oldObject with newObject
+}
+