pyrowave_c_test.cpp

// Copyright (c) 2026 Hans-Kristian Arntzen
// Copyright (c) 2026 AlkaidLab contributors
// SPDX-License-Identifier: MIT AND GPL-3.0-only

#include "pyrowave.h"
#include <stdio.h>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <vector>
#include <memory>

// Smoke test the C API.

#define ASSERT_THAT(x) do { \
	if (!(x)) { fprintf(stderr, "Fatal error executing %s at line %d.\n", #x, __LINE__); std::terminate(); } \
} while(false)

#define CHECKED(x) do { \
	pyrowave_result res = x; \
	if (res != PYROWAVE_SUCCESS) { fprintf(stderr, "Got pyrowave result %d while executing %s at line %d.\n", res, #x, __LINE__); std::terminate(); } \
} while(false)

static void test_encoder_create_validation()
{
	pyrowave_encoder_create_info info = {};
	info.width = 64;
	info.height = 64;

	// No device.
	pyrowave_encoder encoder, dummy;
	ASSERT_THAT(pyrowave_encoder_create(&info, &encoder) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	pyrowave_device device;
	CHECKED(pyrowave_create_default_device(&device));

	info.device = device;
	CHECKED(pyrowave_encoder_create(&info, &encoder));
	ASSERT_THAT(pyrowave_encoder_set_frame_context(encoder, -1) == PYROWAVE_ERROR_INVALID_ARGUMENT);
	ASSERT_THAT(pyrowave_encoder_set_frame_context(encoder, PYROWAVE_MAX_FRAME_CONTEXTS) ==
	            PYROWAVE_ERROR_INVALID_ARGUMENT);

	pyrowave_color_metadata invalid_metadata = {};
	invalid_metadata.transfer = static_cast<pyrowave_transfer_function>(-1);
	ASSERT_THAT(pyrowave_encoder_set_color_metadata(encoder, &invalid_metadata) == PYROWAVE_ERROR_INVALID_ARGUMENT);
	invalid_metadata = {};
	invalid_metadata.transform = static_cast<pyrowave_ycbcr_transform>(2);
	ASSERT_THAT(pyrowave_encoder_set_color_metadata(encoder, &invalid_metadata) == PYROWAVE_ERROR_INVALID_ARGUMENT);
	invalid_metadata = {};
	invalid_metadata.range = static_cast<pyrowave_ycbcr_range>(2);
	ASSERT_THAT(pyrowave_encoder_set_color_metadata(encoder, &invalid_metadata) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// 0 size not allowed.
	info.width = 0;
	info.height = 0;
	ASSERT_THAT(pyrowave_encoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Odd size not allowed.
	info.width = 65;
	info.height = 64;
	ASSERT_THAT(pyrowave_encoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	info.width = 64;
	info.height = 65;
	ASSERT_THAT(pyrowave_encoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Odd size allowed for 444.
	info.chroma = PYROWAVE_CHROMA_SUBSAMPLING_444;
	CHECKED(pyrowave_encoder_create(&info, &dummy));
	pyrowave_encoder_destroy(dummy);

	pyrowave_encoder_destroy(encoder);
	pyrowave_device_destroy(device);
}

static void test_decoder_create_validation()
{
	pyrowave_decoder_create_info info = {};
	info.width = 64;
	info.height = 64;

	// No device.
	pyrowave_decoder decoder, dummy;
	ASSERT_THAT(pyrowave_decoder_create(&info, &decoder) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	pyrowave_device device;
	CHECKED(pyrowave_create_default_device(&device));

	info.device = device;
	CHECKED(pyrowave_decoder_create(&info, &decoder));

	// 0 size not allowed.
	info.width = 0;
	info.height = 0;
	ASSERT_THAT(pyrowave_decoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Odd size not allowed.
	info.width = 65;
	info.height = 64;
	ASSERT_THAT(pyrowave_decoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	info.width = 64;
	info.height = 65;
	ASSERT_THAT(pyrowave_decoder_create(&info, &dummy) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Odd size allowed for 444.
	info.chroma = PYROWAVE_CHROMA_SUBSAMPLING_444;
	CHECKED(pyrowave_decoder_create(&info, &dummy));
	pyrowave_decoder_destroy(dummy);

	// Smoke test that creating device on fragment path doesn't explode.
	info.fragment_path = true;
	CHECKED(pyrowave_decoder_create(&info, &dummy));
	pyrowave_decoder_destroy(dummy);

	pyrowave_decoder_destroy(decoder);
	pyrowave_device_destroy(device);
}

static void test_decode_cpu_buffer_validation(bool fragment_path)
{
	pyrowave_decoder_create_info info = {};
	info.width = 16;
	info.height = 16;
	info.fragment_path = fragment_path;

	pyrowave_decoder decoder;
	CHECKED(pyrowave_create_default_device(&info.device));
	CHECKED(pyrowave_decoder_create(&info, &decoder));

	// This shouldn't be ready.
	ASSERT_THAT(!pyrowave_decoder_decode_is_ready(decoder, false));
	ASSERT_THAT(!pyrowave_decoder_decode_is_ready(decoder, true));

	pyrowave_cpu_buffer cpu_buffer = {};

	uint8_t luma[16][16] = {};
	uint8_t cb[8][16] = {}; // Test strided readback while we're at it.
	uint8_t cr[8][16] = {};

	cpu_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
	cpu_buffer.row_stride_in_bytes[0] = 16;
	cpu_buffer.row_stride_in_bytes[1] = 16;
	cpu_buffer.row_stride_in_bytes[2] = 16;
	cpu_buffer.plane_size_in_bytes[0] = 16 * 16;
	cpu_buffer.plane_size_in_bytes[1] = 16 * 8;
	cpu_buffer.plane_size_in_bytes[2] = 16 * 8;
	cpu_buffer.data[0] = &luma[0][0];
	cpu_buffer.data[1] = &cb[0][0];
	cpu_buffer.data[2] = &cr[0][0];

	cpu_buffer.width = 16;
	cpu_buffer.height = 16;

	// NV12 is banned for decode.
	cpu_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_NV12;
	ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &cpu_buffer) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	cpu_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
	CHECKED(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &cpu_buffer));

	// Assert that we do indeed decode a gray image.
	for (uint32_t y = 0; y < 16; y++)
		for (uint32_t x = 0; x < 16; x++)
			ASSERT_THAT(luma[y][x] == 0x7f || luma[y][x] == 0x80);

	for (uint32_t y = 0; y < 8; y++)
		for (uint32_t x = 0; x < 8; x++)
			ASSERT_THAT(cb[y][x] == 0x7f || cb[y][x] == 0x80);

	for (uint32_t y = 0; y < 8; y++)
		for (uint32_t x = 0; x < 8; x++)
			ASSERT_THAT(cr[y][x] == 0x7f || cr[y][x] == 0x80);

	pyrowave_decoder_destroy(decoder);
	pyrowave_device_destroy(info.device);
}

static void test_high_precision_cpu_buffer_readback(bool fragment_path)
{
	pyrowave_decoder_create_info info = {};
	info.width = 16;
	info.height = 16;
	info.fragment_path = fragment_path;
	pyrowave_decoder decoder;
	CHECKED(pyrowave_create_default_device(&info.device));
	CHECKED(pyrowave_decoder_create(&info, &decoder));

	pyrowave_decoder_set_ycbcr_scaling_factor(decoder, 2.0f);
	for (bool padded : {false, true})
	{
		pyrowave_cpu_buffer buffer = {};
		buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P16;
		buffer.width = info.width;
		buffer.height = info.height;
		std::vector<uint16_t> planes[3];
		for (int plane = 0; plane < 3; plane++)
		{
			const size_t width = plane == 0 ? 16 : 8;
			const size_t height = plane == 0 ? 16 : 8;
			buffer.row_stride_in_bytes[plane] = (width + (padded ? 2 : 0)) * sizeof(uint16_t);
			buffer.plane_size_in_bytes[plane] = buffer.row_stride_in_bytes[plane] * height;
			planes[plane].assign(buffer.plane_size_in_bytes[plane] / sizeof(uint16_t) + 8, 0xa55a);
			buffer.data[plane] = planes[plane].data();
		}
		for (int plane = 0; plane < 3; plane++)
		{
			const size_t height = plane == 0 ? 16 : 8;
			auto invalid = buffer;
			invalid.row_stride_in_bytes[plane]++;
			invalid.plane_size_in_bytes[plane] = invalid.row_stride_in_bytes[plane] * height;
			ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &invalid) ==
			            PYROWAVE_ERROR_INVALID_ARGUMENT);
			invalid.row_stride_in_bytes[plane] = SIZE_MAX - 1;
			invalid.plane_size_in_bytes[plane] = SIZE_MAX;
			ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &invalid) ==
			            PYROWAVE_ERROR_INVALID_ARGUMENT);
		}
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_async(decoder, &buffer, -1) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &buffer, -1) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		if (fragment_path)
		{
			auto unsupported = buffer;
			unsupported.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P10;
			ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_async(decoder, &unsupported, 0) ==
			            PYROWAVE_ERROR_NOT_IMPLEMENTED);
		}
		CHECKED(pyrowave_decoder_decode_cpu_buffer_async(decoder, &buffer, 0));
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_async(decoder, &buffer, 0) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		auto mismatched = buffer;
		mismatched.plane_size_in_bytes[0]++;
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &mismatched, 0) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		mismatched = buffer;
		mismatched.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &mismatched, 0) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		CHECKED(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &buffer, 0));
		ASSERT_THAT(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &buffer, 0) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
		for (int plane = 0; plane < 3; plane++)
		{
			const size_t width = plane == 0 ? 16 : 8;
			const size_t height = plane == 0 ? 16 : 8;
			const size_t pitch = buffer.row_stride_in_bytes[plane] / sizeof(uint16_t);
			for (size_t y = 0; y < height; y++)
				for (size_t x = 0; x < width; x++)
					ASSERT_THAT(planes[plane][y * pitch + x] == 0x7fff ||
					            planes[plane][y * pitch + x] == 0x8000);
			for (size_t index = buffer.plane_size_in_bytes[plane] / sizeof(uint16_t);
			     index < planes[plane].size(); index++)
				ASSERT_THAT(planes[plane][index] == 0xa55a);
		}
	}

	pyrowave_decoder_destroy(decoder);
	pyrowave_device_destroy(info.device);
}

