draw-things-chat

interface FpzipModule {
    _fpzip_get_header(
        inputPtr: number,
        outTypePtr: number,
        outPrecPtr: number,
        outNxPtr: number,
        outNyPtr: number,
        outNzPtr: number,
        outNfPtr: number,
    ): number;
    _fpzip_decompress_f32(inputPtr: number, outputPtr: number): number;
    _fpzip_compress_f32(
        inputPtr: number,
        nValues: number,
        precBits: number,
        outputPtr: number,
        outputCap: number,
    ): number;
    // NNC-compatible compression with proper tensor dimensions
    _fpzip_compress_f32_dims(
        inputPtr: number,
        nx: number,
        ny: number,
        nz: number,
        nf: number,
        precBits: number,
        outputPtr: number,
        outputCap: number,
    ): number;
    _fpzip_errno_code(): number;
    _fpzip_errno_string(): number;
    _malloc(size: number): number;
    _free(ptr: number): void;
    HEAPU8: Uint8Array;
}

/** Tensor dimensions for NNC-compatible fpzip compression */
export interface TensorDims {
    /** Last dimension (typically channels for image tensors) */
    nx: number;
    /** Second-to-last dimension (typically width) */
    ny: number;
    /** Third-to-last dimension (typically height) */
    nz: number;
    /** Remaining dimensions multiplied together (default 1) */
    nf?: number;
}

import fpzipModule from './fpzip_loader.js';

let _fpzip: FpzipModule | null = null;
async function loadFpzip(): Promise<FpzipModule> {
    if (!_fpzip)
        _fpzip = await fpzipModule() as FpzipModule
    return _fpzip!
}

export async function decompress(data: Uint8Array): Promise<Float32Array> {
    const fpzip = await loadFpzip()

    let inputPtr: number | null = null
    let headerPtr: number | null = null
    let outputPtr: number | null = null

    try {
        // Allocate memory for the input data
        inputPtr = fpzip._malloc(data.length);
        fpzip.HEAPU8.set(data, inputPtr);

        // Extract header (type/prec/dims) so we can size the output buffer.
        headerPtr = fpzip._malloc(6 * 4)
        const outTypePtr = headerPtr + 0 * 4
        const outPrecPtr = headerPtr + 1 * 4
        const outNxPtr = headerPtr + 2 * 4
        const outNyPtr = headerPtr + 3 * 4
        const outNzPtr = headerPtr + 4 * 4
        const outNfPtr = headerPtr + 5 * 4

        const headerOk = fpzip._fpzip_get_header(
            inputPtr,
            outTypePtr,
            outPrecPtr,
            outNxPtr,
            outNyPtr,
            outNzPtr,
            outNfPtr,
        )
        if (!headerOk) {
            throw new Error(`Failed to read fpzip header (errno=${fpzip._fpzip_errno_code()})`)
        }

        const heap32 = new Int32Array(fpzip.HEAPU8.buffer)
        const base = headerPtr >> 2
        const type = heap32[base + 0]
        const nx = heap32[base + 2]
        const ny = heap32[base + 3]
        const nz = heap32[base + 4]
        const nf = heap32[base + 5]

        if (type !== 0) {
            throw new Error(`fpzip stream is not float32 (type=${type})`)
        }

        const nValues = nx * ny * nz * nf
        const outBytes = nValues * 4
        outputPtr = fpzip._malloc(outBytes)

        const bytesRead = fpzip._fpzip_decompress_f32(inputPtr, outputPtr)
        if (!bytesRead) {
            throw new Error(`Failed to decompress fpzip stream: ${fpzip._fpzip_errno_code()}`)
        }

        const floats = new Float32Array(nValues)
        floats.set(new Float32Array(fpzip.HEAPU8.buffer, outputPtr, nValues))
        return floats
    } finally {
        if (inputPtr !== null) fpzip._free(inputPtr)
        if (headerPtr !== null) fpzip._free(headerPtr)
        if (outputPtr !== null) fpzip._free(outputPtr)
    }
}

export async function compress(data: Float32Array, precBits = 0, dims?: TensorDims): Promise<Uint8Array> {
    const fpzip = await loadFpzip()

    let inputPtr: number | null = null
    let outputPtr: number | null = null

    try {
        inputPtr = fpzip._malloc(data.length * 4)
        new Float32Array(fpzip.HEAPU8.buffer).set(data, inputPtr >> 2)

        // Start with a conservative cap; fpzip should typically compress, but in worst
        // case it can slightly expand due to headers.
        let cap = Math.max(256, data.length * 4 + 1024)

        for (let attempt = 0; attempt < 8; attempt++) {
            if (outputPtr !== null) fpzip._free(outputPtr)
            outputPtr = fpzip._malloc(cap)

            let bytesWritten: number
            if (dims) {
                // NNC-compatible compression with proper tensor dimensions
                // NNC sets: nx=dims[last], ny=dims[last-1], nz=dims[last-2]
                // For [height, width, channels]: nx=channels, ny=width, nz=height
                bytesWritten = fpzip._fpzip_compress_f32_dims(
                    inputPtr,
                    dims.nx,
                    dims.ny,
                    dims.nz,
                    dims.nf ?? 1,
                    precBits,
                    outputPtr,
                    cap
                )
            } else {
                // Legacy flattened mode
                bytesWritten = fpzip._fpzip_compress_f32(inputPtr, data.length, precBits, outputPtr, cap)
            }

            if (bytesWritten > 0) {
                const out = new Uint8Array(bytesWritten)
                out.set(fpzip.HEAPU8.subarray(outputPtr, outputPtr + bytesWritten))
                return out
            }

            // Retry with a larger buffer if we overflowed.
            cap *= 2
        }

        throw new Error(`Failed to compress fpzip stream after retries: ${fpzip._fpzip_errno_code()}`)
    } finally {
        if (inputPtr !== null) fpzip._free(inputPtr)
        if (outputPtr !== null) fpzip._free(outputPtr)
    }
}