Allow specifying custom environment variables for jobs

Sources/SwiftDriver/Execution/JobExecutor.swift

@@ -333,8 +333,8 @@ class ExecuteJobRule: LLBuildRule {
   private func executeJob(_ engine: LLTaskBuildEngine) {
     let context = engine.jobExecutorContext
     let resolver = context.argsResolver
-    let env = context.env
     let job = key.job
+    let env = context.env.merging(job.extraEnvironment, uniquingKeysWith: { $1 })
 
     let value: DriverBuildValue
     var pid = 0

Sources/SwiftDriver/Jobs/InterpretJob.swift

@@ -36,14 +36,17 @@ extension Driver {
     // The immediate arguments must be last.
     try commandLine.appendLast(.DASHDASH, from: &parsedOptions)
 
-    // FIXME: Set [DY]LD_LIBRARY_PATH, DYLD_FRAMEWORK_PATH if needed.
+    let extraEnvironment = try toolchain.platformSpecificInterpreterEnvironmentVariables(
+      env: self.env, parsedOptions: &parsedOptions, sdkPath: self.sdkPath,
+      targetTriple: self.targetTriple)
 
     return Job(
       kind: .interpret,
       tool: .absolute(try toolchain.getToolPath(.swiftCompiler)),
       commandLine: commandLine,
       inputs:inputs,
-      outputs: []
+      outputs: [],
+      extraEnvironment: extraEnvironment
     )
   }
 }

Sources/SwiftDriver/Jobs/Job.swift

@@ -45,6 +45,9 @@ public struct Job: Codable, Equatable {
 
   /// The outputs produced by the job.
   public var outputs: [TypedVirtualPath]
+
+  /// Any extra environment variables which should be set while running the job.
+  public var extraEnvironment: [String: String]
 
   /// The kind of job.
   public var kind: Kind
@@ -55,14 +58,16 @@ public struct Job: Codable, Equatable {
     commandLine: [ArgTemplate],
     displayInputs: [TypedVirtualPath]? = nil,
     inputs: [TypedVirtualPath],
-    outputs: [TypedVirtualPath]
+    outputs: [TypedVirtualPath],
+    extraEnvironment: [String: String] = [:]
   ) {
     self.kind = kind
     self.tool = tool
     self.commandLine = commandLine
     self.displayInputs = displayInputs ?? []
     self.inputs = inputs
     self.outputs = outputs
+    self.extraEnvironment = extraEnvironment
   }
 }
 
@@ -80,6 +85,13 @@ extension Job: CustomStringConvertible {
       }
     }
 
+    if !self.extraEnvironment.isEmpty {
+      result += " #"
+      for (envVar, val) in extraEnvironment {
+        result += " \(envVar)=\(val)"
+      }
+    }
+
     return result
   }
 }

Sources/SwiftDriver/Jobs/Toolchain+InterpreterSupport.swift

@@ -0,0 +1,81 @@
+//===----- Toolchain+InterpreterSupport.swift - Swift Interpreter Support -===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+
+extension Toolchain {
+  func addPathEnvironmentVariableIfNeeded(
+    _ environmentVariable: String,
+    to env: inout [String : String],
+    currentEnv: [String: String],
+    option: Option,
+    parsedOptions: inout ParsedOptions,
+    extraPaths: [String] = []
+    ) {
+    let argPaths = parsedOptions.arguments(for: option)
+    guard !argPaths.isEmpty || !extraPaths.isEmpty else { return }
+
+    env[environmentVariable] = (argPaths.map({ $0.argument.asSingle }) +
+      extraPaths + [currentEnv[environmentVariable]]).compactMap{ $0 }.joined(separator: ":")
+  }
+}
+
+extension DarwinToolchain {
+  public func platformSpecificInterpreterEnvironmentVariables(
+    env: [String : String],
+    parsedOptions: inout ParsedOptions,
+    sdkPath: String?,
+    targetTriple: Triple) throws -> [String: String] {
+    var envVars: [String: String] = [:]
+
+    let runtimePaths = try runtimeLibraryPaths(
+      for: targetTriple,
+      parsedOptions: &parsedOptions,
+      sdkPath: sdkPath,
+      isShared: true
+    ).map { $0.pathString }
+
+    addPathEnvironmentVariableIfNeeded("DYLD_LIBRARY_PATH", to: &envVars,
+                                       currentEnv: env, option: .L,
+                                       parsedOptions: &parsedOptions,
+                                       extraPaths: runtimePaths)
+
+    addPathEnvironmentVariableIfNeeded("DYLD_FRAMEWORK_PATH", to: &envVars,
+                                       currentEnv: env, option: .F,
+                                       parsedOptions: &parsedOptions)
+
+    return envVars
+  }
+}
+
+extension GenericUnixToolchain {
+  public func platformSpecificInterpreterEnvironmentVariables(
+    env: [String : String],
+    parsedOptions: inout ParsedOptions,
+    sdkPath: String?,
+    targetTriple: Triple) throws -> [String: String] {
+    var envVars: [String: String] = [:]
+
+    let runtimePaths = try runtimeLibraryPaths(
+      for: targetTriple,
+      parsedOptions: &parsedOptions,
+      sdkPath: sdkPath,
+      isShared: true
+    ).map { $0.pathString }
+
+    addPathEnvironmentVariableIfNeeded("LD_LIBRARY_PATH", to: &envVars,
+                                       currentEnv: env, option: .L,
+                                       parsedOptions: &parsedOptions,
+                                       extraPaths: runtimePaths)
+
+    return envVars
+  }
+}

