Quantizer.swift

import Accelerate
import AppKit
import CoreGraphics

struct QuantizationOptions: Equatable {
    /// libimagequant's speed, 1 (slowest, best) to 10; Preferences falls back to this.
    static let defaultSpeed = 3

    var numberOfColors: Int = 256
    var dithering: Bool = false
    var speed: Int = Self.defaultSpeed
    /// 0–100 target; when set, libimagequant uses the fewest colors (up to
    /// `numberOfColors`) that still reach this quality.
    var qualityTarget: Int?
}

struct QuantizationResult {
    var image: NSImage
    var pngData: Data
    /// Input for re-encoding pngData without quantizing again; cleared (to
    /// free the pixel buffer) once the maximum-effort encode has run.
    var bitmap: IndexedBitmap?
    var paletteCount: Int
    /// libimagequant's 0–100 estimate of how well the palette matches the
    /// source; nil when the library doesn't compute it (high speed settings).
    var quality: Int?
}

enum QuantizationError: Error, LocalizedError {
    case failedToCreateAttr
    case failedToCreateImage
    case failedToQuantize(liq_error)
    case failedToRemap(liq_error)
    case failedToGetPixelData
    case failedToCreatePNG

    var errorDescription: String? {
        switch self {
        case .failedToCreateAttr: return "Failed to create quantization attributes"
        case .failedToCreateImage: return "Failed to create quantization image"
        case .failedToQuantize(let e): return "Quantization failed (\(e.rawValue))"
        case .failedToRemap(let e): return "Remapping failed (\(e.rawValue))"
        case .failedToGetPixelData: return "Failed to get pixel data from image"
        case .failedToCreatePNG: return "Failed to create PNG data"
        }
    }
}

