More support for ivars in class extension.

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

5 files changed, 85 insertions, 6 deletions

diff --git a/include/clang/AST/DeclObjC.h b/include/clang/AST/DeclObjC.h
index 686c1eb2f8..26656bf30a 100644
--- a/include/clang/AST/DeclObjC.h
+++ b/include/clang/AST/DeclObjC.h
@@ -527,6 +527,8 @@ public:
     CategoryList = category;
   }
 
+  ObjCCategoryDecl* getClassExtension() const;
+  
   /// isSuperClassOf - Return true if this class is the specified class or is a
   /// super class of the specified interface class.
   bool isSuperClassOf(const ObjCInterfaceDecl *I) const {
@@ -949,6 +951,20 @@ public:
 
   bool IsClassExtension() const { return getIdentifier() == 0; }
   
+  typedef specific_decl_iterator<ObjCIvarDecl> ivar_iterator;
+  ivar_iterator ivar_begin() const {
+    return ivar_iterator(decls_begin());
+  }
+  ivar_iterator ivar_end() const {
+    return ivar_iterator(decls_end());
+  }
+  unsigned ivar_size() const {
+    return std::distance(ivar_begin(), ivar_end());
+  }
+  bool ivar_empty() const {
+    return ivar_begin() == ivar_end();
+  }
+  
   SourceLocation getAtLoc() const { return AtLoc; }
   void setAtLoc(SourceLocation At) { AtLoc = At; }
 
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);
diff --git a/test/SemaObjC/ivar-in-class-extension.m b/test/SemaObjC/ivar-in-class-extension.m
new file mode 100644
index 0000000000..683a78feef
--- /dev/null
+++ b/test/SemaObjC/ivar-in-class-extension.m
@@ -0,0 +1,42 @@
+// RUN: %clang_cc1 -fsyntax-only -fobjc-nonfragile-abi2 -verify %s
+
+@interface SomeClass  @end
+
+int fn1(SomeClass *obj) {
+        obj->privateIvar = 1; // expected-error {{'SomeClass' does not have a member named 'privateIvar}}
+        return obj->publicIvar;     // expected-error {{'SomeClass' does not have a member named 'publicIvar'}}
+}
+
+@interface SomeClass () {
+// @private by default
+        int privateIvar;
+@public
+        int publicIvar;
+}
+@end
+
+int fn2(SomeClass *obj) {
+	obj->publicIvar = 1;
+        return obj->publicIvar    // ok
+             + obj->privateIvar;  // expected-error {{instance variable 'privateIvar' is private}}
+}
+
+@implementation SomeClass
+
+int fn3(SomeClass *obj) {
+	obj->privateIvar = 2;
+        return obj->publicIvar    // ok
+             + obj->privateIvar;  // ok
+    }
+@end
+
+@interface SomeClass (Category)
+    {	// expected-error {{ivar may be placed in a class extension}}
+        int categoryIvar;
+    }
+@end
+
+@interface SomeClass (Category1)
+    {	// expected-error {{ivar may be placed in a class extension}}
+    }
+@end