punktfunk/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletDecoder.swift

// PyroWave native Metal decoder — the Apple twin of pf-client-core's Vulkan decoder
// (crates/pf-client-core/src/video_pyrowave.rs), reimplemented on the presenter's own MTLDevice
// so decode + CSC + present share one device with zero interop (design/pyrowave-codec-plan.md
// §4.7). No upstream C/C++ ships in the app: the bitstream parse below reimplements
// pyrowave_decoder.cpp's push_packet/decode_packet walk, and the two compute kernels
// (MetalWaveletShaders.swift) are hand-ported from the vendored GLSL. The §4.2 upstream pin
// covers this hand-port: a vendored bump means re-diffing two decode shaders and the two 8-byte
// header structs, and it is already a protocol-version event.
//
// Wire shape (all fixed by the host encoder, punktfunk-host encode/linux/pyrowave.rs):
// • One AU = one frame = a self-delimiting stream of packets. Each packet is one 32x32
//   coefficient block for one (component, level, band), self-sized by its 8-byte
//   BitstreamHeader; a per-frame START_OF_FRAME sequence header carries dims + total block
//   count + the VUI bits (chroma 4:2:0, BT.709/BT.2020, limited/full).
// • With `USER_FLAG_CHUNK_ALIGNED` (Phase 4) the AU is a whole number of `shard_payload`-sized
//   windows, each 4-byte-prefixed (used-len u16 LE + kind u16 LE): kind 0 = whole packets,
//   1/2/3 = FRAG chain for a packet bigger than one window. A missing shard of a partial frame
//   arrives as an all-zero window (used = 0) → skipped, its blocks reconstruct as zeros
//   (localized blur, the Phase-4 design intent). The reassembler enables partial delivery
//   core-side automatically for PyroWave sessions.
// • Decode acceptance mirrors upstream decode_is_ready(allow_partial=true): a frame with no
//   SOF or with no more than half its blocks is dropped rather than decoded to garbage.
//
// GPU structure per frame (mirroring pyrowave_decoder.cpp's barriers): one concurrent compute
// encoder with all ~42 dequant dispatches (each writes a distinct band layer — no intra-stage
// hazards), then one concurrent encoder per iDWT level (5) — encoder boundaries provide the
// write→sampled-read synchronization the Vulkan version expresses as pipeline barriers. The
// output is a ring of 4 plane sets (Y full-res + Cb/Cr half-res R8Unorm); ring depth plus
// same-queue hazard tracking keeps a set alive while the presenter still samples it (the same
// scheme as the Vulkan client's ring).

#if canImport(Metal)
import Foundation
import Metal
import os

private let waveletLog = Logger(subsystem: "io.unom.punktfunk", category: "pyrowave")

/// The per-(component, level, band) 32x32-block table — the exact Swift port of
/// `WaveletBuffers::init_block_meta` (pyrowave_common.cpp): the walk order (level 4→0,
/// component 0→2 skipping level-0 chroma in 4:2:0, band (level==4 ? 0 : 1)→3) DEFINES the
/// global `block_index` space the wire packets address, so it must match the encoder exactly.
struct WaveletLayout {
    static let decompositionLevels = 5
    static let alignment = 32
    static let minimumImageSize = 128

    let width: Int
    let height: Int
    let alignedWidth: Int
    let alignedHeight: Int
    /// blockMeta[component][level][band] = (blockOffset32x32, blockStride32x32); -1 offset =
    /// band not coded (level-0 chroma in 4:2:0).
    let blockMeta: [[[(offset: Int, stride: Int)]]]
    let blockCount32: Int

    /// Band-image extent at `level` — mip `level` of the (aligned/2)-sized coefficient image.
    /// Exact halving: the aligned dims are 32-aligned, so /2 is 16-aligned and survives 4 shifts.
    func levelWidth(_ level: Int) -> Int { (alignedWidth / 2) >> level }
    func levelHeight(_ level: Int) -> Int { (alignedHeight / 2) >> level }

