package render_ir_tests

import ir "../../render_ir"
import "core:testing"

@(test)
test_add_compute_dispatch_sets_payload :: proc(t: ^testing.T) {
	frame := ir.init_frame_ir()
	defer ir.destroy_frame_ir(&frame)

	pass := ir.add_pass(&frame, {kind = .Compute, name = "compute"})
	pipeline := ir.add_compute_pipeline(
		&frame,
		"particles",
		ir.default_compute_pipeline_state(ir.Shader_Handle(1)),
	)
	buffer_a := ir.add_buffer(&frame, "a", 64, {.Storage}, {}, .Imported)
	buffer_b := ir.add_buffer(&frame, "b", 128, {.Storage}, {}, .Imported)
	bindings := []ir.Descriptor_Resource_Binding {
		{binding = 0, type = .Storage_Buffer, resource = buffer_a, size = 64},
		{binding = 1, type = .Storage_Buffer, resource = buffer_b, size = 128},
	}
	push_constants := [12]u8{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
	cmd := ir.add_compute_dispatch(
		&frame,
		pass,
		pipeline,
		3,
		{4, 2, 1},
		bindings,
		raw_data(push_constants[:]),
		u32(len(push_constants)),
	)
	push_constants[0] = 99

	dispatch, ok := ir.get_dispatch(&frame, cmd)
	testing.expect(t, ok)
	testing.expect(t, dispatch.pass == pass)
	testing.expect(t, dispatch.pipeline == pipeline)
	testing.expect_value(t, dispatch.groups, [3]u32{4, 2, 1})
	testing.expect_value(t, dispatch.shader_id, u32(3))
	testing.expect_value(t, dispatch.storage_binding_count, u8(2))
	testing.expect_value(t, dispatch.storage_bindings[0].binding, u32(0))
	testing.expect_value(t, dispatch.storage_bindings[0].resource, buffer_a)
	testing.expect_value(t, dispatch.storage_bindings[1].binding, u32(1))
	testing.expect_value(t, dispatch.storage_bindings[1].resource, buffer_b)
	testing.expect_value(t, len(dispatch.push_constants), 12)
	testing.expect_value(t, dispatch.push_constants[0], u8(1))
	testing.expect_value(t, dispatch.push_constants[11], u8(12))
	testing.expect_value(t, dispatch.barrier_after, true)
}

@(test)
test_add_compute_dispatch_clamps_storage_buffers :: proc(t: ^testing.T) {
	frame := ir.init_frame_ir()
	defer ir.destroy_frame_ir(&frame)

	pass := ir.add_pass(&frame, {kind = .Compute, name = "compute"})
	pipeline := ir.add_compute_pipeline(
		&frame,
		"particles",
		ir.default_compute_pipeline_state(ir.Shader_Handle(1)),
	)
	bindings: [ir.MAX_COMPUTE_BINDINGS + 2]ir.Descriptor_Resource_Binding
	for i in 0 ..< len(bindings) {
		resource := ir.add_buffer(&frame, "buffer", 64, {.Storage}, {}, .Imported)
		bindings[i] = {
			binding  = u32(i),
			type     = .Storage_Buffer,
			resource = resource,
			size     = 64,
		}
	}

	cmd := ir.add_compute_dispatch(&frame, pass, pipeline, 3, {1, 1, 1}, bindings[:])

	dispatch, ok := ir.get_dispatch(&frame, cmd)
	testing.expect(t, ok)
	testing.expect_value(t, dispatch.storage_binding_count, u8(ir.MAX_COMPUTE_BINDINGS))
	testing.expect_value(t, dispatch.storage_bindings[0].binding, u32(0))
	testing.expect_value(
		t,
		dispatch.storage_bindings[ir.MAX_COMPUTE_BINDINGS - 1].binding,
		u32(ir.MAX_COMPUTE_BINDINGS - 1),
	)
}

@(test)
test_trace_compute_dispatch :: proc(t: ^testing.T) {
	frame := ir.init_frame_ir()
	defer ir.destroy_frame_ir(&frame)

	pass := ir.add_pass(&frame, {kind = .Compute, name = "compute"})
	pipeline := ir.add_compute_pipeline(
		&frame,
		"particles",
		ir.default_compute_pipeline_state(ir.Shader_Handle(1)),
	)
	push_constants := [12]u8{}
	buffer_a := ir.add_buffer(&frame, "a", 64, {.Storage}, {}, .Imported)
	buffer_b := ir.add_buffer(&frame, "b", 128, {.Storage}, {}, .Imported)
	_ = ir.add_compute_dispatch(
		&frame,
		pass,
		pipeline,
		3,
		{4, 2, 1},
		[]ir.Descriptor_Resource_Binding {
			{binding = 0, type = .Storage_Buffer, resource = buffer_a, size = 64},
			{binding = 1, type = .Storage_Buffer, resource = buffer_b, size = 128},
		},
		raw_data(push_constants[:]),
		u32(len(push_constants)),
	)

	trace := ir.trace_frame_ir(&frame)
	defer delete(trace)

	expected :=
		"Trace\n" +
		"begin_pass #1 Compute name=\"compute\"\n" +
		"bind_pipeline #1 Compute name=\"particles\"\n" +
		"dispatch groups=4,2,1 shader_id=3 buffers=0:1|1:2 push_constants_size=12 barrier_after=true\n" +
		"end_pass #1"
	testing.expect_value(t, trace, expected)
}

@(test)
test_dump_compute_dispatch :: proc(t: ^testing.T) {
	frame := ir.init_frame_ir()
	defer ir.destroy_frame_ir(&frame)

	pass := ir.add_pass(&frame, {kind = .Compute, name = "compute"})
	pipeline := ir.add_compute_pipeline(
		&frame,
		"particles",
		ir.default_compute_pipeline_state(ir.Shader_Handle(1)),
	)
	push_constants := [12]u8{}
	buffer_a := ir.add_buffer(&frame, "a", 64, {.Storage}, {}, .Imported)
	buffer_b := ir.add_buffer(&frame, "b", 128, {.Storage}, {}, .Imported)
	_ = ir.add_compute_dispatch(
		&frame,
		pass,
		pipeline,
		3,
		{4, 2, 1},
		[]ir.Descriptor_Resource_Binding {
			{binding = 0, type = .Storage_Buffer, resource = buffer_a, size = 64},
			{binding = 1, type = .Storage_Buffer, resource = buffer_b, size = 128},
		},
		raw_data(push_constants[:]),
		u32(len(push_constants)),
	)

	dump := ir.dump_frame_ir(&frame)
	defer delete(dump)

	expected :=
		"Frame_IR\n" +
		"resources 2\n" +
		"  #1 Buffer name=\"a\" lifetime=Imported size=64 usage=Storage memory=None\n" +
		"  #2 Buffer name=\"b\" lifetime=Imported size=128 usage=Storage memory=None\n" +
		"passes 1\n" +
		"  #1 Compute name=\"compute\"\n" +
		"shaders 0\n" +
		"descriptor_layouts 0\n" +
		"descriptor_sets 0\n" +
		"pipelines 1\n" +
		"  #1 Compute name=\"particles\" key=\"compute:shader=1;layouts=None;pc=0;pc_stages=None\"\n" +
		"commands 1\n" +
		"  #1 Dispatch pass=1 pipeline=1 groups=4,2,1 shader_id=3 buffers=0:1|1:2 push_constants_size=12 barrier_after=true"
	testing.expect_value(t, dump, expected)
}