Serialize Fingerprints in Swift Modules

Author: CodaFi

include/swift/AST/LazyResolver.h

@@ -84,6 +84,11 @@ class alignas(void*) LazyMemberLoader {
   loadNamedMembers(const IterableDeclContext *IDC, DeclBaseName N,
                    uint64_t contextData) = 0;
 
+  /// Returns the fingerprint associated with the given iterable decl context,
+  /// or \c None if no such fingerprint is available.
+  virtual Optional<std::string>
+  loadFingerprint(const IterableDeclContext *IDC) = 0;
+
   /// Populates the given vector with all conformances for \p D.
   ///
   /// The implementation should \em not call setConformances on \p D.

lib/AST/DeclContext.cpp

@@ -1012,12 +1012,16 @@ IterableDeclContext::castDeclToIterableDeclContext(const Decl *D) {
 }
 
 Optional<std::string> IterableDeclContext::getBodyFingerprint() const {
-  // Only makes sense for contexts in a source file
-  if (!getAsGenericContext()->getParentSourceFile())
-    return None;
+  auto &ctx = getASTContext();
+  // If this decl comes from a serialized module, grab its fingerprint from
+  // the file.
+  if (!getAsGenericContext()->getParentSourceFile()) {
+    auto ci = ctx.getOrCreateLazyIterableContextData(this,
+                                                     /*lazyLoader=*/nullptr);
+    return ci->loader->loadFingerprint(this);
+  }
   auto mutableThis = const_cast<IterableDeclContext *>(this);
-  return evaluateOrDefault(getASTContext().evaluator,
-                           ParseMembersRequest{mutableThis},
+  return evaluateOrDefault(ctx.evaluator, ParseMembersRequest{mutableThis},
                            FingerprintAndMembers())
       .fingerprint;
 }

lib/ClangImporter/ClangImporter.cpp

@@ -3831,6 +3831,12 @@ void ClangImporter::Implementation::lookupAllObjCMembers(
   }
 }
 
+Optional<std::string>
+ClangImporter::Implementation::loadFingerprint(const IterableDeclContext *) {
+  // Clang decls are not fingerprinted in Swift.
+  return None;
+}
+
 TinyPtrVector<ValueDecl *>
 ClangImporter::Implementation::loadNamedMembers(
     const IterableDeclContext *IDC, DeclBaseName N, uint64_t contextData) {

lib/ClangImporter/ImporterImpl.h

@@ -1247,6 +1247,9 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   loadNamedMembers(const IterableDeclContext *IDC, DeclBaseName N,
                    uint64_t contextData) override;
 
+  virtual Optional<std::string>
+  loadFingerprint(const IterableDeclContext *IDC) override;
+
 private:
   void
   loadAllMembersOfObjcContainer(Decl *D,

lib/Serialization/ModuleFile.cpp

@@ -650,6 +650,24 @@ void ModuleFile::loadDerivativeFunctionConfigurations(
   }
 }
 
+Optional<std::string>
+ModuleFile::loadFingerprint(const IterableDeclContext *IDC) {
+  PrettyStackTraceDecl trace("loading fingerprints for", IDC->getDecl());
+
+  assert(IDC->wasDeserialized());
+  assert(IDC->getDeclID() != 0);
+
+  if (!Core->DeclFingerprints) {
+    return None;
+  }
+
+  auto it = Core->DeclFingerprints->find(IDC->getDeclID());
+  if (it == Core->DeclFingerprints->end()) {
+    return None;
+  }
+  return *it;
+}
+
 TinyPtrVector<ValueDecl *>
 ModuleFile::loadNamedMembers(const IterableDeclContext *IDC, DeclBaseName N,
                              uint64_t contextData) {

lib/Serialization/ModuleFile.h

@@ -653,6 +653,9 @@ class ModuleFile
   loadNamedMembers(const IterableDeclContext *IDC, DeclBaseName N,
                    uint64_t contextData) override;
 
+  virtual Optional<std::string>
+  loadFingerprint(const IterableDeclContext *IDC) override;
+
   virtual void
   loadAllConformances(const Decl *D, uint64_t contextData,
                     SmallVectorImpl<ProtocolConformance*> &Conforms) override;

lib/Serialization/ModuleFileCoreTableInfo.h

@@ -566,6 +566,42 @@ class ModuleFileSharedCore::DeclUSRTableInfo {
   }
 };
 
+class ModuleFileSharedCore::DeclFingerprintsTableInfo {
+public:
+  using internal_key_type = uint32_t;
+  using external_key_type = DeclID;
+  using data_type = std::string;
+  using hash_value_type = uint32_t;
+  using offset_type = unsigned;
+
+  internal_key_type GetInternalKey(external_key_type ID) { return ID; }
+
+  hash_value_type ComputeHash(internal_key_type key) {
+    return llvm::hash_value(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) {
+    using namespace llvm::support;
+    unsigned dataLength = endian::readNext<uint16_t, little, unaligned>(data);
+    return {sizeof(uint32_t), dataLength};
+  }
+
+  static internal_key_type ReadKey(const uint8_t *data, unsigned length) {
+    using namespace llvm::support;
+    return endian::readNext<uint32_t, little, unaligned>(data);
+  }
+
+  static data_type ReadData(internal_key_type key, const uint8_t *data,
+                            unsigned length) {
+    using namespace llvm::support;
+    return std::string{reinterpret_cast<const char *>(data), length};
+  }
+};
+
 } // end namespace swift
 
 #endif

lib/Serialization/ModuleFileSharedCore.cpp

@@ -531,6 +531,18 @@ ModuleFileSharedCore::readDeclMembersTable(ArrayRef<uint64_t> fields,
       base + sizeof(uint32_t), base));
 }
 
