Make bitonic sorting on spans

1 files changed, 107 insertions, 64 deletions

diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift
index dff52f2..fa1bee6 100644
--- a/Sources/PixelSort.swift
+++ b/Sources/PixelSort.swift
@@ -24,9 +24,18 @@ struct PixelSort: ParsableCommand {
     @Option(name: .shortAndLong, help: "Upper brightness threshold (0.0–1.0)")
     var upper: Float = 0.9
 
-    @Flag(name: .shortAndLong, help: "Sort descending")
+    @Flag(name: .shortAndLong, help: "Sort descending (reverse)")
     var descending: Bool = false
 
+    @Option(name: .shortAndLong, help: "Gamma applied to sorted pixels (Unity-style composite)")
+    var gamma: Float = 1.0
+
+    @Option(name: .shortAndLong, help: "Clamp maximum sortable span length (default: image width)")
+    var maxSpan: Int?
+
+    @Flag(name: .long, help: "Invert the threshold mask")
+    var invertMask: Bool = false
+
     mutating func run() throws {
         let device = MTLCreateSystemDefaultDevice()!
         let compute = try Compute(device: device)
@@ -44,18 +53,51 @@ struct PixelSort: ParsableCommand {
         let width = cgImage.width
         let height = cgImage.height
 
-        let desc = MTLTextureDescriptor.texture2DDescriptor(
+        let rgbaDesc = MTLTextureDescriptor.texture2DDescriptor(
             pixelFormat: .rgba8Unorm,
             width: width,
             height: height,
             mipmapped: false
         )
-        desc.usage = [.shaderRead, .shaderWrite]
+        rgbaDesc.usage = [.shaderRead, .shaderWrite]
+
+        let texA = device.makeTexture(descriptor: rgbaDesc)!
+        let texB = device.makeTexture(descriptor: rgbaDesc)!
+        texA.label = "original"
+        texB.label = "output"
 
-        let texA = device.makeTexture(descriptor: desc)!
-        let texB = device.makeTexture(descriptor: desc)!
-        texA.label = "texA"
-        texB.label = "texB"
+        let sortedTex = device.makeTexture(descriptor: rgbaDesc)!
+        sortedTex.label = "sorted"
+
+        let maskDesc = MTLTextureDescriptor.texture2DDescriptor(
+            pixelFormat: .r8Uint,
+            width: width,
+            height: height,
+            mipmapped: false
+        )
+        maskDesc.usage = [.shaderRead, .shaderWrite]
+        let maskTex = device.makeTexture(descriptor: maskDesc)!
+        maskTex.label = "mask"
+
+        let spanDesc = MTLTextureDescriptor.texture2DDescriptor(
+            pixelFormat: .r32Uint,
+            width: width,
+            height: height,
+            mipmapped: false
+        )
+        spanDesc.usage = [.shaderRead, .shaderWrite]
+        let spanTex = device.makeTexture(descriptor: spanDesc)!
+        spanTex.label = "spans"
+
+        let valDesc = MTLTextureDescriptor.texture2DDescriptor(
+            pixelFormat: .r16Float,
+            width: width,
+            height: height,
+            mipmapped: false
+        )
+        valDesc.usage = [.shaderRead, .shaderWrite]
+        let valTex = device.makeTexture(descriptor: valDesc)!
+        valTex.label = "sortValues"
 
         // Upload pixel data
         let bytesPerPixel = 4
@@ -86,58 +128,68 @@ struct PixelSort: ParsableCommand {
         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)
+        var createMaskPipeline = try compute.makePipeline(function: library.createMask)
+        var clearSpanPipeline = try compute.makePipeline(function: library.clearSpanBuffer)
+        var identifySpansPipeline = try compute.makePipeline(function: library.identifySpans)
+        var rgbToValPipeline = try compute.makePipeline(function: library.rgbToSortValue)
+        var pixelSortPipeline = try compute.makePipeline(function: library.pixelSortSpan)
+        var compositePipeline = try compute.makePipeline(function: library.composite)
 
-        // 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)
+        var params = Params(
+            width: UInt32(width),
+            height: UInt32(height),
+            sortKey: UInt32(key.metalValue),
+            lowerThreshold: lower,
+            upperThreshold: upper,
+            reverseSorting: descending ? 1 : 0,
+            gamma: gamma,
+            maxSpanLength: UInt32(maxSpan ?? width),
+            invertMask: invertMask ? 1 : 0
+        )
+        let paramBuffer = device.makeBuffer(bytes: &params, length: MemoryLayout<Params>.stride, options: .storageModeShared)!
 
-                // 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
-                )
+        // 1) Mask
+        createMaskPipeline.arguments.colorTex = .texture(texA)
+        createMaskPipeline.arguments.maskTex = .texture(maskTex)
+        createMaskPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: createMaskPipeline, width: width, height: height)
 
-                sortPipeline.arguments.inputTexture = .texture(readTex)
-                sortPipeline.arguments.outputTexture = .texture(writeTex)
+        // 2) Clear span buffer
+        clearSpanPipeline.arguments.spanTex = .texture(spanTex)
+        clearSpanPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: clearSpanPipeline, width: width, height: height)
 
-                let paramBuffer = device.makeBuffer(bytes: &params, length: MemoryLayout<BitonicParams>.stride, options: .storageModeShared)!
-                sortPipeline.arguments.params = .buffer(paramBuffer)
+        // 3) Identify spans (1 thread per row: dispatch width=1)
+        identifySpansPipeline.arguments.maskTex = .texture(maskTex)
+        identifySpansPipeline.arguments.spanTex = .texture(spanTex)
+        identifySpansPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: identifySpansPipeline, width: 1, height: height)
 
-                try compute.run(pipeline: sortPipeline, width: width / 2, height: height)
+        // 4) Sort values
+        rgbToValPipeline.arguments.colorTex = .texture(texA)
+        rgbToValPipeline.arguments.valTex = .texture(valTex)
+        rgbToValPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: rgbToValPipeline, width: width, height: height)
 
-                // Swap
-                let tmp = readTex
-                readTex = writeTex
-                writeTex = tmp
+        // 5) Sort each span into sortedTex
+        pixelSortPipeline.arguments.colorTex = .texture(texA)
+        pixelSortPipeline.arguments.valTex = .texture(valTex)
+        pixelSortPipeline.arguments.spanTex = .texture(spanTex)
+        pixelSortPipeline.arguments.sortedTex = .texture(sortedTex)
+        pixelSortPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: pixelSortPipeline, width: width, height: height)
 
-                subBlockSize /= 2
-            }
-            blockSize *= 2
-        }
+        // 6) Composite only masked pixels into output texB
+        compositePipeline.arguments.maskTex = .texture(maskTex)
+        compositePipeline.arguments.sortedTex = .texture(sortedTex)
+        compositePipeline.arguments.originalTex = .texture(texA)
+        compositePipeline.arguments.outTex = .texture(texB)
+        compositePipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: compositePipeline, width: width, height: height)
 
-        // Read back from readTex (the last write destination after swap)
+        // Read back from output
         var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel)
-        readTex.getBytes(
+        texB.getBytes(
             &outputData,
             bytesPerRow: bytesPerRow,
             from: MTLRegionMake2D(0, 0, width, height),
@@ -173,25 +225,16 @@ struct PixelSort: ParsableCommand {
 
 // MARK: - Helpers
 
-struct BitonicParams {
+struct Params {
     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
+    var reverseSorting: UInt32
+    var gamma: Float
+    var maxSpanLength: UInt32
+    var invertMask: UInt32
 }
 
 enum SortKeyOption: String, ExpressibleByArgument, CaseIterable {