[TBDGen] Split off symbol collection from TBD file generation

hamishknight · hamishknight · commit c5cd19a5cb02 · 2020-07-24T15:49:07.000-07:00
Refactor `TBDGenVisitor` to accept a callback for
when it discovers a symbol, and split off public
symbol gathering into `PublicSymbolsRequest` such
that we don't need to unnecessarily also build a
TBD file which we immediately throw away.
diff --git a/include/swift/AST/TBDGenRequests.h b/include/swift/AST/TBDGenRequests.h
@@ -21,6 +21,10 @@
 #include "swift/AST/SimpleRequest.h"
 
 namespace llvm {
+
+class DataLayout;
+class Triple;
+
 namespace MachO {
 class InterfaceFile;
 } // end namespace MachO
@@ -59,6 +63,9 @@ class TBDGenDescriptor final {
   /// Returns the TBDGen options.
   const TBDGenOptions &getOptions() const { return Opts; }
 
+  const llvm::DataLayout &getDataLayout() const;
+  const llvm::Triple &getTarget() const;
+
   bool operator==(const TBDGenDescriptor &other) const;
   bool operator!=(const TBDGenDescriptor &other) const {
     return !(*this == other);
@@ -77,13 +84,27 @@ llvm::hash_code hash_value(const TBDGenDescriptor &desc);
 void simple_display(llvm::raw_ostream &out, const TBDGenDescriptor &desc);
 SourceLoc extractNearestSourceLoc(const TBDGenDescriptor &desc);
 
-using TBDFileAndSymbols =
-    std::pair<llvm::MachO::InterfaceFile, std::vector<std::string>>;
+using TBDFile = llvm::MachO::InterfaceFile;
 
-/// Computes the TBD file and public symbols for a given module or file.
+/// Computes the TBD file for a given Swift module or file.
 class GenerateTBDRequest
     : public SimpleRequest<GenerateTBDRequest,
-                           TBDFileAndSymbols(TBDGenDescriptor),
+                           TBDFile(TBDGenDescriptor),
+                           RequestFlags::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  TBDFile evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
+};
+
+/// Retrieve the public symbols for a file or module.
+class PublicSymbolsRequest
+    : public SimpleRequest<PublicSymbolsRequest,
+                           std::vector<std::string>(TBDGenDescriptor),
                            RequestFlags::Uncached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -92,7 +113,8 @@ class GenerateTBDRequest
   friend SimpleRequest;
 
   // Evaluation.
-  TBDFileAndSymbols evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
+  std::vector<std::string>
+  evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
 };
 
 /// Report that a request of the given kind is being evaluated, so it
diff --git a/include/swift/AST/TBDGenTypeIDZone.def b/include/swift/AST/TBDGenTypeIDZone.def
@@ -14,5 +14,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-SWIFT_REQUEST(TBDGen, GenerateTBDRequest, TBDFileAndSymbols(TBDGenDescriptor),
+SWIFT_REQUEST(TBDGen, GenerateTBDRequest, TBDFile(TBDGenDescriptor),
+              Uncached, NoLocationInfo)
+SWIFT_REQUEST(TBDGen, PublicSymbolsRequest,
+              std::vector<std::string>(TBDGenDescriptor),
               Uncached, NoLocationInfo)
diff --git a/lib/TBDGen/TBDGen.cpp b/lib/TBDGen/TBDGen.cpp
@@ -61,21 +61,27 @@ static bool isGlobalOrStaticVar(VarDecl *VD) {
   return VD->isStatic() || VD->getDeclContext()->isModuleScopeContext();
 }
 
+TBDGenVisitor::TBDGenVisitor(TBDGenDescriptor desc,
+                             SymbolCallbackFn symbolCallback)
+    : TBDGenVisitor(desc.getTarget(), desc.getDataLayout(),
+                    desc.getParentModule(), desc.getOptions(),
+                    symbolCallback) {}
+
 void TBDGenVisitor::addSymbolInternal(StringRef name,
                                       llvm::MachO::SymbolKind kind,
                                       bool isLinkerDirective) {
   if (!isLinkerDirective && Opts.LinkerDirectivesOnly)
     return;
-  Symbols.addSymbol(kind, name, Targets);
-  if (kind == SymbolKind::GlobalSymbol) {
-    StringSymbols.push_back(name);
+
 #ifndef NDEBUG
+  if (kind == SymbolKind::GlobalSymbol) {
     if (!DuplicateSymbolChecker.insert(name).second) {
       llvm::dbgs() << "TBDGen duplicate symbol: " << name << '\n';
       assert(false && "TBDGen symbol appears twice");
     }
-#endif
   }
+#endif
+  SymbolCallback(name, kind);
 }
 
 static std::vector<OriginallyDefinedInAttr::ActiveVersion>
@@ -1157,21 +1163,15 @@ static bool isApplicationExtensionSafe(const LangOptions &LangOpts) {
          llvm::sys::Process::GetEnv("LD_APPLICATION_EXTENSION_SAFE");
 }
 
-TBDFileAndSymbols
-GenerateTBDRequest::evaluate(Evaluator &evaluator,
-                             TBDGenDescriptor desc) const {
+TBDFile GenerateTBDRequest::evaluate(Evaluator &evaluator,
+                                     TBDGenDescriptor desc) const {
   auto *M = desc.getParentModule();
   auto &opts = desc.getOptions();
-
   auto &ctx = M->getASTContext();
-  const auto &triple = ctx.LangOpts.Target;
-  UniversalLinkageInfo linkInfo(triple, opts.HasMultipleIGMs,
-                                /*forcePublicDecls*/ false);
 
   llvm::MachO::InterfaceFile file;
   file.setFileType(llvm::MachO::FileType::TBD_V4);
-  file.setApplicationExtensionSafe(
-    isApplicationExtensionSafe(M->getASTContext().LangOpts));
+  file.setApplicationExtensionSafe(isApplicationExtensionSafe(ctx.LangOpts));
   file.setInstallName(opts.InstallName);
   file.setTwoLevelNamespace();
   file.setSwiftABIVersion(irgen::getSwiftABIVersion());
@@ -1187,29 +1187,47 @@ GenerateTBDRequest::evaluate(Evaluator &evaluator,
     file.setCompatibilityVersion(*packed);
   }
 
-  llvm::MachO::Target target(triple);
+  llvm::MachO::Target target(ctx.LangOpts.Target);
   file.addTarget(target);
   // Add target variant
   if (ctx.LangOpts.TargetVariant.hasValue()) {
     llvm::MachO::Target targetVar(*ctx.LangOpts.TargetVariant);
     file.addTarget(targetVar);
   }
-  auto *clang = static_cast<ClangImporter *>(ctx.getClangModuleLoader());
-  TBDGenVisitor visitor(file, {target}, clang->getTargetInfo().getDataLayout(),
-                        linkInfo, M, opts);
+
+  llvm::MachO::TargetList targets{target};
+  auto addSymbol = [&](StringRef symbol, SymbolKind kind) {
+    file.addSymbol(kind, symbol, targets);
+  };
+
+  TBDGenVisitor visitor(desc, addSymbol);
+  visitor.visit(desc);
+  return file;
+}
+
+std::vector<std::string>
+PublicSymbolsRequest::evaluate(Evaluator &evaluator,
+                               TBDGenDescriptor desc) const {
+  std::vector<std::string> symbols;
+  auto addSymbol = [&](StringRef symbol, SymbolKind kind) {
+    if (kind == SymbolKind::GlobalSymbol)
+      symbols.push_back(symbol.str());
+  };
+
+  TBDGenVisitor visitor(desc, addSymbol);
   visitor.visit(desc);
-  return std::make_pair(std::move(file), std::move(visitor.StringSymbols));
+  return symbols;
 }
 
 std::vector<std::string> swift::getPublicSymbols(TBDGenDescriptor desc) {
   auto &evaluator = desc.getParentModule()->getASTContext().evaluator;
-  return llvm::cantFail(evaluator(GenerateTBDRequest{desc})).second;
+  return llvm::cantFail(evaluator(PublicSymbolsRequest{desc}));
 }
 void swift::writeTBDFile(ModuleDecl *M, llvm::raw_ostream &os,
                          const TBDGenOptions &opts) {
   auto &evaluator = M->getASTContext().evaluator;
   auto desc = TBDGenDescriptor::forModule(M, opts);
-  auto file = llvm::cantFail(evaluator(GenerateTBDRequest{desc})).first;
+  auto file = llvm::cantFail(evaluator(GenerateTBDRequest{desc}));
   llvm::cantFail(llvm::MachO::TextAPIWriter::writeToStream(os, file),
                  "YAML writing should be error-free");
 }
diff --git a/lib/TBDGen/TBDGenRequests.cpp b/lib/TBDGen/TBDGenRequests.cpp
@@ -11,11 +11,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/AST/TBDGenRequests.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/Evaluator.h"
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/Module.h"
+#include "swift/ClangImporter/ClangImporter.h"
 #include "swift/Subsystems.h"
 #include "swift/TBDGen/TBDGen.h"
+#include "clang/Basic/TargetInfo.h"
 #include "llvm/TextAPI/MachO/InterfaceFile.h"
 
 using namespace swift;
@@ -43,6 +46,16 @@ ModuleDecl *TBDGenDescriptor::getParentModule() const {
   return Input.get<FileUnit *>()->getParentModule();
 }
 
+const llvm::DataLayout &TBDGenDescriptor::getDataLayout() const {
+  auto &ctx = getParentModule()->getASTContext();
+  auto *clang = static_cast<ClangImporter *>(ctx.getClangModuleLoader());
+  return clang->getTargetInfo().getDataLayout();
+}
+
+const llvm::Triple &TBDGenDescriptor::getTarget() const {
+  return getParentModule()->getASTContext().LangOpts.Target;
+}
+
 bool TBDGenDescriptor::operator==(const TBDGenDescriptor &other) const {
   return Input == other.Input && Opts == other.Opts;
 }
diff --git a/lib/TBDGen/TBDGenVisitor.h b/lib/TBDGen/TBDGenVisitor.h
@@ -61,21 +61,21 @@ struct InstallNameStore {
 };
 
 class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
-public:
-  llvm::MachO::InterfaceFile &Symbols;
-  llvm::MachO::TargetList Targets;
-  std::vector<std::string> StringSymbols;
-  const llvm::DataLayout &DataLayout;
-
 #ifndef NDEBUG
   /// Tracks the symbols emitted to ensure we don't emit any duplicates.
   llvm::StringSet<> DuplicateSymbolChecker;
 #endif
 
-  const UniversalLinkageInfo &UniversalLinkInfo;
+  const llvm::DataLayout &DataLayout;
+  UniversalLinkageInfo UniversalLinkInfo;
   ModuleDecl *SwiftModule;
   const TBDGenOptions &Opts;
 
+  using SymbolKind = llvm::MachO::SymbolKind;
+  using SymbolCallbackFn = llvm::function_ref<void(StringRef, SymbolKind)>;
+
+  SymbolCallbackFn SymbolCallback;
+
   /// A set of original function and derivative configuration pairs for which
   /// derivative symbols have been emitted.
   ///
@@ -84,7 +84,6 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   llvm::DenseSet<std::pair<AbstractFunctionDecl *, AutoDiffConfig>>
       AddedDerivatives;
 
-private:
   std::vector<Decl*> DeclStack;
   std::unique_ptr<std::map<std::string, InstallNameStore>>
     previousInstallNameMap;
@@ -141,15 +140,18 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
                                   AutoDiffConfig config);
 
 public:
-  TBDGenVisitor(llvm::MachO::InterfaceFile &symbols,
-                llvm::MachO::TargetList targets,
-                const llvm::DataLayout &dataLayout,
-                const UniversalLinkageInfo &universalLinkInfo,
-                ModuleDecl *swiftModule, const TBDGenOptions &opts)
-      : Symbols(symbols), Targets(targets), DataLayout(dataLayout),
-        UniversalLinkInfo(universalLinkInfo), SwiftModule(swiftModule),
-        Opts(opts),
-        previousInstallNameMap(parsePreviousModuleInstallNameMap())  {}
+  TBDGenVisitor(const llvm::Triple &target, const llvm::DataLayout &dataLayout,
+                ModuleDecl *swiftModule, const TBDGenOptions &opts,
+                SymbolCallbackFn symbolCallback)
+      : DataLayout(dataLayout),
+        UniversalLinkInfo(target, opts.HasMultipleIGMs, /*forcePublic*/ false),
+        SwiftModule(swiftModule), Opts(opts), SymbolCallback(symbolCallback),
+        previousInstallNameMap(parsePreviousModuleInstallNameMap()) {}
+
+  /// Create a new visitor using the target and layout information from a
+  /// TBDGenDescriptor.
+  TBDGenVisitor(TBDGenDescriptor desc, SymbolCallbackFn symbolCallback);
+
   ~TBDGenVisitor() { assert(DeclStack.empty()); }
   void addMainIfNecessary(FileUnit *file) {
     // HACK: 'main' is a special symbol that's always emitted in SILGen if
