Add initial files

1 files changed, 210 insertions, 0 deletions

diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift
new file mode 100644
index 0000000..dff52f2
--- /dev/null
+++ b/Sources/PixelSort.swift
@@ -0,0 +1,210 @@
+import AppKit
+import ArgumentParser
+import Compute
+import Metal
+
+@main
+struct PixelSort: ParsableCommand {
+    static let configuration = CommandConfiguration(
+        abstract: "GPU-accelerated pixel sorting for glitch art"
+    )
+
+    @Argument(help: "Input image path")
+    var input: String
+
+    @Argument(help: "Output image path")
+    var output: String
+
+    @Option(name: .shortAndLong, help: "Sort key: brightness, hue, saturation, red, green, blue")
+    var key: SortKeyOption = .brightness
+
+    @Option(name: .shortAndLong, help: "Lower brightness threshold (0.0–1.0)")
+    var lower: Float = 0.1
+
+    @Option(name: .shortAndLong, help: "Upper brightness threshold (0.0–1.0)")
+    var upper: Float = 0.9
+
+    @Flag(name: .shortAndLong, help: "Sort descending")
+    var descending: Bool = false
+
+    mutating func run() throws {
+        let device = MTLCreateSystemDefaultDevice()!
+        let compute = try Compute(device: device)
+
+        // Load image into a Metal texture
+        let inputURL = URL(fileURLWithPath: input)
+        let outputURL = URL(fileURLWithPath: self.output)
+
+        guard let nsImage = NSImage(contentsOf: inputURL),
+              let cgImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil)
+        else {
+            throw ValidationError("Could not load image at \(input)")
+        }
+
+        let width = cgImage.width
+        let height = cgImage.height
+
+        let desc = MTLTextureDescriptor.texture2DDescriptor(
+            pixelFormat: .rgba8Unorm,
+            width: width,
+            height: height,
+            mipmapped: false
+        )
+        desc.usage = [.shaderRead, .shaderWrite]
+
+        let texA = device.makeTexture(descriptor: desc)!
+        let texB = device.makeTexture(descriptor: desc)!
+        texA.label = "texA"
+        texB.label = "texB"
+
+        // Upload pixel data
+        let bytesPerPixel = 4
+        let bytesPerRow = bytesPerPixel * width
+        var pixelData = [UInt8](repeating: 0, count: width * height * bytesPerPixel)
+        let colorSpace = CGColorSpaceCreateDeviceRGB()
+        guard let ctx = CGContext(
+            data: &pixelData,
+            width: width,
+            height: height,
+            bitsPerComponent: 8,
+            bytesPerRow: bytesPerRow,
+            space: colorSpace,
+            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
+        ) else {
+            throw ValidationError("Failed to create CGContext")
+        }
+        ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
+
+        texA.replace(
+            region: MTLRegionMake2D(0, 0, width, height),
+            mipmapLevel: 0,
+            withBytes: pixelData,
+            bytesPerRow: bytesPerRow
+        )
+
+        // Load shaders
+        let shaderSource = try String(contentsOf: Bundle.module.url(forResource: "Shaders", withExtension: "metal")!, encoding: .utf8)
+        let library = ShaderLibrary.source(shaderSource)
+
+        var copyPipeline = try compute.makePipeline(function: library.copyTexture)
+        var sortPipeline = try compute.makePipeline(function: library.bitonicSortStep)
+
+        // Bitonic sort requires log2(nextPow2(width)) outer passes
+        let n = nextPowerOf2(width)
+
+        // Ping-pong between texA and texB
+        var readTex = texA
+        var writeTex = texB
+
+        var blockSize: Int = 2
+        while blockSize <= n {
+            var subBlockSize = blockSize
+            while subBlockSize >= 2 {
+                // Copy readTex → writeTex so untouched pixels carry forward
+                copyPipeline.arguments.src = .texture(readTex)
+                copyPipeline.arguments.dst = .texture(writeTex)
+                try compute.run(pipeline: copyPipeline, width: width, height: height)
+
+                // Run the bitonic compare-swap step
+                var params = BitonicParams(
+                    width: UInt32(width),
+                    height: UInt32(height),
+                    blockSize: UInt32(blockSize),
+                    subBlockSize: UInt32(subBlockSize),
+                    sortKey: UInt32(key.metalValue),
+                    lowerThreshold: lower,
+                    upperThreshold: upper,
+                    descending: descending ? 1 : 0
+                )
+
+                sortPipeline.arguments.inputTexture = .texture(readTex)
+                sortPipeline.arguments.outputTexture = .texture(writeTex)
+
+                let paramBuffer = device.makeBuffer(bytes: &params, length: MemoryLayout<BitonicParams>.stride, options: .storageModeShared)!
+                sortPipeline.arguments.params = .buffer(paramBuffer)
+
+                try compute.run(pipeline: sortPipeline, width: width / 2, height: height)
+
+                // Swap
+                let tmp = readTex
+                readTex = writeTex
+                writeTex = tmp
+
+                subBlockSize /= 2
+            }
+            blockSize *= 2
+        }
+
+        // Read back from readTex (the last write destination after swap)
+        var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel)
+        readTex.getBytes(
+            &outputData,
+            bytesPerRow: bytesPerRow,
+            from: MTLRegionMake2D(0, 0, width, height),
+            mipmapLevel: 0
+        )
+
+        // Save output
+        guard let outCtx = CGContext(
+            data: &outputData,
+            width: width,
+            height: height,
+            bitsPerComponent: 8,
+            bytesPerRow: bytesPerRow,
+            space: colorSpace,
+            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
+        ),
+        let outCGImage = outCtx.makeImage()
+        else {
+            throw ValidationError("Failed to create output image")
+        }
+
+        let nsOutImage = NSBitmapImageRep(cgImage: outCGImage)
+        let ext = outputURL.pathExtension.lowercased()
+        let fileType: NSBitmapImageRep.FileType = ext == "jpg" || ext == "jpeg" ? .jpeg : .png
+        guard let data = nsOutImage.representation(using: fileType, properties: [:]) else {
+            throw ValidationError("Failed to encode output image")
+        }
+        try data.write(to: outputURL)
+
+        print("Pixel-sorted image saved to \(self.output) (\(width)×\(height))")
+    }
+}
+
+// MARK: - Helpers
+
+struct BitonicParams {
+    var width: UInt32
+    var height: UInt32
+    var blockSize: UInt32
+    var subBlockSize: UInt32
+    var sortKey: UInt32
+    var lowerThreshold: Float
+    var upperThreshold: Float
+    var descending: UInt32
+}
+
+func nextPowerOf2(_ n: Int) -> Int {
+    var v = n - 1
+    v |= v >> 1
+    v |= v >> 2
+    v |= v >> 4
+    v |= v >> 8
+    v |= v >> 16
+    return v + 1
+}
+
+enum SortKeyOption: String, ExpressibleByArgument, CaseIterable {
+    case brightness, hue, saturation, red, green, blue
+
+    var metalValue: Int {
+        switch self {
+        case .brightness:  return 0
+        case .hue:         return 1
+        case .saturation:  return 2
+        case .red:         return 3
+        case .green:       return 4
+        case .blue:        return 5
+        }
+    }
+}