Merge pull request #37105 from bnbarham/serialize-offsets

[Serialization] Store offset of decls in .swiftsourceinfo

Author: bnbarham

Diff for: include/swift/AST/ASTContext.h

@@ -73,6 +73,7 @@ namespace swift {
   class DerivativeAttr;
   class DifferentiableAttr;
   class ExtensionDecl;
+  struct ExternalSourceLocs;
   class ForeignRepresentationInfo;
   class FuncDecl;
   class GenericContext;
@@ -1184,6 +1185,10 @@ class ASTContext final {
 
 private:
   friend Decl;
+
+  Optional<ExternalSourceLocs *> getExternalSourceLocs(const Decl *D);
+  void setExternalSourceLocs(const Decl *D, ExternalSourceLocs *Locs);
+
   Optional<std::pair<RawComment, bool>> getRawComment(const Decl *D);
   void setRawComment(const Decl *D, RawComment RC, bool FromSerialized);
 

Diff for: include/swift/AST/Decl.h

@@ -62,6 +62,7 @@ namespace swift {
   class DynamicSelfType;
   class Type;
   class Expr;
+  struct ExternalSourceLocs;
   class CaptureListExpr;
   class DeclRefExpr;
   class ForeignAsyncConvention;
@@ -688,14 +689,12 @@ class alignas(1 << DeclAlignInBits) Decl {
   void operator=(const Decl&) = delete;
   SourceLoc getLocFromSource() const;
 
-  struct CachedExternalSourceLocs {
-    SourceLoc Loc;
-    SourceLoc StartLoc;
-    SourceLoc EndLoc;
-    SmallVector<CharSourceRange, 4> DocRanges;
-  };
-  mutable CachedExternalSourceLocs const *CachedSerializedLocs = nullptr;
-  const CachedExternalSourceLocs *getSerializedLocs() const;
+  /// Returns the serialized locations of this declaration from the
+  /// corresponding .swiftsourceinfo file. "Empty" (ie. \c BufferID of 0, an
+  /// invalid \c Loc, and empty \c DocRanges) if either there is no
+  /// .swiftsourceinfo or the buffer could not be created, eg. if the file
+  /// no longer exists.
+  const ExternalSourceLocs *getSerializedLocs() const;
 
   /// Directly set the invalid bit
   void setInvalidBit();

Diff for: include/swift/AST/FileUnit.h

@@ -164,7 +164,8 @@ class FileUnit : public DeclContext {
     return None;
   }
 
-  virtual Optional<BasicDeclLocs> getBasicLocsForDecl(const Decl *D) const {
+  virtual Optional<ExternalSourceLocs::RawLocs>
+  getExternalRawLocsForDecl(const Decl *D) const {
     return None;
   }
 

Diff for: include/swift/AST/SourceFile.h

@@ -399,7 +399,8 @@ class SourceFile final : public FileUnit {
 
   Identifier getDiscriminatorForPrivateValue(const ValueDecl *D) const override;
   Identifier getPrivateDiscriminator() const { return PrivateDiscriminator; }
-  Optional<BasicDeclLocs> getBasicLocsForDecl(const Decl *D) const override;
+  Optional<ExternalSourceLocs::RawLocs>
+  getExternalRawLocsForDecl(const Decl *D) const override;
 
   /// Returns the synthesized file for this source file, if it exists.
   SynthesizedFileUnit *getSynthesizedFile() const { return SynthesizedFile; };

Diff for: include/swift/Basic/BasicSourceInfo.h

@@ -15,25 +15,42 @@
 
 #include "swift/Basic/Fingerprint.h"
 #include "swift/Basic/LLVM.h"
+#include "swift/Basic/SourceLoc.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/Support/Chrono.h"
 
 namespace swift {
 
 class SourceFile;
 
-struct SourcePosition {
-  uint32_t Line = 0;
-  uint32_t Column = 0;
-  bool isValid() const { return Line && Column; }
-};
-
-struct BasicDeclLocs {
-  StringRef SourceFilePath;
-  SmallVector<std::pair<SourcePosition, uint32_t>, 4> DocRanges;
-  SourcePosition Loc;
-  SourcePosition StartLoc;
-  SourcePosition EndLoc;
+struct ExternalSourceLocs {
+  struct LocationDirective {
+    uint32_t Offset = 0;
+    int32_t LineOffset = 0;
+    uint32_t Length = 0;
+    StringRef Name;
+
+    bool isValid() const { return Length > 0; }
+  };
+
+  struct RawLoc {
+    uint32_t Offset = 0;
+    uint32_t Line = 0;
+    uint32_t Column = 0;
+    LocationDirective Directive;
+  };
+
+  struct RawLocs {
+    StringRef SourceFilePath;
+    SmallVector<std::pair<RawLoc, uint32_t>, 4> DocRanges;
+    RawLoc Loc;
+    RawLoc StartLoc;
+    RawLoc EndLoc;
+  };
+
+  unsigned BufferID = 0;
+  SourceLoc Loc;
+  SmallVector<CharSourceRange, 4> DocRanges;
 };
 
 class BasicSourceFileInfo {

Diff for: include/swift/Basic/SourceManager.h

@@ -24,6 +24,18 @@ namespace swift {
 
 /// This class manages and owns source buffers.
 class SourceManager {
+public:
+  /// A \c #sourceLocation-defined virtual file region, representing the source
+  /// source after a \c #sourceLocation (or between two). It provides a
+  /// filename and line offset to be applied to \c SourceLoc's within its
+  /// \c Range.
+  struct VirtualFile {
+    CharSourceRange Range;
+    std::string Name;
+    int LineOffset;
+  };
+
+private:
   llvm::SourceMgr LLVMSourceMgr;
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem;
   unsigned CodeCompletionBufferID = 0U;
@@ -52,12 +64,6 @@ class SourceManager {
   };
   ReplacedRangeType ReplacedRange;
 
-  // \c #sourceLocation directive handling.
-  struct VirtualFile {
-    CharSourceRange Range;
-    std::string Name;
-    int LineOffset;
-  };
   std::map<const char *, VirtualFile> VirtualFiles;
   mutable std::pair<const char *, const VirtualFile*> CachedVFile = {nullptr, nullptr};
 
@@ -143,6 +149,10 @@ class SourceManager {
   /// Adds a memory buffer to the SourceManager, taking ownership of it.
   unsigned addNewSourceBuffer(std::unique_ptr<llvm::MemoryBuffer> Buffer);
 
+  /// Add a \c #sourceLocation-defined virtual file region of \p Length.
+  void createVirtualFile(SourceLoc Loc, StringRef Name, int LineOffset,
+                         unsigned Length);
+
   /// Add a \c #sourceLocation-defined virtual file region.
   ///
   /// By default, this region continues to the end of the buffer.
@@ -275,18 +285,17 @@ class SourceManager {
 
   std::string getLineString(unsigned BufferID, unsigned LineNumber);
 
+  /// Retrieve the buffer ID for \p Path, loading if necessary.
+  unsigned getExternalSourceBufferID(StringRef Path);
+
   SourceLoc getLocFromExternalSource(StringRef Path, unsigned Line, unsigned Col);
-private:
+
+  /// Retrieve the virtual file for the given \p Loc, or nullptr if none exists.
   const VirtualFile *getVirtualFile(SourceLoc Loc) const;
-  unsigned getExternalSourceBufferId(StringRef Path);
-  int getLineOffset(SourceLoc Loc) const {
-    if (auto VFile = getVirtualFile(Loc))
-      return VFile->LineOffset;
-    else
-      return 0;
-  }
 
-public:
+  /// Whether or not \p Loc is after a \c #sourceLocation directive, ie. its
+  /// file, line, and column should be reported using the information in the
+  /// directive.
   bool isLocInVirtualFile(SourceLoc Loc) const {
     return getVirtualFile(Loc) != nullptr;
   }
@@ -295,6 +304,14 @@ class SourceManager {
   /// owned by \p otherManager. Returns an invalid SourceLoc if it cannot be
   /// translated.
   SourceLoc getLocForForeignLoc(SourceLoc otherLoc, SourceManager &otherMgr);
+
+private:
+  int getLineOffset(SourceLoc Loc) const {
+    if (auto VFile = getVirtualFile(Loc))
+      return VFile->LineOffset;
+    else
+      return 0;
+  }
 };
 
 } // end namespace swift

Diff for: include/swift/Serialization/SerializedModuleLoader.h

@@ -392,7 +392,8 @@ class SerializedASTFile final : public LoadedFile {
 
   Optional<StringRef> getGroupNameByUSR(StringRef USR) const override;
 
-  Optional<BasicDeclLocs> getBasicLocsForDecl(const Decl *D) const override;
+  Optional<ExternalSourceLocs::RawLocs>
+  getExternalRawLocsForDecl(const Decl *D) const override;
 
   void collectAllGroups(SmallVectorImpl<StringRef> &Names) const override;
 

Diff for: lib/AST/ASTContext.cpp

@@ -287,6 +287,10 @@ struct ASTContext::Implementation {
   /// The module loader used to load Clang modules from DWARF.
   ClangModuleLoader *TheDWARFModuleLoader = nullptr;
 
+  /// Map from Swift declarations to deserialized resolved locations, ie.
+  /// actual \c SourceLocs that require opening their external buffer.
+  llvm::DenseMap<const Decl *, ExternalSourceLocs *> ExternalSourceLocs;
+
   /// Map from Swift declarations to raw comments.
   llvm::DenseMap<const Decl *, std::pair<RawComment, bool>> RawComments;
 
@@ -2048,6 +2052,20 @@ ModuleDecl *ASTContext::getStdlibModule(bool loadIfAbsent) {
   return TheStdlibModule;
 }
 
+Optional<ExternalSourceLocs *>
+ASTContext::getExternalSourceLocs(const Decl *D) {
+  auto Known = getImpl().ExternalSourceLocs.find(D);
+  if (Known == getImpl().ExternalSourceLocs.end())
+    return None;
+
+  return Known->second;
+}
+
+void ASTContext::setExternalSourceLocs(const Decl *D,
+                                       ExternalSourceLocs *Locs) {
+  getImpl().ExternalSourceLocs[D] = Locs;
+}
+
 Optional<std::pair<RawComment, bool>> ASTContext::getRawComment(const Decl *D) {
   auto Known = getImpl().RawComments.find(D);
   if (Known == getImpl().RawComments.end())

Diff for: lib/AST/Decl.cpp

@@ -610,36 +610,50 @@ case DeclKind::ID: return cast<ID##Decl>(this)->getLocFromSource();
   llvm_unreachable("Unknown decl kind");
 }
 
-const Decl::CachedExternalSourceLocs *Decl::getSerializedLocs() const {
-  if (CachedSerializedLocs) {
-    return CachedSerializedLocs;
-  }
+const ExternalSourceLocs *Decl::getSerializedLocs() const {
+  auto &Context = getASTContext();
+  if (auto EL = Context.getExternalSourceLocs(this).getValueOr(nullptr))
+    return EL;
+
+  static ExternalSourceLocs NullLocs{};
+
   auto *File = cast<FileUnit>(getDeclContext()->getModuleScopeContext());
-  assert(File->getKind() == FileUnitKind::SerializedAST &&
-         "getSerializedLocs() should only be called on decls in "
-         "a 'SerializedASTFile'");
-  auto Locs = File->getBasicLocsForDecl(this);
-  if (!Locs.hasValue()) {
-    static const Decl::CachedExternalSourceLocs NullLocs{};
+  if (File->getKind() != FileUnitKind::SerializedAST)
+    return &NullLocs;
+
+  auto RawLocs = File->getExternalRawLocsForDecl(this);
+  if (!RawLocs.hasValue()) {
+    // Don't read .swiftsourceinfo again on failure
+    Context.setExternalSourceLocs(this, &NullLocs);
     return &NullLocs;
   }
-  auto *Result = getASTContext().Allocate<Decl::CachedExternalSourceLocs>();
-  auto &SM = getASTContext().SourceMgr;
-#define CASE(X)                                                               \
-Result->X = SM.getLocFromExternalSource(Locs->SourceFilePath, Locs->X.Line,   \
-                                       Locs->X.Column);
-  CASE(Loc)
-  CASE(StartLoc)
-  CASE(EndLoc)
-#undef CASE
 
-  for (const auto &LineColumnAndLength : Locs->DocRanges) {
-    auto Start = SM.getLocFromExternalSource(Locs->SourceFilePath,
-      LineColumnAndLength.first.Line,
-      LineColumnAndLength.first.Column);
-    Result->DocRanges.push_back({ Start, LineColumnAndLength.second });
+  auto &SM = getASTContext().SourceMgr;
+  unsigned BufferID = SM.getExternalSourceBufferID(RawLocs->SourceFilePath);
+  if (!BufferID) {
+    // Don't try open the file again on failure
+    Context.setExternalSourceLocs(this, &NullLocs);
+    return &NullLocs;
   }
 
+  auto ResolveLoc = [&](const ExternalSourceLocs::RawLoc &Raw) -> SourceLoc {
+    // If the decl had a presumed loc, create its virtual file so that
+    // getPresumedLineAndColForLoc works from serialized locations as well.
+    if (Raw.Directive.isValid()) {
+      auto &LD = Raw.Directive;
+      SourceLoc Loc = SM.getLocForOffset(BufferID, LD.Offset);
+      SM.createVirtualFile(Loc, LD.Name, LD.LineOffset, LD.Length);
+    }
+    return SM.getLocForOffset(BufferID, Raw.Offset);
+  };
+
+  auto *Result = getASTContext().Allocate<ExternalSourceLocs>();
+  Result->BufferID = BufferID;
+  Result->Loc = ResolveLoc(RawLocs->Loc);
+  for (auto &Range : RawLocs->DocRanges) {
+    Result->DocRanges.emplace_back(ResolveLoc(Range.first), Range.second);
+  }
+  Context.setExternalSourceLocs(this, Result);
   return Result;
 }
 

Diff for: lib/AST/Module.cpp

@@ -855,37 +855,50 @@ TypeDecl *SourceFile::lookupLocalType(llvm::StringRef mangledName) const {
   return nullptr;
 }
 
-Optional<BasicDeclLocs>
-SourceFile::getBasicLocsForDecl(const Decl *D) const {
+Optional<ExternalSourceLocs::RawLocs>
+SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
   auto *FileCtx = D->getDeclContext()->getModuleScopeContext();
   assert(FileCtx == this && "D doesn't belong to this source file");
   if (FileCtx != this) {
     // D doesn't belong to this file. This shouldn't happen in practice.
     return None;
   }
-  if (D->getLoc().isInvalid())
+
+  SourceLoc Loc = D->getLoc(/*SerializedOK=*/false);
+  if (Loc.isInvalid())
     return None;
+
   SourceManager &SM = getASTContext().SourceMgr;
-  BasicDeclLocs Result;
-  Result.SourceFilePath = SM.getDisplayNameForLoc(D->getLoc());
+  auto BufferID = SM.findBufferContainingLoc(Loc);
 
-  for (const auto &SRC : D->getRawComment(/*SerializedOK*/false).Comments) {
-    auto LineAndCol = SM.getLineAndColumnInBuffer(SRC.Range.getStart());
-    Result.DocRanges.push_back(
-        std::make_pair(SourcePosition{LineAndCol.first, LineAndCol.second},
-                       SRC.Range.getByteLength()));
-  }
+  ExternalSourceLocs::RawLocs Result;
+  auto setLoc = [&](ExternalSourceLocs::RawLoc &RawLoc, SourceLoc Loc) {
+    if (!Loc.isValid())
+      return;
 
-  auto setLineColumn = [&SM](SourcePosition &Home, SourceLoc Loc) {
-    if (Loc.isValid()) {
-      std::tie(Home.Line, Home.Column) = SM.getPresumedLineAndColumnForLoc(Loc);
-    }
+    RawLoc.Offset = SM.getLocOffsetInBuffer(Loc, BufferID);
+    std::tie(RawLoc.Line, RawLoc.Column) = SM.getLineAndColumnInBuffer(Loc);
+
+    auto *VF = SM.getVirtualFile(Loc);
+    if (!VF)
+      return;
+
+    RawLoc.Directive.Offset =
+        SM.getLocOffsetInBuffer(VF->Range.getStart(), BufferID);
+    RawLoc.Directive.LineOffset = VF->LineOffset;
+    RawLoc.Directive.Length = VF->Range.getByteLength();
+    RawLoc.Directive.Name = StringRef(VF->Name);
   };
-#define SET(X) setLineColumn(Result.X, D->get##X());
-  SET(Loc)
-  SET(StartLoc)
-  SET(EndLoc)
-#undef SET
+
+  Result.SourceFilePath = SM.getIdentifierForBuffer(BufferID);
+  for (const auto &SRC : D->getRawComment(/*SerializedOK=*/false).Comments) {
+    Result.DocRanges.emplace_back(ExternalSourceLocs::RawLoc(),
+                                  SRC.Range.getByteLength());
+    setLoc(Result.DocRanges.back().first, SRC.Range.getStart());
+  }
+  setLoc(Result.Loc, D->getLoc(/*SerializedOK=*/false));
+  setLoc(Result.StartLoc, D->getStartLoc());
+  setLoc(Result.EndLoc, D->getEndLoc());
   return Result;
 }