Introduce SymbolSourceMapRequest

include/swift/AST/ASTTypeIDZone.def

@@ -28,6 +28,7 @@ SWIFT_TYPEID(Requirement)
 SWIFT_TYPEID(ResilienceExpansion)
 SWIFT_TYPEID(FragileFunctionKind)
 SWIFT_TYPEID(TangentPropertyInfo)
+SWIFT_TYPEID(SymbolSourceMap)
 SWIFT_TYPEID(Type)
 SWIFT_TYPEID(TypePair)
 SWIFT_TYPEID(TypeWitnessAndDecl)

include/swift/AST/ASTTypeIDs.h

@@ -60,6 +60,7 @@ class Requirement;
 enum class ResilienceExpansion : unsigned;
 struct FragileFunctionKind;
 class SourceFile;
+class SymbolSourceMap;
 struct TangentPropertyInfo;
 class Type;
 class TypeAliasDecl;

include/swift/AST/TBDGenRequests.h

@@ -19,6 +19,9 @@
 
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/SimpleRequest.h"
+#include "swift/IRGen/Linking.h"
+#include "swift/SIL/SILDeclRef.h"
+#include "swift/TBDGen/TBDGen.h"
 
 namespace llvm {
 
@@ -37,12 +40,11 @@ namespace swift {
 
 class FileUnit;
 class ModuleDecl;
-struct TBDGenOptions;
 
 class TBDGenDescriptor final {
   using FileOrModule = llvm::PointerUnion<FileUnit *, ModuleDecl *>;
   FileOrModule Input;
-  const TBDGenOptions &Opts;
+  TBDGenOptions Opts;
 
   TBDGenDescriptor(FileOrModule input, const TBDGenOptions &opts)
       : Input(input), Opts(opts) {
@@ -62,6 +64,7 @@ class TBDGenDescriptor final {
 
   /// Returns the TBDGen options.
   const TBDGenOptions &getOptions() const { return Opts; }
+  TBDGenOptions &getOptions() { return Opts; }
 
   const llvm::DataLayout &getDataLayout() const;
   const llvm::Triple &getTarget() const;
@@ -117,6 +120,123 @@ class PublicSymbolsRequest
   evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
 };
 
+/// Describes the origin of a particular symbol, including the stage of
+/// compilation it is introduced, as well as information on what decl introduces
+/// it.
+class SymbolSource {
+public:
+  enum class Kind {
+    /// A symbol introduced when emitting a SIL decl.
+    SIL,
+
+    /// A symbol introduced when emitting LLVM IR.
+    IR,
+
+    /// A symbol used to customize linker behavior, introduced by TBDGen.
+    LinkerDirective,
+
+    /// A symbol with an unknown origin.
+    // FIXME: This should be eliminated.
+    Unknown
+  };
+  Kind kind;
+
+private:
+  union {
+    SILDeclRef silDeclRef;
+    irgen::LinkEntity irEntity;
+  };
+
+  explicit SymbolSource(SILDeclRef ref) : kind(Kind::SIL) {
+    silDeclRef = ref;
+  }
+  explicit SymbolSource(irgen::LinkEntity entity) : kind(Kind::IR) {
+    irEntity = entity;
+  }
+  explicit SymbolSource(Kind kind) : kind(kind) {
+    assert(kind == Kind::LinkerDirective || kind == Kind::Unknown);
+  }
+
+public:
+  static SymbolSource forSILDeclRef(SILDeclRef ref) {
+    return SymbolSource{ref};
+  }
+  static SymbolSource forIRLinkEntity(irgen::LinkEntity entity) {
+    return SymbolSource{entity};
+  }
+  static SymbolSource forLinkerDirective() {
+    return SymbolSource{Kind::LinkerDirective};
+  }
+  static SymbolSource forUnknown() {
+    return SymbolSource{Kind::Unknown};
+  }
+
+  bool isLinkerDirective() const {
+    return kind == Kind::LinkerDirective;
+  }
+
+  SILDeclRef getSILDeclRef() const {
+    assert(kind == Kind::SIL);
+    return silDeclRef;
+  }
+  irgen::LinkEntity getIRLinkEntity() const {
+    assert(kind == Kind::IR);
+    return irEntity;
+  }
+};
+
+/// Maps a symbol back to its source for lazy compilation.
+class SymbolSourceMap {
+  friend class SymbolSourceMapRequest;
+
+  using Storage = llvm::StringMap<SymbolSource>;
+  const Storage *storage;
+
+  explicit SymbolSourceMap(const Storage *storage) : storage(storage) {
+    assert(storage);
+  }
+
+public:
+  Optional<SymbolSource> find(StringRef symbol) const {
+    auto result = storage->find(symbol);
+    if (result == storage->end())
+      return None;
+    return result->second;
+  }
+
+  friend bool operator==(const SymbolSourceMap &lhs,
+                         const SymbolSourceMap &rhs) {
+    return lhs.storage == rhs.storage;
+  }
+  friend bool operator!=(const SymbolSourceMap &lhs,
+                         const SymbolSourceMap &rhs) {
+    return !(lhs == rhs);
+  }
+
+  friend void simple_display(llvm::raw_ostream &out, const SymbolSourceMap &) {
+    out << "(symbol storage map)";
+  }
+};
+
+/// Computes a map of symbols to their SymbolSource for a file or module.
+class SymbolSourceMapRequest
+    : public SimpleRequest<SymbolSourceMapRequest,
+                           SymbolSourceMap(TBDGenDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  SymbolSourceMap evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
 /// Report that a request of the given kind is being evaluated, so it
 /// can be recorded by the stats reporter.
 template <typename Request>

include/swift/AST/TBDGenTypeIDZone.def

@@ -19,3 +19,6 @@ SWIFT_REQUEST(TBDGen, GenerateTBDRequest, TBDFile(TBDGenDescriptor),
 SWIFT_REQUEST(TBDGen, PublicSymbolsRequest,
               std::vector<std::string>(TBDGenDescriptor),
               Uncached, NoLocationInfo)
+SWIFT_REQUEST(TBDGen, SymbolSourceMapRequest,
+              SymbolSourceMap(TBDGenDescriptor),
+              Cached, NoLocationInfo)

include/swift/TBDGen/TBDGen.h

@@ -14,7 +14,6 @@
 
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/StringSet.h"
-#include "swift/AST/TBDGenRequests.h"
 #include "swift/Basic/Version.h"
 #include <vector>
 
@@ -25,6 +24,7 @@ class raw_ostream;
 namespace swift {
 class FileUnit;
 class ModuleDecl;
+class TBDGenDescriptor;
 
 /// Options for controlling the exact set of symbols included in the TBD
 /// output.
@@ -38,6 +38,9 @@ struct TBDGenOptions {
   /// Only collect linker directive symbols.
   bool LinkerDirectivesOnly = false;
 
+  /// Whether to include only symbols with public linkage.
+  bool PublicSymbolsOnly = true;
+
   /// The install_name to use in the TBD file.
   std::string InstallName;
 
@@ -66,6 +69,7 @@ struct TBDGenOptions {
     return lhs.HasMultipleIGMs == rhs.HasMultipleIGMs &&
            lhs.IsInstallAPI == rhs.IsInstallAPI &&
            lhs.LinkerDirectivesOnly == rhs.LinkerDirectivesOnly &&
+           lhs.PublicSymbolsOnly == rhs.PublicSymbolsOnly &&
            lhs.InstallName == rhs.InstallName &&
            lhs.ModuleLinkName == rhs.ModuleLinkName &&
            lhs.CurrentVersion == rhs.CurrentVersion &&
@@ -82,8 +86,9 @@ struct TBDGenOptions {
     using namespace llvm;
     return hash_combine(
         opts.HasMultipleIGMs, opts.IsInstallAPI, opts.LinkerDirectivesOnly,
-        opts.InstallName, opts.ModuleLinkName, opts.CurrentVersion,
-        opts.CompatibilityVersion, opts.ModuleInstallNameMapPath,
+        opts.PublicSymbolsOnly, opts.InstallName, opts.ModuleLinkName,
+        opts.CurrentVersion, opts.CompatibilityVersion,
+        opts.ModuleInstallNameMapPath,
         hash_combine_range(opts.embedSymbolsFromModules.begin(),
                            opts.embedSymbolsFromModules.end()));
   }

lib/FrontendTool/FrontendTool.cpp

@@ -35,6 +35,7 @@
 #include "swift/AST/IRGenRequests.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/TBDGenRequests.h"
 #include "swift/AST/TypeRefinementContext.h"
 #include "swift/Basic/Dwarf.h"
 #include "swift/Basic/Edit.h"

lib/FrontendTool/TBD.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/TBDGenRequests.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Demangling/Demangle.h"
 #include "swift/Frontend/FrontendOptions.h"

lib/IRGen/IRGenRequests.cpp

@@ -16,8 +16,8 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/AST/TBDGenRequests.h"
 #include "swift/Subsystems.h"
-#include "swift/TBDGen/TBDGen.h"
 #include "llvm/IR/Module.h"
 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"