1 files changed, 65 insertions, 26 deletions

diff --git a/lib/Rewrite/RewriteObjC.cpp b/lib/Rewrite/RewriteObjC.cpp
index ba518069eb..645883606a 100644
--- a/lib/Rewrite/RewriteObjC.cpp
+++ b/lib/Rewrite/RewriteObjC.cpp
@@ -247,6 +247,10 @@ namespace {
     // Syntactic Rewriting.
     void RewriteInclude();
     void RewriteForwardClassDecl(DeclGroupRef D);
+    void RewriteForwardClassDecl(const llvm::SmallVector<Decl*, 8> &DG);
+    void RewriteForwardClassEpilogue(ObjCClassDecl *ClassDecl, 
+                                     const std::string &typedefString);
+      
     void RewritePropertyImplDecl(ObjCPropertyImplDecl *PID,
                                  ObjCImplementationDecl *IMD,
                                  ObjCCategoryImplDecl *CID);
@@ -722,8 +726,23 @@ void RewriteObjC::HandleTopLevelSingleDecl(Decl *D) {
     // Recurse into linkage specifications
     for (DeclContext::decl_iterator DI = LSD->decls_begin(),
                                  DIEnd = LSD->decls_end();
-         DI != DIEnd; ++DI)
+         DI != DIEnd; ) {
+      if (isa<ObjCClassDecl>((*DI))) {
+        SmallVector<Decl *, 8> DG;
+        Decl *D = (*DI);
+        SourceLocation Loc = D->getLocation();
+        while (DI != DIEnd &&
+               isa<ObjCClassDecl>(D) && D->getLocation() == Loc) {
+          DG.push_back(D);
+          ++DI;
+          D = (*DI);
+        }
+        RewriteForwardClassDecl(DG);
+        continue;
+      }
       HandleTopLevelSingleDecl(*DI);
+      ++DI;
+    }
   }
   // If we have a decl in the main file, see if we should rewrite it.
   if (SM->isFromMainFile(Loc))
@@ -891,41 +910,61 @@ void RewriteObjC::RewritePropertyImplDecl(ObjCPropertyImplDecl *PID,
   InsertText(onePastSemiLoc, Setr);
 }
 
+static void RewriteOneForwardClassDecl(ObjCInterfaceDecl *ForwardDecl,
+                                       std::string &typedefString) {
+  typedefString += "#ifndef _REWRITER_typedef_";
+  typedefString += ForwardDecl->getNameAsString();
+  typedefString += "\n";
+  typedefString += "#define _REWRITER_typedef_";
+  typedefString += ForwardDecl->getNameAsString();
+  typedefString += "\n";
+  typedefString += "typedef struct objc_object ";
+  typedefString += ForwardDecl->getNameAsString();
+  typedefString += ";\n#endif\n";
+}
+
+void RewriteObjC::RewriteForwardClassEpilogue(ObjCClassDecl *ClassDecl,
+                                              const std::string &typedefString) {
+    SourceLocation startLoc = ClassDecl->getLocation();
+    const char *startBuf = SM->getCharacterData(startLoc);
+    const char *semiPtr = strchr(startBuf, ';'); 
+    // Replace the @class with typedefs corresponding to the classes.
+    ReplaceText(startLoc, semiPtr-startBuf+1, typedefString);  
+}
+
 void RewriteObjC::RewriteForwardClassDecl(DeclGroupRef D) {
-  SourceLocation startLoc;
   std::string typedefString;
-  const char *startBuf = 0;
-  const char *semiPtr = 0;
   for (DeclGroupRef::iterator I = D.begin(), E = D.end(); I != E; ++I) {
     ObjCClassDecl *ClassDecl = cast<ObjCClassDecl>(*I);
     ObjCInterfaceDecl *ForwardDecl = ClassDecl->getForwardInterfaceDecl();
     if (I == D.begin()) {
-      // Get the start location and compute the semi location.
-      startLoc = ClassDecl->getLocation();
-      startBuf = SM->getCharacterData(startLoc);
-      semiPtr = strchr(startBuf, ';');
+      // Translate to typedef's that forward reference structs with the same name
+      // as the class. As a convenience, we include the original declaration
+      // as a comment.
       typedefString += "// @class ";
       typedefString += ForwardDecl->getNameAsString();
       typedefString += ";\n";
     }
-    // Translate to typedef's that forward reference structs with the same name
-    // as the class. As a convenience, we include the original declaration
-    // as a comment.
-  
-  
-    typedefString += "#ifndef _REWRITER_typedef_";
-    typedefString += ForwardDecl->getNameAsString();
-    typedefString += "\n";
-    typedefString += "#define _REWRITER_typedef_";
-    typedefString += ForwardDecl->getNameAsString();
-    typedefString += "\n";
-    typedefString += "typedef struct objc_object ";
-    typedefString += ForwardDecl->getNameAsString();
-    typedefString += ";\n#endif\n";
-  }
-
-  // Replace the @class with typedefs corresponding to the classes.
-  ReplaceText(startLoc, semiPtr-startBuf+1, typedefString);
+    RewriteOneForwardClassDecl(ForwardDecl, typedefString);
+  }
+  DeclGroupRef::iterator I = D.begin();
+  RewriteForwardClassEpilogue(cast<ObjCClassDecl>(*I), typedefString);
+}
+
+void RewriteObjC::RewriteForwardClassDecl(
+                                const llvm::SmallVector<Decl*, 8> &D) {
+  std::string typedefString;
+  for (unsigned i = 0; i < D.size(); i++) {
+    ObjCClassDecl *ClassDecl = cast<ObjCClassDecl>(D[i]);
+    ObjCInterfaceDecl *ForwardDecl = ClassDecl->getForwardInterfaceDecl();
+    if (i == 0) {
+      typedefString += "// @class ";
+      typedefString += ForwardDecl->getNameAsString();
+      typedefString += ";\n";
+    }
+    RewriteOneForwardClassDecl(ForwardDecl, typedefString);
+  }
+  RewriteForwardClassEpilogue(cast<ObjCClassDecl>(D[0]), typedefString);
 }
 
 void RewriteObjC::RewriteMethodDeclaration(ObjCMethodDecl *Method) {