bleam 0.0.7 → 0.0.9

package/templates/native/scripts/build/mlx-frameworks.ts

@@ -0,0 +1,421 @@
+import { createHash } from 'node:crypto'
+import { execFileSync } from 'node:child_process'
+import {
+  cpSync,
+  existsSync,
+  mkdirSync,
+  readdirSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from 'node:fs'
+import path from 'node:path'
+import { pathToFileURL } from 'node:url'
+
+type MLXFrameworkAsset = {
+  name: string
+  url: string
+  sha256: string
+}
+
+export const MLX_SWIFT_VERSION = '0.31.6'
+export const MLX_CATALYST_LIBRARY_IDENTIFIER = 'ios-arm64_x86_64-maccatalyst'
+export const MLX_MACOS_LIBRARY_IDENTIFIER = 'macos-arm64_x86_64'
+export const SWIFT_NUMERICS_VERSION = '1.1.1'
+
+const SWIFT_NUMERICS_ASSET: MLXFrameworkAsset = {
+  name: `swift-numerics-${SWIFT_NUMERICS_VERSION}.zip`,
+  url: `https://github.com/apple/swift-numerics/archive/refs/tags/${SWIFT_NUMERICS_VERSION}.zip`,
+  sha256: 'f37ab2a80b6507ad2af7f564d7481ab72eb1f7a70692238de12ba442443462e0',
+}
+
+const SWIFT_NUMERICS_SWIFT_MODULES = ['RealModule', 'ComplexModule', 'Numerics']
+const SWIFT_NUMERICS_ARCHES = [
+  { arch: 'arm64', triple: 'arm64-apple-ios-macabi' },
+  { arch: 'x86_64', triple: 'x86_64-apple-ios-macabi' },
+]
+const SWIFT_NUMERICS_MACOS_ARCHES = [
+  { arch: 'arm64', triple: 'arm64-apple-macos' },
+  { arch: 'x86_64', triple: 'x86_64-apple-macos' },
+]
+
+export const MLX_SWIFT_ASSETS: MLXFrameworkAsset[] = [
+  {
+    name: 'Cmlx.xcframework.zip',
+    url: 'https://github.com/ml-explore/mlx-swift/releases/download/0.31.6/Cmlx.xcframework.zip',
+    sha256: 'a202bf1dcfe1e64404adabfeb5eb363332e3a6221d18e4289ca0663fa3ab86c9',
+  },
+  {
+    name: 'MLX.xcframework.zip',
+    url: 'https://github.com/ml-explore/mlx-swift/releases/download/0.31.6/MLX.xcframework.zip',
+    sha256: '697f45a2ec2c0c312d9aa041a2b21f9c3706f18276823854dd96396118004c25',
+  },
+  {
+    name: 'MLXNN.xcframework.zip',
+    url: 'https://github.com/ml-explore/mlx-swift/releases/download/0.31.6/MLXNN.xcframework.zip',
+    sha256: 'afe3c8fa7c95e343266897b42bc7b312c3b7f4baac50ead955161e288eb51d79',
+  },
+  {
+    name: 'MLXOptimizers.xcframework.zip',
+    url: 'https://github.com/ml-explore/mlx-swift/releases/download/0.31.6/MLXOptimizers.xcframework.zip',
+    sha256: '6273556851fdd0c869b1b245917062acbec73f50a08944b77c4f6206400df5c1',
+  },
+]
+
+export function mlxFrameworkCacheRoot(rootDir: string) {
+  return path.join(rootDir, 'build', 'Vendor', 'MLXSwift', MLX_SWIFT_VERSION)
+}
+
+export function mlxCatalystFrameworksRoot(rootDir: string) {
+  return path.join(
+    rootDir,
+    'ios',
+    'Vendor',
+    'MLXSwift',
+    MLX_SWIFT_VERSION,
+  )
+}
+
+export function mlxMacOSFrameworksRoot(rootDir: string) {
+  return path.join(
+    rootDir,
+    'ios',
+    'Vendor',
+    'MLXSwift',
+    `${MLX_SWIFT_VERSION}-macos`,
+  )
+}
+
+function assetFrameworkName(asset: MLXFrameworkAsset) {
+  return asset.name.replace(/\.xcframework\.zip$/, '')
+}
+
+function sha256File(filePath: string) {
+  return createHash('sha256').update(readFileSync(filePath)).digest('hex')
+}
+
+function downloadAsset(asset: MLXFrameworkAsset, destination: string) {
+  if (existsSync(destination) && sha256File(destination) === asset.sha256) {
+    return
+  }
+
+  rmSync(destination, { force: true })
+  execFileSync('curl', ['--fail', '--location', '--output', destination, asset.url], {
+    stdio: 'inherit',
+  })
+
+  const actual = sha256File(destination)
+  if (actual !== asset.sha256) {
+    rmSync(destination, { force: true })
+    throw new Error(
+      `Checksum mismatch for ${asset.name}: expected ${asset.sha256}, got ${actual}`,
+    )
+  }
+}
+
+function findExtractedXCFramework(extractRoot: string, frameworkName: string) {
+  const direct = path.join(extractRoot, `${frameworkName}.xcframework`)
+  if (existsSync(direct)) {
+    return direct
+  }
+
+  for (const entry of readdirSync(extractRoot)) {
+    const candidate = path.join(extractRoot, entry, `${frameworkName}.xcframework`)
+    if (existsSync(candidate)) {
+      return candidate
+    }
+  }
+
+  throw new Error(`Could not find ${frameworkName}.xcframework in ${extractRoot}`)
+}
+
+function extractCatalystFramework(asset: MLXFrameworkAsset, zipPath: string, rootDir: string) {
+  extractFramework(asset, zipPath, rootDir, MLX_CATALYST_LIBRARY_IDENTIFIER, mlxCatalystFrameworksRoot(rootDir))
+}
+
+function extractMacOSFramework(asset: MLXFrameworkAsset, zipPath: string, rootDir: string) {
+  extractFramework(asset, zipPath, rootDir, MLX_MACOS_LIBRARY_IDENTIFIER, mlxMacOSFrameworksRoot(rootDir))
+}
+
+function extractFramework(
+  asset: MLXFrameworkAsset,
+  zipPath: string,
+  rootDir: string,
+  libraryIdentifier: string,
+  destinationRoot: string,
+) {
+  const frameworkName = assetFrameworkName(asset)
+  const cacheRoot = mlxFrameworkCacheRoot(rootDir)
+  const extractRoot = path.join(cacheRoot, 'Extracted', frameworkName)
+  const destination = path.join(destinationRoot, `${frameworkName}.framework`)
+
+  rmSync(extractRoot, { recursive: true, force: true })
+  mkdirSync(extractRoot, { recursive: true })
+  execFileSync('ditto', ['-x', '-k', zipPath, extractRoot])
+
+  const xcframework = findExtractedXCFramework(extractRoot, frameworkName)
+  const framework = path.join(
+    xcframework,
+    libraryIdentifier,
+    `${frameworkName}.framework`,
+  )
+  if (!existsSync(framework)) {
+    throw new Error(
+      `${asset.name} does not contain ${libraryIdentifier}/${frameworkName}.framework`,
+    )
+  }
+
+  mkdirSync(destinationRoot, { recursive: true })
+  rmSync(destination, { recursive: true, force: true })
+  cpSync(framework, destination, { recursive: true })
+  normalizeVersionedFrameworkSymlinks(destination, frameworkName)
+}
+
+function findSwiftNumericsSourceRoot(extractRoot: string) {
+  const direct = path.join(extractRoot, `swift-numerics-${SWIFT_NUMERICS_VERSION}`)
+  if (existsSync(path.join(direct, 'Package.swift'))) {
+    return direct
+  }
+
+  for (const entry of readdirSync(extractRoot)) {
+    const candidate = path.join(extractRoot, entry)
+    if (existsSync(path.join(candidate, 'Package.swift'))) {
+      return candidate
+    }
+  }
+
+  throw new Error(`Could not find swift-numerics Package.swift in ${extractRoot}`)
+}
+
+function copySwiftNumericsModuleFramework(
+  moduleName: string,
+  moduleBuildRoot: string,
+  destinationRoot: string,
+  arches = SWIFT_NUMERICS_ARCHES,
+) {
+  const frameworkRoot = path.join(destinationRoot, `${moduleName}.framework`)
+  const moduleRoot = path.join(frameworkRoot, 'Modules', `${moduleName}.swiftmodule`)
+  rmSync(frameworkRoot, { recursive: true, force: true })
+  mkdirSync(moduleRoot, { recursive: true })
+
+  for (const { arch, triple } of arches) {
+    const sourceBase = path.join(
+      moduleBuildRoot,
+      `${moduleName}.build`,
+      'Objects-normal',
+      arch,
+      moduleName,
+    )
+
+    for (const extension of [
+      'swiftmodule',
+      'swiftinterface',
+      'private.swiftinterface',
+      'swiftdoc',
+      'abi.json',
+    ]) {
+      const source = `${sourceBase}.${extension}`
+      if (existsSync(source)) {
+        cpSync(source, path.join(moduleRoot, `${triple}.${extension}`))
+      }
+    }
+  }
+}
+
+function copySwiftNumericsShimsFramework(sourceRoot: string, destinationRoot: string) {
+  const frameworkRoot = path.join(destinationRoot, '_NumericsShims.framework')
+  const headersRoot = path.join(frameworkRoot, 'Headers')
+  const modulesRoot = path.join(frameworkRoot, 'Modules')
+  rmSync(frameworkRoot, { recursive: true, force: true })
+  mkdirSync(headersRoot, { recursive: true })
+  mkdirSync(modulesRoot, { recursive: true })
+  cpSync(
+    path.join(sourceRoot, 'Sources', '_NumericsShims', 'include', '_NumericsShims.h'),
+    path.join(headersRoot, '_NumericsShims.h'),
+  )
+  writeFileSync(
+    path.join(modulesRoot, 'module.modulemap'),
+    [
+      'framework module _NumericsShims {',
+      '  umbrella header "_NumericsShims.h"',
+      '  export *',
+      '}',
+      '',
+    ].join('\n'),
+  )
+}
+
+function ensureSwiftNumericsMacOSModules(rootDir: string) {
+  const macOSRoot = mlxMacOSFrameworksRoot(rootDir)
+  const numericsModule = path.join(
+    macOSRoot,
+    'Numerics.framework',
+    'Modules',
+    'Numerics.swiftmodule',
+    'arm64-apple-macos.swiftmodule',
+  )
+  if (existsSync(numericsModule)) {
+    return
+  }
+
+  const cacheRoot = mlxFrameworkCacheRoot(rootDir)
+  const downloadRoot = path.join(cacheRoot, 'Downloads')
+  const zipPath = path.join(downloadRoot, SWIFT_NUMERICS_ASSET.name)
+  downloadAsset(SWIFT_NUMERICS_ASSET, zipPath)
+
+  const extractRoot = path.join(cacheRoot, 'Extracted', `swift-numerics-${SWIFT_NUMERICS_VERSION}`)
+  rmSync(extractRoot, { recursive: true, force: true })
+  mkdirSync(extractRoot, { recursive: true })
+  execFileSync('ditto', ['-x', '-k', zipPath, extractRoot])
+
+  const sourceRoot = findSwiftNumericsSourceRoot(extractRoot)
+  const buildRoot = path.join(cacheRoot, 'Build', `swift-numerics-${SWIFT_NUMERICS_VERSION}-macos`)
+  const archivePath = path.join(buildRoot, 'Numerics-macos.xcarchive')
+  const derivedDataPath = path.join(buildRoot, 'DerivedData')
+  rmSync(buildRoot, { recursive: true, force: true })
+  mkdirSync(buildRoot, { recursive: true })
+
+  execFileSync(
+    'xcodebuild',
+    [
+      'archive',
+      '-scheme',
+      'Numerics',
+      '-destination',
+      'generic/platform=macOS',
+      '-archivePath',
+      archivePath,
+      '-derivedDataPath',
+      derivedDataPath,
+      'SKIP_INSTALL=NO',
+      'BUILD_LIBRARY_FOR_DISTRIBUTION=YES',
+    ],
+    { cwd: sourceRoot, stdio: 'inherit' },
+  )
+
+  const moduleBuildRoot = path.join(
+    derivedDataPath,
+    'Build',
+    'Intermediates.noindex',
+    'ArchiveIntermediates',
+    'Numerics',
+    'IntermediateBuildFilesPath',
+    'swift-numerics.build',
+    'Release',
+  )
+
+  mkdirSync(macOSRoot, { recursive: true })
+  copySwiftNumericsShimsFramework(sourceRoot, macOSRoot)
+  for (const moduleName of SWIFT_NUMERICS_SWIFT_MODULES) {
+    copySwiftNumericsModuleFramework(moduleName, moduleBuildRoot, macOSRoot, SWIFT_NUMERICS_MACOS_ARCHES)
+  }
+}
+
+function ensureSwiftNumericsCatalystModules(rootDir: string) {
+  const catalystRoot = mlxCatalystFrameworksRoot(rootDir)
+  const numericsModule = path.join(
+    catalystRoot,
+    'Numerics.framework',
+    'Modules',
+    'Numerics.swiftmodule',
+    'arm64-apple-ios-macabi.swiftmodule',
+  )
+  if (existsSync(numericsModule)) {
+    return
+  }
+
+  const cacheRoot = mlxFrameworkCacheRoot(rootDir)
+  const downloadRoot = path.join(cacheRoot, 'Downloads')
+  const zipPath = path.join(downloadRoot, SWIFT_NUMERICS_ASSET.name)
+  downloadAsset(SWIFT_NUMERICS_ASSET, zipPath)
+
+  const extractRoot = path.join(cacheRoot, 'Extracted', `swift-numerics-${SWIFT_NUMERICS_VERSION}`)
+  rmSync(extractRoot, { recursive: true, force: true })
+  mkdirSync(extractRoot, { recursive: true })
+  execFileSync('ditto', ['-x', '-k', zipPath, extractRoot])
+
+  const sourceRoot = findSwiftNumericsSourceRoot(extractRoot)
+  const buildRoot = path.join(cacheRoot, 'Build', `swift-numerics-${SWIFT_NUMERICS_VERSION}`)
+  const archivePath = path.join(buildRoot, 'Numerics-maccatalyst.xcarchive')
+  const derivedDataPath = path.join(buildRoot, 'DerivedData')
+  rmSync(buildRoot, { recursive: true, force: true })
+  mkdirSync(buildRoot, { recursive: true })
+
+  execFileSync(
+    'xcodebuild',
+    [
+      'archive',
+      '-scheme',
+      'Numerics',
+      '-destination',
+      'generic/platform=macOS,variant=Mac Catalyst',
+      '-archivePath',
+      archivePath,
+      '-derivedDataPath',
+      derivedDataPath,
+      'SKIP_INSTALL=NO',
+      'BUILD_LIBRARY_FOR_DISTRIBUTION=YES',
+    ],
+    { cwd: sourceRoot, stdio: 'inherit' },
+  )
+
+  const moduleBuildRoot = path.join(
+    derivedDataPath,
+    'Build',
+    'Intermediates.noindex',
+    'ArchiveIntermediates',
+    'Numerics',
+    'IntermediateBuildFilesPath',
+    'swift-numerics.build',
+    'Release-maccatalyst',
+  )
+
+  mkdirSync(catalystRoot, { recursive: true })
+  copySwiftNumericsShimsFramework(sourceRoot, catalystRoot)
+  for (const moduleName of SWIFT_NUMERICS_SWIFT_MODULES) {
+    copySwiftNumericsModuleFramework(moduleName, moduleBuildRoot, catalystRoot)
+  }
+}
+
+function replaceSymlink(linkPath: string, target: string) {
+  rmSync(linkPath, { recursive: true, force: true })
+  execFileSync('ln', ['-sfn', target, linkPath])
+}
+
+function normalizeVersionedFrameworkSymlinks(frameworkPath: string, frameworkName: string) {
+  const versionA = path.join(frameworkPath, 'Versions', 'A')
+  if (!existsSync(versionA)) {
+    return
+  }
+
+  replaceSymlink(path.join(frameworkPath, 'Versions', 'Current'), 'A')
+
+  for (const entry of [frameworkName, 'Modules', 'Resources', 'Headers', 'PrivateHeaders']) {
+    if (existsSync(path.join(versionA, entry))) {
+      replaceSymlink(path.join(frameworkPath, entry), path.join('Versions', 'Current', entry))
+    }
+  }
+}
+
+export function ensureMLXSwiftCatalystFrameworks(rootDir: string = process.cwd()) {
+  const cacheRoot = mlxFrameworkCacheRoot(rootDir)
+  const downloadRoot = path.join(cacheRoot, 'Downloads')
+  mkdirSync(downloadRoot, { recursive: true })
+
+  for (const asset of MLX_SWIFT_ASSETS) {
+    const zipPath = path.join(downloadRoot, asset.name)
+    downloadAsset(asset, zipPath)
+    extractCatalystFramework(asset, zipPath, rootDir)
+    extractMacOSFramework(asset, zipPath, rootDir)
+  }
+
+  ensureSwiftNumericsCatalystModules(rootDir)
+  ensureSwiftNumericsMacOSModules(rootDir)
+}
+
+if (
+  process.argv[1] &&
+  import.meta.url === pathToFileURL(path.resolve(process.argv[1])).href
+) {
+  ensureMLXSwiftCatalystFrameworks(process.cwd())
+}