    init(width: Int, height: Int) {
        self.width = width
        self.height = height
        let align = { (v: Int) in
            max((v + Self.alignment - 1) & ~(Self.alignment - 1), Self.minimumImageSize)
        }
        alignedWidth = align(width)
        alignedHeight = align(height)

        var meta = [[[(offset: Int, stride: Int)]]](
            repeating: [[(offset: Int, stride: Int)]](
                repeating: [(offset: Int, stride: Int)](repeating: (-1, 0), count: 4),
                count: Self.decompositionLevels),
            count: 3)
        var count32 = 0
        let aw = alignedWidth
        let ah = alignedHeight
        for level in stride(from: Self.decompositionLevels - 1, through: 0, by: -1) {
            for component in 0..<3 {
                if level == 0 && component != 0 { continue } // 4:2:0: no top-level chroma
                for band in (level == Self.decompositionLevels - 1 ? 0 : 1)..<4 {
                    let levelW = (aw / 2) >> level
                    let levelH = (ah / 2) >> level
                    let blocksX8 = (levelW + 7) / 8
                    let blocksY8 = (levelH + 7) / 8
                    let blocksX32 = (levelW + 31) / 32
                    meta[component][level][band] = (count32, blocksX32)
                    // accumulate_block_mapping's 32x32 count.
                    count32 += ((blocksX8 + 3) / 4) * ((blocksY8 + 3) / 4)
                }
            }
        }
        blockMeta = meta
        blockCount32 = count32
    }
}

/// One parsed frame, CPU side: the per-block payload offset table + the flat payload words the
/// dequant kernel consumes (packet words INCLUDING each 8-byte header, as upstream uploads
/// them), plus the sequence header's facts.
struct ParsedWaveletFrame {
    var layout: WaveletLayout
    /// Per 32x32 block: u32 word offset into `payload`, or UInt32.max = block missing.
    var offsets: [UInt32]
    var payload: [UInt32]
    var totalBlocks: Int
    var decodedBlocks: Int
    /// VUI bits from the sequence header (BitstreamSequenceHeader).
    var bt2020: Bool
    var fullRange: Bool

    /// The frame's Y′CbCr→RGB signal for the presenter's planar CSC. PyroWave today is always
    /// BT.709 limited (the host's fixed contract), but the sequence header signals it, so honor
    /// what it says.
    var cscSignal: CscRows.Signal {
        CscRows.Signal(matrix: bt2020 ? 9 : 1, fullRange: fullRange)
    }
}

enum WaveletBitstream {
    /// Window kinds of the chunk-aligned framing (host WIN_* constants).
    private static let winPacked: UInt16 = 0
    private static let winFragFirst: UInt16 = 1
    private static let winFragCont: UInt16 = 2
    private static let winFragLast: UInt16 = 3

