[Source manager] Start tracking generated-source info in the source manager

Establish the relationship for generated sources, whether for macro
expansions or (via a small stretch) replacing function bodies with
other bodies, in the source manager itself. This makes the information
available for diagnostic rendering, and unifies a little bit of the
representation, although it isn't used for much yet.

Author: DougGregor

include/swift/AST/ASTNode.h

@@ -72,7 +72,13 @@ namespace swift {
     FUNC(Decl)
     FUNC(Pattern)
 #undef FUNC
-    
+
+    static inline ASTNode getFromOpaqueValue(void *ptr) {
+      ASTNode result;
+      result.Val = decltype(result.Val)::getFromOpaqueValue(ptr);
+      return result;
+    }
+
     SWIFT_DEBUG_DUMP;
     void dump(llvm::raw_ostream &OS, unsigned Indent = 0) const;
 

include/swift/AST/SourceFile.h

@@ -194,15 +194,6 @@ class SourceFile final : public FileUnit {
   friend ASTContext;
 
 public:
-  /// For source files created to hold the source code created by expanding
-  /// a macro, this is the AST node that describes the macro expansion.
-  ///
-  /// The source location of this AST node is the place in the source that
-  /// triggered the creation of the macro expansion whose resulting source
-  /// code is in this source file. This field is only valid when
-  /// the \c SourceFileKind is \c MacroExpansion.
-  const ASTNode macroExpansion;
-
   /// Appends the given declaration to the end of the top-level decls list. Do
   /// not add any additional uses of this function.
   void addTopLevelDecl(Decl *d);
@@ -327,8 +318,7 @@ class SourceFile final : public FileUnit {
   llvm::StringMap<SourceFilePathInfo> getInfoForUsedFilePaths() const;
 
   SourceFile(ModuleDecl &M, SourceFileKind K, Optional<unsigned> bufferID,
-             ParsingOptions parsingOpts = {}, bool isPrimary = false,
-             ASTNode macroExpansion = ASTNode());
+             ParsingOptions parsingOpts = {}, bool isPrimary = false);
 
   ~SourceFile();
 
@@ -489,6 +479,15 @@ class SourceFile final : public FileUnit {
     return BufferID;
   }
 
+  /// For source files created to hold the source code created by expanding
+  /// a macro, this is the AST node that describes the macro expansion.
+  ///
+  /// The source location of this AST node is the place in the source that
+  /// triggered the creation of the macro expansion whose resulting source
+  /// code is in this source file. This will only produce a non-null value when
+  /// the \c SourceFileKind is \c MacroExpansion.
+  ASTNode getMacroExpansion() const;
+
   /// When this source file is enclosed within another source file, for example
   /// because it describes a macro expansion, return the source file it was
   /// enclosed in.

include/swift/Basic/SourceManager.h

@@ -23,6 +23,40 @@
 
 namespace swift {
 
+/// Augments a buffer that was created specifically to hold generated source
+/// code with the reasons for it being generated.
+class GeneratedSourceInfo {
+public:
+  /// The kind of generated source code.
+  enum Kind {
+    /// The expansion of a macro.
+    MacroExpansion,
+
+    /// A new function body that is replacing an existing function body.
+    ReplacedFunctionBody,
+  } kind;
+
+  /// The buffer ID for the enclosing buffer, in which originalSourceRange
+  /// resides.
+  unsigned originalBufferID;
+
+  /// The source range in the enclosing buffer where this source was generated.
+  ///
+  /// This could point at a macro expansion or at the implicit location at
+  /// which source code was generated. Conceptually, one can think of the
+  /// buffer described by a \c GeneratedSource instance as replacing the
+  /// code in the \c originalSourceRange.
+  SourceRange originalSourceRange;
+
+  /// The source range in the generated-source buffer where the generated
+  /// code exists. This might be a subrange of the buffer containing the
+  /// generated source, but it will never be from a different buffer.
+  SourceRange generatedSourceRange;
+
+  /// The opaque pointer for an ASTNode.
+  void *astNode;
+};
+
 /// This class manages and owns source buffers.
 class SourceManager {
 public:
@@ -51,7 +85,10 @@ class SourceManager {
   /// to speed up stats.
   mutable llvm::DenseMap<StringRef, llvm::vfs::Status> StatusCache;
 
-  /// Holds replaced ranges. Keys are orignal ranges, and values are new ranges
+  /// Holds generated source information, indexed by the buffer ID.
+  llvm::DenseMap<unsigned, GeneratedSourceInfo> GeneratedSourceInfos;
+
+  /// Holds replaced ranges. Keys are original ranges, and values are new ranges
   /// in different buffers. This is used for code completion and ASTContext
   /// reusing compilation.
   llvm::DenseMap<SourceRange, SourceRange> ReplacedRanges;
@@ -108,9 +145,12 @@ class SourceManager {
   const llvm::DenseMap<SourceRange, SourceRange> &getReplacedRanges() const {
     return ReplacedRanges;
   }
-  void setReplacedRange(SourceRange Orig, SourceRange New) {
-    ReplacedRanges[Orig] = New;
-  }
+
+  /// Set the generated source information associated with a given buffer.
+  void setGeneratedSourceInfo(unsigned bufferID, GeneratedSourceInfo);
+
+  /// Retrieve the generated source information for the given buffer.
+  Optional<GeneratedSourceInfo> getGeneratedSourceInfo(unsigned bufferID) const;
 
   /// Record the starting source location of a regex literal.
   void recordRegexLiteralStartLoc(SourceLoc loc) {

lib/AST/ASTScopeCreation.cpp

@@ -248,7 +248,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
                                        ScopeCreator *scopeCreator)
     : SF(SF), scopeCreator(scopeCreator) {
   if (auto enclosingSF = SF->getEnclosingSourceFile()) {
-    SourceLoc parentLoc = SF->macroExpansion.getStartLoc();
+    SourceLoc parentLoc = SF->getMacroExpansion().getStartLoc();
     if (auto parentScope = findStartingScopeForLookup(enclosingSF, parentLoc)) {
       parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
     }

lib/AST/Module.cpp

@@ -877,10 +877,20 @@ void SourceFile::lookupClassMembers(ImportPath::Access accessPath,
   cache.lookupClassMembers(accessPath, consumer);
 }
 
-SourceFile *SourceFile::getEnclosingSourceFile() const {
+ASTNode SourceFile::getMacroExpansion() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;
 
+  auto genInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  return ASTNode::getFromOpaqueValue(genInfo.astNode);
+}
+
+SourceFile *SourceFile::getEnclosingSourceFile() const {
+  auto macroExpansion = getMacroExpansion();
+  if (!macroExpansion)
+    return nullptr;
+
   auto sourceLoc = macroExpansion.getStartLoc();
   return getParentModule()->getSourceFileContainingLocation(sourceLoc);
 }
@@ -3263,17 +3273,13 @@ ModuleDecl::computeFileIDMap(bool shouldDiagnose) const {
 
 SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
                        Optional<unsigned> bufferID,
-                       ParsingOptions parsingOpts, bool isPrimary,
-                       ASTNode macroExpansion)
+                       ParsingOptions parsingOpts, bool isPrimary)
     : FileUnit(FileUnitKind::Source, M), BufferID(bufferID ? *bufferID : -1),
-      ParsingOpts(parsingOpts), IsPrimary(isPrimary),
-      macroExpansion(macroExpansion), Kind(K) {
+      ParsingOpts(parsingOpts), IsPrimary(isPrimary), Kind(K) {
   M.getASTContext().addDestructorCleanup(*this);
 
   assert(!IsPrimary || M.isMainModule() &&
          "A primary cannot appear outside the main module");
-  assert(macroExpansion.isNull() == (K != SourceFileKind::MacroExpansion) &&
-         "Macro expansions always need an expansion node");
 
   if (isScriptMode()) {
     bool problem = M.registerEntryPointFile(this, SourceLoc(), None);

lib/Basic/SourceLoc.cpp

@@ -252,6 +252,32 @@ StringRef SourceManager::extractText(CharSourceRange Range,
                        Range.getByteLength());
 }
 
+void SourceManager::setGeneratedSourceInfo(
+    unsigned bufferID, GeneratedSourceInfo info
+) {
+  assert(GeneratedSourceInfos.count(bufferID) == 0);
+  GeneratedSourceInfos[bufferID] = info;
+
+  switch (info.kind) {
+  case GeneratedSourceInfo::MacroExpansion:
+    break;
+
+  case GeneratedSourceInfo::ReplacedFunctionBody:
+    // Keep track of the replaced range.
+    ReplacedRanges[info.originalSourceRange] = info.generatedSourceRange;
+    break;
+  }
+}
+
+Optional<GeneratedSourceInfo> SourceManager::getGeneratedSourceInfo(
+    unsigned bufferID
+) const {
+  auto known = GeneratedSourceInfos.find(bufferID);
+  if (known == GeneratedSourceInfos.end())
+    return None;
+  return known->second;
+}
+
 Optional<unsigned>
 SourceManager::findBufferContainingLocInternal(SourceLoc Loc) const {
   assert(Loc.isValid());

lib/IDETool/CompileInstance.cpp

@@ -189,7 +189,16 @@ void retypeCheckFunctionBody(AbstractFunctionDecl *func,
   SourceRange newRange{rangeStartLoc, rangeEndLoc};
 
   // Reset the body range of the function decl, and re-typecheck it.
-  origSM.setReplacedRange(func->getOriginalBodySourceRange(), newRange);
+  origSM.setGeneratedSourceInfo(
+      sliceBufferID,
+      GeneratedSourceInfo{
+        GeneratedSourceInfo::ReplacedFunctionBody,
+        *func->getParentSourceFile()->getBufferID(),
+        func->getOriginalBodySourceRange(),
+        newRange,
+        func
+      }
+  );
   func->setBodyToBeReparsed(newRange);
   (void)func->getTypecheckedBody();
 }

lib/IDETool/IDEInspectionInstance.cpp

@@ -342,7 +342,17 @@ bool IDEInspectionInstance::performCachedOperationIfPossible(
                              newBufferStart.getAdvancedLoc(newInfo.EndOffset));
 
     auto *AFD = cast<AbstractFunctionDecl>(DC);
-    SM.setReplacedRange(AFD->getOriginalBodySourceRange(), newBodyRange);
+    SM.setGeneratedSourceInfo(
+        newBufferID,
+        GeneratedSourceInfo{
+          GeneratedSourceInfo::ReplacedFunctionBody,
+          *AFD->getParentSourceFile()->getBufferID(),
+          AFD->getOriginalBodySourceRange(),
+          newBodyRange,
+          AFD
+        }
+    );
+
     AFD->setBodyToBeReparsed(newBodyRange);
     oldSF->clearScope();
 

lib/Sema/TypeCheckMacros.cpp

@@ -220,13 +220,22 @@ Expr *swift::expandMacroExpr(
   auto macroBuffer =
       llvm::MemoryBuffer::getMemBufferCopy(evaluatedSource, bufferName);
   unsigned macroBufferID = sourceMgr.addNewSourceBuffer(std::move(macroBuffer));
+  auto macroBufferRange = sourceMgr.getRangeForBuffer(macroBufferID);
+  GeneratedSourceInfo sourceInfo{
+    GeneratedSourceInfo::MacroExpansion,
+    *sourceFile->getBufferID(),
+    expr->getSourceRange(),
+    SourceRange(macroBufferRange.getStart(), macroBufferRange.getEnd()),
+    ASTNode(expr).getOpaqueValue()
+  };
+  sourceMgr.setGeneratedSourceInfo(macroBufferID, sourceInfo);
   free((void*)evaluatedSource.data());
 
   // Create a source file to hold the macro buffer. This is automatically
   // registered with the enclosing module.
   auto macroSourceFile = new (ctx) SourceFile(
       *dc->getParentModule(), SourceFileKind::MacroExpansion, macroBufferID,
-      /*parsingOpts=*/{}, /*isPrimary=*/false, expr);
+      /*parsingOpts=*/{}, /*isPrimary=*/false);
 
   // Parse the expression.
   Parser parser(macroBufferID, *macroSourceFile, &ctx.Diags, nullptr, nullptr);