Handle LLVM LTO option to perform thin and full link-time-optimization

Sources/SwiftDriver/Driver/Driver.swift

@@ -165,6 +165,9 @@ public struct Driver {
   /// The type of the primary output generated by the compiler.
   @_spi(Testing) public let compilerOutputType: FileType?
 
+  /// The type of the link-time-optimization we expect to perform.
+  @_spi(Testing) public let lto: LTOKind?
+
   /// The type of the primary output generated by the linker.
   @_spi(Testing) public let linkerOutputType: LinkOutputType?
 
@@ -376,6 +379,7 @@ public struct Driver {
       self.allSourcesFileList = nil
     }
 
+    self.lto = Self.ltoKind(&parsedOptions, diagnosticsEngine: diagnosticsEngine)
     // Figure out the primary outputs from the driver.
     (self.compilerOutputType, self.linkerOutputType) = Self.determinePrimaryOutputs(&parsedOptions, driverKind: driverKind, diagnosticsEngine: diagnosticEngine)
 
@@ -589,6 +593,18 @@ extension Driver {
   }
 }
 
+extension Driver {
+  private static func ltoKind(_ parsedOptions: inout ParsedOptions,
+                              diagnosticsEngine: DiagnosticsEngine) -> LTOKind? {
+    guard let arg = parsedOptions.getLastArgument(.lto)?.asSingle else { return nil }
+    guard let kind = LTOKind(rawValue: arg) else {
+      diagnosticsEngine.emit(.error_invalid_arg_value(arg: .lto, value: arg))
+      return nil
+    }
+    return kind
+  }
+}
+
 // MARK: - Response files.
 extension Driver {
   /// Tokenize a single line in a response file.
@@ -1034,6 +1050,7 @@ extension Driver {
     // By default, the driver does not link its output. However, this will be updated below.
     var compilerOutputType: FileType? = (driverKind == .interactive ? nil : .object)
     var linkerOutputType: LinkOutputType? = nil
+    let objectLikeFileType: FileType = parsedOptions.getLastArgument(.lto) != nil ? .llvmBitcode : .object
 
     if let outputOption = parsedOptions.getLast(in: .modes) {
       switch outputOption.option {
@@ -1042,11 +1059,11 @@ extension Driver {
           diagnosticsEngine.emit(.error_static_emit_executable_disallowed)
         }
         linkerOutputType = .executable
-        compilerOutputType = .object
+        compilerOutputType = objectLikeFileType
 
       case .emitLibrary:
         linkerOutputType = parsedOptions.hasArgument(.static) ? .staticLibrary : .dynamicLibrary
-        compilerOutputType = .object
+        compilerOutputType = objectLikeFileType
 
       case .emitObject, .c:
         compilerOutputType = .object

Sources/SwiftDriver/Driver/LinkKind.swift

@@ -21,3 +21,11 @@ public enum LinkOutputType {
   /// A static library (e.g., .a or .lib)
   case staticLibrary
 }
+
+/// Describes the kind of link-time-optimization we expect to perform.
+public enum LTOKind: String, Hashable {
+  /// Perform LLVM ThinLTO.
+  case llvmThin = "llvm-thin"
+  /// Perform LLVM full LTO.
+  case llvmFull = "llvm-full"
+}

Sources/SwiftDriver/Jobs/AutolinkExtractJob.swift

@@ -17,7 +17,7 @@ extension Driver {
     // pull the info we need from the .o files directly and pass them as an
     // argument input file to the linker.
     // FIXME: Also handle Cygwin and MinGW
-    guard inputs.count > 0 && targetTriple.objectFormat == .elf else {
+    guard inputs.count > 0 && targetTriple.objectFormat == .elf && lto == nil else {
       return nil
     }
 

Sources/SwiftDriver/Jobs/DarwinToolchain+LinkerSupport.swift

@@ -14,11 +14,11 @@ import TSCUtility
 import SwiftOptions
 
 extension DarwinToolchain {
-  internal func findARCLiteLibPath() throws -> AbsolutePath? {
+  internal func findXcodeClangLibPath(_ additionalPath: String) throws -> AbsolutePath? {
     let path = try getToolPath(.swiftCompiler)
       .parentDirectory // 'swift'
       .parentDirectory // 'bin'
-      .appending(components: "lib", "arc")
+      .appending(components: "lib", additionalPath)
 
     if fileSystem.exists(path) { return path }
 
@@ -28,11 +28,22 @@ extension DarwinToolchain {
       return clangPath
         .parentDirectory // 'clang'
         .parentDirectory // 'bin'
-        .appending(components: "lib", "arc")
+        .appending(components: "lib", additionalPath)
     }
     return nil
   }
 
+  internal func findARCLiteLibPath() throws -> AbsolutePath? {
+    return try findXcodeClangLibPath("arc")
+  }
+
+  internal func addLTOLibArgs(to commandLine: inout [Job.ArgTemplate]) throws {
+    if let path = try findXcodeClangLibPath("libLTO.dylib") {
+      commandLine.appendFlag("-lto_library")
+      commandLine.appendPath(path)
+    }
+  }
+
   func addLinkRuntimeLibraryRPath(
     to commandLine: inout [Job.ArgTemplate],
     parsedOptions: inout ParsedOptions,
@@ -183,6 +194,7 @@ extension DarwinToolchain {
     inputs: [TypedVirtualPath],
     outputFile: VirtualPath,
     shouldUseInputFileList: Bool,
+    lto: LTOKind?,
     sdkPath: String?,
     sanitizers: Set<Sanitizer>,
     targetInfo: FrontendTargetInfo
@@ -260,7 +272,7 @@ extension DarwinToolchain {
       try commandLine.appendAllArguments(.Xlinker, from: &parsedOptions)
 
     case .staticLibrary:
-      linkerTool = .staticLinker
+      linkerTool = .staticLinker(lto)
       commandLine.appendFlag(.static)
     }
 
@@ -269,6 +281,9 @@ extension DarwinToolchain {
       parsedOptions: &parsedOptions,
       targetTriple: targetTriple
     )
+
+    try addLTOLibArgs(to: &commandLine)
+
     let targetVariantTriple = targetInfo.targetVariant?.triple
     addDeploymentTargetArgs(
       to: &commandLine,
@@ -320,6 +335,8 @@ extension DarwinToolchain {
           inputModules.append(input.file)
         } else if input.type == .object {
           inputPaths.append(input.file)
+        } else if input.type == .llvmBitcode {
+          inputPaths.append(input.file)
         }
       }
       commandLine.appendPath(.fileList(path, .list(inputPaths)))
@@ -337,6 +354,8 @@ extension DarwinToolchain {
           return [.flag("-add_ast_path"), .path(path.file)]
         } else if path.type == .object {
           return [.path(path.file)]
+        } else if path.type == .llvmBitcode {
+          return [.path(path.file)]
         } else {
           return []
         }

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -177,6 +177,7 @@ extension Driver {
     try commandLine.appendLast(.localizationPath, from: &parsedOptions)
     try commandLine.appendLast(.requireExplicitAvailability, from: &parsedOptions)
     try commandLine.appendLast(.requireExplicitAvailabilityTarget, from: &parsedOptions)
+    try commandLine.appendLast(.lto, from: &parsedOptions)
     try commandLine.appendAll(.D, from: &parsedOptions)
     try commandLine.appendAll(.sanitizeEQ, from: &parsedOptions)
     try commandLine.appendAll(.debugPrefixMap, from: &parsedOptions)

Sources/SwiftDriver/Jobs/GenericUnixToolchain+LinkerSupport.swift

@@ -50,6 +50,7 @@ extension GenericUnixToolchain {
     inputs: [TypedVirtualPath],
     outputFile: VirtualPath,
     shouldUseInputFileList: Bool,
+    lto: LTOKind?,
     sdkPath: String?,
     sanitizers: Set<Sanitizer>,
     targetInfo: FrontendTargetInfo
@@ -70,6 +71,8 @@ extension GenericUnixToolchain {
       var linker: String?
       if let arg = parsedOptions.getLastArgument(.useLd) {
         linker = arg.asSingle
+      } else if lto != nil {
+        linker = "lld"
       } else {
         linker = defaultLinker(for: targetTriple)
       }
@@ -247,6 +250,15 @@ extension GenericUnixToolchain {
         commandLine.appendFlag("-u__llvm_profile_runtime")
       }
 
+      if let lto = lto {
+        switch lto {
+        case .llvmFull:
+          commandLine.appendFlag("-flto=full")
+        case .llvmThin:
+          commandLine.appendFlag("-flto=thin")
+        }
+      }
+
       // Run clang++ in verbose mode if "-v" is set
       try commandLine.appendLast(.v, from: &parsedOptions)
 
@@ -268,7 +280,7 @@ extension GenericUnixToolchain {
       commandLine.appendPath(outputFile)
 
       commandLine.append(contentsOf: inputs.map { .path($0.file) })
-      return try getToolPath(.staticLinker)
+      return try getToolPath(.staticLinker(lto))
     }
 
   }

Sources/SwiftDriver/Jobs/LinkJob.swift

@@ -45,6 +45,7 @@ extension Driver {
       inputs: inputs,
       outputFile: outputFile,
       shouldUseInputFileList: shouldUseInputFileList,
+      lto: lto,
       sdkPath: sdkPath,
       sanitizers: enabledSanitizers,
       targetInfo: frontendTargetInfo

Sources/SwiftDriver/Jobs/Planning.swift

@@ -103,6 +103,7 @@ extension Driver {
     func addLinkerInput(_ li: TypedVirtualPath) { linkerInputs.append(li) }
 
     var moduleInputs = [TypedVirtualPath]()
+    let acceptBitcodeAsLinkerInput = lto == .llvmThin || lto == .llvmFull
     func addModuleInput(_ mi: TypedVirtualPath) { moduleInputs.append(mi) }
     var moduleInputsFromJobOutputs = [TypedVirtualPath]()
     func addModuleInputFromJobOutputs(_ mis: TypedVirtualPath) {
@@ -113,7 +114,8 @@ extension Driver {
         switch jobOutput.type {
           case .object, .autolink:
             addLinkerInput(jobOutput)
-
+          case .llvmBitcode where acceptBitcodeAsLinkerInput:
+            addLinkerInput(jobOutput)
           case .swiftModule:
             addModuleInputFromJobOutputs(jobOutput)
 
@@ -245,7 +247,7 @@ extension Driver {
         addJob(job)
       }
 
-    case .object, .autolink:
+    case .object, .autolink, .llvmBitcode:
       if linkerOutputType != nil {
         addLinkerInput(input)
       } else {

Sources/SwiftDriver/Toolchains/GenericUnixToolchain.swift

@@ -54,8 +54,11 @@ import TSCBasic
     switch tool {
     case .swiftCompiler:
       return try lookup(executable: "swift-frontend")
-    case .staticLinker:
+    case .staticLinker(nil):
       return try lookup(executable: "ar")
+    case .staticLinker(.llvmFull),
+         .staticLinker(.llvmThin):
+      return try lookup(executable: "llvm-ar")
     case .dynamicLinker:
       // FIXME: This needs to look in the tools_directory first.
       return try lookup(executable: "clang")

Sources/SwiftDriver/Toolchains/Toolchain.swift

@@ -13,9 +13,9 @@ import Foundation
 import TSCBasic
 import SwiftOptions
 
-public enum Tool {
+public enum Tool: Hashable {
   case swiftCompiler
-  case staticLinker
+  case staticLinker(LTOKind?)
   case dynamicLinker
   case clang
   case swiftAutolinkExtract
@@ -66,6 +66,7 @@ public enum Tool {
     inputs: [TypedVirtualPath],
     outputFile: VirtualPath,
     shouldUseInputFileList: Bool,
+    lto: LTOKind?,
     sdkPath: String?,
     sanitizers: Set<Sanitizer>,
     targetInfo: FrontendTargetInfo

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -248,6 +248,11 @@ final class SwiftDriverTests: XCTestCase {
     let driver3 = try Driver(args: ["swiftc", "-static", "foo.swift", "-emit-library"])
     XCTAssertEqual(driver3.compilerOutputType, .object)
     XCTAssertEqual(driver3.linkerOutputType, .staticLibrary)
+
+    let driver4 = try Driver(args: ["swiftc", "-lto=llvm-thin", "foo.swift", "-emit-library"])
+    XCTAssertEqual(driver4.compilerOutputType, .llvmBitcode)
+    let driver5 = try Driver(args: ["swiftc", "-lto=llvm-full", "foo.swift", "-emit-library"])
+    XCTAssertEqual(driver5.compilerOutputType, .llvmBitcode)
   }
 
   func testPrimaryOutputKindsDiagnostics() throws {
@@ -851,6 +856,34 @@ final class SwiftDriverTests: XCTestCase {
       XCTAssertFalse(cmd.contains(.flag("-dylib")))
       XCTAssertFalse(cmd.contains(.flag("-shared")))
     }
+
+    do {
+      // lto linking
+      var driver1 = try Driver(args: commonArgs + ["-emit-executable", "-target", "x86_64-apple-macosx10.15", "-lto=llvm-thin"], env: env)
+      let plannedJobs1 = try driver1.planBuild()
+      XCTAssertFalse(plannedJobs1.contains(where: { $0.kind == .autolinkExtract }))
+      let linkJob1 = plannedJobs1.first(where: { $0.kind == .link })
+      XCTAssertTrue(linkJob1?.tool.name.contains("ld"))
+      XCTAssertTrue(linkJob1?.commandLine.contains(.flag("-lto_library")))
+
+      var driver2 = try Driver(args: commonArgs + ["-emit-executable", "-target", "x86_64-unknown-linux", "-lto=llvm-thin"], env: env)
+      let plannedJobs2 = try driver2.planBuild()
+      XCTAssertFalse(plannedJobs2.contains(where: { $0.kind == .autolinkExtract }))
+      let linkJob2 = plannedJobs2.first(where: { $0.kind == .link })
+      XCTAssertTrue(linkJob2?.tool.name.contains("clang"))
+      XCTAssertTrue(linkJob2?.commandLine.contains(.flag("-flto=thin")))
+
+      var driver3 = try Driver(args: commonArgs + ["-emit-executable", "-target", "x86_64-unknown-linux", "-lto=llvm-full"], env: env)
+      let plannedJobs3 = try driver3.planBuild()
+      XCTAssertFalse(plannedJobs3.contains(where: { $0.kind == .autolinkExtract }))
+
+      let compileJob3 = try XCTUnwrap(plannedJobs3.first(where: { $0.kind == .compile }))
+      XCTAssertTrue(compileJob3.outputs.contains { $0.file.basename.hasSuffix(".bc") })
+
+      let linkJob3 = try XCTUnwrap(plannedJobs3.first(where: { $0.kind == .link }))
+      XCTAssertTrue(linkJob3.tool.name.contains("clang"))
+      XCTAssertTrue(linkJob3.commandLine.contains(.flag("-flto=full")))
+    }
 
     #if os(macOS)
     // dsymutil won't be found on Linux.
@@ -2436,6 +2469,18 @@ final class SwiftDriverTests: XCTestCase {
     }
   }
 
+  func testLTOOption() throws {
+    XCTAssertEqual(try Driver(args: ["swiftc"]).lto, nil)
+
+    XCTAssertEqual(try Driver(args: ["swiftc", "-lto=llvm-thin"]).lto, .llvmThin)
+
+    XCTAssertEqual(try Driver(args: ["swiftc", "-lto=llvm-full"]).lto, .llvmFull)
+
+    try assertDriverDiagnostics(args: ["swiftc", "-lto=nop"]) { driver, verify in
+      verify.expect(.error("invalid value 'nop' in '-lto='"))
+    }
+  }
+
   func testScanDependenciesOption() throws {
     do {
       var driver = try Driver(args: ["swiftc", "-scan-dependencies", "foo.swift"])