TBDGen: add a flag for embedding external symbols in emitted tbd file

include/swift/AST/DiagnosticsFrontend.def

@@ -284,6 +284,9 @@ REMARK(platform_previous_install_name, none,
 ERROR(unknown_platform_name, none,
       "unkown platform name %0", (StringRef))
 
+ERROR(unknown_swift_module_name, none,
+      "cannot find Swift module with name %0", (StringRef))
+
 ERROR(cannot_find_install_name, none,
       "cannot find previous install name for module %0 in %1", (StringRef, StringRef))
 

include/swift/Option/Options.td

@@ -344,6 +344,9 @@ def emit_tbd_path : Separate<["-"], "emit-tbd-path">,
 def emit_tbd_path_EQ : Joined<["-"], "emit-tbd-path=">,
   Flags<[FrontendOption, NoInteractiveOption, ArgumentIsPath]>,
   Alias<emit_tbd_path>;
+def embed_tbd_for_module : Separate<["-"], "embed-tbd-for-module">,
+  Flags<[FrontendOption]>,
+  HelpText<"Embed symbols from the module in the emitted tbd file">;
 
 def serialize_diagnostics : Flag<["-"], "serialize-diagnostics">,
   Flags<[FrontendOption, NoInteractiveOption, DoesNotAffectIncrementalBuild]>,

include/swift/TBDGen/TBDGen.h

@@ -53,6 +53,12 @@ struct TBDGenOptions {
   /// The path to a Json file indicating the module name to install-name map
   /// used by @_originallyDefinedIn
   std::string ModuleInstallNameMapPath;
+
+  /// For these modules, TBD gen should embed their symbols in the emitted tbd
+  /// file.
+  /// Typically, these modules are static linked libraries. Thus their symbols
+  /// are embeded in the current dylib.
+  std::vector<std::string> embedSymbolsFromModules;
 };
 
 void enumeratePublicSymbols(FileUnit *module, llvm::StringSet<> &symbols,

lib/Driver/ToolChains.cpp

@@ -264,6 +264,9 @@ static void addCommonFrontendArgs(const ToolChain &TC, const OutputInfo &OI,
   // Pass through the values passed to -Xfrontend.
   inputArgs.AddAllArgValues(arguments, options::OPT_Xfrontend);
 
+  // Pass on module names whose symbols should be embeded in tbd.
+  inputArgs.AddAllArgs(arguments, options::OPT_embed_tbd_for_module);
+
   if (auto *A = inputArgs.getLastArg(options::OPT_working_directory)) {
     // Add -Xcc -working-directory before any other -Xcc options to ensure it is
     // overridden by an explicit -Xcc -working-directory, although having a

lib/Frontend/CompilerInvocation.cpp

@@ -1105,6 +1105,9 @@ static bool ParseTBDGenArgs(TBDGenOptions &Opts, ArgList &Args,
   if (const Arg *A = Args.getLastArg(OPT_previous_module_installname_map_file)) {
     Opts.ModuleInstallNameMapPath = A->getValue();
   }
+  for (auto A : Args.getAllArgValues(OPT_embed_tbd_for_module)) {
+    Opts.embedSymbolsFromModules.push_back(StringRef(A).str());
+  }
   return false;
 }
 

lib/FrontendTool/FrontendTool.cpp

@@ -1435,11 +1435,14 @@ static bool validateTBDIfNeeded(const CompilerInvocation &Invocation,
   }
 
   const bool allSymbols = mode == FrontendOptions::TBDValidationMode::All;
+  // We should ignore embeded symbols from external modules for validation.
+  TBDGenOptions Opts = Invocation.getTBDGenOptions();
+  Opts.embedSymbolsFromModules.clear();
   return MSF.is<SourceFile *>()
              ? validateTBD(MSF.get<SourceFile *>(), IRModule,
-                           Invocation.getTBDGenOptions(), allSymbols)
+                           Opts, allSymbols)
              : validateTBD(MSF.get<ModuleDecl *>(), IRModule,
-                           Invocation.getTBDGenOptions(), allSymbols);
+                           Opts, allSymbols);
 }
 
 static bool generateCode(const CompilerInvocation &Invocation,

lib/TBDGen/TBDGen.cpp

@@ -1005,9 +1005,25 @@ static void enumeratePublicSymbolsAndWrite(ModuleDecl *M, FileUnit *singleFile,
     assert(M == singleFile->getParentModule() && "mismatched file and module");
     visitFile(singleFile);
   } else {
-    for (auto *file : M->getFiles()) {
-      visitFile(file);
+    llvm::SmallVector<ModuleDecl*, 4> Modules;
+    Modules.push_back(M);
+    for (auto Name: opts.embedSymbolsFromModules) {
+      if (auto *MD = ctx.getModuleByName(Name)) {
+        // If it is a clang module, the symbols should be collected by TAPI.
+        if (!MD->isNonSwiftModule()) {
+          Modules.push_back(MD);
+          continue;
+        }
+      }
+      // Diagnose module name that cannot be found
+      ctx.Diags.diagnose(SourceLoc(), diag::unknown_swift_module_name, Name);
     }
+    // Collect symbols in each module.
+    llvm::for_each(Modules, [&](ModuleDecl *M) {
+      for (auto *file : M->getFiles()) {
+        visitFile(file);
+      }
+    });
   }
 
   if (os) {

test/TBD/embed-symbol.swift

@@ -0,0 +1,21 @@
+// REQUIRES: VENDOR=apple
+// RUN: %empty-directory(%t)
+
+// RUN: echo 'public class Foo {}' > %t/foo.swift
+// RUN: echo 'public class Bar {}' > %t/bar.swift
+// RUN: %target-swift-frontend -emit-module %t/foo.swift -emit-module-path %t/foo.swiftmodule
+// RUN: %target-swift-frontend -emit-module %t/bar.swift -emit-module-path %t/bar.swiftmodule
+// RUN: %target-swift-frontend -typecheck %s -emit-tbd -emit-tbd-path %t/flag-not-provided.tbd -I %t -module-name main
+// RUN: %FileCheck %s --check-prefix FLAG-NOT-PROVIDED < %t/flag-not-provided.tbd
+
+// RUN: %target-swift-frontend -typecheck %s -emit-tbd -emit-tbd-path %t/flag-provided.tbd -I %t -embed-tbd-for-module foo -embed-tbd-for-module bar -module-name main
+// RUN: %FileCheck %s --check-prefix FLAG-PROVIDED < %t/flag-provided.tbd
+
+import foo
+import bar
+
+// FLAG-NOT-PROVIDED-NOT: $s3bar3BarCMa
+// FLAG-NOT-PROVIDED-NOT: $s3foo3FooCMa
+
+// FLAG-PROVIDED: $s3bar3BarCMa
+// FLAG-PROVIDED: $s3foo3FooCMa