    /// Parse one AU into the dequant kernel's inputs. `windowSize` > 0 with `chunkAligned`
    /// walks the Phase-4 shard-window framing first; otherwise the AU is one packet stream.
    /// nil = drop the frame (malformed, no SOF, or not enough blocks survived loss to be worth
    /// decoding — upstream's `decoded_blocks > total/2` partial rule).
    static func parse(au: Data, chunkAligned: Bool, windowSize: Int) -> ParsedWaveletFrame? {
        var state = ParseState()
        let ok = au.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> Bool in
            guard let base = raw.baseAddress?.assumingMemoryBound(to: UInt8.self) else {
                return false
            }
            let count = raw.count
            if chunkAligned, windowSize >= 8 {
                // Whole windows only; a trailing partial window would be a framing bug.
                guard count % windowSize == 0 else { return false }
                var frag: [UInt8] = []
                var fragLive = false
                var pos = 0
                while pos < count {
                    let win = UnsafeBufferPointer(start: base + pos, count: windowSize)
                    pos += windowSize
                    let used = Int(win[0]) | (Int(win[1]) << 8)
                    let kind = UInt16(win[2]) | (UInt16(win[3]) << 8)
                    // A zeroed (missing) shard or an overrun drops the window AND breaks any
                    // fragment chain riding across it (mirrors video_pyrowave.rs push_window).
                    guard used > 0, 4 + used <= windowSize else {
                        frag.removeAll(keepingCapacity: true)
                        fragLive = false
                        continue
                    }
                    let body = UnsafeBufferPointer(start: win.baseAddress! + 4, count: used)
                    switch kind {
                    case winPacked:
                        frag.removeAll(keepingCapacity: true)
                        fragLive = false
                        guard state.pushPackets(body) else { return false }
                    case winFragFirst:
                        frag.removeAll(keepingCapacity: true)
                        frag.append(contentsOf: body)
                        fragLive = true
                    case winFragCont:
                        if fragLive { frag.append(contentsOf: body) }
                    case winFragLast:
                        if fragLive {
                            frag.append(contentsOf: body)
                            let ok = frag.withUnsafeBufferPointer { state.pushPackets($0) }
                            guard ok else { return false }
                        }
                        frag.removeAll(keepingCapacity: true)
                        fragLive = false
                    default:
                        frag.removeAll(keepingCapacity: true)
                        fragLive = false
                    }
                }
                return true
            }
            return state.pushPackets(UnsafeBufferPointer(start: base, count: count))
        }
        guard ok, let frame = state.finish() else { return nil }
        // Upstream decode_is_ready(allow_partial=true): with no SOF the frame is undecodable;
        // at half the blocks or fewer it is presumed garbage.
        guard frame.totalBlocks > 0, frame.decodedBlocks > frame.totalBlocks / 2 else {
            return nil
        }
        return frame
    }

    /// Streaming packet-walk state (pyrowave_decoder.cpp push_packet + decode_packet). The
    /// SOF sequence header arrives first in every host AU, which fixes the dims → layout →
    /// offset-table size before any coefficient packet lands; a coefficient packet before the
    /// SOF (its window was lost) is skipped — its block just stays missing.
    private struct ParseState {
        var layout: WaveletLayout?
        var offsets: [UInt32] = []
        var payload: [UInt32] = []
        var totalBlocks = 0
        var decodedBlocks = 0
        var bt2020 = false
        var fullRange = false
        var sawSOF = false

        mutating func pushPackets(_ buf: UnsafeBufferPointer<UInt8>) -> Bool {
            guard let base = buf.baseAddress else { return true }
            var pos = 0
            let count = buf.count
            while count - pos >= 8 {
                let word0 = loadWord(base, pos)
                let word1 = loadWord(base, pos + 4)
                let extended = (word0 >> 31) & 1
                if extended != 0 {
                    // BitstreamSequenceHeader: w-1[0:14] h-1[14:28] seq[28:31] ext[31];
                    // total[0:24] code[24:26] chroma[26] prim[27] trc[28] mtx[29] range[30]
                    // siting[31].
                    let code = (word1 >> 24) & 0x3
                    guard code == 0 else { return false } // only START_OF_FRAME is defined
                    let chromaRes = (word1 >> 26) & 1
                    guard chromaRes == 0 else { return false } // host contract: 4:2:0
                    let w = Int(word0 & 0x3fff) + 1
                    let h = Int((word0 >> 14) & 0x3fff) + 1
                    guard w >= 2, h >= 2, w % 2 == 0, h % 2 == 0 else { return false }
                    if sawSOF {
                        // One frame, one geometry — a second SOF must agree.
                        guard layout?.width == w, layout?.height == h else { return false }
                    } else {
                        sawSOF = true
                        let l = WaveletLayout(width: w, height: h)
                        layout = l
                        offsets = [UInt32](repeating: .max, count: l.blockCount32)
                        payload.reserveCapacity(64 * 1024 / 4)
                        totalBlocks = Int(word1 & 0xff_ffff)
                        bt2020 = (word1 >> 29) & 1 != 0
                        fullRange = (word1 >> 30) & 1 == 0 // YCBCR_RANGE_FULL = 0
                    }
                    pos += 8
                    continue
                }
                // BitstreamHeader: ballot[0:16] payload_words[16:28] seq[28:31] ext[31];
                // quant_code[0:8] block_index[8:32]. payload_words counts u32s INCLUDING the
                // 8-byte header.
                let payloadWords = Int((word0 >> 16) & 0xfff)
                guard payloadWords >= 2, pos + payloadWords * 4 <= count else { return false }
                let blockIndex = Int(word1 >> 8)
                if let layout, blockIndex < layout.blockCount32 {
                    // First write wins (duplicate packets are ignored, like upstream).
                    if offsets[blockIndex] == .max {
                        offsets[blockIndex] = UInt32(payload.count)
                        decodedBlocks += 1
                        payload.reserveCapacity(payload.count + payloadWords)
                        for w in 0..<payloadWords {
                            payload.append(loadWord(base, pos + w * 4))
                        }
                    }
                } else if layout != nil {
                    return false // out-of-bounds block index — corrupt stream
                }
                // No layout yet (SOF lost): skip the packet, the block stays missing.
                pos += payloadWords * 4
            }
            // In the windowed framing, `used` delimits exactly; dense AUs must also consume
            // fully (upstream errors on trailing bytes).
            return pos == count
        }

        private func loadWord(_ base: UnsafePointer<UInt8>, _ offset: Int) -> UInt32 {
            UInt32(base[offset])
                | (UInt32(base[offset + 1]) << 8)
                | (UInt32(base[offset + 2]) << 16)
                | (UInt32(base[offset + 3]) << 24)
        }

        func finish() -> ParsedWaveletFrame? {
            guard let layout else { return nil }
            return ParsedWaveletFrame(
                layout: layout, offsets: offsets, payload: payload,
                totalBlocks: totalBlocks, decodedBlocks: decodedBlocks,
                bt2020: bt2020, fullRange: fullRange)
        }
    }
}