static void test_encode_cpu_buffer_validation(bool nv12)
{
	pyrowave_encoder_create_info info = {};
	info.width = 16;
	info.height = 16;

	pyrowave_encoder encoder;
	CHECKED(pyrowave_create_default_device(&info.device));
	CHECKED(pyrowave_encoder_create(&info, &encoder));

	const pyrowave_rate_control rate_control = { 1024 };
	pyrowave_cpu_buffer cpu_buffer = {};

	uint8_t y[16][16] = {};
	uint8_t cb[8][8] = {};
	uint8_t cr[8][8] = {};

	cpu_buffer.format = nv12 ? PYROWAVE_CPU_BUFFER_FORMAT_NV12 : PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
	cpu_buffer.row_stride_in_bytes[0] = 16;
	cpu_buffer.row_stride_in_bytes[1] = nv12 ? 16 : 8;
	cpu_buffer.row_stride_in_bytes[2] = nv12 ? 0 : 8;
	cpu_buffer.plane_size_in_bytes[0] = 16 * 16;
	cpu_buffer.plane_size_in_bytes[1] = (nv12 ? 2 : 1) * 8 * 8;
	cpu_buffer.plane_size_in_bytes[2] = nv12 ? 0 : 8 * 8;
	cpu_buffer.data[0] = &y[0][0];
	cpu_buffer.data[1] = &cb[0][0];
	cpu_buffer.data[2] = &cr[0][0];

	cpu_buffer.width = 16;
	cpu_buffer.height = 16;
	CHECKED(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control));

	if (nv12)
	{
		auto misaligned = cpu_buffer;
		misaligned.row_stride_in_bytes[1]++;
		misaligned.plane_size_in_bytes[1] = misaligned.row_stride_in_bytes[1] * (info.height / 2);
		ASSERT_THAT(pyrowave_encoder_encode_cpu(encoder, &misaligned, &rate_control) ==
		            PYROWAVE_ERROR_INVALID_ARGUMENT);
	}

	// Mismatching width/height against encoder.
	cpu_buffer.width = 15;
	cpu_buffer.height = 16;
	ASSERT_THAT(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	cpu_buffer.width = 16;
	cpu_buffer.height = 15;
	ASSERT_THAT(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Too small row strides.
	cpu_buffer.width = 16;
	cpu_buffer.height = 16;
	cpu_buffer.row_stride_in_bytes[1] = nv12 ? 15 : 7;
	ASSERT_THAT(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	// Too small plane size.
	cpu_buffer.row_stride_in_bytes[1] = nv12 ? 16 : 8;
	cpu_buffer.plane_size_in_bytes[1] = (nv12 ? 2 : 1) * 8 * 8 - 1;
	ASSERT_THAT(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control) == PYROWAVE_ERROR_INVALID_ARGUMENT);

	pyrowave_encoder_destroy(encoder);
	pyrowave_device_destroy(info.device);
}

static void test_error_correction_api()
{
	pyrowave_device device;
	CHECKED(pyrowave_create_default_device(&device));

	constexpr int Width = 1920;
	constexpr int Height = 1080;

	pyrowave_decoder_create_info decoder_info = {};
	decoder_info.device = device;
	decoder_info.width = Width; // Test somewhat odd size. Quite relevant for fragment path as well.
	decoder_info.height = Height;
	decoder_info.chroma = PYROWAVE_CHROMA_SUBSAMPLING_420;

	pyrowave_encoder_create_info encoder_info = {};
	encoder_info.device = device;
	encoder_info.width = Width;
	encoder_info.height = Height;
	encoder_info.chroma = PYROWAVE_CHROMA_SUBSAMPLING_420;

	pyrowave_decoder decoder;
	pyrowave_encoder encoder;
	CHECKED(pyrowave_decoder_create(&decoder_info, &decoder));
	CHECKED(pyrowave_encoder_create(&encoder_info, &encoder));

	std::unique_ptr<uint8_t[]> luma(new uint8_t[Width * Height]);
	pyrowave_cpu_buffer cpu_buffer = {};

	cpu_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
	cpu_buffer.row_stride_in_bytes[0] = Width;
	cpu_buffer.row_stride_in_bytes[1] = Width / 2;
	cpu_buffer.row_stride_in_bytes[2] = Width / 2;
	cpu_buffer.plane_size_in_bytes[0] = Width * Height;
	cpu_buffer.plane_size_in_bytes[1] = Width * Height / 4;
	cpu_buffer.plane_size_in_bytes[2] = Width * Height / 4;
	cpu_buffer.data[0] = luma.get();
	cpu_buffer.data[1] = luma.get();
	cpu_buffer.data[2] = luma.get();

	const auto mirror = [](int value)
	{
		value &= 511;
		if (value >= 256)
			value = 511 - value;
		return value;
	};

	for (int y = 0; y < Height; y++)
		for (int x = 0; x < Width; x++)
			luma[y * Width + x] = uint8_t(mirror(13 * x + 17 * y));

	cpu_buffer.width = Width;
	cpu_buffer.height = Height;
	const pyrowave_rate_control rate_control = { 256 * 1024 };
	CHECKED(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control));

	size_t num_packets;
	CHECKED(pyrowave_encoder_compute_num_packets_with_padding(encoder, 4 * 1024, 2000, &num_packets));

	std::vector<uint8_t> bitstream(256 * 1024);
	std::vector<pyrowave_packet> packets(num_packets);
	CHECKED(pyrowave_encoder_packetize_with_padding(encoder, packets.data(), 4 * 1024, 2000, &num_packets, bitstream.data(), bitstream.size()));
	ASSERT_THAT(num_packets > 1);
	ASSERT_THAT(packets[0].offset == 0);
	ASSERT_THAT(packets[0].size < 4 * 1024 - 2000);
	for (size_t i = 1; i < num_packets; i++)
		ASSERT_THAT(packets[i].size <= 4 * 1024);

	{
		const void *mapped_bitstream = nullptr;
		const void *mapped_meta = nullptr;
		size_t mapped_bitstream_size = 0;
		size_t mapped_metadata_size = 0;
		CHECKED(pyrowave_encoder_get_mapped_raw_bitstream(encoder, &mapped_bitstream, &mapped_bitstream_size, &mapped_meta, &mapped_metadata_size));
		ASSERT_THAT(mapped_bitstream);
		ASSERT_THAT(mapped_meta);
		ASSERT_THAT(mapped_bitstream_size == rate_control.maximum_bitstream_size + mapped_metadata_size);
	}

	// Test a theoretical situation.
	for (int bands = 1; bands < 4; bands++)
	{
		size_t num_critical_packets;
		size_t num_active_blocks;
		CHECKED(pyrowave_encoder_get_num_active_blocks(encoder, bands, &num_active_blocks));

		std::vector<uint32_t> active_words((num_active_blocks + 31) / 32);
		CHECKED(pyrowave_encoder_compute_block_active_words(encoder, bands, active_words.data(), active_words.size()));

		CHECKED(pyrowave_encoder_compute_num_critical_packets(encoder, bands, 4 * 1024, 2000, &num_critical_packets));

		pyrowave_decoder_clear(decoder);

		for (size_t i = 0; i < num_critical_packets; i++)
		{
			CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packets[i].offset, packets[i].size));
			bool should_be_ready = i + 1 == num_critical_packets;
			ASSERT_THAT(pyrowave_decoder_decode_is_ready_with_sideband(
				decoder, true, bands, 0.0f,
				active_words.data(), active_words.size()) == should_be_ready);

			// If minimum packet ratio is large, we will fail due to that too.
			ASSERT_THAT(!pyrowave_decoder_decode_is_ready_with_sideband(
				decoder, true, bands, 0.5f,
				active_words.data(), active_words.size()));
		}
	}

	pyrowave_decoder_destroy(decoder);
	pyrowave_encoder_destroy(encoder);
	pyrowave_device_destroy(device);
}

