Teach the Driver to Read Fine-Grained Dependency Graphs in Swiftdeps Files

Author: CodaFi

include/swift/AST/FineGrainedDependencies.h

@@ -851,6 +851,8 @@ class SourceFileDepGraph {
   /// Read a swiftdeps file from \p buffer and return a SourceFileDepGraph if
   /// successful.
   Optional<SourceFileDepGraph> static loadFromBuffer(llvm::MemoryBuffer &);
+  Optional<SourceFileDepGraph> static loadFromSwiftModuleBuffer(
+      llvm::MemoryBuffer &);
 
   void verifySame(const SourceFileDepGraph &other) const;
 

include/swift/AST/FineGrainedDependencyFormat.h

@@ -48,6 +48,7 @@ using NodeKindField = BCFixed<3>;
 using DeclAspectField = BCFixed<1>;
 
 const unsigned RECORD_BLOCK_ID = llvm::bitc::FIRST_APPLICATION_BLOCKID;
+const unsigned INCREMENTAL_INFORMATION_BLOCK_ID = 196;
 
 /// The swiftdeps file format consists of a METADATA record, followed by zero or more
 /// IDENTIFIER_NODE records.
@@ -113,10 +114,16 @@ namespace record_block {
 }
 
 /// Tries to read the dependency graph from the given buffer.
-/// Returns true if there was an error.
+/// Returns \c true if there was an error.
 bool readFineGrainedDependencyGraph(llvm::MemoryBuffer &buffer,
                                     SourceFileDepGraph &g);
 
+/// Tries to read the dependency graph from the given buffer, assuming that it
+/// is in the format of a swiftmodule file.
+/// Returns \c true if there was an error.
+bool readFineGrainedDependencyGraphFromSwiftModule(llvm::MemoryBuffer &buffer,
+                                                   SourceFileDepGraph &g);
+
 /// Tries to read the dependency graph from the given path name.
 /// Returns true if there was an error.
 bool readFineGrainedDependencyGraph(llvm::StringRef path,

include/swift/Driver/FineGrainedDependencyDriverGraph.h

@@ -336,6 +336,8 @@ class ModuleDepGraph {
                                      const SourceFileDepGraph &,
                                      DiagnosticEngine &);
 
+  Changes loadFromSwiftModuleBuffer(const driver::Job *, llvm::MemoryBuffer &,
+                                    DiagnosticEngine &);
 
 private:
   /// Read a SourceFileDepGraph belonging to \p job from \p buffer

lib/AST/FineGrainedDependencies.cpp

@@ -58,6 +58,15 @@ SourceFileDepGraph::loadFromBuffer(llvm::MemoryBuffer &buffer) {
   return Optional<SourceFileDepGraph>(std::move(fg));
 }
 
+Optional<SourceFileDepGraph>
+SourceFileDepGraph::loadFromSwiftModuleBuffer(llvm::MemoryBuffer &buffer) {
+  SourceFileDepGraph fg;
+  if (swift::fine_grained_dependencies::
+          readFineGrainedDependencyGraphFromSwiftModule(buffer, fg))
+    return None;
+  return Optional<SourceFileDepGraph>(std::move(fg));
+}
+
 //==============================================================================
 // MARK: SourceFileDepGraph access
 //==============================================================================

lib/AST/FineGrainedDependencyFormat.cpp

@@ -45,6 +45,7 @@ class Deserializer {
 public:
   Deserializer(llvm::MemoryBufferRef Data) : Cursor(Data) {}
   bool readFineGrainedDependencyGraph(SourceFileDepGraph &g);
+  bool readFineGrainedDependencyGraphFromSwiftModule(SourceFileDepGraph &g);
 };
 
 } // end namespace
@@ -485,3 +486,106 @@ bool swift::fine_grained_dependencies::writeFineGrainedDependencyGraphToPath(
     return false;
   });
 }
