[Serialization] Store offset of decls in .swiftsourceinfo

Rather than calculate the offset from the line and column by reading the
source file, just store the offset to begin with.

Since the location is used through `getLoc()`, also store the *original*
location rather than the presumed. Add line directive information so
that the presumed location can be retrieved if needed.

Move the serialized location cache out of `Decl` and into `ASTContext`
since very few will actually have their locations deserialized. Actually
use that cache - the old one was never actually written to.

Author: bnbarham

include/swift/AST/ASTContext.h

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

include/swift/AST/Decl.h

@@ -62,6 +62,7 @@ namespace swift {
   class DynamicSelfType;
   class Type;
   class Expr;
+  struct ExternalDeclLocs;
   class CaptureListExpr;
   class DeclRefExpr;
   class ForeignAsyncConvention;
@@ -686,16 +687,8 @@ class alignas(1 << DeclAlignInBits) Decl {
 
   Decl(const Decl&) = delete;
   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;
+  SourceLoc getLocFromSource() const;
 
   /// Directly set the invalid bit
   void setInvalidBit();
@@ -807,6 +800,15 @@ class alignas(1 << DeclAlignInBits) Decl {
   /// in diagnostics.
   SourceLoc getLoc(bool SerializedOK = true) const;
 
+  /// Returns the serialized locations of this declaration from the
+  /// corresponding .swiftsourceinfo file. Empty if it is within the current
+  /// module and thus does not have any .swiftsourceinfo.
+  ///
+  /// By default the locations are simple file, offset, line, and column. To
+  /// resolve these into \c SourceLocs, set \p Resolve to \c true. Note that
+  /// this must load the external file, hence being off by default.
+  const ExternalDeclLocs &getSerializedLocs(bool Resolve = false) const;
+
   /// Returns the source range of the entire declaration.
   SourceRange getSourceRange() const;
 

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<BasicDeclPositions>
+  getBasicPositionsForDecl(const Decl *D) const {
     return None;
   }
 

include/swift/AST/SourceFile.h

@@ -407,7 +407,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<BasicDeclPositions>
+  getBasicPositionsForDecl(const Decl *D) const override;
 
   /// Returns the synthesized file for this source file, if it exists.
   SynthesizedFileUnit *getSynthesizedFile() const { return SynthesizedFile; };

include/swift/Basic/BasicSourceInfo.h

@@ -15,27 +15,48 @@
 
 #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 LocationDirective {
+  uint32_t Offset = 0;
+  int32_t LineOffset = 0;
+  uint32_t Length = 0;
+  StringRef Name;
+
+  bool isValid() const { return Length > 0; }
+};
+
 struct SourcePosition {
+  uint32_t Offset = 0;
   uint32_t Line = 0;
   uint32_t Column = 0;
+  LocationDirective Directive;
+
   bool isValid() const { return Line && Column; }
 };
 
-struct BasicDeclLocs {
+struct BasicDeclPositions {
   StringRef SourceFilePath;
   SmallVector<std::pair<SourcePosition, uint32_t>, 4> DocRanges;
   SourcePosition Loc;
   SourcePosition StartLoc;
   SourcePosition EndLoc;
 };
 
+struct ExternalDeclLocs {
+  unsigned BufferID = 0;
+  SourceLoc Loc;
+  SmallVector<CharSourceRange, 4> DocRanges;
+  BasicDeclPositions *Positions = nullptr;
+  bool Resolved = false;
+};
+
 class BasicSourceFileInfo {
   /// If this is non-null, fields other than 'FilePath' hasn't been populated.
   /// The 'getInt()' part indicates this instance is constructed with a

include/swift/Basic/SourceManager.h

@@ -22,6 +22,13 @@
 
 namespace swift {
 
+// \c #sourceLocation directive handling.
+struct VirtualFile {
+  CharSourceRange Range;
+  std::string Name;
+  int LineOffset;
+};
+
 /// This class manages and owns source buffers.
 class SourceManager {
   llvm::SourceMgr LLVMSourceMgr;
@@ -52,12 +59,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 +144,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 +280,21 @@ 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.
+  ///
+  /// A virtual file represents the source after a \c #sourceLocation (or
+  /// between two), providing a filename and line offset to be applied to the
+  /// given \p Loc.
   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 +303,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

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<BasicDeclPositions>
+  getBasicPositionsForDecl(const Decl *D) const override;
 
   void collectAllGroups(SmallVectorImpl<StringRef> &Names) const override;
 

lib/AST/ASTContext.cpp

@@ -284,6 +284,9 @@ struct ASTContext::Implementation {
   /// The module loader used to load Clang modules from DWARF.
   ClangModuleLoader *TheDWARFModuleLoader = nullptr;
 
+  /// Map from Swift declarations to deserialized locations
+  llvm::DenseMap<const Decl *, ExternalDeclLocs *> ExternalLocs;
+
   /// Map from Swift declarations to raw comments.
   llvm::DenseMap<const Decl *, std::pair<RawComment, bool>> RawComments;
 
@@ -2020,6 +2023,18 @@ ModuleDecl *ASTContext::getStdlibModule(bool loadIfAbsent) {
   return TheStdlibModule;
 }
 
+Optional<ExternalDeclLocs *> ASTContext::getExternalLocs(const Decl *D) {
+  auto Known = getImpl().ExternalLocs.find(D);
+  if (Known == getImpl().ExternalLocs.end())
+    return None;
+
+  return Known->second;
+}
+
+void ASTContext::setExternalLocs(const Decl *D, ExternalDeclLocs *Locs) {
+  getImpl().ExternalLocs[D] = Locs;
+}
+
 Optional<std::pair<RawComment, bool>> ASTContext::getRawComment(const Decl *D) {
   auto Known = getImpl().RawComments.find(D);
   if (Known == getImpl().RawComments.end())

lib/AST/Decl.cpp

@@ -610,37 +610,58 @@ case DeclKind::ID: return cast<ID##Decl>(this)->getLocFromSource();
   llvm_unreachable("Unknown decl kind");
 }
 
-const Decl::CachedExternalSourceLocs *Decl::getSerializedLocs() const {
-  if (CachedSerializedLocs) {
-    return CachedSerializedLocs;
-  }
-  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{};
-    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
+const ExternalDeclLocs &Decl::getSerializedLocs(bool Resolve) const {
+  auto &Context = getASTContext();
+  ExternalDeclLocs *Result = nullptr;
+  if (auto EL = Context.getExternalLocs(this)) {
+    Result = *EL;
+  } else {
+    static ExternalDeclLocs NullLocs{};
+    auto *File = cast<FileUnit>(getDeclContext()->getModuleScopeContext());
+    if (File->getKind() != FileUnitKind::SerializedAST)
+      return NullLocs;
+
+    auto Positions = File->getBasicPositionsForDecl(this);
+    if (!Positions.hasValue()) {
+      Result = &NullLocs;
+    } else {
+      Result = getASTContext().Allocate<ExternalDeclLocs>();
+      Result->Positions = getASTContext().AllocateObjectCopy(*Positions);
+    }
+    Context.setExternalLocs(this, Result);
+  }
+
+  // Not all callers need SourceLoc's. Creating them requires loading the file,
+  // so avoid that where we can.
+  if (Result->Positions && Resolve && !Result->Resolved) {
+    auto &SM = getASTContext().SourceMgr;
+    auto *Positions = Result->Positions;
+
+    auto ResolveLoc = [&](const SourcePosition &Pos) -> SourceLoc {
+      // If the decl had a presumed loc, create its virtual file so that
+      // getPresumedLineAndColForLoc works from serialized locations as well.
+      if (Pos.Directive.isValid()) {
+        auto &SL = Pos.Directive;
+        SourceLoc Loc = SM.getLocForOffset(Result->BufferID, SL.Offset);
+        SM.createVirtualFile(Loc, SL.Name, SL.LineOffset, SL.Length);
+      }
+      return SM.getLocForOffset(Result->BufferID, Pos.Offset);
+    };
 
-  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 });
-  }
+    Result->BufferID = SM.getExternalSourceBufferID(Positions->SourceFilePath);
 
-  return Result;
+    // File might not exist since it could have moved since .swiftsourceinfo
+    // was created.
+    if (Result->BufferID) {
+      Result->Loc = ResolveLoc(Positions->Loc);
+      for (auto &Range : Positions->DocRanges) {
+        Result->DocRanges.emplace_back(ResolveLoc(Range.first), Range.second);
+      }
+    }
+
+    Result->Resolved = true;
+  }
+  return *Result;
 }
 
 StringRef Decl::getAlternateModuleName() const {
@@ -683,7 +704,7 @@ static_assert(sizeof(checkSourceLocType(&ID##Decl::getLoc)) == 2, \
   case FileUnitKind::SerializedAST: {
     if (!SerializedOK)
       return SourceLoc();
-    return getSerializedLocs()->Loc;
+    return getSerializedLocs(/*Resolve=*/true).Loc;
   }
   case FileUnitKind::Builtin:
   case FileUnitKind::Synthesized: