[IDE] Avoid relying on after-the-fact use of addFile

We shouldn't be attempting to append SourceFiles
to the module after-the-fact for syntactic macro
expansion, refactor things such that the SourceFile
is created alongside the ModuleDecl.

Author: hamishknight

include/swift/IDETool/SyntacticMacroExpansion.h

@@ -45,35 +45,32 @@ class SyntacticMacroExpansionInstance {
   DiagnosticEngine Diags{SourceMgr};
   std::unique_ptr<ASTContext> Ctx;
   ModuleDecl *TheModule = nullptr;
+  SourceFile *SF = nullptr;
   llvm::StringMap<MacroDecl *> MacroDecls;
 
-  /// Create 'SourceFile' using the buffer.
-  swift::SourceFile *getSourceFile(llvm::MemoryBuffer *inputBuf);
-
   /// Synthesize 'MacroDecl' AST object to use the expansion.
   swift::MacroDecl *
   getSynthesizedMacroDecl(swift::Identifier name,
                           const MacroExpansionSpecifier &expansion);
 
-  /// Expand single 'expansion' in SF.
-  void expand(swift::SourceFile *SF,
-                    const MacroExpansionSpecifier &expansion,
-                    SourceEditConsumer &consumer);
+  /// Expand single 'expansion'.
+  void expand(const MacroExpansionSpecifier &expansion,
+              SourceEditConsumer &consumer);
 
 public:
   SyntacticMacroExpansionInstance() {}
 
-  /// Setup the instance with \p args .
+  /// Setup the instance with \p args and a given \p inputBuf.
   bool setup(StringRef SwiftExecutablePath, ArrayRef<const char *> args,
+             llvm::MemoryBuffer *inputBuf,
              std::shared_ptr<PluginRegistry> plugins, std::string &error);
 
   ASTContext &getASTContext() { return *Ctx; }
 
-  /// Expand all macros in \p inputBuf and send the edit results to \p consumer.
-  /// Expansions are specified by \p expansions .
-  void expandAll(llvm::MemoryBuffer *inputBuf,
-                     ArrayRef<MacroExpansionSpecifier> expansions,
-                     SourceEditConsumer &consumer);
+  /// Expand all macros and send the edit results to \p consumer. Expansions are
+  /// specified by \p expansions .
+  void expandAll(ArrayRef<MacroExpansionSpecifier> expansions,
+                 SourceEditConsumer &consumer);
 };
 
 /// Manager object to vend 'SyntacticMacroExpansionInstance'.
@@ -86,9 +83,11 @@ class SyntacticMacroExpansion {
                           std::shared_ptr<PluginRegistry> Plugins)
       : SwiftExecutablePath(SwiftExecutablePath), Plugins(Plugins) {}
 
-  /// Get instance configured with the specified compiler arguments.
+  /// Get instance configured with the specified compiler arguments and
+  /// input buffer.
   std::shared_ptr<SyntacticMacroExpansionInstance>
-  getInstance(ArrayRef<const char *> args, std::string &error);
+  getInstance(ArrayRef<const char *> args, llvm::MemoryBuffer *inputBuf,
+              std::string &error);
 };
 
 } // namespace ide

lib/IDETool/SyntacticMacroExpansion.cpp

@@ -26,11 +26,13 @@ using namespace ide;
 
 std::shared_ptr<SyntacticMacroExpansionInstance>
 SyntacticMacroExpansion::getInstance(ArrayRef<const char *> args,
+                                     llvm::MemoryBuffer *inputBuf,
                                      std::string &error) {
   // Create and configure a new instance.
   auto instance = std::make_shared<SyntacticMacroExpansionInstance>();
 
-  bool failed = instance->setup(SwiftExecutablePath, args, Plugins, error);
+  bool failed =
+      instance->setup(SwiftExecutablePath, args, inputBuf, Plugins, error);
   if (failed)
     return nullptr;
 
@@ -39,7 +41,8 @@ SyntacticMacroExpansion::getInstance(ArrayRef<const char *> args,
 
 bool SyntacticMacroExpansionInstance::setup(
     StringRef SwiftExecutablePath, ArrayRef<const char *> args,
-    std::shared_ptr<PluginRegistry> plugins, std::string &error) {
+    llvm::MemoryBuffer *inputBuf, std::shared_ptr<PluginRegistry> plugins,
+    std::string &error) {
   SmallString<256> driverPath(SwiftExecutablePath);
   llvm::sys::path::remove_filename(driverPath);
   llvm::sys::path::append(driverPath, "swiftc");
@@ -73,37 +76,16 @@ bool SyntacticMacroExpansionInstance::setup(
   pluginLoader->setRegistry(plugins.get());
   Ctx->setPluginLoader(std::move(pluginLoader));
 
-  // Create a module where SourceFiles reside.
+  // Create the ModuleDecl and SourceFile.
   Identifier ID = Ctx->getIdentifier(invocation.getModuleName());
   TheModule = ModuleDecl::create(ID, *Ctx);
 
-  return false;
-}
-
-SourceFile *
-SyntacticMacroExpansionInstance::getSourceFile(llvm::MemoryBuffer *inputBuf) {
-
-  // If there is a SourceFile with the same name and the content, use it.
-  // Note that this finds the generated source file that was created in the
-  // previous expansion requests.
-  if (auto bufID =
-          SourceMgr.getIDForBufferIdentifier(inputBuf->getBufferIdentifier())) {
-    if (inputBuf->getBuffer() == SourceMgr.getEntireTextForBuffer(*bufID)) {
-      SourceLoc bufLoc = SourceMgr.getLocForBufferStart(*bufID);
-      if (SourceFile *existing =
-              TheModule->getSourceFileContainingLocation(bufLoc)) {
-        return existing;
-      }
-    }
-  }
-
-  // Otherwise, create a new SourceFile.
-  SourceFile *SF = new (getASTContext()) SourceFile(
-      *TheModule, SourceFileKind::Main, SourceMgr.addMemBufferCopy(inputBuf));
+  SF = new (*Ctx) SourceFile(*TheModule, SourceFileKind::Main,
+                             SourceMgr.addMemBufferCopy(inputBuf));
   SF->setImports({});
   TheModule->addFile(*SF);
 
-  return SF;
+  return false;
 }
 
 MacroDecl *SyntacticMacroExpansionInstance::getSynthesizedMacroDecl(
@@ -425,8 +407,7 @@ class MacroExpansionFinder : public ASTWalker {
 } // namespace
 
 void SyntacticMacroExpansionInstance::expand(
-    SourceFile *SF, const MacroExpansionSpecifier &expansion,
-    SourceEditConsumer &consumer) {
+    const MacroExpansionSpecifier &expansion, SourceEditConsumer &consumer) {
 
   // Find the expansion at 'expansion.offset'.
   MacroExpansionFinder expansionFinder(
@@ -475,13 +456,9 @@ void SyntacticMacroExpansionInstance::expand(
 }
 
 void SyntacticMacroExpansionInstance::expandAll(
-    llvm::MemoryBuffer *inputBuf, ArrayRef<MacroExpansionSpecifier> expansions,
+    ArrayRef<MacroExpansionSpecifier> expansions,
     SourceEditConsumer &consumer) {
-
-  // Create a source file.
-  SourceFile *SF = getSourceFile(inputBuf);
-
   for (const auto &expansion : expansions) {
-    expand(SF, expansion, consumer);
+    expand(expansion, consumer);
   }
 }

tools/SourceKit/lib/SwiftLang/SwiftSyntacticMacroExpansion.cpp

@@ -30,7 +30,7 @@ void SwiftLangSupport::expandMacroSyntactically(
     CategorizedEditsReceiver receiver) {
 
   std::string error;
-  auto instance = SyntacticMacroExpansions->getInstance(args, error);
+  auto instance = SyntacticMacroExpansions->getInstance(args, inputBuf, error);
   if (!instance) {
     return receiver(
         RequestResult<ArrayRef<CategorizedEdits>>::fromError(error));
@@ -85,6 +85,6 @@ void SwiftLangSupport::expandMacroSyntactically(
   }
 
   RequestRefactoringEditConsumer consumer(receiver);
-  instance->expandAll(inputBuf, expansions, consumer);
+  instance->expandAll(expansions, consumer);
   // consumer automatically send the results on destruction.
 }