static void test_basic_encoder_roundtrip(bool fragment_decode, bool nv12_encode, pyrowave_chroma_subsampling chroma)
{
	if (chroma == PYROWAVE_CHROMA_SUBSAMPLING_444 && nv12_encode)
		return;

	if (fragment_decode)
		return;

	pyrowave_device device;
	CHECKED(pyrowave_create_default_device(&device));

	constexpr int Width = 34;
	constexpr int Height = 30;

	pyrowave_decoder_create_info decoder_info = {};
	decoder_info.device = device;
	decoder_info.width = Width; // Test somewhat odd size. Quite relevant for fragment path as well.
	decoder_info.height = Height;
	decoder_info.fragment_path = fragment_decode;
	decoder_info.chroma = chroma;

	pyrowave_encoder_create_info encoder_info = {};
	encoder_info.device = device;
	encoder_info.width = Width;
	encoder_info.height = Height;
	encoder_info.chroma = chroma;

	pyrowave_decoder decoder;
	pyrowave_encoder encoder;
	CHECKED(pyrowave_decoder_create(&decoder_info, &decoder));
	CHECKED(pyrowave_encoder_create(&encoder_info, &encoder));

	const pyrowave_color_metadata expected_color_metadata = {
		PYROWAVE_COLOR_PRIMARIES_BT2020,
		PYROWAVE_TRANSFER_HLG,
		PYROWAVE_YCBCR_BT2020,
		PYROWAVE_YCBCR_FULL,
		1,
	};
	CHECKED(pyrowave_encoder_set_color_metadata(encoder, &expected_color_metadata));

	uint8_t luma[Height][Width] = {};
	uint8_t cb[Height][Width] = {};
	uint8_t cr[Height][Width] = {};
	uint8_t cbcr[Height][Width][2] = {};

	uint8_t decode_luma[Height][Width] = {};
	uint8_t decode_cb[Height][Width] = {};
	uint8_t decode_cr[Height][Width] = {};

	pyrowave_cpu_buffer cpu_buffer = {};

	cpu_buffer.format =
			nv12_encode
				? PYROWAVE_CPU_BUFFER_FORMAT_NV12
				: (chroma == PYROWAVE_CHROMA_SUBSAMPLING_444
					   ? PYROWAVE_CPU_BUFFER_FORMAT_YUV444P
					   : PYROWAVE_CPU_BUFFER_FORMAT_YUV420P);

	cpu_buffer.row_stride_in_bytes[0] = Width;
	cpu_buffer.row_stride_in_bytes[1] = nv12_encode ? sizeof(cbcr[0]) : sizeof(cb[0]);
	cpu_buffer.row_stride_in_bytes[2] = nv12_encode ? 0 : sizeof(cr[0]);
	cpu_buffer.plane_size_in_bytes[0] = sizeof(luma);
	cpu_buffer.plane_size_in_bytes[1] = nv12_encode ? sizeof(cbcr) : sizeof(cb);
	cpu_buffer.plane_size_in_bytes[2] = nv12_encode ? 0 : sizeof(cr);
	cpu_buffer.data[0] = &luma[0][0];

	if (nv12_encode)
	{
		cpu_buffer.data[1] = &cbcr[0][0][0];
	}
	else
	{
		cpu_buffer.data[1] = &cb[0][0];
		cpu_buffer.data[2] = &cr[0][0];
	}

	for (int y = 0; y < Height; y++)
	{
		for (int x = 0; x < Width; x++)
		{
			luma[y][x] = uint8_t(3 * x + 5 * y);

			uint8_t cb_signal = 7 * x + 3 * y;
			uint8_t cr_signal = 3 * x + 5 * y;

			if (nv12_encode)
			{
				cbcr[y][x][0] = cb_signal;
				cbcr[y][x][1] = cr_signal;
			}
			else
			{
				cb[y][x] = cb_signal;
				cr[y][x] = cr_signal;
			}
		}
	}

	cpu_buffer.width = Width;
	cpu_buffer.height = Height;
	const pyrowave_rate_control rate_control = { 64 * 1024 }; // Just give it something massive.
	CHECKED(pyrowave_encoder_encode_cpu(encoder, &cpu_buffer, &rate_control));

	size_t num_packets;
	CHECKED(pyrowave_encoder_compute_num_packets(encoder, 64 * 1024, &num_packets));
	ASSERT_THAT(num_packets == 1);

	CHECKED(pyrowave_encoder_compute_num_packets_with_padding(encoder, 64 * 1024, 64 * 1024 - 4, &num_packets));
	ASSERT_THAT(num_packets == 2);

	std::vector<uint8_t> bitstream(64 * 1024);
	std::vector<uint8_t> bitstream_padded(64 * 1024);
	pyrowave_packet packet = {};
	CHECKED(pyrowave_encoder_packetize(encoder, &packet, 64 * 1024, &num_packets, bitstream.data(), bitstream.size()));
	ASSERT_THAT(num_packets == 1);
	ASSERT_THAT(packet.offset == 0);
	ASSERT_THAT(packet.size != 0);
	ASSERT_THAT(packet.size <= bitstream.size());
	bitstream.resize(packet.size);

	// Just padding on its own should not change the bitstream in any way if the splits don't happen.
	CHECKED(pyrowave_encoder_packetize_with_padding(encoder, &packet, 64 * 1024, 16, &num_packets,
		bitstream_padded.data(), bitstream_padded.size()));
	ASSERT_THAT(num_packets == 1);
	ASSERT_THAT(packet.offset == 0);
	ASSERT_THAT(packet.size != 0);
	ASSERT_THAT(packet.size <= bitstream_padded.size());
	bitstream_padded.resize(packet.size);
	ASSERT_THAT(packet.size == bitstream.size());
	ASSERT_THAT(std::memcmp(bitstream.data(), bitstream_padded.data(), packet.size) == 0);

	CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packet.offset, packet.size));
	ASSERT_THAT(pyrowave_decoder_decode_is_ready(decoder, false));
	pyrowave_color_metadata decoded_color_metadata = {};
	ASSERT_THAT(pyrowave_decoder_get_color_metadata(decoder, &decoded_color_metadata));
	ASSERT_THAT(decoded_color_metadata.primaries == expected_color_metadata.primaries);
	ASSERT_THAT(decoded_color_metadata.transfer == expected_color_metadata.transfer);
	ASSERT_THAT(decoded_color_metadata.transform == expected_color_metadata.transform);
	ASSERT_THAT(decoded_color_metadata.range == expected_color_metadata.range);
	ASSERT_THAT(decoded_color_metadata.chroma_siting == expected_color_metadata.chroma_siting);
	ASSERT_THAT(pyrowave_decoder_decode_is_ready_with_sideband(decoder, false, 4, 0.0f, nullptr, 0));
	pyrowave_decoder_clear(decoder);
	pyrowave_color_metadata cleared_color_metadata = {};
	ASSERT_THAT(!pyrowave_decoder_get_color_metadata(decoder, &cleared_color_metadata));
	ASSERT_THAT(cleared_color_metadata.primaries == PYROWAVE_COLOR_PRIMARIES_BT709);
	ASSERT_THAT(cleared_color_metadata.transfer == PYROWAVE_TRANSFER_BT709);
	ASSERT_THAT(cleared_color_metadata.transform == PYROWAVE_YCBCR_BT709);
	ASSERT_THAT(cleared_color_metadata.range == PYROWAVE_YCBCR_FULL);
	ASSERT_THAT(cleared_color_metadata.chroma_siting == 0);
	ASSERT_THAT(!pyrowave_decoder_decode_is_ready(decoder, false));
	ASSERT_THAT(!pyrowave_decoder_decode_is_ready_with_sideband(decoder, false, 4, 0.0f, nullptr, 0));
	CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packet.offset, packet.size));
	ASSERT_THAT(pyrowave_decoder_decode_is_ready(decoder, false));
	ASSERT_THAT(pyrowave_decoder_decode_is_ready_with_sideband(decoder, false, 4, 0.0f, nullptr, 0));

	cpu_buffer.data[0] = &decode_luma[0][0];
	cpu_buffer.data[1] = &decode_cb[0][0];
	cpu_buffer.data[2] = &decode_cr[0][0];
	cpu_buffer.row_stride_in_bytes[1] = sizeof(decode_cb[0]);
	cpu_buffer.row_stride_in_bytes[2] = sizeof(decode_cr[0]);
	cpu_buffer.plane_size_in_bytes[1] = sizeof(decode_cb);
	cpu_buffer.plane_size_in_bytes[2] = sizeof(decode_cr);
	cpu_buffer.format = chroma == PYROWAVE_CHROMA_SUBSAMPLING_444 ?
			PYROWAVE_CPU_BUFFER_FORMAT_YUV444P : PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;

	CHECKED(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &cpu_buffer));

	for (int y = 0; y < Height; y++)
	{
		for (int x = 0; x < Width; x++)
		{
			int d = std::abs(int(decode_luma[y][x]) - int(luma[y][x]));
			// With the "infinite" bitrate we get here,
			// accept a maximum 1 ULP error.
			ASSERT_THAT(d <= 1);

			if (chroma == PYROWAVE_CHROMA_SUBSAMPLING_444 || (!nv12_encode && y < Height / 2 && x < Width / 2))
			{
				// Allow more error for chroma.
				d = std::abs(int(decode_cb[y][x]) - int(cb[y][x]));
				ASSERT_THAT(d <= 1);
				d = std::abs(int(decode_cr[y][x]) - int(cr[y][x]));
				ASSERT_THAT(d <= 1);
			}
		}
	}

	if (nv12_encode)
	{
		for (int y = 0; y < Height / 2; y++)
		{
			for (int x = 0; x < Width / 2; x++)
			{
				int d = std::abs(int(decode_cb[y][x]) - int(cbcr[y][x][0]));
				ASSERT_THAT(d <= 1);
				d = std::abs(int(decode_cr[y][x]) - int(cbcr[y][x][1]));
				ASSERT_THAT(d <= 1);
			}
		}
	}

	pyrowave_decoder_destroy(decoder);
	pyrowave_encoder_destroy(encoder);
	pyrowave_device_destroy(device);
}

