[darwin] add a check for a matching python architecture

Author: charles-zablit

Sources/SwiftDriver/Driver/Driver.swift

@@ -873,12 +873,10 @@ public struct Driver {
       throw Error.subcommandPassedToDriver
     }
 
-    #if os(Windows)
-      if case .normal(isRepl: true) = invocationMode {
-        checkIfMatchingPythonArch(
-          cwd: ProcessEnv.cwd, envBlock: envBlock, diagnosticsEngine: diagnosticsEngine)
-      }
-    #endif
+    // if case .normal(isRepl: true) = invocationMode {
+      checkIfMatchingPythonArch(
+        cwd: ProcessEnv.cwd, envBlock: envBlock, diagnosticsEngine: diagnosticsEngine)
+    // }
 
     var args = args
     if let additional = env["ADDITIONAL_SWIFT_DRIVER_FLAGS"] {

Sources/SwiftDriver/Utilities/PythonArchitecture.swift

@@ -15,28 +15,40 @@ import func TSCBasic.lookupExecutablePath
 /// of the Python installation.
 ///
 /// When installing the x86 toolchain on ARM64 Windows, if the user does not
-/// install an x86 version of Python, they will get acryptic error message
+/// install an x86 version of Python, they will get a cryptic error message
 /// when running lldb (`0xC000007B`). Calling this function before invoking
-/// lldb gives them a warning to help troublshoot the issue.
+/// lldb gives them a warning to help troubleshoot the issue.
 ///
 /// - Parameters:
 ///   - cwd: The current working directory.
-///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
+///   - env: The parent shell's ProcessEnvironmentBlock.
 ///   - diagnosticsEngine: DiagnosticsEngine instance to use for printing the warning.
 public func checkIfMatchingPythonArch(
   cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, diagnosticsEngine: DiagnosticsEngine
 ) {
   #if arch(arm64)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.arm64
+    let toolchainArchitecture = ExecutableArchitecture.arm64
   #elseif arch(x86_64)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.x64
+    let toolchainArchitecture = ExecutableArchitecture.x64
   #elseif arch(x86)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.x86
+    let toolchainArchitecture = ExecutableArchitecture.x86
   #else
     return
   #endif
-  let pythonArchitecture = COFFBinaryExecutableArchitecture.readWindowsExecutableArchitecture(
-    cwd: cwd, envBlock: envBlock, filename: "python.exe")
+
+  #if os(Windows)
+    let pythonArchitecture = ExecutableArchitecture.readWindowsExecutableArchitecture(
+      cwd: cwd, envBlock: envBlock, filename: "python.exe")
+  #elseif os(macOS)
+    let pythonArchitecture = ExecutableArchitecture.readDarwinExecutableArchitecture(
+      cwd: cwd, envBlock: envBlock, filename: "python3")
+  #else
+    return
+  #endif
+
+  if pythonArchitecture == .universal {
+    return
+  }
 
   if toolchainArchitecture != pythonArchitecture {
     diagnosticsEngine.emit(
@@ -50,10 +62,11 @@ public func checkIfMatchingPythonArch(
 }
 
 /// Some of the architectures that can be stored in a COFF header.
-enum COFFBinaryExecutableArchitecture: String {
+enum ExecutableArchitecture: String {
   case x86 = "X86"
   case x64 = "X64"
   case arm64 = "ARM64"
+  case universal = "Universal"
   case unknown = "Unknown"
 
   static func fromPEMachineByte(machine: UInt16) -> Self {
@@ -66,59 +79,111 @@ enum COFFBinaryExecutableArchitecture: String {
     }
   }
 
-  /// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture
-  /// of the binary.
-  ///
-  /// - Parameters:
-  ///   - cwd: The current working directory.
-  ///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
-  ///   - filename: The name of the file we are resolving the architecture of.
-  /// - Returns: The architecture of the file which was found in the `Path`.
-  static func readWindowsExecutableArchitecture(
-    cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
-  ) -> Self {
-    let searchPaths = getEnvSearchPaths(pathString: envBlock["Path"], currentWorkingDirectory: cwd)
-    guard
-      let filePath = lookupExecutablePath(
-        filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
-    else {
-      return .unknown
-    }
-    guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
-      return .unknown
+  static func fromMachoCPUType(cpuType: Int32) -> Self {
+    // https://en.wikipedia.org/wiki/Mach-O
+    switch cpuType {
+    case 0x0100_0007: return .x86
+    case 0x0100_000c: return .arm64
+    default: return .unknown
     }
+  }
 
-    defer { fileHandle.closeFile() }
-
-    // Infering the architecture of a Windows executable from its COFF header involves the following:
-    // 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long).
-    // 2. Jump to that offset and read the next 6 bytes.
-    // 3. The first 4 are the signature which should be equal to 0x50450000.
-    // 4. The last 2 are the machine architecture which can be infered from the value we get.
-    //
-    // The link below provides a visualization of the COFF header and the process to get to it.
-    // https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg
-    fileHandle.seek(toFileOffset: 0x3C)
-    guard let offsetPointer = try? fileHandle.read(upToCount: 4),
-      offsetPointer.count == 4
-    else {
-      return .unknown
-    }
+  #if os(Windows)
+    /// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture
+    /// of the binary.
+    ///
+    /// - Parameters:
+    ///   - cwd: The current working directory.
+    ///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
+    ///   - filename: The name of the file we are resolving the architecture of.
+    /// - Returns: The architecture of the file which was found in the `Path`.
+    static func readWindowsExecutableArchitecture(
+      cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
+    ) -> Self {
+      let searchPaths = getEnvSearchPaths(
+        pathString: envBlock["Path"], currentWorkingDirectory: cwd)
+      guard
+        let filePath = lookupExecutablePath(
+          filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
+      else {
+        return .unknown
+      }
+      guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
+        return .unknown
+      }
 
-    let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
+      defer { fileHandle.closeFile() }
 
-    fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
-    guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
-      return .unknown
-    }
+      // Infering the architecture of a Windows executable from its COFF header involves the following:
+      // 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long).
+      // 2. Jump to that offset and read the next 6 bytes.
+      // 3. The first 4 are the signature which should be equal to 0x50450000.
+      // 4. The last 2 are the machine architecture which can be infered from the value we get.
+      //
+      // The link below provides a visualization of the COFF header and the process to get to it.
+      // https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg
+      fileHandle.seek(toFileOffset: 0x3C)
+      guard let offsetPointer = try? fileHandle.read(upToCount: 4),
+        offsetPointer.count == 4
+      else {
+        return .unknown
+      }
 
-    let signature = coffHeader.prefix(4)
-    let machineBytes = coffHeader.suffix(2)
+      let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
 
-    guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
-      return .unknown
+      fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
+      guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
+        return .unknown
+      }
+
+      let signature = coffHeader.prefix(4)
+      let machineBytes = coffHeader.suffix(2)
+
+      guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
+        return .unknown
+      }
+
+      return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
     }
+  #endif
 
-    return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
-  }
+  #if os(macOS)
+    static func readDarwinExecutableArchitecture(
+      cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
+    ) -> Self {
+      let magicNumber: UInt32 = 0xcafe_babe
+
+      let searchPaths = getEnvSearchPaths(
+        pathString: envBlock["PATH"], currentWorkingDirectory: cwd)
+      guard
+        let filePath = lookupExecutablePath(
+          filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
+      else {
+        return .unknown
+      }
+      guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
+        return .unknown
+      }
+
+      defer {
+        try? fileHandle.close()
+      }
+
+      // The first 4 bytes of a Mach-O header contain the magic number. We use it to determine if the binary is
+      // universal.
+      // https://github.com/apple/darwin-xnu/blob/main/EXTERNAL_HEADERS/mach-o/loader.h
+      let magicData = fileHandle.readData(ofLength: 4)
+      let magic = magicData.withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }
+
+      if magic == magicNumber {
+        return .universal
+      }
+
+      // If the binary is not universal, the next 4 bytes contain the CPU type.
+      fileHandle.seek(toFileOffset: 4)
+      let cpuTypeData = fileHandle.readData(ofLength: 4)
+      let cpuType = cpuTypeData.withUnsafeBytes { $0.load(as: Int32.self) }
+      return Self.fromMachoCPUType(cpuType: cpuType)
+    }
+  #endif
 }