diff --git a/Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift b/Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift
index 7e1ac8eab..aae05f9b4 100644
--- a/Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift
+++ b/Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift
@@ -585,7 +585,10 @@ extension Driver {
         // To match the legacy driver behavior, make sure we add the first input file
         // to the output file map if compiling without primary inputs (WMO), even
         // if there aren't any corresponding outputs.
-        entries[inputFiles[0].fileHandle] = [:]
+        guard let firstSourceInputHandle = inputFiles.first(where:{ $0.type == .swift })?.fileHandle  else {
+          fatalError("Formulating swift-frontend invocation without any input .swift files")
+        }
+        entries[firstSourceInputHandle] = [:]
       }
 
       for flaggedPair in flaggedInputOutputPairs {
diff --git a/Sources/SwiftDriver/Jobs/Planning.swift b/Sources/SwiftDriver/Jobs/Planning.swift
index 51a5e312e..bf74de97b 100644
--- a/Sources/SwiftDriver/Jobs/Planning.swift
+++ b/Sources/SwiftDriver/Jobs/Planning.swift
@@ -329,10 +329,11 @@ extension Driver {
     emitModuleTrace: Bool
   ) throws -> Job? {
     guard case .singleCompile = compilerMode,
-          inputFiles.allSatisfy({ $0.type.isPartOfSwiftCompilation })
+          inputFiles.contains(where: { $0.type.isPartOfSwiftCompilation })
     else { return nil }
 
-    if parsedOptions.hasArgument(.embedBitcode) {
+    if parsedOptions.hasArgument(.embedBitcode),
+       inputFiles.allSatisfy({ $0.type.isPartOfSwiftCompilation }) {
       let compile = try compileJob(primaryInputs: [],
                                    outputType: .llvmBitcode,
                                    addJobOutputs: addJobOutputs,
diff --git a/Tests/SwiftDriverTests/SwiftDriverTests.swift b/Tests/SwiftDriverTests/SwiftDriverTests.swift
index 30807a43a..24241ed2f 100644
--- a/Tests/SwiftDriverTests/SwiftDriverTests.swift
+++ b/Tests/SwiftDriverTests/SwiftDriverTests.swift
@@ -2824,6 +2824,27 @@ final class SwiftDriverTests: XCTestCase {
     
   }
 
+  func testWMOWithNonSourceInput() throws {
+    var driver1 = try Driver(args: [
+      "swiftc", "-whole-module-optimization", "danger.o", "foo.swift", "bar.swift", "wibble.swift", "-module-name", "Test",
+      "-driver-filelist-threshold=0"
+    ])
+    let plannedJobs = try driver1.planBuild().removingAutolinkExtractJobs()
+    XCTAssertEqual(plannedJobs.count, 2)
+    let compileJob = plannedJobs[0]
+    XCTAssertEqual(compileJob.kind, .compile)
+    XCTAssert(compileJob.commandLine.contains(.flag("-supplementary-output-file-map")))
+    let argIdx = try XCTUnwrap(compileJob.commandLine.firstIndex(where: { $0 == .flag("-supplementary-output-file-map") }))
+    let supplOutputs = compileJob.commandLine[argIdx+1]
+    guard case let .path(path) = supplOutputs,
+          case let .fileList(_, fileList) = path,
+          case let .outputFileMap(outFileMap) = fileList else {
+      throw StringError("Unexpected argument for output file map")
+    }
+    let firstKey: String = try VirtualPath.lookup(XCTUnwrap(outFileMap.entries.keys.first)).description
+    XCTAssertEqual(firstKey, "foo.swift")
+  }
+
   func testDashDashPassingDownInput() throws {
     do {
       var driver = try Driver(args: ["swiftc", "-module-name=ThisModule", "-wmo", "-num-threads", "4", "-emit-module", "-o", "test.swiftmodule", "--", "main.swift", "multi-threaded.swift"])