Support importing of ObjC categories from modules.

The initial incentive was to fix a crash when PCH chaining categories
to an interface, but the fix was done in the "modules way" that I hear
is popular with the kids these days.

Each module stores the local chain of categories and we combine them
when the interface is loaded. We also warn if non-dependent modules
introduce duplicate named categories.

llvm-svn: 138926

clang/include/clang/AST/ASTMutationListener.h

@@ -22,6 +22,8 @@ namespace clang {
   class ClassTemplateSpecializationDecl;
   class FunctionDecl;
   class FunctionTemplateDecl;
+  class ObjCCategoryDecl;
+  class ObjCInterfaceDecl;
 
 /// \brief An abstract interface that should be implemented by listeners
 /// that want to be notified when an AST entity gets modified after its
@@ -54,6 +56,10 @@ public:
 
   /// \brief A static data member was implicitly instantiated.
   virtual void StaticDataMemberInstantiated(const VarDecl *D) {}
+
+  /// \brief A new objc category class was added for an interface.
+  virtual void AddedObjCCategoryToInterface(const ObjCCategoryDecl *CatD,
+                                            const ObjCInterfaceDecl *IFD) {}
 };
 
 } // end namespace clang

clang/include/clang/Serialization/ASTBitCodes.h

@@ -64,6 +64,11 @@ namespace clang {
     /// \brief a Decl::Kind/DeclID pair.
     typedef std::pair<uint32_t, DeclID> KindDeclIDPair;
 
+    // FIXME: Turn these into classes so we can have some type safety when
+    // we go from local ID to global and vice-versa.
+    typedef DeclID LocalDeclID;
+    typedef DeclID GlobalDeclID;
+
     /// \brief An ID number that refers to a type in an AST file.
     ///
     /// The ID of a type is partitioned into two parts: the lower
@@ -402,7 +407,11 @@ namespace clang {
 
       /// \brief Record code for the source manager line table information,
       /// which stores information about #line directives.
-      SOURCE_MANAGER_LINE_TABLE = 48
+      SOURCE_MANAGER_LINE_TABLE = 48,
+
+      /// \brief Record code for ObjC categories in a module that are chained to
+      /// an interface.
+      OBJC_CHAINED_CATEGORIES
     };
 
     /// \brief Record types used within a source manager block.

clang/include/clang/Serialization/ASTReader.h

@@ -36,6 +36,7 @@
 #include "llvm/ADT/OwningPtr.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Bitcode/BitstreamReader.h"
 #include "llvm/Support/DataTypes.h"
 #include <deque>
@@ -315,6 +316,10 @@ private:
   /// most recent declarations in another AST file.
   FirstLatestDeclIDMap FirstLatestDeclIDs;
 
+  /// \brief Set of ObjC interfaces that have categories chained to them in
+  /// other modules.
+  llvm::DenseSet<serialization::GlobalDeclID> ObjCChainedCategoriesInterfaces;
+
   /// \brief Read the records that describe the contents of declcontexts.
   bool ReadDeclContextStorage(Module &M, 
                               llvm::BitstreamCursor &Cursor,
@@ -681,6 +686,8 @@ private:
   Decl *ReadDeclRecord(serialization::DeclID ID);
   RecordLocation DeclCursorForID(serialization::DeclID ID);
   void loadDeclUpdateRecords(serialization::DeclID ID, Decl *D);
+  void loadObjCChainedCategories(serialization::GlobalDeclID ID,
+                                 ObjCInterfaceDecl *D);
   
   RecordLocation getLocalBitOffset(uint64_t GlobalOffset);
   uint64_t getGlobalBitOffset(Module &M, uint32_t LocalOffset);
@@ -899,6 +906,10 @@ public:
   /// \brief Map from a local declaration ID within a given module to a 
   /// global declaration ID.
   serialization::DeclID getGlobalDeclID(Module &F, unsigned LocalID) const;
+
+  /// \brief Returns true if global DeclID \arg ID originated from module
+  /// \arg M.
+  bool isDeclIDFromModule(serialization::GlobalDeclID ID, Module &M) const;
   
   /// \brief Resolve a declaration ID into a declaration, potentially
   /// building a new declaration.

clang/include/clang/Serialization/ASTWriter.h

@@ -258,6 +258,19 @@ private:
   /// \brief Decls that will be replaced in the current dependent AST file.
   DeclsToRewriteTy DeclsToRewrite;
 
+  struct ChainedObjCCategoriesData {
+    /// \brief The interface in the imported module.
+    const ObjCInterfaceDecl *Interface;
+    /// \brief ID of the interface.
+    serialization::DeclID InterfaceID;
+    /// \brief ID of the locally tail category ID that got chained to the
+    /// imported interface.
+    serialization::DeclID TailCatID;
+  };
+  /// \brief ObjC categories that got chained to an interface imported from
+  /// another module.
+  SmallVector<ChainedObjCCategoriesData, 16> LocalChainedObjCCategories;
+
   /// \brief Decls that have been replaced in the current dependent AST file.
   ///
   /// When a decl changes fundamentally after being deserialized (this shouldn't
@@ -350,6 +363,7 @@ private:
   void WriteAttributes(const AttrVec &Attrs, RecordDataImpl &Record);
   void WriteDeclUpdatesBlocks();
   void WriteDeclReplacementsBlock();
+  void WriteChainedObjCCategories();
   void WriteDeclContextVisibleUpdate(const DeclContext *DC);
   void WriteFPPragmaOptions(const FPOptions &Opts);
   void WriteOpenCLExtensions(Sema &SemaRef);
@@ -620,6 +634,8 @@ public:
                                               const FunctionDecl *D);
   virtual void CompletedImplicitDefinition(const FunctionDecl *D);
   virtual void StaticDataMemberInstantiated(const VarDecl *D);
+  virtual void AddedObjCCategoryToInterface(const ObjCCategoryDecl *CatD,
+                                            const ObjCInterfaceDecl *IFD);
 };
 
 /// \brief AST and semantic-analysis consumer that generates a