+
+static bool checkModuleSignature(llvm::BitstreamCursor &cursor,
+                                 ArrayRef<unsigned char> signature) {
+  for (unsigned char byte : signature) {
+    if (cursor.AtEndOfStream())
+      return false;
+    if (llvm::Expected<llvm::SimpleBitstreamCursor::word_t> maybeRead =
+            cursor.Read(8)) {
+      if (maybeRead.get() != byte)
+        return false;
+    } else {
+      consumeError(maybeRead.takeError());
+      return false;
+    }
+  }
+  return true;
+}
+
+static bool enterTopLevelModuleBlock(llvm::BitstreamCursor &cursor, unsigned ID,
+                                     bool shouldReadBlockInfo = true) {
+  llvm::Expected<llvm::BitstreamEntry> maybeNext = cursor.advance();
+  if (!maybeNext) {
+    consumeError(maybeNext.takeError());
+    return false;
+  }
+  llvm::BitstreamEntry next = maybeNext.get();
+
+  if (next.Kind != llvm::BitstreamEntry::SubBlock)
+    return false;
+
+  if (next.ID == RECORD_BLOCK_ID) {
+    if (shouldReadBlockInfo) {
+      if (!cursor.ReadBlockInfoBlock())
+        return false;
+    } else {
+      if (cursor.SkipBlock())
+        return false;
+    }
+    return enterTopLevelModuleBlock(cursor, ID, false);
+  }
+
+  if (next.ID != ID)
+    return false;
+
+  if (llvm::Error Err = cursor.EnterSubBlock(ID)) {
+    // FIXME this drops the error on the floor.
+    consumeError(std::move(Err));
+    return false;
+  }
+
+  return true;
+}
+
+bool swift::fine_grained_dependencies::
+    readFineGrainedDependencyGraphFromSwiftModule(llvm::MemoryBuffer &buffer,
+                                                  SourceFileDepGraph &g) {
+  Deserializer deserializer(buffer.getMemBufferRef());
+  return deserializer.readFineGrainedDependencyGraphFromSwiftModule(g);
+}
+
+bool Deserializer::readFineGrainedDependencyGraphFromSwiftModule(
+    SourceFileDepGraph &g) {
+  if (!checkModuleSignature(Cursor, {0xE2, 0x9C, 0xA8, 0x0E}) ||
+      !enterTopLevelModuleBlock(Cursor, RECORD_BLOCK_ID, false)) {
+    return false;
+  }
+
+  llvm::BitstreamEntry topLevelEntry;
+
+  while (!Cursor.AtEndOfStream()) {
+    llvm::Expected<llvm::BitstreamEntry> maybeEntry =
+        Cursor.advance(llvm::BitstreamCursor::AF_DontPopBlockAtEnd);
+    if (!maybeEntry) {
+      consumeError(maybeEntry.takeError());
+      return false;
+    }
+    topLevelEntry = maybeEntry.get();
+    if (topLevelEntry.Kind != llvm::BitstreamEntry::SubBlock)
+      break;
+
+    switch (topLevelEntry.ID) {
+    case INCREMENTAL_INFORMATION_BLOCK_ID: {
+      if (llvm::Error Err =
+              Cursor.EnterSubBlock(INCREMENTAL_INFORMATION_BLOCK_ID)) {
+        consumeError(std::move(Err));
+        return false;
+      }
+      readFineGrainedDependencyGraph(g);
+      break;
+    }
+
+    default:
+      // Unknown top-level block, possibly for use by a future version of the
+      // module format.
+      if (Cursor.SkipBlock()) {
+        return false;
+      }
+      break;
+    }
+  }
+
+  return false;
+}

lib/Driver/FineGrainedDependencyDriverGraph.cpp

@@ -108,6 +108,27 @@ ModuleDepGraph::Changes ModuleDepGraph::loadFromSourceFileDepGraph(
   return changes;
 }
 
+ModuleDepGraph::Changes ModuleDepGraph::loadFromSwiftModuleBuffer(
+    const Job *Cmd, llvm::MemoryBuffer &buffer, DiagnosticEngine &diags) {
+  FrontendStatsTracer tracer(
+      stats, "fine-grained-dependencies-loadFromSwiftModuleBuffer");
+  PrettyStackTraceStringAction stackTrace(
+      "loading fine-grained dependency graph from swiftmodule",
+      buffer.getBufferIdentifier());
+
+   Optional<SourceFileDepGraph> sourceFileDepGraph =
+      SourceFileDepGraph::loadFromSwiftModuleBuffer(buffer);
+  if (!sourceFileDepGraph)
+    return None;
+  registerJob(Cmd);
+  auto changes = integrate(*sourceFileDepGraph, buffer.getBufferIdentifier());
+  if (verifyFineGrainedDependencyGraphAfterEveryImport)
+    verify();
+  if (emitFineGrainedDependencyDotFileAfterEveryImport)
+    emitDotFileForJob(diags, Cmd);
+  return changes;
+}
+
 bool ModuleDepGraph::haveAnyNodesBeenTraversedIn(const Job *cmd) const {
   std::string swiftDeps = getSwiftDeps(cmd).str();
 

lib/Serialization/ModuleFormat.h

@@ -20,10 +20,11 @@
 #define SWIFT_SERIALIZATION_MODULEFORMAT_H
 
 #include "swift/AST/Decl.h"
+#include "swift/AST/FineGrainedDependencyFormat.h"
 #include "swift/AST/Types.h"
+#include "llvm/ADT/PointerEmbeddedInt.h"
 #include "llvm/Bitcode/RecordLayout.h"
 #include "llvm/Bitstream/BitCodes.h"
-#include "llvm/ADT/PointerEmbeddedInt.h"
 
 namespace swift {
 namespace serialization {
@@ -724,16 +725,19 @@ enum BlockID {
   /// This is part of a stable format and should not be renumbered.
   ///
   /// Though we strive to keep the format stable, breaking the format of
-  /// .swiftsourceinfo doesn't have consequences as serious as breaking the format
-  /// of .swiftdoc because .swiftsourceinfo file is for local development use only.
+  /// .swiftsourceinfo doesn't have consequences as serious as breaking the
+  /// format
+  /// of .swiftdoc because .swiftsourceinfo file is for local development use
+  /// only.
   ///
   /// \sa decl_locs_block
   DECL_LOCS_BLOCK_ID,
 
   /// The incremental dependency information block.
   ///
   /// This is part of a stable format and should not be renumbered.
-  INCREMENTAL_INFORMATION_BLOCK_ID = 196,
+  INCREMENTAL_INFORMATION_BLOCK_ID =
+      fine_grained_dependencies::INCREMENTAL_INFORMATION_BLOCK_ID,
 };
 
 /// The record types within the control block.