Merge pull request swiftlang#27746 from harlanhaskins/away-in-a-mangler

[TBDGen] Mangle symbols before putting them in the TBD

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/Process.h"
 #include "llvm/Support/YAMLTraits.h"
@@ -53,10 +56,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");
   }
@@ -648,7 +656,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
+}