static void test_extended_cpu_formats(bool subsampled)
{
	pyrowave_device device;
	CHECKED(pyrowave_create_default_device(&device));

	constexpr int Width = 16;
	constexpr int Height = 16;

	pyrowave_encoder_create_info encoder_info = {};
	encoder_info.device = device;
	encoder_info.width = Width;
	encoder_info.height = Height;
	encoder_info.chroma = subsampled ? PYROWAVE_CHROMA_SUBSAMPLING_420 : PYROWAVE_CHROMA_SUBSAMPLING_444;

	pyrowave_decoder_create_info decoder_info = {};
	decoder_info.device = device;
	decoder_info.width = Width;
	decoder_info.height = Height;
	decoder_info.chroma = subsampled ? PYROWAVE_CHROMA_SUBSAMPLING_420 : PYROWAVE_CHROMA_SUBSAMPLING_444;

	pyrowave_decoder decoder;
	pyrowave_encoder encoder;
	CHECKED(pyrowave_decoder_create(&decoder_info, &decoder));
	CHECKED(pyrowave_encoder_create(&encoder_info, &encoder));

	pyrowave_encoder_set_ycbcr_scaling_factor(encoder, 0.5f);
	pyrowave_decoder_set_ycbcr_scaling_factor(decoder, 2.0f);

	const struct
	{
		pyrowave_cpu_buffer_format format;
		int y, cb, cr;
	} tests[] = {
		{ subsampled ? PYROWAVE_CPU_BUFFER_FORMAT_YUV420P10 : PYROWAVE_CPU_BUFFER_FORMAT_YUV444P10, 0x3ff, 0x200, 0 },
		{ subsampled ? PYROWAVE_CPU_BUFFER_FORMAT_YUV420P16 : PYROWAVE_CPU_BUFFER_FORMAT_YUV444P16, 0xffff, 0x8000, 0 },
	};

	for (auto &test : tests)
	{
		// 10-bit encode path.
		uint16_t luma[Width * Height];
		uint16_t cb[Width * Height];
		uint16_t cr[Width * Height];

		for (int i = 0; i < Width * Height; i++)
		{
			luma[i] = test.y;
			cb[i] = test.cb;
			cr[i] = test.cr;
		}

		pyrowave_cpu_buffer encode_buffer = {}, decode_buffer = {};

		encode_buffer.data[0] = luma;
		encode_buffer.data[1] = cb;
		encode_buffer.data[2] = cr;
		encode_buffer.format = test.format;
		encode_buffer.width = Width;
		encode_buffer.height = Height;
		for (int i = 0; i < 3; i++)
		{
			encode_buffer.row_stride_in_bytes[i] = Width * sizeof(uint16_t);
			encode_buffer.plane_size_in_bytes[i] = Width * Height * sizeof(uint16_t);
			if (i && subsampled)
				encode_buffer.plane_size_in_bytes[i] /= 2;
		}

		decode_buffer = encode_buffer;

		const pyrowave_rate_control rate_control = { 64 * 1024 };

		// Reuse this encoder after 8-bit planar/NV12 input so cached upload
		// images must change format before reading the 10/16-bit input.
		uint8_t warmup_planes[3][Width * Height];
		memset(warmup_planes, 128, sizeof(warmup_planes));
		pyrowave_cpu_buffer warmup = {};
		warmup.width = Width;
		warmup.height = Height;
		warmup.format = subsampled ? PYROWAVE_CPU_BUFFER_FORMAT_YUV420P : PYROWAVE_CPU_BUFFER_FORMAT_YUV444P;
		for (int plane = 0; plane < 3; plane++)
		{
			warmup.data[plane] = warmup_planes[plane];
			warmup.row_stride_in_bytes[plane] = Width;
			warmup.plane_size_in_bytes[plane] = Width * (subsampled && plane != 0 ? Height / 2 : Height);
		}
		size_t warmup_packets = 0;
		CHECKED(pyrowave_encoder_encode_cpu(encoder, &warmup, &rate_control));
		CHECKED(pyrowave_encoder_compute_num_packets(encoder, rate_control.maximum_bitstream_size, &warmup_packets));
		if (subsampled)
		{
			warmup.format = PYROWAVE_CPU_BUFFER_FORMAT_NV12;
			warmup.data[2] = nullptr;
			CHECKED(pyrowave_encoder_encode_cpu(encoder, &warmup, &rate_control));
			CHECKED(pyrowave_encoder_compute_num_packets(encoder, rate_control.maximum_bitstream_size, &warmup_packets));
		}

		std::vector<uint8_t> bitstream;
		bitstream.reserve(rate_control.maximum_bitstream_size);
		CHECKED(pyrowave_encoder_encode_cpu(encoder, &encode_buffer, &rate_control));

		size_t after_packets = 0;
		pyrowave_packet packet;
		bitstream.resize(rate_control.maximum_bitstream_size);
		CHECKED(pyrowave_encoder_packetize(encoder, &packet, rate_control.maximum_bitstream_size,
				&after_packets, bitstream.data(), bitstream.size()));

		CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packet.offset, packet.size));
		decode_buffer.format = subsampled ? PYROWAVE_CPU_BUFFER_FORMAT_YUV420P16 : PYROWAVE_CPU_BUFFER_FORMAT_YUV444P16;
		CHECKED(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &decode_buffer));

		const auto abs_diff = [](int a, int b) { return std::abs(a - b); };

		int ref_y, ref_cb, ref_cr;

		if (test.format == PYROWAVE_CPU_BUFFER_FORMAT_YUV444P10 ||
		    test.format == PYROWAVE_CPU_BUFFER_FORMAT_YUV420P10)
		{
			ref_y = (test.y * 0xffff + 0x1ff) / 0x3ff;
			ref_cb = (test.cb * 0xffff + 0x1ff) / 0x3ff;
			ref_cr = (test.cr * 0xffff + 0x1ff) / 0x3ff;
		}
		else
		{
			ref_y = test.y;
			ref_cb = test.cb;
			ref_cr = test.cr;
		}

		// The quantizers and FP math isn't accurate enough to resolve full 16-bit accuracy unless perhaps
		// using full 32-bit precision everywhere.
		for (int y = 0; y < Height; y++)
		{
			for (int x = 0; x < Width; x++)
			{
				ASSERT_THAT(abs_diff(luma[y * Width + x], ref_y) < 0x40);
				if (!subsampled || (x < Width / 2 && y < Height / 2))
				{
					ASSERT_THAT(abs_diff(cb[y * Width + x], ref_cb) < 0x40);
					ASSERT_THAT(abs_diff(cr[y * Width + x], ref_cr) < 0x40);
				}
			}
		}

		CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packet.offset, packet.size));
		decode_buffer.format = subsampled ? PYROWAVE_CPU_BUFFER_FORMAT_YUV420P10 : PYROWAVE_CPU_BUFFER_FORMAT_YUV444P10;
		CHECKED(pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &decode_buffer));

		if (test.format == PYROWAVE_CPU_BUFFER_FORMAT_YUV444P16 ||
		    test.format == PYROWAVE_CPU_BUFFER_FORMAT_YUV420P16)
		{
			ref_y = (test.y * 0x3ff + 0x7fff) / 0xffff;
			ref_cb = (test.cb * 0x3ff + 0x7fff) / 0xffff;
			ref_cr = (test.cr * 0x3ff + 0x7fff) / 0xffff;
		}
		else
		{
			ref_y = test.y;
			ref_cb = test.cb;
			ref_cr = test.cr;
		}

		for (int y = 0; y < Height; y++)
		{
			for (int x = 0; x < Width; x++)
			{
				ASSERT_THAT(abs_diff(luma[y * Width + x], ref_y) <= 1);
				if (!subsampled || (x < Width / 2 && y < Height / 2))
				{
					ASSERT_THAT(abs_diff(cb[y * Width + x], ref_cb) <= 1);
					ASSERT_THAT(abs_diff(cr[y * Width + x], ref_cr) <= 1);
				}
			}
		}
	}
}

