3 files changed, 93 insertions, 3 deletions

diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index 05627b47eb..304799cf91 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -824,6 +824,36 @@ unsigned ASTContext::CountSynthesizedIvars(const ObjCInterfaceDecl *OI)
   return count;
 }
 
+/// \brief Get the implementation of ObjCInterfaceDecl,or NULL if none exists.
+ObjCImplementationDecl *ASTContext::getObjCImplementation(ObjCInterfaceDecl *D) {
+  llvm::DenseMap<ObjCContainerDecl*, ObjCImplDecl*>::iterator
+    I = ObjCImpls.find(D);
+  if (I != ObjCImpls.end())
+    return cast<ObjCImplementationDecl>(I->second);
+  return 0;
+}
+/// \brief Get the implementation of ObjCCategoryDecl, or NULL if none exists.
+ObjCCategoryImplDecl *ASTContext::getObjCImplementation(ObjCCategoryDecl *D) {
+  llvm::DenseMap<ObjCContainerDecl*, ObjCImplDecl*>::iterator
+    I = ObjCImpls.find(D);
+  if (I != ObjCImpls.end())
+    return cast<ObjCCategoryImplDecl>(I->second);
+  return 0;
+}
+
+/// \brief Set the implementation of ObjCInterfaceDecl.
+void ASTContext::setObjCImplementation(ObjCInterfaceDecl *IFaceD,
+                           ObjCImplementationDecl *ImplD) {
+  assert(IFaceD && ImplD && "Passed null params");
+  ObjCImpls[IFaceD] = ImplD;
+}
+/// \brief Set the implementation of ObjCCategoryDecl.
+void ASTContext::setObjCImplementation(ObjCCategoryDecl *CatD,
+                           ObjCCategoryImplDecl *ImplD) {
+  assert(CatD && ImplD && "Passed null params");
+  ObjCImpls[CatD] = ImplD;
+}
+
 /// getInterfaceLayoutImpl - Get or compute information about the
 /// layout of the given interface.
 ///
diff --git a/lib/AST/DeclObjC.cpp b/lib/AST/DeclObjC.cpp
index 613dc58a57..4eee28b468 100644
--- a/lib/AST/DeclObjC.cpp
+++ b/lib/AST/DeclObjC.cpp
@@ -363,6 +363,15 @@ void ObjCInterfaceDecl::Destroy(ASTContext &C) {
   Decl::Destroy(C);
 }
 
+ObjCImplementationDecl *ObjCInterfaceDecl::getImplementation() const {
+  return getASTContext().getObjCImplementation(
+                                          const_cast<ObjCInterfaceDecl*>(this));
+}
+
+void ObjCInterfaceDecl::setImplementation(ObjCImplementationDecl *ImplD) {
+  getASTContext().setObjCImplementation(this, ImplD);
+}
+
 
 /// FindCategoryDeclaration - Finds category declaration in the list of
 /// categories for this class and returns it. Name of the category is passed
@@ -529,6 +538,16 @@ ObjCCategoryDecl *ObjCCategoryDecl::Create(ASTContext &C, DeclContext *DC,
   return new (C) ObjCCategoryDecl(DC, L, Id);
 }
 
+ObjCCategoryImplDecl *ObjCCategoryDecl::getImplementation() const {
+  return getASTContext().getObjCImplementation(
+                                           const_cast<ObjCCategoryDecl*>(this));
+}
+
+void ObjCCategoryDecl::setImplementation(ObjCCategoryImplDecl *ImplD) {
+  getASTContext().setObjCImplementation(this, ImplD);
+}
+
+
 //===----------------------------------------------------------------------===//
 // ObjCCategoryImplDecl
 //===----------------------------------------------------------------------===//
@@ -547,6 +566,23 @@ void ObjCImplDecl::addPropertyImplementation(ObjCPropertyImplDecl *property) {
   addDecl(property);
 }
 
+void ObjCImplDecl::setClassInterface(ObjCInterfaceDecl *IFace) {
+  ASTContext &Ctx = getASTContext();
+
+  if (ObjCImplementationDecl *ImplD
+        = dyn_cast_or_null<ObjCImplementationDecl>(this))
+    if (IFace)
+      Ctx.setObjCImplementation(IFace, ImplD);
+
+  else if (ObjCCategoryImplDecl *ImplD =
+             dyn_cast_or_null<ObjCCategoryImplDecl>(this)) {
+    if (ObjCCategoryDecl *CD = IFace->FindCategoryDeclaration(getIdentifier()))
+      Ctx.setObjCImplementation(CD, ImplD);
+  }
+
+  ClassInterface = IFace;
+}
+
 /// FindPropertyImplIvarDecl - This method lookup the ivar in the list of
 /// properties implemented in this category @implementation block and returns
 /// the implemented property that uses it.
diff --git a/lib/Sema/SemaDeclObjC.cpp b/lib/Sema/SemaDeclObjC.cpp
index 8dc14eebe4..440b0fba6a 100644
--- a/lib/Sema/SemaDeclObjC.cpp
+++ b/lib/Sema/SemaDeclObjC.cpp
@@ -605,6 +605,19 @@ Sema::DeclPtrTy Sema::ActOnStartCategoryImplementation(
                       IdentifierInfo *ClassName, SourceLocation ClassLoc,
                       IdentifierInfo *CatName, SourceLocation CatLoc) {
   ObjCInterfaceDecl *IDecl = getObjCInterfaceDecl(ClassName);
+  ObjCCategoryDecl *CatIDecl = 0;
+  if (IDecl) {
+    CatIDecl = IDecl->FindCategoryDeclaration(CatName);
+    if (!CatIDecl) {
+      // Category @implementation with no corresponding @interface.
+      // Create and install one.
+      CatIDecl = ObjCCategoryDecl::Create(Context, CurContext, SourceLocation(),
+                                          CatName);
+      CatIDecl->setClassInterface(IDecl);
+      CatIDecl->insertNextClassCategory();
+    }
+  }
+
   ObjCCategoryImplDecl *CDecl = 
     ObjCCategoryImplDecl::Create(Context, CurContext, AtCatImplLoc, CatName,
                                  IDecl);
@@ -615,8 +628,17 @@ Sema::DeclPtrTy Sema::ActOnStartCategoryImplementation(
   // FIXME: PushOnScopeChains?
   CurContext->addDecl(CDecl);
 
-  /// TODO: Check that CatName, category name, is not used in another
-  // implementation.
+  /// Check that CatName, category name, is not used in another implementation.
+  if (CatIDecl) {
+    if (CatIDecl->getImplementation()) {
+      Diag(ClassLoc, diag::err_dup_implementation_category) << ClassName
+        << CatName;
+      Diag(CatIDecl->getImplementation()->getLocation(),
+           diag::note_previous_definition);
+    } else
+      CatIDecl->setImplementation(CDecl);
+  }
+
   ObjCCategoryImpls.push_back(CDecl);
   
   CheckObjCDeclScope(CDecl);
@@ -697,8 +719,10 @@ Sema::DeclPtrTy Sema::ActOnStartClassImplementation(
   if (LookupObjCImplementation(ClassName))
     // FIXME: Don't leak everything!
     Diag(ClassLoc, diag::err_dup_implementation_class) << ClassName;
-  else // add it to the list.
+  else { // add it to the list.
+    IDecl->setImplementation(IMPDecl);
     PushOnScopeChains(IMPDecl, TUScope);
+  }
   return DeclPtrTy::make(IMPDecl);
 }