/// One decoded frame's output planes, handed to the presenter's planar render path. The
/// textures belong to the decoder's ring — ring depth (4) plus same-queue hazard tracking keep
/// them valid while referenced. Public because it rides inside `ReadyImage`.
public struct WaveletPlanes: @unchecked Sendable {
    public let y: MTLTexture
    public let cb: MTLTexture
    public let cr: MTLTexture
    public let csc: CscUniform
    public var width: Int { y.width }
    public var height: Int { y.height }
}

public final class MetalWaveletDecoder {
    /// Matches the Vulkan client's ring: deep enough that a slot is never rewritten while the
    /// presenter still samples it in practice; same-queue hazard tracking is the hard backstop.
    private static let ringDepth = 4

    /// Device-capability gate for advertisement (SessionModel) and the settings picker: the
    /// dequant kernel needs simdgroup prefix sums with its 16 header lanes inside one
    /// simdgroup, so compile the real kernels once and check the pipeline facts. Apple6 (A13)
    /// and every Mac2 device pass the family check; the compile probe is authoritative.
    public static let supported: Bool = {
        guard let device = MTLCreateSystemDefaultDevice() else { return false }
        guard device.supportsFamily(.apple6) || device.supportsFamily(.mac2) else { return false }
        do {
            let lib = try device.makeLibrary(source: waveletShaderSource, options: nil)
            guard let dequant = lib.makeFunction(name: "wavelet_dequant") else { return false }
            let p = try device.makeComputePipelineState(function: dequant)
            var shift = false
            let fc = MTLFunctionConstantValues()
            fc.setConstantValue(&shift, type: .bool, index: 0)
            _ = try lib.makeFunction(name: "idwt", constantValues: fc)
            return p.threadExecutionWidth >= 16 && p.maxTotalThreadsPerThreadgroup >= 128
        } catch {
            waveletLog.info("pyrowave probe: kernels rejected (\(error, privacy: .public))")
            return false
        }
    }()

    private let device: MTLDevice
    private let queue: MTLCommandQueue
    private let dequantPipeline: MTLComputePipelineState
    private let idwtPipeline: MTLComputePipelineState
    private let idwtShiftPipeline: MTLComputePipelineState
    private let mirrorSampler: MTLSamplerState

    // Size-dependent state, rebuilt when the SOF dims change (this is also the mid-stream
    // Reconfigure/resize path — the wavelet decoder is fixed-size per geometry).
    private var layout: WaveletLayout?
    /// coefficients[component][level]: 4-slice R16Float (levels 0–1) / R32Float (levels 2–4)
    /// texture2d_array — the band images (precision-1 split, see MetalWaveletShaders).
    private var coefficients: [[MTLTexture]] = []
    /// llViews[component][level]: slice-0 (LL band) 2D write view of `coefficients` — the iDWT
    /// output target chaining level L+1 into level L.
    private var llViews: [[MTLTexture]] = []