actor Quantizer {

    private static let srgb = CGColorSpace(name: CGColorSpace.sRGB)

    func quantize(cgImage: CGImage, options: QuantizationOptions) throws -> QuantizationResult {
        // Slider scrubbing queues multiple requests on this actor; skip any
        // whose task was already cancelled by a newer request.
        try Task.checkCancellation()

        let width = cgImage.width
        let height = cgImage.height
        let pixelCount = width * height

        // Owned here, not by a context, because libimagequant keeps pointing
        // at these bytes until the remap below is done.
        let pixelData = UnsafeMutableRawPointer.allocate(byteCount: pixelCount * 4, alignment: 16)
        defer { pixelData.deallocate() }
        guard Self.readStraightRGBA(cgImage, into: pixelData) else {
            throw QuantizationError.failedToGetPixelData
        }

        guard let attr = liq_attr_create() else {
            throw QuantizationError.failedToCreateAttr
        }
        defer { liq_attr_destroy(attr) }

        Self.configure(attr, with: options)

        guard let liqImage = liq_image_create_rgba(attr, pixelData, Int32(width), Int32(height), 0) else {
            throw QuantizationError.failedToCreateImage
        }
        defer { liq_image_destroy(liqImage) }

        try Task.checkCancellation()
        var resultPtr: OpaquePointer?
        let quantErr = liq_image_quantize(liqImage, attr, &resultPtr)
        guard quantErr == LIQ_OK, let result = resultPtr else {
            throw QuantizationError.failedToQuantize(quantErr)
        }
        defer { liq_result_destroy(result) }

        liq_set_dithering_level(result, options.dithering ? 1.0 : 0.0)

        let remapped = UnsafeMutablePointer<UInt8>.allocate(capacity: pixelCount)
        defer { remapped.deallocate() }

        let remapErr = liq_write_remapped_image(result, liqImage, remapped, pixelCount)
        guard remapErr == LIQ_OK else {
            throw QuantizationError.failedToRemap(remapErr)
        }

        guard let palettePtr = liq_get_palette(result) else {
            throw QuantizationError.failedToQuantize(LIQ_OK)
        }
        let palette = palettePtr.pointee
        let colorCount = Int(palette.count)

        let paletteEntries: [IndexedPNGEncoder.PaletteEntry] = withUnsafePointer(to: palette.entries) {
            entriesPtr in
            entriesPtr.withMemoryRebound(to: liq_color.self, capacity: 256) { colors in
                (0..<colorCount).map { i in
                    IndexedPNGEncoder.PaletteEntry(
                        red: colors[i].r, green: colors[i].g, blue: colors[i].b, alpha: colors[i].a)
                }
            }
        }
        let indices = [UInt8](UnsafeBufferPointer(start: remapped, count: pixelCount))
        let bitmap = IndexedBitmap(width: width, height: height, palette: paletteEntries, pixels: indices)

        // Encode a real indexed PNG (PLTE + tRNS) by hand; see IndexedPNGEncoder
        // for why ImageIO isn't used.
        guard let pngData = IndexedPNGEncoder.encode(bitmap) else {
            throw QuantizationError.failedToCreatePNG
        }

        guard
            let nsImage = Self.makeDisplayImage(
                palette: paletteEntries, indices: indices, width: width, height: height
            )
        else {
            throw QuantizationError.failedToCreatePNG
        }

        let measuredQuality = liq_get_quantization_quality(result)
        return QuantizationResult(
            image: nsImage,
            pngData: pngData,
            bitmap: bitmap,
            paletteCount: colorCount,
            quality: measuredQuality >= 0 ? Int(measuredQuality) : nil
        )
    }

    /// Writes the image into `destination` as straight-alpha sRGB RGBA, rows
    /// packed at width × 4 bytes, the layout libimagequant reads. vImage
    /// converts directly; a CGBitmapContext premultiplies, and undoing that
    /// loses the colour of nearly transparent pixels.
    /// /documentation/accelerate/vimagebuffer_initwithcgimage(_:_:_:_:_:)
    /// sRGB, not device RGB, which is "not recommended when color preservation
    /// is important". /documentation/coregraphics/cgcolorspacecreatedevicergb()
    private static func readStraightRGBA(_ cgImage: CGImage, into destination: UnsafeMutableRawPointer)
        -> Bool
    {
        guard let srgb,
            var format = vImage_CGImageFormat(
                bitsPerComponent: 8,
                bitsPerPixel: 32,
                colorSpace: srgb,
                bitmapInfo: CGBitmapInfo(
                    rawValue: CGImageAlphaInfo.last.rawValue | CGBitmapInfo.byteOrder32Big.rawValue)
            )
        else {
            return false
        }
        var buffer = vImage_Buffer(
            data: destination,
            height: vImagePixelCount(cgImage.height),
            width: vImagePixelCount(cgImage.width),
            rowBytes: cgImage.width * 4
        )
        // kvImageNoAllocate: fill the caller's packed buffer rather than one
        // vImage sizes, whose rows may be padded. /documentation/accelerate/kvimagenoallocate
        return vImageBuffer_InitWithCGImage(&buffer, &format, nil, cgImage, vImage_Flags(kvImageNoAllocate))
            == kvImageNoError
    }

    private static func configure(_ attr: OpaquePointer, with options: QuantizationOptions) {
        liq_set_max_colors(attr, Int32(min(options.numberOfColors, 256)))
        liq_set_speed(attr, Int32(options.speed))
        if let target = options.qualityTarget {
            // Min 0 so quantization never fails outright; max is the target.
            liq_set_quality(attr, 0, Int32(max(0, min(target, 100))))
        }
    }

    /// Expands palette + indices into an RGBA CGImage for on-screen display.
    private static func makeDisplayImage(
        palette: [IndexedPNGEncoder.PaletteEntry],
        indices: [UInt8],
        width: Int,
        height: Int
    ) -> NSImage? {
        let pixelCount = width * height
        var outputPixels = [UInt8](repeating: 0, count: pixelCount * 4)
        for i in 0..<pixelCount {
            let color = palette[Int(indices[i])]
            let offset = i * 4
            outputPixels[offset + 0] = color.red
            outputPixels[offset + 1] = color.green
            outputPixels[offset + 2] = color.blue
            outputPixels[offset + 3] = color.alpha
        }

        // sRGB, like the palette it shows, so the preview matches the saved file.
        guard let srgb,
            let dataProvider = CGDataProvider(data: Data(outputPixels) as CFData),
            let outputCGImage = CGImage(
                width: width,
                height: height,
                bitsPerComponent: 8,
                bitsPerPixel: 32,
                bytesPerRow: width * 4,
                space: srgb,
                bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.last.rawValue),
                provider: dataProvider,
                decode: nil,
                shouldInterpolate: false,
                intent: .defaultIntent
            )
        else {
            return nil
        }

        return NSImage(cgImage: outputCGImage, size: NSSize(width: width, height: height))
    }
}