clang/include/clang/Serialization/Module.h

@@ -273,6 +273,15 @@ public:
   /// \brief Information about the lexical and visible declarations
   /// for each DeclContext.
   DeclContextInfosMap DeclContextInfos;
+
+  typedef llvm::DenseMap<serialization::GlobalDeclID,
+             std::pair<serialization::LocalDeclID, serialization::LocalDeclID> >
+    ChainedObjCCategoriesMap;
+  /// \brief ObjC categories that got chained to an interface from another
+  /// module.
+  /// Key is the ID of the interface.
+  /// Value is a pair of linked category DeclIDs (head category, tail category).
+  ChainedObjCCategoriesMap ChainedObjCCategories;
   
   // === Types ===
   

clang/lib/AST/DeclObjC.cpp

@@ -14,6 +14,7 @@
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Stmt.h"
+#include "clang/AST/ASTMutationListener.h"
 #include "llvm/ADT/STLExtras.h"
 using namespace clang;
 
@@ -918,7 +919,8 @@ ObjCCategoryDecl *ObjCCategoryDecl::Create(ASTContext &C, DeclContext *DC,
     // Link this category into its class's category list.
     CatDecl->NextClassCategory = IDecl->getCategoryList();
     IDecl->setCategoryList(CatDecl);
-    IDecl->setChangedSinceDeserialization(true);
+    if (ASTMutationListener *L = C.getASTMutationListener())
+      L->AddedObjCCategoryToInterface(CatDecl, IDecl);
   }
 
   return CatDecl;

clang/lib/Serialization/ASTReader.cpp

@@ -2391,6 +2391,20 @@ ASTReader::ReadASTBlock(Module &F) {
           = std::make_pair(&F, Record[I+1]);
       break;
     }