    private struct Slot {
        var y: MTLTexture
        var cb: MTLTexture
        var cr: MTLTexture
        var offsets: MTLBuffer
        var payload: MTLBuffer
    }

    private var slots: [Slot] = []
    private var nextSlot = 0

    /// The current geometry (from the last SOF that built the resources) — the pump reports
    /// decoded-size changes to the resize overlay from this. PUMP THREAD.
    var decodedSize: (width: Int, height: Int)? {
        layout.map { ($0.width, $0.height) }
    }

    /// The pump thread owns `decode`; everything mutable is confined to it.
    init?(device: MTLDevice, queue: MTLCommandQueue) {
        self.device = device
        self.queue = queue
        do {
            let lib = try device.makeLibrary(source: waveletShaderSource, options: nil)
            guard let dequantFn = lib.makeFunction(name: "wavelet_dequant") else { return nil }
            dequantPipeline = try device.makeComputePipelineState(function: dequantFn)
            var shift = false
            let fcOff = MTLFunctionConstantValues()
            fcOff.setConstantValue(&shift, type: .bool, index: 0)
            idwtPipeline = try device.makeComputePipelineState(
                function: try lib.makeFunction(name: "idwt", constantValues: fcOff))
            shift = true
            let fcOn = MTLFunctionConstantValues()
            fcOn.setConstantValue(&shift, type: .bool, index: 0)
            idwtShiftPipeline = try device.makeComputePipelineState(
                function: try lib.makeFunction(name: "idwt", constantValues: fcOn))
        } catch {
            waveletLog.error("pyrowave: pipeline build failed (\(error, privacy: .public))")
            return nil
        }
        guard dequantPipeline.threadExecutionWidth >= 16,
              dequantPipeline.maxTotalThreadsPerThreadgroup >= 128
        else { return nil }
        // Upstream's mirror_repeat_sampler: mirrored repeat, NEAREST everything, normalized
        // coords — the idwt gather footprint + coordinate nudge depend on exactly this.
        let samp = MTLSamplerDescriptor()
        samp.sAddressMode = .mirrorRepeat
        samp.tAddressMode = .mirrorRepeat
        samp.minFilter = .nearest
        samp.magFilter = .nearest
        samp.mipFilter = .notMipmapped
        samp.normalizedCoordinates = true
        guard let sampler = device.makeSamplerState(descriptor: samp) else { return nil }
        mirrorSampler = sampler
    }

