qoi decoder

qoi.ts

export async function decode(stream: AsyncIterable<Uint8Array>): Promise<ImageData> {
	let smallBuffer = new Uint8Array(14);

	let off = 0;
	let bufferEnd = 0;
	let buffer = smallBuffer;
	
	let needed = 14;
	let decode = decodeHeader;

	let width: number;
	let height: number;
	let channels: number;
	let colorspace: number;
	let output: Uint8Array;
	let output32: Uint32Array;
	let outOff = 0;

	let prevPixelOffset = 0;
	let hashTable = new Uint32Array(64);

	let runSize = 0;

nextChunk:
	for await (let chunk of stream) {
		if (needed > smallBuffer.length) {
			throw new Error(`trailing data after qoi`);
		}
		let chunkOff = -(bufferEnd - off);
		for (; chunkOff < 0; ) {
			let inBuf = bufferEnd - off;

			if (inBuf < needed) {
				for (let i = 0; i < inBuf; i++) {
					smallBuffer[i] = buffer[i + off];
				}
				buffer = smallBuffer;
				off = 0;
				bufferEnd = inBuf;

				for (;bufferEnd < needed; ) {
					smallBuffer[bufferEnd++] = chunk[chunkOff + bufferEnd];
				}
			}

			if (bufferEnd < needed) {
				continue nextChunk;
			}
			
			let startOff = off;
			decode();
			chunkOff += off - startOff;
		}

		off = chunkOff
		buffer = chunk;
		bufferEnd = chunk.length;

		for (; bufferEnd - off >= needed ; ) {
			decode();
		}
	}

	if (bufferEnd !== off) {
		throw new Error(`trailing data after qoi`);
	}
	if (needed !== Infinity) {
		throw new Error(`truncated qoi`);
	}

	return new ImageData(new Uint8ClampedArray(output!.buffer, output!.byteOffset, output!.byteLength), width!, height!);

	function u32(): number {
		return (buffer[off++] << 24 | buffer[off++] << 16 | buffer[off++] << 8 | buffer[off++]) >>> 0;
	}

	function decodeHeader() {
		let header = u32();
		if (header !== 0x716f6966) {
			throw new Error(`Invalid magic, got ${header.toString(16)}, not 'qoif'`);
		}

		width = u32();
		height = u32();
		channels = buffer[off++];
		colorspace = buffer[off++];

		output = new Uint8Array(width * height * 4);
		output[3] = 0xFF;
		output32 = new Uint32Array(output.buffer);

		decode = decodePixels;
		needed = 8;
	}

	function decodePixels() {
		for (; bufferEnd - off >= 8; ) {
			if (outOff == output.length) {
				decode = decodeTrailer;
				return;
			}

			let tag = buffer[off++];
			let tag2 = tag >> 6;

			if (tag2 == 3 && tag < 0xFE) {
				runSize += (tag & 0x3f) + 1;

				if (outOff + (runSize * 4) < output.length) {
					continue;
				}
			}

			if (runSize > 0) {
				if (runSize < 12) {
					// yes, the overhead of calling fill high enough on firefox that doing this manually is faster for small fills
					for (let i = 0; i < runSize; i++) {
						output32[(outOff >> 2) + i] = output32[prevPixelOffset >> 2];
					}
				} else {
					output32.fill(output32[prevPixelOffset >> 2], outOff >> 2, (outOff >> 2) + runSize);
				}
				outOff += runSize * 4;
				hashTable[((output[outOff - 4] * 3) + (output[outOff - 3] * 5) + (output[outOff - 2] * 7) + (output[outOff - 1] * 11)) & 0x3f] = output32[(outOff - 4) >> 2];
				runSize = 0;
			}

			if (tag === 0xFE) {
				output[outOff++] = buffer[off++];
				output[outOff++] = buffer[off++];
				output[outOff++] = buffer[off++];
				output[outOff++] = output[prevPixelOffset + 3];
			} else if (tag === 0xFF) {
				output[outOff++] = buffer[off++];
				output[outOff++] = buffer[off++];
				output[outOff++] = buffer[off++];
				output[outOff++] = buffer[off++];
			} else if (tag2 === 0) {
				let value = hashTable[tag & 0x3f];
				output32[outOff >> 2] = value;
				outOff += 4;
			} else if (tag2 === 1) {
				output[outOff++] = output[prevPixelOffset++] + ((tag >> 4) & 3) - 2;
				output[outOff++] = output[prevPixelOffset++] + ((tag >> 2) & 3) - 2;
				output[outOff++] = output[prevPixelOffset++] + (tag & 3) - 2;
				output[outOff++] = output[prevPixelOffset++];
			} else if (tag2 === 2) {
				let dg = (tag & 0x3f) - 32;
				let drb = buffer[off++];
				output[outOff++] = output[prevPixelOffset++] + dg + (drb >> 4) - 8;
				output[outOff++] = output[prevPixelOffset++] + dg;
				output[outOff++] = output[prevPixelOffset++] + dg + (drb & 0xF) - 8;
				output[outOff++] = output[prevPixelOffset++];
			}

			hashTable[((output[outOff - 4] * 3) + (output[outOff - 3] * 5) + (output[outOff - 2] * 7) + (output[outOff - 1] * 11)) & 0x3f] = output32[(outOff - 4) >> 2];
			prevPixelOffset = outOff - 4;
		}
	}

	function decodeTrailer() {
		for (let want of [0, 0, 0, 0, 0, 0, 0, 1]) {
			let v = buffer[off++]
			if (v !== want) {
				throw new Error(`Invalid trailer byte ${v}, wanted ${want}`);
			}
		}
		needed = Infinity;
	}
}