IndexedPNGEncoder.swift

import Compression
import Foundation

/// Palette + per-pixel indices: everything IndexedPNGEncoder needs to produce
/// a PNG, independent of how the quantization was done.
struct IndexedBitmap {
    var width: Int
    var height: Int
    var palette: [IndexedPNGEncoder.PaletteEntry]
    var pixels: [UInt8]
}

/// Encodes palette-quantized pixels as a color-type-3 (indexed) PNG, including
/// PLTE/tRNS chunks and minimal 1/2/4/8 bit depth. Apple documents no way to
/// have ImageIO write a palette: its PNG properties have no palette or bit-depth
/// key (/documentation/imageio/png-image-properties), and an indexed
/// CGColorSpace's table has no alpha for tRNS. A truecolor file would discard
/// the size benefit of quantization.
enum IndexedPNGEncoder {

    struct PaletteEntry: Equatable {
        var red: UInt8
        var green: UInt8
        var blue: UInt8
        var alpha: UInt8
    }

    enum CompressionEffort {
        /// Apple's Compression framework: fastest, and "the zlib encoder at
        /// level 5 only". /documentation/compression/compression_zlib
        /// Used for interactive previews.
        case fast
        /// System zlib at level 9: a few percent smaller, several times
        /// slower. Used for the file that actually gets saved.
        case maximum
    }

    /// Convenience for callers that carry the encoder's input as one value.
    static func encode(_ bitmap: IndexedBitmap, effort: CompressionEffort = .fast) -> Data? {
        encode(
            width: bitmap.width, height: bitmap.height,
            palette: bitmap.palette, pixels: bitmap.pixels,
            effort: effort
        )
    }

    /// Smallest PNG-legal bit depth (1, 2, 4 or 8) that can index the palette.
    static func bitDepth(forPaletteCount count: Int) -> Int {
        switch count {
        case ...2: return 1
        case ...4: return 2
        case ...16: return 4
        default: return 8
        }
    }

    /// `pixels` are palette indices, one byte per pixel, row-major.
    /// Returns nil for structurally invalid input or if compression fails.
    static func encode(
        width: Int, height: Int, palette: [PaletteEntry], pixels: [UInt8],
        effort: CompressionEffort = .fast
    ) -> Data? {
        guard width > 0, height > 0,
            (1...256).contains(palette.count),
            pixels.count == width * height,
            !pixels.contains(where: { Int($0) >= palette.count })
        else {
            return nil
        }

        let depth = bitDepth(forPaletteCount: palette.count)
        let scanlines = packScanlines(pixels: pixels, width: width, height: height, bitDepth: depth)
        let compressed = effort == .maximum ? zlibCompressMax(scanlines) : zlibCompress(scanlines)
        guard let idat = compressed else { return nil }

        var ihdr = Data()
        ihdr.appendBigEndian(UInt32(width))
        ihdr.appendBigEndian(UInt32(height))
        ihdr.append(UInt8(depth))
        ihdr.append(3)  // color type: indexed
        ihdr.append(0)  // compression
        ihdr.append(0)  // filter
        ihdr.append(0)  // interlace

        var plte = Data(capacity: palette.count * 3)
        for entry in palette {
            plte.append(entry.red)
            plte.append(entry.green)
            plte.append(entry.blue)
        }

        var png = Data([137, 80, 78, 71, 13, 10, 26, 10])
        png.append(chunk("IHDR", ihdr))
        png.append(chunk("PLTE", plte))

        // tRNS holds palette alphas, truncated after the last non-opaque entry
        // and omitted entirely for fully opaque palettes.
        if let lastTransparent = palette.lastIndex(where: { $0.alpha < 255 }) {
            let trns = Data(palette[...lastTransparent].map(\.alpha))
            png.append(chunk("tRNS", trns))
        }

        png.append(chunk("IDAT", idat))
        png.append(chunk("IEND", Data()))
        return png
    }