+std::unique_ptr<ModuleFileSharedCore::SerializedDeclFingerprintsTable>
+ModuleFileSharedCore::readDeclFingerprintsTable(ArrayRef<uint64_t> fields,
+                                                StringRef blobData) const {
+  uint32_t tableOffset;
+  index_block::DeclFingerprintsLayout::readRecord(fields, tableOffset);
+  auto base = reinterpret_cast<const uint8_t *>(blobData.data());
+
+  using OwnedTable = std::unique_ptr<SerializedDeclFingerprintsTable>;
+  return OwnedTable(SerializedDeclFingerprintsTable::Create(
+      base + tableOffset, base + sizeof(uint32_t), base));
+}
+
 std::unique_ptr<ModuleFileSharedCore::SerializedObjCMethodTable>
 ModuleFileSharedCore::readObjCMethodTable(ArrayRef<uint64_t> fields,
                                     StringRef blobData) const {
@@ -685,6 +697,9 @@ bool ModuleFileSharedCore::readIndexBlock(llvm::BitstreamCursor &cursor) {
       case index_block::DECL_MEMBER_NAMES:
         DeclMemberNames = readDeclMemberNamesTable(scratch, blobData);
         break;
+      case index_block::DECL_FINGERPRINTS:
+        DeclFingerprints = readDeclFingerprintsTable(scratch, blobData);
+        break;
       case index_block::LOCAL_DECL_CONTEXT_OFFSETS:
         assert(blobData.empty());
         allocateBuffer(LocalDeclContexts, scratch);

lib/Serialization/ModuleFileSharedCore.h

@@ -240,6 +240,10 @@ class ModuleFileSharedCore {
   using SerializedDeclMembersTable =
       llvm::OnDiskIterableChainedHashTable<DeclMembersTableInfo>;
 
+  class DeclFingerprintsTableInfo;
+  using SerializedDeclFingerprintsTable =
+      llvm::OnDiskIterableChainedHashTable<DeclFingerprintsTableInfo>;
+
   std::unique_ptr<SerializedDeclTable> TopLevelDecls;
   std::unique_ptr<SerializedDeclTable> OperatorDecls;
   std::unique_ptr<SerializedDeclTable> PrecedenceGroupDecls;
@@ -250,6 +254,7 @@ class ModuleFileSharedCore {
   std::unique_ptr<SerializedOpaqueReturnTypeDeclTable> OpaqueReturnTypeDecls;
   std::unique_ptr<SerializedNestedTypeDeclsTable> NestedTypeDecls;
   std::unique_ptr<SerializedDeclMemberNamesTable> DeclMemberNames;
+  std::unique_ptr<SerializedDeclFingerprintsTable> DeclFingerprints;
 
   class ObjCMethodTableInfo;
   using SerializedObjCMethodTable =
@@ -399,6 +404,12 @@ class ModuleFileSharedCore {
   std::unique_ptr<SerializedDeclMembersTable>
   readDeclMembersTable(ArrayRef<uint64_t> fields, StringRef blobData) const;
 
+  /// Read an on-disk local declid-string hash table stored in
+  /// index_block::DeclFingerprintsLayout format.
+  std::unique_ptr<SerializedDeclFingerprintsTable>
+  readDeclFingerprintsTable(ArrayRef<uint64_t> fields,
+                            StringRef blobData) const;
+
   /// Read an on-disk derivative function configuration table stored in
   /// index_block::DerivativeFunctionConfigTableLayout format.
   std::unique_ptr<ModuleFileSharedCore::SerializedDerivativeFunctionConfigTable>

lib/Serialization/ModuleFormat.h

@@ -56,7 +56,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 586; // allow errors in modules
+const uint16_t SWIFTMODULE_VERSION_MINOR = 587; // fingerprints in modules
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2010,14 +2010,15 @@ namespace index_block {
     PRECEDENCE_GROUPS,
     NESTED_TYPE_DECLS,
     DECL_MEMBER_NAMES,
+    DECL_FINGERPRINTS,
 
     ORDERED_TOP_LEVEL_DECLS,
 
     SUBSTITUTION_MAP_OFFSETS,
     CLANG_TYPE_OFFSETS,
     LastRecordKind = CLANG_TYPE_OFFSETS,
   };
-  
+
   constexpr const unsigned RecordIDFieldWidth = 5;
   static_assert(LastRecordKind < (1 << RecordIDFieldWidth),
                 "not enough bits for all record kinds");
@@ -2078,6 +2079,12 @@ namespace index_block {
     RecordIDField,        // record ID
     BCArray<DeclIDField>  // list of decls by ID
   >;
+
+  using DeclFingerprintsLayout = BCRecordLayout<
+    DECL_FINGERPRINTS, // record ID
+    BCVBR<16>,   // table offset within the blob (see below)
+    BCBlob       // map from member DeclIDs to strings
+  >;
 }
 
 /// \sa DECL_MEMBER_TABLES_BLOCK_ID

lib/Serialization/Serialization.cpp

@@ -434,6 +434,43 @@ namespace {
       }
     }
   };
+
+  class DeclFingerprintsTableInfo {
+  public:
+    using key_type = DeclID;
+    using key_type_ref = const key_type &;
+    using data_type = std::string;
+    using data_type_ref = llvm::StringRef;
+    using hash_value_type = uint32_t;
+    using offset_type = unsigned;
+
+    hash_value_type ComputeHash(key_type_ref key) {
+      return llvm::hash_value(static_cast<uint32_t>(key));
+    }
+
+    std::pair<unsigned, unsigned>
+    EmitKeyDataLength(raw_ostream &out, key_type_ref key, data_type_ref data) {
+      assert((data.size() == 32) && "Fingerprint must be 32-bytes long!");
+      endian::Writer writer(out, little);
+      // No need to write the key length; it's constant.
+      writer.write<uint16_t>(data.size());
+      return {sizeof(uint32_t), data.size()};
+    }
+
+    void EmitKey(raw_ostream &out, key_type_ref key, unsigned len) {
+      static_assert(declIDFitsIn32Bits(), "DeclID too large");
+      assert(len == sizeof(uint32_t));
+      endian::Writer writer(out, little);
+      writer.write<uint32_t>(key);
+    }
+
+    void EmitData(raw_ostream &out, key_type_ref key, data_type_ref data,
+                  unsigned len) {
+      static_assert(declIDFitsIn32Bits(), "DeclID too large");
+      endian::Writer writer(out, little);
+      out << data;
+    }
+  };
 } // end anonymous namespace
 
 static ModuleDecl *getModule(ModuleOrSourceFile DC) {
@@ -799,6 +836,7 @@ void Serializer::writeBlockInfoBlock() {
   BLOCK_RECORD(index_block, PRECEDENCE_GROUPS);
   BLOCK_RECORD(index_block, NESTED_TYPE_DECLS);
   BLOCK_RECORD(index_block, DECL_MEMBER_NAMES);
+  BLOCK_RECORD(index_block, DECL_FINGERPRINTS);
   BLOCK_RECORD(index_block, ORDERED_TOP_LEVEL_DECLS);
 
   BLOCK(DECL_MEMBER_TABLES_BLOCK);
@@ -4806,6 +4844,27 @@ writeDeclMembersTable(const decl_member_tables_block::DeclMembersLayout &mems,
   mems.emit(scratch, tableOffset, hashTableBlob);
 }
 
+static void
+writeDeclFingerprintsTable(const index_block::DeclFingerprintsLayout &fpl,
+                           const Serializer::DeclFingerprintsTable &table) {
+  SmallVector<uint64_t, 8> scratch;
+  llvm::SmallString<4096> hashTableBlob;
+  uint32_t tableOffset;
+  {
+    llvm::OnDiskChainedHashTableGenerator<DeclFingerprintsTableInfo> generator;
+    for (auto &entry : table) {
+      generator.insert(entry.first, entry.second);
+    }
+
+    llvm::raw_svector_ostream blobStream(hashTableBlob);
+    // Make sure that no bucket is at offset 0
+    endian::write<uint32_t>(blobStream, 0, little);
+    tableOffset = generator.Emit(blobStream);
+  }
+
+  fpl.emit(scratch, tableOffset, hashTableBlob);
+}
+
 namespace {
   /// Used to serialize the on-disk Objective-C method hash table.
   class ObjCMethodTableInfo {
@@ -5078,6 +5137,7 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
   LocalTypeHashTableGenerator localTypeGenerator, opaqueReturnTypeGenerator;
   ExtensionTable extensionDecls;
   UniquedDerivativeFunctionConfigTable uniquedDerivativeConfigs;
+  DeclFingerprintsTable declFingerprints;
   bool hasLocalTypes = false;
   bool hasOpaqueReturnTypes = false;
 
@@ -5137,6 +5197,9 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
       // derived conformance (for example, ==), force them to be
       // serialized.
       if (auto IDC = dyn_cast<IterableDeclContext>(D)) {
+        if (auto bodyFP = IDC->getBodyFingerprint()) {
+          declFingerprints[addDeclRef(D)] = *bodyFP;
+        }
         collectInterestingNestedDeclarations(*this, IDC->getMembers(),
                                              operatorMethodDecls, objcMethods,
                                              nestedTypeDecls,
@@ -5167,6 +5230,9 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
       localTypeGenerator.insert(MangledName, addDeclRef(TD));
 
       if (auto IDC = dyn_cast<IterableDeclContext>(TD)) {
+        if (auto bodyFP = IDC->getBodyFingerprint()) {
+          declFingerprints[addDeclRef(TD)] = *bodyFP;
+        }
         collectInterestingNestedDeclarations(*this, IDC->getMembers(),
                                              operatorMethodDecls, objcMethods,
                                              nestedTypeDecls,
@@ -5233,6 +5299,12 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
       writeNestedTypeDeclsTable(NestedTypeDeclsTable, nestedTypeDecls);
     }
 
+    if (!declFingerprints.empty()) {
+      // Write decl fingerprints.
+      index_block::DeclFingerprintsLayout DeclsFingerprints(Out);
+      writeDeclFingerprintsTable(DeclsFingerprints, declFingerprints);
+    }
+
     // Convert uniqued derivative function config table to serialization-
     // ready format: turn `GenericSignature` to `GenericSignatureID`.
     DerivativeFunctionConfigTable derivativeConfigs;

lib/Serialization/Serialization.h

@@ -277,6 +277,8 @@ class Serializer : public SerializerBase {
       Identifier,
       llvm::SmallSetVector<std::pair<Identifier, GenericSignature>, 4>>;
 
+  using DeclFingerprintsTable = llvm::MapVector<uint32_t, std::string>;
+
 private:
   /// A map from identifiers to methods and properties with the given name.
   ///