static void test_basic_system_stability(bool realtime_prio)
{
	pyrowave_device device;

	if (realtime_prio)
	{
		CHECKED(pyrowave_create_device_by_compat(0, 0, nullptr, nullptr, nullptr,
			VK_QUEUE_GLOBAL_PRIORITY_REALTIME, &device));
	}
	else
	{
		CHECKED(pyrowave_create_default_device(&device));
	}

	fprintf(stderr, "Got global priority: expected %u, got %u\n",
	        realtime_prio ? VK_QUEUE_GLOBAL_PRIORITY_REALTIME : VK_QUEUE_GLOBAL_PRIORITY_MEDIUM,
	        pyrowave_device_get_global_priority(device));

	// 4K, upper bound of normal usage.
	constexpr int Width = 3840;
	constexpr int Height = 2160;

	pyrowave_decoder_create_info decoder_info = {};
	decoder_info.device = device;
	decoder_info.width = Width; // Test somewhat odd size. Quite relevant for fragment path as well.
	decoder_info.height = Height;

	pyrowave_encoder_create_info encoder_info = {};
	encoder_info.device = device;
	encoder_info.width = Width;
	encoder_info.height = Height;

	pyrowave_decoder decoder;
	pyrowave_encoder encoder;
	CHECKED(pyrowave_decoder_create(&decoder_info, &decoder));
	CHECKED(pyrowave_encoder_create(&encoder_info, &encoder));

	std::vector<uint8_t> luma(Width * Height);
	std::vector<uint16_t> cbcr(Width * Height / 4);
	std::vector<uint8_t> decode_luma(Width * Height);
	std::vector<uint8_t> decode_cb(Width * Height / 4);
	std::vector<uint8_t> decode_cr(Width * Height / 4);

	pyrowave_cpu_buffer encode_buffer = {}, decode_buffer = {};

	encode_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_NV12;
	encode_buffer.row_stride_in_bytes[0] = Width;
	encode_buffer.row_stride_in_bytes[1] = Width;
	encode_buffer.plane_size_in_bytes[0] = Width * Height;
	encode_buffer.plane_size_in_bytes[1] = Width * Height / 2;
	encode_buffer.data[0] = luma.data();
	encode_buffer.data[1] = cbcr.data();
	encode_buffer.width = Width;
	encode_buffer.height = Height;

	decode_buffer.format = PYROWAVE_CPU_BUFFER_FORMAT_YUV420P;
	decode_buffer.row_stride_in_bytes[0] = Width;
	decode_buffer.row_stride_in_bytes[1] = Width / 2;
	decode_buffer.row_stride_in_bytes[2] = Width / 2;
	decode_buffer.plane_size_in_bytes[0] = Width * Height;
	decode_buffer.plane_size_in_bytes[1] = Width * Height / 4;
	decode_buffer.plane_size_in_bytes[2] = Width * Height / 4;
	decode_buffer.data[0] = decode_luma.data();
	decode_buffer.data[1] = decode_cb.data();
	decode_buffer.data[2] = decode_cr.data();
	decode_buffer.width = Width;
	decode_buffer.height = Height;

	const auto mirror = [](int v) -> uint8_t
	{
		v &= 511;
		if (v > 255)
			v = 511 - v;
		ASSERT_THAT(v >= 0 && v <= 255);
		return uint8_t(v);
	};

	// Just generate a synthetic signal.
	for (int y = 0; y < Height; y++)
		for (int x = 0; x < Width; x++)
			luma[y * Width + x] = mirror(3 * x + 5 * y);

	for (int y = 0; y < Height / 2; y++)
		for (int x = 0; x < Width / 2; x++)
			cbcr[y * Width / 2 + x] = mirror(7 * x + y * 3) * 0x100 + mirror(y * 5 + x * 7);

	std::vector<uint8_t> bitstream;
	std::vector<pyrowave_packet> packets;

	for (int iter = 0; iter < 100; iter++)
	{
		// 240mbit equivalent for 60 fps.
		const pyrowave_rate_control rate_control = { 500000 };

		// Get some test coverage for async compute path.
		CHECKED(pyrowave_device_set_queue_type(device, iter % 2 ? VK_QUEUE_COMPUTE_BIT : VK_QUEUE_GRAPHICS_BIT));

		// Very basic testing.
		CHECKED(pyrowave_encoder_set_frame_context(encoder, iter % 2));

		bitstream.reserve(rate_control.maximum_bitstream_size);
		CHECKED(pyrowave_encoder_encode_cpu(encoder, &encode_buffer, &rate_control));

		size_t num_packets, after_packets;
		CHECKED(pyrowave_encoder_compute_num_packets(encoder, 8 * 1024, &num_packets));
		ASSERT_THAT(num_packets > 1);

		packets.resize(num_packets);
		bitstream.resize(rate_control.maximum_bitstream_size);
		CHECKED(pyrowave_encoder_packetize(encoder, packets.data(), 8 * 1024, &after_packets, bitstream.data(), bitstream.size()));
		ASSERT_THAT(num_packets == after_packets);

		// Verify that the bitstream is sound. We should be able to decode it.
		size_t total_bitstream_size = 0;

		for (auto &packet : packets)
		{
			CHECKED(pyrowave_decoder_push_packet(decoder, bitstream.data() + packet.offset, packet.size));
			// When we push the last packet, we should get a complete frame.
			ASSERT_THAT(pyrowave_decoder_decode_is_ready(decoder, false) == (&packet == &packets.back()));
			total_bitstream_size += packet.size;
		}

		// Verify that we tightly hit our rate control budget.
		ASSERT_THAT(total_bitstream_size <= rate_control.maximum_bitstream_size);
		ASSERT_THAT(total_bitstream_size >= 95 * rate_control.maximum_bitstream_size / 100);

		CHECKED(pyrowave_decoder_decode_cpu_buffer_async(decoder, &decode_buffer, iter % 2));
		CHECKED(pyrowave_decoder_decode_cpu_buffer_complete(decoder, &decode_buffer, iter % 2));
	}

	pyrowave_decoder_destroy(decoder);
	pyrowave_encoder_destroy(encoder);

	// Verify that PSNR is under control.
	double y_error = 0.0;
	double cb_error = 0.0;
	double cr_error = 0.0;

	for (int i = 0; i < Width * Height; i++)
	{
		double y_d = double(luma[i]) - double(decode_luma[i]);
		y_error += y_d * y_d;
	}

	for (int i = 0; i < Width * Height / 4; i++)
	{
		double cb_d = double(cbcr[i] & 0xff) - double(decode_cb[i]);
		cb_error += cb_d * cb_d;
		double cr_d = double(cbcr[i] >> 8) - double(decode_cr[i]);
		cr_error += cr_d * cr_d;
	}

	double y_signal = 255.0 * 255.0 * Width * Height;
	double chroma_signal = 255.0 * 255.0 * (Width * Height / 4);
	double y_psnr = y_signal / y_error;
	double cb_psnr = chroma_signal / cb_error;
	double cr_psnr = chroma_signal / cr_error;

	// 40 dB, arbitrarily chosen for testing purposes.
	ASSERT_THAT(y_psnr > 10000.0);
	ASSERT_THAT(cb_psnr > 10000.0);
	ASSERT_THAT(cr_psnr > 10000.0);

	// We're not hitting 60 dB. That'd mean we cheated or something or added bugs in the test code.
	ASSERT_THAT(y_psnr < 1000000.0);
	ASSERT_THAT(cb_psnr < 1000000.0);
	ASSERT_THAT(cr_psnr < 1000000.0);

	int blah = 0;
	pyrowave_device_report_performance_stats(device, [](void *userdata, const char *msg)
	{
		*static_cast<int *>(userdata) = 42;
		printf("performance cb: %s\n", msg);
	}, &blah, true);

	// Verify that userdata gets passed correctly.
	ASSERT_THAT(blah == 42);

	pyrowave_device_destroy(device);
}

int main()
{
	printf("Running system stability test ...\n");
	test_basic_system_stability(false);
	printf("Running system stability test ... (REALTIME) \n");
	test_basic_system_stability(true);

	// Correctness tests for small-ish outputs.
	for (int variant = 0; variant < 8; variant++)
	{
		printf("Running roundtrip variant %d test ...\n", variant);
		test_basic_encoder_roundtrip(
			(variant & 1) != 0, (variant & 2) != 0,
			(variant & 4) != 0 ? PYROWAVE_CHROMA_SUBSAMPLING_444 : PYROWAVE_CHROMA_SUBSAMPLING_420);
	}

	test_extended_cpu_formats(false);
	test_extended_cpu_formats(true);

	test_error_correction_api();

	// Validate that we handle error inputs gracefully.
	printf("Running error handling tests ...\n");
	test_decode_cpu_buffer_validation(false);
	test_decode_cpu_buffer_validation(true);
	test_high_precision_cpu_buffer_readback(false);
	test_high_precision_cpu_buffer_readback(true);
	test_encode_cpu_buffer_validation(false);
	test_encode_cpu_buffer_validation(true);
	test_encoder_create_validation();
	test_decoder_create_validation();

	printf("Passed all tests :)\n");
}