Use on-disk hash table to serialize and deserialize availability domain

information

rdar://147122406

Author: ahatanaka

clang/include/clang/AST/ExternalASTSource.h

@@ -142,6 +142,11 @@ class ExternalASTSource : public RefCountedBase<ExternalASTSource> {
   /// Update an out-of-date identifier.
   virtual void updateOutOfDateIdentifier(const IdentifierInfo &II) {}
 
+  // Retrieve the variable declaration that has the information for the domain.
+  virtual Decl *getAvailabilityDomainDecl(StringRef DomainName) {
+    return nullptr;
+  }
+
   /// Find all declarations with the given name in the given context,
   /// and add them to the context by calling SetExternalVisibleDeclsForName
   /// or SetNoExternalVisibleDeclsForName.

clang/include/clang/Serialization/ASTBitCodes.h

@@ -738,6 +738,9 @@ enum ASTRecordTypes {
   /// Record code for Sema's vector of functions/blocks with effects to
   /// be verified.
   DECLS_WITH_EFFECTS_TO_VERIFY = 72,
+
+  /// Record code for availability domain table.
+  AVAILABILITY_DOMAIN_TABLE = 73,
 };
 
 /// Record types used within a source manager block.

clang/include/clang/Serialization/ASTReader.h

@@ -1445,6 +1445,7 @@ class ASTReader
                                ASTReaderListener *Listener,
                                bool ValidateDiagnosticOptions);
 
+  Decl *getAvailabilityDomainDecl(StringRef DomainName) override;
   llvm::Error ReadASTBlock(ModuleFile &F, unsigned ClientLoadCapabilities);
   llvm::Error ReadExtensionBlock(ModuleFile &F);
   void ReadModuleOffsetMap(ModuleFile &F) const;

clang/include/clang/Serialization/ASTWriter.h

@@ -569,6 +569,7 @@ class ASTWriter : public ASTDeserializationListener,
   uint64_t WriteDeclContextLexicalBlock(ASTContext &Context,
                                         const DeclContext *DC);
   uint64_t WriteDeclContextVisibleBlock(ASTContext &Context, DeclContext *DC);
+  void writeAvailabilityDomainDecls(ASTContext &Context);
   void WriteTypeDeclOffsets();
   void WriteFileDeclIDsMap();
   void WriteComments(ASTContext &Context);

clang/include/clang/Serialization/ModuleFile.h