+
+    case OBJC_CHAINED_CATEGORIES: {
+      if (Record.size() % 3 != 0) {
+        Error("invalid OBJC_CHAINED_CATEGORIES block in AST file");
+        return Failure;
+      }
+      for (unsigned I = 0, N = Record.size(); I != N; I += 3) {
+        serialization::GlobalDeclID GlobID = getGlobalDeclID(F, Record[I]);
+        F.ChainedObjCCategories[GlobID] = std::make_pair(Record[I+1],
+                                                         Record[I+2]);
+        ObjCChainedCategoriesInterfaces.insert(GlobID);
+      }
+      break;
+    }
         
     case CXX_BASE_SPECIFIER_OFFSETS: {
       if (F.LocalNumCXXBaseSpecifiers != 0) {
@@ -4075,6 +4089,13 @@ ASTReader::getGlobalDeclID(Module &F, unsigned LocalID) const {
   return LocalID + I->second;
 }
 
+bool ASTReader::isDeclIDFromModule(serialization::GlobalDeclID ID,
+                                   Module &M) const {
+  GlobalDeclMapType::const_iterator I = GlobalDeclMap.find(ID);
+  assert(I != GlobalDeclMap.end() && "Corrupted global declaration map");
+  return &M == I->second;
+}
+
 Decl *ASTReader::GetDecl(DeclID ID) {
   if (ID < NUM_PREDEF_DECL_IDS) {    
     switch ((PredefinedDeclIDs)ID) {

clang/lib/Serialization/ASTReaderDecl.cpp

@@ -14,6 +14,7 @@
 
 #include "ASTCommon.h"
 #include "clang/Serialization/ASTReader.h"
+#include "clang/Sema/SemaDiagnostic.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclVisitor.h"
@@ -104,6 +105,11 @@ namespace clang {
     void UpdateDecl(Decl *D, Module &Module,
                     const RecordData &Record);
 
+    static void setNextObjCCategory(ObjCCategoryDecl *Cat,
+                                    ObjCCategoryDecl *Next) {
+      Cat->NextClassCategory = Next;
+    }
+
     void VisitDecl(Decl *D);
     void VisitTranslationUnitDecl(TranslationUnitDecl *TU);
     void VisitNamedDecl(NamedDecl *ND);
@@ -1714,6 +1720,9 @@ Decl *ASTReader::ReadDeclRecord(DeclID ID) {
   // Load any relevant update records.
   loadDeclUpdateRecords(ID, D);
   
+  if (ObjCChainedCategoriesInterfaces.count(ID))
+    loadObjCChainedCategories(ID, cast<ObjCInterfaceDecl>(D));
+  
   // If we have deserialized a declaration that has a definition the
   // AST consumer might need to know about, queue it.
   // We don't pass it to the consumer immediately because we may be in recursive
@@ -1751,6 +1760,131 @@ void ASTReader::loadDeclUpdateRecords(serialization::DeclID ID, Decl *D) {
   }
 }
 
+namespace {
+  /// \brief Given an ObjC interface, goes through the modules and links to the
+  /// interface all the categories for it.
+  class ObjCChainedCategoriesVisitor {
+    ASTReader &Reader;
+    serialization::GlobalDeclID InterfaceID;
+    ObjCInterfaceDecl *Interface;
+    ObjCCategoryDecl *GlobHeadCat, *GlobTailCat;
+    llvm::DenseMap<DeclarationName, ObjCCategoryDecl *> NameCategoryMap;
+
+  public:
+    ObjCChainedCategoriesVisitor(ASTReader &Reader,
+                                 serialization::GlobalDeclID InterfaceID,
+                                 ObjCInterfaceDecl *Interface)
+      : Reader(Reader), InterfaceID(InterfaceID), Interface(Interface),
+        GlobHeadCat(0), GlobTailCat(0) { }
+
+    static bool visit(Module &M, void *UserData) {
+      return static_cast<ObjCChainedCategoriesVisitor *>(UserData)->visit(M);
+    }
+
+    bool visit(Module &M) {
+      if (Reader.isDeclIDFromModule(InterfaceID, M))
+        return true; // We reached the module where the interface originated
+                    // from. Stop traversing the imported modules.
+
+      Module::ChainedObjCCategoriesMap::iterator
+        I = M.ChainedObjCCategories.find(InterfaceID);
+      if (I == M.ChainedObjCCategories.end())
+        return false;
+
+      ObjCCategoryDecl *
+        HeadCat = Reader.GetLocalDeclAs<ObjCCategoryDecl>(M, I->second.first);
+      ObjCCategoryDecl *
+        TailCat = Reader.GetLocalDeclAs<ObjCCategoryDecl>(M, I->second.second);
+
+      addCategories(HeadCat, TailCat);
+      return false;
+    }
+
+    void addCategories(ObjCCategoryDecl *HeadCat,
+                       ObjCCategoryDecl *TailCat = 0) {
+      if (!HeadCat) {
+        assert(!TailCat);
+        return;
+      }
+
+      if (!TailCat) {
+        TailCat = HeadCat;
+        while (TailCat->getNextClassCategory())
+          TailCat = TailCat->getNextClassCategory();
+      }
+
+      if (!GlobHeadCat) {
+        GlobHeadCat = HeadCat;
+        GlobTailCat = TailCat;
+      } else {
+        ASTDeclReader::setNextObjCCategory(GlobTailCat, HeadCat);
+        GlobTailCat = TailCat;
+      }
+
+      llvm::DenseSet<DeclarationName> Checked;
+      for (ObjCCategoryDecl *Cat = HeadCat,
+                            *CatEnd = TailCat->getNextClassCategory();
+             Cat != CatEnd; Cat = Cat->getNextClassCategory()) {
+        if (Checked.count(Cat->getDeclName()))
+          continue;
+        Checked.insert(Cat->getDeclName());
+        checkForDuplicate(Cat);
+      }
+    }
+
+    /// \brief Warns for duplicate categories that come from different modules.
+    void checkForDuplicate(ObjCCategoryDecl *Cat) {
+      DeclarationName Name = Cat->getDeclName();
+      // Find the top category with the same name. We do not want to warn for
+      // duplicates along the established chain because there were already
+      // warnings for them when the module was created. We only want to warn for
+      // duplicates between non-dependent modules:
+      //
+      //   MT
+      //  /  \
+      // ML  MR
+      //
+      // We want to warn for duplicates between ML and MR,not between ML and MT.
+      //
+      // FIXME: We should not warn for duplicates in diamond:
+      //
+      //   MT
+      //  /  \
+      // ML  MR
+      //  \  /
+      //   MB
+      //
+      // If there are duplicates in ML/MR, there will be warning when creating
+      // MB *and* when importing MB. We should not warn when importing.
+      for (ObjCCategoryDecl *Next = Cat->getNextClassCategory(); Next;
+             Next = Next->getNextClassCategory()) {
+        if (Next->getDeclName() == Name)
+          Cat = Next;
+      }
+
+      ObjCCategoryDecl *&PrevCat = NameCategoryMap[Name];
+      if (!PrevCat)
+        PrevCat = Cat;
+
+      if (PrevCat != Cat) {
+        Reader.Diag(Cat->getLocation(), diag::warn_dup_category_def)
+          << Interface->getDeclName() << Name;
+        Reader.Diag(PrevCat->getLocation(), diag::note_previous_definition);
+      }
+    }
+
+    ObjCCategoryDecl *getHeadCategory() const { return GlobHeadCat; }
+  };
+}
+
+void ASTReader::loadObjCChainedCategories(serialization::GlobalDeclID ID,
+                                          ObjCInterfaceDecl *D) {
+  ObjCChainedCategoriesVisitor Visitor(*this, ID, D);
+  ModuleMgr.visit(ObjCChainedCategoriesVisitor::visit, &Visitor);
+  // Also add the categories that the interface already links to.
+  Visitor.addCategories(D->getCategoryList());
+  D->setCategoryList(Visitor.getHeadCategory());
+}
 
 void ASTDeclReader::UpdateDecl(Decl *D, Module &Module,
                                const RecordData &Record) {

clang/lib/Serialization/ASTWriter.cpp

@@ -3188,6 +3188,7 @@ void ASTWriter::WriteASTCore(Sema &SemaRef, MemorizeStatCalls *StatCalls,
 
   WriteDeclUpdatesBlocks();
   WriteDeclReplacementsBlock();
+  WriteChainedObjCCategories();
 
   // Some simple statistics
   Record.clear();
@@ -3236,6 +3237,26 @@ void ASTWriter::WriteDeclReplacementsBlock() {
   Stream.EmitRecord(DECL_REPLACEMENTS, Record);
 }
 
+void ASTWriter::WriteChainedObjCCategories() {
+  if (LocalChainedObjCCategories.empty())
+    return;
+
+  RecordData Record;
+  for (SmallVector<ChainedObjCCategoriesData, 16>::iterator
+         I = LocalChainedObjCCategories.begin(),
+         E = LocalChainedObjCCategories.end(); I != E; ++I) {
+    ChainedObjCCategoriesData &Data = *I;
+    serialization::DeclID
+        HeadCatID = getDeclID(Data.Interface->getCategoryList());
+    assert(HeadCatID != 0 && "Category not written ?");
+
+    Record.push_back(Data.InterfaceID);
+    Record.push_back(HeadCatID);
+    Record.push_back(Data.TailCatID);
+  }
+  Stream.EmitRecord(OBJC_CHAINED_CATEGORIES, Record);
+}
+
 void ASTWriter::AddSourceLocation(SourceLocation Loc, RecordDataImpl &Record) {
   Record.push_back(Loc.getRawEncoding());
 }
@@ -4037,4 +4058,17 @@ void ASTWriter::StaticDataMemberInstantiated(const VarDecl *D) {
       D->getMemberSpecializationInfo()->getPointOfInstantiation(), Record);
 }
 
+void ASTWriter::AddedObjCCategoryToInterface(const ObjCCategoryDecl *CatD,
+                                             const ObjCInterfaceDecl *IFD) {
+  if (IFD->getPCHLevel() == 0)
+    return; // Declaration not imported from PCH.
+  if (CatD->getNextClassCategory() &&
+      CatD->getNextClassCategory()->getPCHLevel() == 0)
+    return; // We already recorded that the tail of a category chain should be
+            // attached to an interface.
+
+  ChainedObjCCategoriesData Data =  { IFD, GetDeclRef(IFD), GetDeclRef(CatD) };
+  LocalChainedObjCCategories.push_back(Data);
+}
+
 ASTSerializationListener::~ASTSerializationListener() { }

clang/test/Modules/objc-categories.m

@@ -0,0 +1,85 @@
+// RUN: mkdir -p %t
+// RUN: %clang_cc1 -emit-module -o %t/diamond_top.pcm %s -D MODULE_TOP
+// RUN: %clang_cc1 -I %t -emit-module -o %t/diamond_left.pcm %s -D MODULE_LEFT
+// RUN: %clang_cc1 -I %t -emit-module -o %t/diamond_right.pcm %s -D MODULE_RIGHT
+// RUN: %clang_cc1 -I %t -emit-module -o %t/diamond_bottom.pcm %s -D MODULE_BOTTOM
+// RUN: %clang_cc1 -I %t %s -verify
+
+/*============================================================================*/
+#ifdef MODULE_TOP
+
+@interface Foo
+-(void)top;
+@end
+
+/*============================================================================*/
+#elif defined(MODULE_LEFT)
+
+__import_module__ diamond_top;
+
+@interface Foo(Left)
+-(void)left;
+@end
+
+@interface LeftFoo
+-(void)left;
+@end
+
+@interface Foo(Duplicate) // expected-note {{previous definition}}
+@end
+
+@interface Foo(Duplicate)
+@end
+
+/*============================================================================*/
+#elif defined(MODULE_RIGHT)
+
+__import_module__ diamond_top;
+
+@interface Foo(Right1)
+-(void)right1;
+@end
+
+@interface Foo(Right2)
+-(void)right2;
+@end
+
+@interface Foo(Duplicate) // expected-warning {{duplicate definition of category}}
+@end
+
+/*============================================================================*/
+#elif defined(MODULE_BOTTOM)
+
+__import_module__ diamond_left;
+__import_module__ diamond_right;
+
+@interface Foo(Bottom)
+-(void)bottom;
+@end
+
+@interface LeftFoo(Bottom)
+-(void)bottom;
+@end
+
+/*============================================================================*/
+#else
+
+__import_module__ diamond_bottom;
+
+@interface Foo(Source)
+-(void)source;
+@end
+
+void test(Foo *foo, LeftFoo *leftFoo) {
+  [foo source];
+  [foo bottom];
+  [foo left];
+  [foo right1];
+  [foo right2];
+  [foo top];
+
+  [leftFoo left];
+  [leftFoo bottom];
+}
+
+#endif

clang/test/PCH/chain-categories.m

@@ -0,0 +1,51 @@
+// Without PCH
+// RUN: %clang_cc1 -fsyntax-only -verify -include %s -include %s %s
+
+// With PCH
+// RUN: %clang_cc1 -fsyntax-only -verify %s -chain-include %s -chain-include %s
+
+#ifndef HEADER1
+#define HEADER1
+//===----------------------------------------------------------------------===//
+// Primary header
+
+@interface NSObject
+- (id)init;
+- (void)finalize;
+@end
+
+//===----------------------------------------------------------------------===//
+#elif !defined(HEADER2)
+#define HEADER2
+#if !defined(HEADER1)
+#error Header inclusion order messed up
+#endif
+
+//===----------------------------------------------------------------------===//
+// Dependent header
+
+@interface MyClass : NSObject
++(void)meth;
+@end
+
+@interface NSObject(ObjExt)
+-(void)extMeth;
+@end
+
+//===----------------------------------------------------------------------===//
+#else
+//===----------------------------------------------------------------------===//
+
+@implementation MyClass
++(void)meth {}
+-(void)finalize {
+  [super finalize];
+}
+@end
+
+void test(NSObject *o) {
+  [o extMeth];
+}
+
+//===----------------------------------------------------------------------===//
+#endif