    /// Decode one AU. Synchronous CPU parse + async GPU decode: returns false when the frame
    /// was dropped (malformed / SOF lost / not enough blocks); on true, `completion` fires on a
    /// Metal callback thread once the planes are decoded (nil = the GPU pass errored).
    /// PUMP THREAD only.
    func decode(
        au: Data, chunkAligned: Bool, windowSize: Int,
        completion: @escaping @Sendable (WaveletPlanes?) -> Void
    ) -> Bool {
        guard
            let frame = WaveletBitstream.parse(
                au: au, chunkAligned: chunkAligned, windowSize: windowSize)
        else { return false }

        if layout?.width != frame.layout.width || layout?.height != frame.layout.height {
            guard rebuild(layout: frame.layout) else { return false }
        }
        guard let layout, !slots.isEmpty else { return false }

        var slot = slots[nextSlot]
        // Grow the payload buffer to the frame (+16-byte zeroed guard: the kernel's 64-bit
        // sign-window load and eager plane-byte prefetch may read past the payload end —
        // upstream pads its Vulkan buffer for exactly this).
        let payloadBytes = frame.payload.count * 4
        if slot.payload.length < payloadBytes + 16 {
            guard
                let grown = device.makeBuffer(
                    length: max(64 * 1024, (payloadBytes + 16) * 2), options: .storageModeShared)
            else { return false }
            slot.payload = grown
            slots[nextSlot] = slot
        }
        frame.offsets.withUnsafeBytes { src in
            slot.offsets.contents().copyMemory(
                from: src.baseAddress!, byteCount: min(src.count, slot.offsets.length))
        }
        frame.payload.withUnsafeBytes { src in
            slot.payload.contents().copyMemory(from: src.baseAddress!, byteCount: src.count)
        }
        memset(slot.payload.contents() + payloadBytes, 0, 16)

        guard let cmd = queue.makeCommandBuffer() else { return false }

        // Stage 1: dequant — every (component, level, band) block grid in one concurrent
        // encoder (each dispatch writes its own band layer; no intra-stage hazards, exactly
        // like the barrier-free Vulkan dispatch loop).
        guard let dequant = cmd.makeComputeCommandEncoder(dispatchType: .concurrent) else {
            return false
        }
        dequant.label = "pyrowave dequant"
        dequant.setComputePipelineState(dequantPipeline)
        dequant.setBuffer(slot.offsets, offset: 0, index: 0)
        dequant.setBuffer(slot.payload, offset: 0, index: 1)
        for level in 0..<WaveletLayout.decompositionLevels {
            for component in 0..<3 {
                if level == 0 && component != 0 { continue } // 4:2:0
                for band in (level == WaveletLayout.decompositionLevels - 1 ? 0 : 1)..<4 {
                    let meta = layout.blockMeta[component][level][band]
                    let w = layout.levelWidth(level)
                    let h = layout.levelHeight(level)
                    var regs = DequantRegisters(
                        resolution: SIMD2(Int32(w), Int32(h)),
                        outputLayer: Int32(band),
                        blockOffset32x32: Int32(meta.offset),
                        blockStride32x32: Int32(meta.stride))
                    dequant.setTexture(coefficients[component][level], index: 0)
                    dequant.setBytes(
                        &regs, length: MemoryLayout<DequantRegisters>.stride, index: 2)
                    dequant.dispatchThreadgroups(
                        MTLSize(width: (w + 31) / 32, height: (h + 31) / 32, depth: 1),
                        threadsPerThreadgroup: MTLSize(width: 128, height: 1, depth: 1))
                }
            }
        }
        dequant.endEncoding()

        // Stage 2: iDWT, coarsest level in — one encoder per level; the encoder boundary is
        // the write→sampled-read barrier chaining each level's LL into the next.
        for inputLevel in stride(from: WaveletLayout.decompositionLevels - 1, through: 0, by: -1) {
            guard let idwt = cmd.makeComputeCommandEncoder(dispatchType: .concurrent) else {
                return false
            }
            idwt.label = "pyrowave idwt L\(inputLevel)"
            idwt.setSamplerState(mirrorSampler, index: 0)
            // Resolution rides TRANSPOSED (the kernel transposes on load and store).
            let rx = layout.levelHeight(inputLevel)
            let ry = layout.levelWidth(inputLevel)
            var regs = IdwtRegisters(
                resolution: SIMD2(Int32(rx), Int32(ry)),
                invResolution: SIMD2(1.0 / Float(rx), 1.0 / Float(ry)))
            idwt.setBytes(&regs, length: MemoryLayout<IdwtRegisters>.stride, index: 0)
            let grid = MTLSize(width: (rx + 15) / 16, height: (ry + 15) / 16, depth: 1)
            let group = MTLSize(width: 64, height: 1, depth: 1)
            if inputLevel == 0 {
                // 4:2:0: the final full-res pass is luma only (chroma finished at level 1).
                idwt.setComputePipelineState(idwtShiftPipeline)
                idwt.setTexture(coefficients[0][0], index: 0)
                idwt.setTexture(slot.y, index: 1)
                idwt.dispatchThreadgroups(grid, threadsPerThreadgroup: group)
            } else {
                for component in 0..<3 {
                    idwt.setTexture(coefficients[component][inputLevel], index: 0)
                    if component != 0 && inputLevel == 1 {
                        // 4:2:0 chroma emits its final half-res plane one level early.
                        idwt.setComputePipelineState(idwtShiftPipeline)
                        idwt.setTexture(component == 1 ? slot.cb : slot.cr, index: 1)
                    } else {
                        idwt.setComputePipelineState(idwtPipeline)
                        idwt.setTexture(llViews[component][inputLevel - 1], index: 1)
                    }
                    idwt.dispatchThreadgroups(grid, threadsPerThreadgroup: group)
                }
            }
            idwt.endEncoding()
        }

        let planes = WaveletPlanes(
            y: slot.y, cb: slot.cb, cr: slot.cr,
            csc: CscRows.rows(frame.cscSignal, depth: 8, msbPacked: false))
        cmd.addCompletedHandler { buffer in
            completion(buffer.error == nil ? planes : nil)
        }
        cmd.commit()
        nextSlot = (nextSlot + 1) % Self.ringDepth
        return true
    }