@@ -405,6 +405,10 @@ class ModuleFile {
   /// the header files.
   void *HeaderFileInfoTable = nullptr;
 
+  /// The on-disk hash table that contains information about availability
+  /// domains.
+  void *AvailabilityDomainTable = nullptr;
+
   // === Submodule information ===
 
   /// The number of submodules in this module.

clang/lib/AST/ASTContext.cpp

@@ -908,9 +908,15 @@ ASTContext::getFeatureAvailInfo(StringRef FeatureName) const {
   if (Kind != FeatureAvailKind::None)
     return {Kind};
 
+  Decl *D = nullptr;
   auto I = AvailabilityDomainMap.find(FeatureName);
   if (I != AvailabilityDomainMap.end())
-    return getFeatureAvailInfo(I->second).second;
+    D = I->second;
+  else if (auto *Source = getExternalSource())
+    D = Source->getAvailabilityDomainDecl(FeatureName);
+
+  if (D)
+    return getFeatureAvailInfo(D).second;
 
   return {};
 }
@@ -13214,10 +13220,6 @@ bool ASTContext::DeclMustBeEmitted(const Decl *D) {
   const auto *VD = cast<VarDecl>(D);
   assert(VD->isFileVarDecl() && "Expected file scoped var");
 
-  // Variables annotated with availability_domain need to be deserialized.
-  if (VD->hasAttr<AvailabilityDomainAttr>())
-    return true;
-
   // If the decl is marked as `declare target to`, it should be emitted for the
   // host and for the device.
   if (LangOpts.OpenMP &&

clang/lib/Serialization/ASTReader.cpp

@@ -3279,6 +3279,74 @@ ASTReader::ReadControlBlock(ModuleFile &F,
   }
 }
 
+AvailabilityDomainsTableReaderTrait::hash_value_type
+AvailabilityDomainsTableReaderTrait::ComputeHash(internal_key_type Key) {
+  return llvm::djbHash(Key);
+}
+
+std::pair<unsigned, unsigned>
+AvailabilityDomainsTableReaderTrait::ReadKeyDataLength(const uint8_t *&Data) {
+  unsigned KeyLength =
+      llvm::support::endian::readNext<uint16_t, llvm::endianness::little>(Data);
+  unsigned DataLength =
+      llvm::support::endian::readNext<uint16_t, llvm::endianness::little>(Data);
+  return {KeyLength, DataLength};
+}
+
+AvailabilityDomainsTableReaderTrait::internal_key_type
+AvailabilityDomainsTableReaderTrait::ReadKey(const uint8_t *Data,
+                                             unsigned Length) {
+  return llvm::StringRef(reinterpret_cast<const char *>(Data), Length);
+}
+
+AvailabilityDomainsTableReaderTrait::data_type
+AvailabilityDomainsTableReaderTrait::ReadData(internal_key_type Key,
+                                              const uint8_t *Data,
+                                              unsigned Length) {
+  return llvm::support::endian::readNext<data_type, llvm::endianness::little>(
+      Data);
+}
+
+class AvailabilityDomainVisitor {
+  std::optional<Decl *> D;
+  StringRef DomainName;
+  ASTReader &Reader;
+
+public:
+  explicit AvailabilityDomainVisitor(StringRef DomainName, ASTReader &Reader)
+      : DomainName(DomainName), Reader(Reader) {}
+
+  bool operator()(ModuleFile &M) {
+    AvailabilityDomainLookupTable *Table =
+        static_cast<AvailabilityDomainLookupTable *>(M.AvailabilityDomainTable);
+
+    if (!Table)
+      return false;
+
+    // Look in the on-disk hash table for an entry for this domain name.
+    AvailabilityDomainLookupTable::iterator Pos = Table->find(DomainName);
+    if (Pos == Table->end())
+      return false;
+
+    LocalDeclID LocalID = LocalDeclID::get(Reader, M, *Pos);
+    GlobalDeclID ID = Reader.getGlobalDeclID(M, LocalID);
+    D = Reader.GetDecl(ID);
+    return true;
+  }
+
+  std::optional<Decl *> getDecl() const { return D; }
+};
+
+Decl *ASTReader::getAvailabilityDomainDecl(StringRef DomainName) {
+  AvailabilityDomainVisitor Visitor(DomainName, *this);
+  ModuleMgr.visit(Visitor);
+
+  if (std::optional<Decl *> D = Visitor.getDecl())
+    return *D;
+
+  return nullptr;
+}
+
 llvm::Error ASTReader::ReadASTBlock(ModuleFile &F,
                                     unsigned ClientLoadCapabilities) {
   BitstreamCursor &Stream = F.Stream;
@@ -3565,6 +3633,14 @@ llvm::Error ASTReader::ReadASTBlock(ModuleFile &F,
         UnusedFileScopedDecls.push_back(ReadDeclID(F, Record, I));
       break;
 
+    case AVAILABILITY_DOMAIN_TABLE: {
+      auto Data = Blob.data();
+      F.AvailabilityDomainTable = AvailabilityDomainLookupTable::Create(
+          (const unsigned char *)Data + Record[0], (const unsigned char *)Data,
+          AvailabilityDomainsTableReaderTrait());
+      break;
+    }
+
     case DELEGATING_CTORS:
       for (unsigned I = 0, N = Record.size(); I != N; /*in loop*/)
         DelegatingCtorDecls.push_back(ReadDeclID(F, Record, I));

clang/lib/Serialization/ASTReaderDecl.cpp

@@ -1683,10 +1683,6 @@ ASTDeclReader::RedeclarableResult ASTDeclReader::VisitVarDeclImpl(VarDecl *VD) {
   }
   }
 
-  if (auto *Attr = VD->getAttr<AvailabilityDomainAttr>())
-    Reader.getContext().addAvailabilityDomainMap(Attr->getName()->getName(),
-                                                 VD);
-
   return Redecl;
 }
 

clang/lib/Serialization/ASTReaderInternals.h

@@ -282,6 +282,38 @@ class HeaderFileInfoTrait {
 using HeaderFileInfoLookupTable =
     llvm::OnDiskChainedHashTable<HeaderFileInfoTrait>;
 
+/// This trait class is used for searching the on-disk hash table that stores
+/// feature availability information. The hash table maps feature names to Decl
+/// IDs that corresponds to the variables declared in header files.
+class AvailabilityDomainsTableReaderTrait {
+public:
+  using internal_key_type = llvm::StringRef;
+  using external_key_type = internal_key_type;
+  using data_type = clang::serialization::DeclID;
+  using hash_value_type = uint32_t;
+  using offset_type = unsigned;
+
+  internal_key_type GetInternalKey(external_key_type Key) { return Key; }
+  external_key_type GetExternalKey(internal_key_type Key) { return Key; }
+
+  hash_value_type ComputeHash(internal_key_type Key);
+
+  static bool EqualKey(internal_key_type LHS, internal_key_type RHS) {
+    return LHS == RHS;
+  }
+
+  static std::pair<unsigned, unsigned> ReadKeyDataLength(const uint8_t *&Data);
+
+  static internal_key_type ReadKey(const uint8_t *Data, unsigned Length);
+
+  static data_type ReadData(internal_key_type Key, const uint8_t *Data,
+                            unsigned Length);
+};
+
+/// The on-disk hash table used for availability domain lookup.
+using AvailabilityDomainLookupTable =
+    llvm::OnDiskChainedHashTable<AvailabilityDomainsTableReaderTrait>;
+
 } // namespace reader
 
 } // namespace serialization

clang/lib/Serialization/ASTWriter.cpp

@@ -941,6 +941,7 @@ void ASTWriter::WriteBlockInfoBlock() {
   RECORD(PP_ASSUME_NONNULL_LOC);
   RECORD(PP_UNSAFE_BUFFER_USAGE);
   RECORD(VTABLES_TO_EMIT);
+  RECORD(AVAILABILITY_DOMAIN_TABLE);
 
   // SourceManager Block.
   BLOCK(SOURCE_MANAGER_BLOCK);
@@ -3392,6 +3393,76 @@ uint64_t ASTWriter::WriteDeclContextLexicalBlock(ASTContext &Context,
   return Offset;
 }
 
+namespace {
+class AvailabilityDomainsTableTrait {
+public:
+  using key_type = StringRef;
+  using key_type_ref = key_type;
+  using data_type = clang::serialization::DeclID;
+  using data_type_ref = const data_type &;
+  using hash_value_type = uint32_t;
+  using offset_type = unsigned;
+
+  hash_value_type ComputeHash(key_type_ref Key) { return llvm::djbHash(Key); }
+
+  std::pair<unsigned, unsigned>
+  EmitKeyDataLength(raw_ostream &Out, key_type_ref Key, data_type_ref) {
+    uint32_t KeyLength = Key.size();
+    uint32_t DataLength = sizeof(data_type);
+
+    llvm::support::endian::Writer Writer(Out, llvm::endianness::little);
+    Writer.write<uint16_t>(KeyLength);
+    Writer.write<uint16_t>(DataLength);
+    return {KeyLength, DataLength};
+  }
+
+  void EmitKey(raw_ostream &Out, key_type_ref Key, unsigned) { Out << Key; }
+
+  void EmitData(raw_ostream &Out, key_type_ref, data_type_ref Data, unsigned) {
+    llvm::support::endian::Writer writer(Out, llvm::endianness::little);
+    writer.write<data_type>(Data);
+  }
+};
+} // namespace
+
+void ASTWriter::writeAvailabilityDomainDecls(ASTContext &Context) {
+  using namespace llvm;
+  if (Context.AvailabilityDomainMap.empty())
+    return;
+
+  AvailabilityDomainsTableTrait GeneratorTrait;
+  llvm::OnDiskChainedHashTableGenerator<AvailabilityDomainsTableTrait>
+      Generator;
+
+  for (auto &P : Context.AvailabilityDomainMap)
+    Generator.insert(P.first, GetDeclRef(P.second).getRawValue(),
+                     GeneratorTrait);
+
+  // Create the on-disk hash table in a buffer.
+  SmallString<4096> AvailDomainTable;
+  uint32_t BucketOffset;
+  {
+    using namespace llvm::support;
+    llvm::raw_svector_ostream Out(AvailDomainTable);
+    // Make sure that no bucket is at offset 0
+    endian::write<uint32_t>(Out, 0, llvm::endianness::little);
+    BucketOffset = Generator.Emit(Out, GeneratorTrait);
+  }
+
+  // Create a blob abbreviation.
+  auto Abbrev = std::make_shared<BitCodeAbbrev>();
+  Abbrev->Add(BitCodeAbbrevOp(AVAILABILITY_DOMAIN_TABLE));
+  Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Fixed, 32));
+  Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob));
+  unsigned MethodPoolAbbrev = Stream.EmitAbbrev(std::move(Abbrev));
+
+  // Write the table.
+  {
+    RecordData::value_type Record[] = {AVAILABILITY_DOMAIN_TABLE, BucketOffset};
+    Stream.EmitRecordWithBlob(MethodPoolAbbrev, Record, AvailDomainTable);
+  }
+}
+
 void ASTWriter::WriteTypeDeclOffsets() {
   using namespace llvm;
 
@@ -5798,6 +5869,8 @@ void ASTWriter::WriteDeclAndTypes(ASTContext &Context) {
   // Write the visible updates to DeclContexts.
   for (auto *DC : UpdatedDeclContexts)
     WriteDeclContextVisibleUpdate(Context, DC);
+
+  writeAvailabilityDomainDecls(Context);
 }
 
 void ASTWriter::WriteDeclUpdatesBlocks(ASTContext &Context,

clang/test/CodeGen/feature-availability.c

@@ -4,6 +4,9 @@
 // RUN: %clang_cc1 -triple arm64-apple-macosx -fblocks -ffeature-availability=feature1:on -ffeature-availability=feature2:off -emit-pch -o %t %s
 // RUN: %clang_cc1 -triple arm64-apple-macosx -fblocks -ffeature-availability=feature1:on -ffeature-availability=feature2:off -include-pch %t -emit-llvm -o - %s | FileCheck %s
 
+// RUN: %clang_cc1 -triple arm64-apple-macosx -fblocks -emit-pch -o %t -DUSE_DOMAIN %s
+// RUN: %clang_cc1 -triple arm64-apple-macosx -fblocks -include-pch %t -emit-llvm -o - -DUSE_DOMAIN %s | FileCheck --check-prefixes=CHECK,DOMAIN %s
+
 // CHECK: %[[STRUCT_S0:.*]] = type { i32 }
 // CHECK: @g0 = external global i32, align 4
 // CHECK-NOT: @g1

clang/test/Modules/Inputs/feature-availability/feature1.h

@@ -0,0 +1,7 @@
+#ifndef FEATURE1_H
+#define FEATURE1_H
+#include <feature-availability.h>
+
+static struct __AvailabilityDomain feature1 __attribute__((availability_domain(feature1))) = {__AVAILABILITY_DOMAIN_ENABLED, 0};
+
+#endif

clang/test/Modules/Inputs/feature-availability/feature2.h

@@ -0,0 +1,8 @@
+#ifndef FEATURE2_H
+#define FEATURE2_H
+#include <feature-availability.h>
+#include "feature1.h"
+
+static struct __AvailabilityDomain feature2 __attribute__((availability_domain(feature2))) = {__AVAILABILITY_DOMAIN_DISABLED, 0};
+
+#endif

clang/test/Modules/Inputs/feature-availability/module.modulemap

@@ -0,0 +1,9 @@
+module Feature1 {
+  header "feature1.h"
+  export *
+}
+
+module Feature2 {
+  header "feature2.h"
+  export *
+}