Sources/SwiftDriver/Toolchains/Toolchain.swift

@@ -59,6 +59,12 @@ public protocol Toolchain {
     targetTriple: Triple,
     isShared: Bool
   ) throws -> String
+
+  func platformSpecificInterpreterEnvironmentVariables(
+    env: [String: String],
+    parsedOptions: inout ParsedOptions,
+    sdkPath: String?,
+    targetTriple: Triple) throws -> [String: String]
 }
 
 extension Toolchain {

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -890,6 +890,12 @@ final class SwiftDriverTests: XCTestCase {
       XCTAssertTrue(job.commandLine.contains(.flag("foo")))
 
       XCTAssertFalse(job.commandLine.contains(.flag("--")))
+
+      #if os(macOS)
+      XCTAssertTrue(job.extraEnvironment.keys.contains("DYLD_LIBRARY_PATH"))
+      #elseif os(Linux)
+      XCTAssertTrue(job.extraEnvironment.keys.contains("LD_LIBRARY_PATH"))
+      #endif
     }
 
     do {
@@ -909,6 +915,30 @@ final class SwiftDriverTests: XCTestCase {
       XCTAssertTrue(job.commandLine.contains(.flag("args")))
       XCTAssertTrue(job.commandLine.contains(.flag("-for=foo")))
     }
+    #if os(macOS)
+    do {
+      var driver = try Driver(args: ["swift", "-L/path/to/lib", "-F/path/to/framework", "foo.swift"])
+      let plannedJobs = try driver.planBuild()
+      XCTAssertEqual(plannedJobs.count, 1)
+      let job = plannedJobs[0]
+      XCTAssertEqual(job.inputs.count, 1)
+      XCTAssertEqual(job.inputs[0].file, .relative(RelativePath("foo.swift")))
+      XCTAssertEqual(job.outputs.count, 0)
+      XCTAssertTrue(job.extraEnvironment.contains { $0 == "DYLD_LIBRARY_PATH" && $1.contains("/path/to/lib") })
+      XCTAssertTrue(job.extraEnvironment.contains { $0 == "DYLD_FRAMEWORK_PATH" && $1.contains("/path/to/framework") })
+    }
+    #elseif os(Linux)
+    do {
+      var driver = try Driver(args: ["swift", "-L/path/to/lib", "foo.swift"])
+      let plannedJobs = try driver.planBuild()
+      XCTAssertEqual(plannedJobs.count, 1)
+      let job = plannedJobs[0]
+      XCTAssertEqual(job.inputs.count, 1)
+      XCTAssertEqual(job.inputs[0].file, .relative(RelativePath("foo.swift")))
+      XCTAssertEqual(job.outputs.count, 0)
+      XCTAssertTrue(job.extraEnvironment.contains { $0 == "LD_LIBRARY_PATH" && $1.contains("/path/to/lib") })
+    }
+    #endif
   }
 
   func testTargetTriple() throws {