1 files changed, 14 insertions, 7 deletions
diff --git a/clang/lib/Rewrite/RewriteModernObjC.cpp b/clang/lib/Rewrite/RewriteModernObjC.cpp
index abdb09b2790..4ba55403cb8 100644
--- a/clang/lib/Rewrite/RewriteModernObjC.cpp
+++ b/clang/lib/Rewrite/RewriteModernObjC.cpp
@@ -934,14 +934,17 @@ void RewriteModernObjC::RewriteCategoryDecl(ObjCCategoryDecl *CatDecl) {
 
   // FIXME: handle category headers that are declared across multiple lines.
   ReplaceText(LocStart, 0, "// ");
-
+  if (CatDecl->getIvarLBraceLoc().isValid())
+    InsertText(CatDecl->getIvarLBraceLoc(), "// ");
   for (ObjCCategoryDecl::ivar_iterator
        I = CatDecl->ivar_begin(), E = CatDecl->ivar_end(); I != E; ++I) {
     ObjCIvarDecl *Ivar = (*I);
     SourceLocation LocStart = Ivar->getLocStart();
     ReplaceText(LocStart, 0, "// ");
   } 
-
+  if (CatDecl->getIvarRBraceLoc().isValid())
+    InsertText(CatDecl->getIvarRBraceLoc(), "// ");
+  
   for (ObjCCategoryDecl::prop_iterator I = CatDecl->prop_begin(),
        E = CatDecl->prop_end(); I != E; ++I)
     RewriteProperty(*I);
@@ -1153,12 +1156,16 @@ void RewriteModernObjC::RewriteImplementationDecl(Decl *OID) {
 
   if (IMD) {
     InsertText(IMD->getLocStart(), "// ");
-      for (ObjCImplementationDecl::ivar_iterator
-           I = IMD->ivar_begin(), E = IMD->ivar_end(); I != E; ++I) {
-        ObjCIvarDecl *Ivar = (*I);
-        SourceLocation LocStart = Ivar->getLocStart();
-        ReplaceText(LocStart, 0, "// ");
+    if (IMD->getIvarLBraceLoc().isValid())
+      InsertText(IMD->getIvarLBraceLoc(), "// ");
+    for (ObjCImplementationDecl::ivar_iterator
+         I = IMD->ivar_begin(), E = IMD->ivar_end(); I != E; ++I) {
+      ObjCIvarDecl *Ivar = (*I);
+      SourceLocation LocStart = Ivar->getLocStart();
+      ReplaceText(LocStart, 0, "// ");
     }
+    if (IMD->getIvarRBraceLoc().isValid())
+      InsertText(IMD->getIvarRBraceLoc(), "// ");
   }
   else
     InsertText(CID->getLocStart(), "// ");