    /// Prefixes each row with a None filter byte and packs indices big-endian
    /// (leftmost pixel in the most significant bits), rows padded to whole bytes.
    static func packScanlines(pixels: [UInt8], width: Int, height: Int, bitDepth: Int) -> [UInt8] {
        let pixelsPerByte = 8 / bitDepth
        let rowBytes = (width + pixelsPerByte - 1) / pixelsPerByte
        var out = [UInt8](repeating: 0, count: (rowBytes + 1) * height)

        out.withUnsafeMutableBufferPointer { dst in
            pixels.withUnsafeBufferPointer { src in
                // Empty only for a zero-row or zero-column image, whose scanlines
                // are the all-zero filter bytes `out` already holds.
                guard let dstBase = dst.baseAddress, let srcBase = src.baseAddress else { return }
                for row in 0..<height {
                    var write = row * (rowBytes + 1)
                    dst[write] = 0  // filter: None
                    write += 1
                    if bitDepth == 8 {
                        UnsafeMutableRawPointer(dstBase + write)
                            .copyMemory(from: srcBase + row * width, byteCount: width)
                        continue
                    }
                    var byte: UInt8 = 0
                    var bitsUsed = 0
                    for col in 0..<width {
                        byte = (byte << bitDepth) | src[row * width + col]
                        bitsUsed += bitDepth
                        if bitsUsed == 8 {
                            dst[write] = byte
                            write += 1
                            byte = 0
                            bitsUsed = 0
                        }
                    }
                    if bitsUsed > 0 {
                        dst[write] = byte << (8 - bitsUsed)
                    }
                }
            }
        }
        return out
    }

    /// Raw DEFLATE from the Compression framework wrapped in a zlib stream
    /// (header + Adler-32 trailer), as PNG's IDAT requires.
    static func zlibCompress(_ source: [UInt8]) -> Data? {
        guard !source.isEmpty else { return nil }
        let capacity = source.count + source.count / 4 + 256
        let destination = UnsafeMutablePointer<UInt8>.allocate(capacity: capacity)
        defer { destination.deallocate() }

        let written = compression_encode_buffer(
            destination, capacity,
            source, source.count,
            nil, COMPRESSION_ZLIB
        )
        guard written > 0 else { return nil }

        var out = Data([0x78, 0x9C])
        out.append(destination, count: written)
        out.appendBigEndian(adler32(source))
        return out
    }

    /// Complete zlib stream from system libz at maximum compression.
    static func zlibCompressMax(_ source: [UInt8]) -> Data? {
        guard !source.isEmpty else { return nil }
        var destLen = compressBound(uLong(source.count))
        var dest = Data(count: Int(destLen))
        let status = dest.withUnsafeMutableBytes { destBuf in
            source.withUnsafeBufferPointer { src -> Int32 in
                guard let destBase = destBuf.bindMemory(to: UInt8.self).baseAddress,
                    let srcBase = src.baseAddress
                else { return Z_BUF_ERROR }
                return compress2(destBase, &destLen, srcBase, uLong(source.count), 9)
            }
        }
        guard status == Z_OK else { return nil }
        dest.count = Int(destLen)
        return dest
    }

    // Checksums come from libz (linked for zlibCompressMax anyway); these
    // wrappers only adapt Swift types to the C signatures.

    static func adler32(_ bytes: [UInt8]) -> UInt32 {
        bytes.withUnsafeBufferPointer { buf in
            UInt32(ImageAlpha.adler32(1, buf.baseAddress, uInt(buf.count)))
        }
    }

    static func crc32(_ data: Data) -> UInt32 {
        data.withUnsafeBytes { (buf: UnsafeRawBufferPointer) in
            UInt32(ImageAlpha.crc32(0, buf.bindMemory(to: UInt8.self).baseAddress, uInt(buf.count)))
        }
    }

    private static func chunk(_ type: String, _ payload: Data) -> Data {
        var out = Data()
        out.appendBigEndian(UInt32(payload.count))
        var typed = Data(type.utf8)
        typed.append(payload)
        out.append(typed)
        out.appendBigEndian(crc32(typed))
        return out
    }
}

private extension Data {
    mutating func appendBigEndian(_ value: UInt32) {
        append(UInt8((value >> 24) & 0xFF))
        append(UInt8((value >> 16) & 0xFF))
        append(UInt8((value >> 8) & 0xFF))
        append(UInt8(value & 0xFF))
    }
}