Allow messaging expression as foreach's collection expression.

llvm-svn: 45793

clang/AST/StmtSerialization.cpp

@@ -921,16 +921,18 @@ ObjCEncodeExpr* ObjCEncodeExpr::CreateImpl(Deserializer& D) {
 
 void ObjCForCollectionStmt::EmitImpl(Serializer& S) const {
   S.Emit(ForLoc);
+  S.Emit(RParenLoc);
   S.BatchEmitOwnedPtrs(getElement(),getCollection(),getBody());
 }
 
 ObjCForCollectionStmt* ObjCForCollectionStmt::CreateImpl(Deserializer& D) {
   SourceLocation ForLoc = SourceLocation::ReadVal(D);
+  SourceLocation RParenLoc = SourceLocation::ReadVal(D);
   Stmt* Element;
   Expr* Collection;
   Stmt* Body;
   D.BatchReadOwnedPtrs(Element,Collection,Body);  
-  return new ObjCForCollectionStmt(Element,Collection,Body,ForLoc);
+  return new ObjCForCollectionStmt(Element,Collection,Body,ForLoc, RParenLoc);
 }
 
 void ObjCIvarRefExpr::EmitImpl(Serializer& S) const {

clang/Driver/RewriteTest.cpp

@@ -822,9 +822,7 @@ void RewriteTest::SynthCountByEnumWithState(std::string &buf) {
 ///
 Stmt *RewriteTest::RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S) {
   SourceLocation startLoc = S->getLocStart();
-  SourceLocation collectionLoc = S->getCollection()->getLocStart();
   const char *startBuf = SM->getCharacterData(startLoc);
-  const char *startCollectionBuf = SM->getCharacterData(collectionLoc);
   const char *elementName;
   std::string elementTypeAsString;
   std::string buf;
@@ -852,14 +850,26 @@ Stmt *RewriteTest::RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S) {
   buf += "id items[16];\n\t";
   // id l_collection = (id)
   buf += "id l_collection = (id)";
+  // Find start location of 'collection' the hard way!
+  const char *startCollectionBuf = startBuf;
+  startCollectionBuf += 3;  // skip 'for'
+  startCollectionBuf = strchr(startCollectionBuf, '(');
+  startCollectionBuf++; // skip '('
+  // find 'in' and skip it.
+  while (*startCollectionBuf != ' ' ||
+         *(startCollectionBuf+1) != 'i' || *(startCollectionBuf+2) != 'n' ||
+         (*(startCollectionBuf+3) != ' ' &&
+          *(startCollectionBuf+3) != '[' && *(startCollectionBuf+3) != '('))
+    startCollectionBuf++;
+  startCollectionBuf += 3;
+  
   // Replace: "for (type element in" with string constructed thus far. 
   Rewrite.ReplaceText(startLoc, startCollectionBuf - startBuf,
                       buf.c_str(), buf.size());
   // Replace ')' in for '(' type elem in collection ')' with ';'
-  SourceLocation endCollectionLoc = S->getCollection()->getLocEnd();
-  const char *endCollectionBuf = SM->getCharacterData(endCollectionLoc);
-  const char *lparenBuf = strchr(endCollectionBuf+1, ')');
-  SourceLocation lparenLoc = startLoc.getFileLocWithOffset(lparenBuf-startBuf);
+  SourceLocation rightParenLoc = S->getRParenLoc();
+  const char *rparenBuf = SM->getCharacterData(rightParenLoc);
+  SourceLocation lparenLoc = startLoc.getFileLocWithOffset(rparenBuf-startBuf);
   buf = ";\n\t";
   
   // unsigned long limit = [l_collection countByEnumeratingWithState:&enumState

clang/Sema/SemaStmt.cpp

@@ -531,7 +531,7 @@ Sema::ActOnForStmt(SourceLocation ForLoc, SourceLocation LParenLoc,
 }
 
 Action::StmtResult 
-Sema::ActOnObjCForCollectionStmt(SourceLocation ForColLoc, 
+Sema::ActOnObjCForCollectionStmt(SourceLocation ForLoc, 
                                  SourceLocation LParenLoc, 
                                  StmtTy *first, ExprTy *second,
                                  SourceLocation RParenLoc, StmtTy *body) {
@@ -553,16 +553,16 @@ Sema::ActOnObjCForCollectionStmt(SourceLocation ForColLoc,
   else
     FirstType = static_cast<Expr*>(first)->getType();
   if (!isObjCObjectPointerType(FirstType))
-      Diag(ForColLoc, diag::err_selector_element_type,
+      Diag(ForLoc, diag::err_selector_element_type,
            FirstType.getAsString(), First->getSourceRange());
   if (Second) {
     DefaultFunctionArrayConversion(Second);
     QualType SecondType = Second->getType();
     if (!isObjCObjectPointerType(SecondType))
-      Diag(ForColLoc, diag::err_collection_expr_type,
+      Diag(ForLoc, diag::err_collection_expr_type,
            SecondType.getAsString(), Second->getSourceRange());
   }
-  return new ObjCForCollectionStmt(First, Second, Body, ForColLoc);
+  return new ObjCForCollectionStmt(First, Second, Body, ForLoc, RParenLoc);
 }
 
 Action::StmtResult 

clang/include/clang/AST/Stmt.h

@@ -580,14 +580,16 @@ class ObjCForCollectionStmt : public Stmt {
   enum { ELEM, COLLECTION, BODY, END_EXPR };
   Stmt* SubExprs[END_EXPR]; // SubExprs[ELEM] is an expression or declstmt.
   SourceLocation ForLoc;
+  SourceLocation RParenLoc;
 public:
   ObjCForCollectionStmt(Stmt *Elem, Expr *Collect, Stmt *Body, 
-                        SourceLocation FCL) 
+                        SourceLocation FCL, SourceLocation RPL) 
   : Stmt(ObjCForCollectionStmtClass) {
     SubExprs[ELEM] = Elem;
     SubExprs[COLLECTION] = reinterpret_cast<Stmt*>(Collect);
     SubExprs[BODY] = Body;
     ForLoc = FCL;
+    RParenLoc = RPL;
   }
     
   Stmt *getElement() { return SubExprs[ELEM]; }
@@ -601,7 +603,9 @@ public:
     return reinterpret_cast<Expr*>(SubExprs[COLLECTION]);
   }
   const Stmt *getBody() const { return SubExprs[BODY]; }
-    
+  
+  SourceLocation getRParenLoc() const { return RParenLoc; }
+  
   virtual SourceRange getSourceRange() const { 
     return SourceRange(ForLoc, SubExprs[BODY]->getLocEnd()); 
   }

clang/test/Sema/rewrite-foreach-4.m

@@ -0,0 +1,32 @@
+// RUN: clang -rewrite-test %s
+
+@interface MyList
+- (id) allKeys;
+@end
+    
+@implementation MyList
+- (unsigned int)countByEnumeratingWithState:  (struct __objcFastEnumerationState *)state objects:  (id *)items count:(unsigned int)stackcount
+{
+        return 0;
+}
+- (id) allKeys {}
+@end
+
+@interface MyList (BasicTest)
+- (void)compilerTestAgainst;
+@end
+
+int LOOP();
+@implementation MyList (BasicTest)
+- (void)compilerTestAgainst {
+  MyList * el;
+        for (el in [el allKeys]) { LOOP(); 
+	  }
+
+        for (id el1 in[el allKeys]) { LOOP(); 
+	  }
+        for (el in([el allKeys])) { LOOP(); 
+	  }
+}
+@end
+