More support for ivars in class extension.

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

3 files changed, 27 insertions, 6 deletions

diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index e8b84c65dd..c8caeb62b3 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -900,6 +900,14 @@ void ASTContext::CollectProtocolSynthesizedIvars(const ObjCProtocolDecl *PD,
 ///
 void ASTContext::CollectNonClassIvars(const ObjCInterfaceDecl *OI,
                                 llvm::SmallVectorImpl<ObjCIvarDecl*> &Ivars) {
+  // Find ivars declared in class extension.
+  if (const ObjCCategoryDecl *CDecl = OI->getClassExtension()) {
+    for (ObjCCategoryDecl::ivar_iterator I = CDecl->ivar_begin(),
+         E = CDecl->ivar_end(); I != E; ++I) {
+      Ivars.push_back(*I);
+    }
+  }
+  
   for (ObjCInterfaceDecl::prop_iterator I = OI->prop_begin(),
        E = OI->prop_end(); I != E; ++I) {
     if (ObjCIvarDecl *Ivar = (*I)->getPropertyIvarDecl())
diff --git a/lib/AST/DeclObjC.cpp b/lib/AST/DeclObjC.cpp
index 131e098d04..8decafa35e 100644
--- a/lib/AST/DeclObjC.cpp
+++ b/lib/AST/DeclObjC.cpp
@@ -202,6 +202,17 @@ void ObjCInterfaceDecl::mergeClassExtensionProtocolList(
   setProtocolList(ProtocolRefs.data(), NumProtoRefs, ProtocolLocs.data(), C);
 }
 
+/// getClassExtension - Find class extension of the given class.
+// FIXME. can speed it up, if need be.
+ObjCCategoryDecl* ObjCInterfaceDecl::getClassExtension() const {
+  const ObjCInterfaceDecl* ClassDecl = this;
+  for (ObjCCategoryDecl *CDecl = ClassDecl->getCategoryList(); CDecl;
+       CDecl = CDecl->getNextClassCategory())
+    if (CDecl->IsClassExtension())
+      return CDecl;
+  return 0;
+}
+
 ObjCIvarDecl *ObjCInterfaceDecl::lookupInstanceVariable(IdentifierInfo *ID,
                                               ObjCInterfaceDecl *&clsDeclared) {
   ObjCInterfaceDecl* ClassDecl = this;
@@ -210,6 +221,12 @@ ObjCIvarDecl *ObjCInterfaceDecl::lookupInstanceVariable(IdentifierInfo *ID,
       clsDeclared = ClassDecl;
       return I;
     }
+    if (const ObjCCategoryDecl *CDecl = ClassDecl->getClassExtension())
+      if (ObjCIvarDecl *I = CDecl->getIvarDecl(ID)) {
+        clsDeclared = ClassDecl;
+        return I;
+      }
+      
     ClassDecl = ClassDecl->getSuperClass();
   }
   return NULL;
diff --git a/lib/Sema/SemaDeclObjC.cpp b/lib/Sema/SemaDeclObjC.cpp
index 5f69dc4462..e40e8fedc0 100644
--- a/lib/Sema/SemaDeclObjC.cpp
+++ b/lib/Sema/SemaDeclObjC.cpp
@@ -599,13 +599,9 @@ ActOnStartCategoryInterface(SourceLocation AtInterfaceLoc,
                             SourceLocation EndProtoLoc) {
   ObjCCategoryDecl *CDecl = 0;
   ObjCInterfaceDecl *IDecl = getObjCInterfaceDecl(ClassName, ClassLoc);
-  if (!CategoryName) {
+  if (!CategoryName)
     // Class extensions require a special treatment. Use an existing one.
-    for (CDecl = IDecl->getCategoryList(); CDecl;
-         CDecl = CDecl->getNextClassCategory())
-      if (CDecl->IsClassExtension()) 
-        break;  
-  }
+    CDecl = IDecl->getClassExtension();
   if (!CDecl) {
     CDecl = ObjCCategoryDecl::Create(Context, CurContext, AtInterfaceLoc, ClassLoc,
                                      CategoryLoc, CategoryName);