[Serialization] Store offset of decls in .swiftsourceinfo

The locations stored in .swiftsourceinfo included the presumed file,
line, and column. When a location is requested it would read these, open
the external file, create a line map, and find the offset corresponding
to that line/column.

The offset is known during serialization though, so output it as well to
avoid having to read the file and generate the line map.

Since the serialized location is returned from `Decl::getLoc()`, it
should not be the presumed location. Instead, also output the line
directives so that the presumed location can be built as per normal
locations.

Finally, move the cache out of `Decl` and into `ASTContext`, since very
few declarations will actually have their locations deserialized. Make
sure to actually write to that cache so it's used - the old cache was
never written to.

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

@@ -864,37 +864,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;
 }