[TBDGen] Mangle symbols before putting them in the TBD

The linker expects to see mangled symbols in the TBD, otherwise it won't
be able to link anything. Use LLVM's mangler to mangle them.

Fixes rdar://54055049

Author: Harlan Haskins

lib/FrontendTool/TBD.cpp

@@ -25,6 +25,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/Mangler.h"
 #include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/Support/FileSystem.h"
 #include <vector>
@@ -85,7 +86,14 @@ static bool validateSymbolSet(DiagnosticEngine &diags,
   std::vector<StringRef> irNotTBD;
 
   for (auto &nameValue : IRModule.getValueSymbolTable()) {
-    auto name = nameValue.getKey();
+    // TBDGen inserts mangled names (usually with a leading '_') into its
+    // symbol table, so make sure to mangle IRGen names before comparing them
+    // with what TBDGen created.
+    auto unmangledName = nameValue.getKey();
+    SmallString<128> name;
+    llvm::Mangler::getNameWithPrefix(name, unmangledName,
+                                     IRModule.getDataLayout());
+
     auto value = nameValue.getValue();
     if (auto GV = dyn_cast<llvm::GlobalValue>(value)) {
       // Is this a symbol that should be listed?

lib/TBDGen/TBDGen.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/Basic/LLVM.h"
+#include "swift/ClangImporter/ClangImporter.h"
 #include "swift/IRGen/IRGenPublic.h"
 #include "swift/IRGen/Linking.h"
 #include "swift/SIL/FormalLinkage.h"
@@ -32,7 +33,9 @@
 #include "swift/SIL/SILWitnessTable.h"
 #include "swift/SIL/SILWitnessVisitor.h"
 #include "swift/SIL/TypeLowering.h"
+#include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/IR/Mangler.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/TextAPI/MachO/InterfaceFile.h"
@@ -52,10 +55,15 @@ static bool isGlobalOrStaticVar(VarDecl *VD) {
 }
 
 void TBDGenVisitor::addSymbol(StringRef name, SymbolKind kind) {
-  Symbols.addSymbol(kind, name, Archs);
+  // The linker expects to see mangled symbol names in TBD files, so make sure
+  // to mangle before inserting the symbol.
+  SmallString<32> mangled;
+  llvm::Mangler::getNameWithPrefix(mangled, name, DataLayout);
+
+  Symbols.addSymbol(kind, mangled, Archs);
 
   if (StringSymbols && kind == SymbolKind::GlobalSymbol) {
-    auto isNewValue = StringSymbols->insert(name).second;
+    auto isNewValue = StringSymbols->insert(mangled).second;
     (void)isNewValue;
     assert(isNewValue && "symbol appears twice");
   }
@@ -640,7 +648,10 @@ static void enumeratePublicSymbolsAndWrite(ModuleDecl *M, FileUnit *singleFile,
   file.addArch(arch);
   file.setPlatform(getPlatformKind(target));
 
-  TBDGenVisitor visitor(file, arch, symbols, linkInfo, M, opts);
+  auto *clang = static_cast<ClangImporter *>(ctx.getClangModuleLoader());
+  TBDGenVisitor visitor(file, arch, symbols,
+                        clang->getTargetInfo().getDataLayout(),
+                        linkInfo, M, opts);
 
   auto visitFile = [&](FileUnit *file) {
     if (file == M->getFiles()[0]) {

lib/TBDGen/TBDGenVisitor.h

@@ -34,6 +34,10 @@
 using namespace swift::irgen;
 using StringSet = llvm::StringSet<>;
 
+namespace llvm {
+class DataLayout;
+}
+
 namespace swift {
 
 struct TBDGenOptions;
@@ -45,6 +49,7 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   llvm::MachO::InterfaceFile &Symbols;
   llvm::MachO::ArchitectureSet Archs;
   StringSet *StringSymbols;
+  const llvm::DataLayout &DataLayout;
 
   const UniversalLinkageInfo &UniversalLinkInfo;
   ModuleDecl *SwiftModule;
@@ -72,11 +77,12 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
 public:
   TBDGenVisitor(llvm::MachO::InterfaceFile &symbols,
                 llvm::MachO::ArchitectureSet archs, StringSet *stringSymbols,
+                const llvm::DataLayout &dataLayout,
                 const UniversalLinkageInfo &universalLinkInfo,
                 ModuleDecl *swiftModule, const TBDGenOptions &opts)
       : Symbols(symbols), Archs(archs), StringSymbols(stringSymbols),
-        UniversalLinkInfo(universalLinkInfo), SwiftModule(swiftModule),
-        Opts(opts) {}
+        DataLayout(dataLayout), UniversalLinkInfo(universalLinkInfo),
+        SwiftModule(swiftModule), Opts(opts) {}
 
   void addMainIfNecessary(FileUnit *file) {
     // HACK: 'main' is a special symbol that's always emitted in SILGen if

test/TBD/Inputs/api_grab_bag.swift

@@ -0,0 +1,31 @@
+public class PublicClass {
+  public func method() {
+  }
+
+  public init() {
+  }
+}
+
+public class PublicSubclass: PublicClass {
+  public override func method() {
+  }
+}
+
+public protocol PublicProtocol {
+  var publicStruct: PublicStruct { get }
+}
+
+public struct PublicStruct {
+  public init() {}
+}
+
+extension PublicStruct: PublicProtocol {
+  public var publicStruct: PublicStruct { return self }
+}
+
+public enum PublicEnum: PublicProtocol {
+  case caseOne
+  case caseTwo
+
+  public var publicStruct: PublicStruct { return PublicStruct() }
+}

test/TBD/linking-with-tbd.swift

@@ -0,0 +1,40 @@
+// REQUIRES: VENDOR=apple
+
+// 1. Create a skeleton of a framework
+// RUN: %empty-directory(%t/APIGrabBag.framework/Modules/APIGrabBag.swiftmodule)
+// RUN: %empty-directory(%t/APIGrabBag.framework/Headers)
+
+// 1. Compile api_grab_bag.swift to a .tbd and put it in %t
+
+// RUN: %target-swift-frontend -emit-module -o %t/APIGrabBag.framework/Modules/APIGrabBag.swiftmodule/%target-cpu.swiftmodule -emit-tbd-path %t/APIGrabBag.framework/APIGrabBag.tbd %S/Inputs/api_grab_bag.swift -module-name APIGrabBag -tbd-install_name %t/APIGrabBag.framework/APIGrabBag
+
+// 2. Compile the current file against the TBD
+
+// RUN: %target-build-swift -emit-executable %s -o %t/executable -F %t -framework APIGrabBag
+
+// 3. Install the actual dylib into the framework
+
+// RUN: %target-build-swift -emit-library %S/Inputs/api_grab_bag.swift -module-name APIGrabBag -o %t/APIGrabBag.framework/APIGrabBag
+
+// 4. Codesign the executable and run it
+
+// RUN: %target-codesign %t/executable %t/APIGrabBag.framework/APIGrabBag
+// RUN: %target-run %t/executable
+
+import APIGrabBag
+
+func useAPIs() {
+  let c = PublicClass()
+  c.method()
+
+  let sub = PublicSubclass()
+  sub.method()
+
+  let s = PublicStruct()
+  let t = s.publicStruct
+
+  var e = PublicEnum.caseOne
+  e = .caseTwo
+
+  _ = e.publicStruct
+}