    /// (Re)allocate every size-dependent resource for `layout`'s geometry. Also the mid-stream
    /// resize path: a Reconfigure shows up here as new SOF dims.
    private func rebuild(layout newLayout: WaveletLayout) -> Bool {
        waveletLog.info(
            "pyrowave: building decoder \(newLayout.width)x\(newLayout.height) (aligned \(newLayout.alignedWidth)x\(newLayout.alignedHeight), \(newLayout.blockCount32) blocks)")
        var coeff: [[MTLTexture]] = []
        var lls: [[MTLTexture]] = []
        for component in 0..<3 {
            var perLevel: [MTLTexture] = []
            var perLevelLL: [MTLTexture] = []
            for level in 0..<WaveletLayout.decompositionLevels {
                let desc = MTLTextureDescriptor()
                desc.textureType = .type2DArray
                desc.arrayLength = 4
                // Upstream precision 1: fp16 storage for the two finest levels, fp32 for the
                // coarse levels whose values feed every later reconstruction step.
                desc.pixelFormat = level < 2 ? .r16Float : .r32Float
                desc.width = newLayout.levelWidth(level)
                desc.height = newLayout.levelHeight(level)
                desc.usage = [.shaderRead, .shaderWrite]
                desc.storageMode = .private
                guard let tex = device.makeTexture(descriptor: desc) else { return false }
                tex.label = "pyrowave coeff c\(component) L\(level)"
                guard
                    let ll = tex.makeTextureView(
                        pixelFormat: desc.pixelFormat, textureType: .type2D,
                        levels: 0..<1, slices: 0..<1)
                else { return false }
                ll.label = "pyrowave LL c\(component) L\(level)"
                perLevel.append(tex)
                perLevelLL.append(ll)
            }
            coeff.append(perLevel)
            lls.append(perLevelLL)
        }

        var newSlots: [Slot] = []
        for i in 0..<Self.ringDepth {
            let plane = { (w: Int, h: Int, name: String) -> MTLTexture? in
                let desc = MTLTextureDescriptor.texture2DDescriptor(
                    pixelFormat: .r8Unorm, width: w, height: h, mipmapped: false)
                desc.usage = [.shaderRead, .shaderWrite]
                desc.storageMode = .private
                let t = self.device.makeTexture(descriptor: desc)
                t?.label = name
                return t
            }
            guard
                let y = plane(newLayout.width, newLayout.height, "pyrowave Y[\(i)]"),
                let cb = plane(newLayout.width / 2, newLayout.height / 2, "pyrowave Cb[\(i)]"),
                let cr = plane(newLayout.width / 2, newLayout.height / 2, "pyrowave Cr[\(i)]"),
                let offsets = device.makeBuffer(
                    length: max(newLayout.blockCount32 * 4, 4), options: .storageModeShared),
                let payload = device.makeBuffer(length: 64 * 1024, options: .storageModeShared)
            else { return false }
            newSlots.append(Slot(y: y, cb: cb, cr: cr, offsets: offsets, payload: payload))
        }

        coefficients = coeff
        llViews = lls
        slots = newSlots
        nextSlot = 0
        layout = newLayout
        return true
    }

    // MSL-side layouts (MetalWaveletShaders.swift) — keep in lockstep.
    private struct DequantRegisters {
        var resolution: SIMD2<Int32>
        var outputLayer: Int32
        var blockOffset32x32: Int32
        var blockStride32x32: Int32
    }

    private struct IdwtRegisters {
        var resolution: SIMD2<Int32>
        var invResolution: SIMD2<Float>
    }
}
#endif