/*
 * Copyright 2022 Philip Rebohle for Valve Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include "dxil_mesh.hpp"
#include "dxil_common.hpp"
#include "logging.hpp"
#include "opcodes/converter_impl.hpp"
#include "spirv_module.hpp"
#include "dxil_geometry.hpp"

namespace dxil_spv
{
bool emit_set_mesh_output_counts_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	Operation *op = impl.allocate(spv::OpSetMeshOutputsEXT);

	spv::Id num_vertex_id, num_prim_id;

	// If we have a degenerate case where either of these is 0, we have to declare 0 as max output.
	if (impl.execution_mode_meta.stage_output_num_vertex == 0)
		num_vertex_id = impl.builder().makeUintConstant(0);
	else
		num_vertex_id = impl.get_id_for_value(instruction->getOperand(1));
	op->add_id(num_vertex_id);

	if (impl.execution_mode_meta.stage_output_num_primitive == 0)
		num_prim_id = impl.builder().makeUintConstant(0);
	else
		num_prim_id = impl.get_id_for_value(instruction->getOperand(2));
	op->add_id(num_prim_id);

	if (impl.options.quirks.mesh_outputs_bounds_check)
	{
		// Technically the first thread wins so we need to broadcast the results,
		// but in all cases where we care about this quirk, the inputs are threadgroup uniform, so this is fine.
		auto *assert_vertex = impl.allocate(spv::OpUGreaterThan, impl.builder().makeBoolType());
		assert_vertex->add_id(num_vertex_id);
		assert_vertex->add_id(impl.builder().makeUintConstant(impl.execution_mode_meta.stage_output_num_vertex));

		auto *assert_prim = impl.allocate(spv::OpUGreaterThan, impl.builder().makeBoolType());
		assert_prim->add_id(num_prim_id);
		assert_prim->add_id(impl.builder().makeUintConstant(impl.execution_mode_meta.stage_output_num_primitive));

		auto *should_exit = impl.allocate(spv::OpLogicalOr, impl.builder().makeBoolType());
		should_exit->add_id(assert_vertex->id);
		should_exit->add_id(assert_prim->id);

		impl.add(assert_vertex);
		impl.add(assert_prim);
		impl.add(should_exit);

		auto *return_op = impl.allocate(spv::PseudoOpReturnCond);
		return_op->add_id(should_exit->id);
		impl.add(return_op);

		// OpSampledImage must be consumed in same block.
		// We'll split blocks here, so just recreate the combined sampler image if needed.
		impl.combined_image_sampler_cache.clear();
	}

	if (impl.options.instruction_instrumentation.enabled &&
	    impl.options.instruction_instrumentation.type == InstructionInstrumentationType::ExpectAssume)
	{
		auto *assert_vertex = impl.allocate(spv::OpULessThanEqual, impl.builder().makeBoolType());
		assert_vertex->add_id(num_vertex_id);
		assert_vertex->add_id(impl.builder().makeUintConstant(impl.execution_mode_meta.stage_output_num_vertex));
		auto *assume_vertex = impl.allocate(spv::OpAssumeTrueKHR);
		assume_vertex->add_id(assert_vertex->id);

		auto *assert_prim = impl.allocate(spv::OpULessThanEqual, impl.builder().makeBoolType());
		assert_prim->add_id(num_prim_id);
		assert_prim->add_id(impl.builder().makeUintConstant(impl.execution_mode_meta.stage_output_num_primitive));
		auto *assume_prim = impl.allocate(spv::OpAssumeTrueKHR);
		assume_prim->add_id(assert_prim->id);

		impl.add(assert_vertex);
		impl.add(assume_vertex);
		impl.add(assert_prim);
		impl.add(assume_prim);
	}

	impl.add(op);

	// Workaround shader compiler bugs by emitting a conditional return.
	if (!impl.shader_analysis.has_side_effects)
	{
		auto &builder = impl.builder();

		builder.addCapability(spv::CapabilityGroupNonUniformBallot);
		spv::Id num_ids = impl.spirv_module.get_builtin_shader_input(spv::BuiltInNumSubgroups);
		auto *load_op = impl.allocate(spv::OpLoad, builder.makeUintType(32));
		load_op->add_id(num_ids);
		impl.add(load_op);

		auto *one_subgroup = impl.allocate(spv::OpIEqual, builder.makeBoolType());
		one_subgroup->add_id(load_op->id);
		one_subgroup->add_id(builder.makeUintConstant(1));
		impl.add(one_subgroup);

		auto *broadcast_vert = impl.allocate(spv::OpGroupNonUniformBroadcastFirst, builder.makeUintType(32));
		broadcast_vert->add_id(builder.makeUintConstant(spv::ScopeSubgroup));
		broadcast_vert->add_id(num_vertex_id);
		impl.add(broadcast_vert);

		auto *broadcast_prim = impl.allocate(spv::OpGroupNonUniformBroadcastFirst, builder.makeUintType(32));
		broadcast_prim->add_id(builder.makeUintConstant(spv::ScopeSubgroup));
		broadcast_prim->add_id(num_prim_id);
		impl.add(broadcast_prim);

		auto *vert_zero = impl.allocate(spv::OpIEqual, builder.makeBoolType());
		vert_zero->add_id(broadcast_vert->id);
		vert_zero->add_id(builder.makeUintConstant(0));
		impl.add(vert_zero);

		auto *prim_zero = impl.allocate(spv::OpIEqual, builder.makeBoolType());
		prim_zero->add_id(broadcast_prim->id);
		prim_zero->add_id(builder.makeUintConstant(0));
		impl.add(prim_zero);

		auto *degenerate_meshlet = impl.allocate(spv::OpLogicalOr, builder.makeBoolType());
		degenerate_meshlet->add_id(vert_zero->id);
		degenerate_meshlet->add_id(prim_zero->id);
		impl.add(degenerate_meshlet);

		auto *early_return = impl.allocate(spv::OpLogicalAnd, builder.makeBoolType());
		early_return->add_id(one_subgroup->id);
		early_return->add_id(degenerate_meshlet->id);
		impl.add(early_return);

		// This may also be useful in some cases if application does misc shader work after SetMeshOutputsEXT.
		// Use a pseudo-op (arbitrarily chosen as OpLifetimeStop as the name is convenient) for conditional return
		// since we have no easy way of introducing additional control flow
		// in the LLVM -> SPIR-V emitter.
		// This gets resolved during final SPIR-V lowering.
		auto *cond_ret = impl.allocate(spv::PseudoOpReturnCond);
		cond_ret->add_id(early_return->id);
		impl.add(cond_ret);
	}

	emit_view_instancing_fixed_layer_viewport(impl, false, num_prim_id);

	return true;
}

bool emit_emit_indices_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	auto &builder = impl.builder();

	// If we for some reason have max primitives 0 in the execution mode,
	// just ignore any access to index buffer.
	if (!impl.primitive_index_array_id || impl.execution_mode_meta.stage_output_num_primitive == 0)
		return true;

	unsigned index_dim = impl.execution_mode_meta.primitive_index_dimension;
	spv::Id index_type_id = impl.get_type_id(DXIL::ComponentType::U32, 1, index_dim);
	spv::Id index_scalar_type_id = impl.get_type_id(DXIL::ComponentType::U32, 1, 1);

	spv::Id components[3];
	for (unsigned i = 0; i < index_dim; i++)
		components[i] = impl.get_id_for_value(instruction->getOperand(2 + i));
	spv::Id index_id = impl.build_vector(index_scalar_type_id, components, index_dim);

	Operation *op = impl.allocate(spv::OpAccessChain,
	                              builder.makePointer(spv::StorageClassOutput, index_type_id));
	spv::Id ptr_id = op->id;

	op->add_id(impl.primitive_index_array_id);
	op->add_id(impl.get_id_for_value(instruction->getOperand(1)));
	impl.add(op);

	op = impl.allocate(spv::OpStore);
	op->add_ids({ ptr_id, index_id });
	impl.add(op);
	return true;
}

bool emit_store_vertex_output_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	// If we for some reason have max vertices 0 in the execution mode,
	// just ignore any access to vertex output.
	if (impl.execution_mode_meta.stage_output_num_vertex == 0)
		return true;

	auto &builder = impl.builder();
	uint32_t output_element_index;
	if (!get_constant_operand(instruction, 1, &output_element_index))
		return false;

	// Need special handling for clip distance.
	auto *clip_cull_meta = output_clip_cull_distance_meta(impl, output_element_index);
	if (clip_cull_meta)
		return emit_store_clip_cull_distance(impl, instruction, *clip_cull_meta);

	const auto &meta = impl.output_elements_meta[output_element_index];

	uint32_t var_id = meta.id;
	uint32_t ptr_id;

	spv::Id output_type_id = builder.getContainedTypeId(builder.getDerefTypeId(var_id));

	bool row_index = false;
	if (builder.isArrayType(output_type_id))
	{
		row_index = true;
		output_type_id = builder.getContainedTypeId(output_type_id);
	}
	uint32_t num_cols = builder.getNumTypeComponents(output_type_id);

	Operation *op = impl.allocate(
			spv::OpAccessChain, builder.makePointer(spv::StorageClassOutput, builder.getScalarTypeId(output_type_id)));
	ptr_id = op->id;

	op->add_id(var_id);
	op->add_id(impl.get_id_for_value(instruction->getOperand(5)));
	if (row_index)
		op->add_id(impl.get_id_for_value(instruction->getOperand(2)));
	if (num_cols > 1)
		op->add_id(impl.get_id_for_value(instruction->getOperand(3), 32));

	impl.add(op);

	impl.register_externally_visible_write(instruction->getOperand(4));
	spv::Id store_value = impl.get_id_for_value(instruction->getOperand(4));

	op = impl.allocate(spv::OpStore);
	op->add_ids({ ptr_id, impl.fixup_store_type_io(meta.component_type, 1, store_value) });
	impl.add(op);
	return true;
}

bool emit_store_primitive_output_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	// If we for some reason have max primitives 0 in the execution mode,
	// just ignore any access to primitive output.
	if (impl.execution_mode_meta.stage_output_num_primitive == 0)
		return true;

	auto &builder = impl.builder();
	uint32_t output_element_index;
	if (!get_constant_operand(instruction, 1, &output_element_index))
		return false;

	const auto &meta = impl.patch_elements_meta[output_element_index];
	uint32_t var_id = meta.id;
	spv::Id ptr_id;

	spv::Id output_type_id = builder.getContainedTypeId(builder.getDerefTypeId(var_id));

	bool row_index = false;
	if (builder.isArrayType(output_type_id))
	{
		row_index = true;
		output_type_id = builder.getContainedTypeId(output_type_id);
	}
	uint32_t num_cols = builder.getNumTypeComponents(output_type_id);

	impl.register_externally_visible_write(instruction->getOperand(4));
	spv::Id store_value = impl.fixup_store_type_io(meta.component_type, 1,
	                                               impl.get_id_for_value(instruction->getOperand(4)));

	Operation *op;
	spv::BuiltIn builtin;
	if (impl.spirv_module.query_builtin_shader_output(var_id, &builtin) && builtin == spv::BuiltInCullPrimitiveEXT)
	{
		// Special case where the output variable is bool, but the value we get is uint32
		op = impl.allocate(spv::OpINotEqual, builder.makeBoolType());
		op->add_id(store_value);
		op->add_id(builder.makeUintConstant(0));
		impl.add(op);
		store_value = op->id;
	}

	if (impl.options.multiview.enable && impl.options.multiview.last_pre_rasterization_stage)
	{
		if (meta.semantic == DXIL::Semantic::RenderTargetArrayIndex)
		{
			auto *add = impl.allocate(spv::OpIAdd, builder.makeUintType(32));
			add->add_id(store_value);
			add->add_id(build_layer_offset_id(impl));
			impl.add(add);

			store_value = add->id;
		}
		else if (meta.semantic == DXIL::Semantic::ViewPortArrayIndex &&
		         impl.options.multiview.view_instance_to_viewport_spec_id != UINT32_MAX)
		{
			auto *add = impl.allocate(spv::OpIAdd, builder.makeUintType(32));
			add->add_id(store_value);
			add->add_id(build_viewport_offset_id(impl));
			impl.add(add);

			store_value = add->id;
		}
	}

	op = impl.allocate(spv::OpAccessChain,
	                   builder.makePointer(spv::StorageClassOutput, builder.getScalarTypeId(output_type_id)));
	ptr_id = op->id;
	op->add_id(var_id);
	op->add_id(impl.get_id_for_value(instruction->getOperand(5)));
	if (row_index)
		op->add_id(impl.get_id_for_value(instruction->getOperand(2)));
	if (num_cols > 1)
		op->add_id(impl.get_id_for_value(instruction->getOperand(3), 32));

	impl.add(op);

	op = impl.allocate(spv::OpStore);
	op->add_ids({ ptr_id, store_value });
	impl.add(op);
	return true;
}

bool emit_dispatch_mesh_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	Operation *op = impl.allocate(spv::OpEmitMeshTasksEXT);

	for (unsigned i = 1; i <= 4; i++)
		op->add_id(impl.get_id_for_value(instruction->getOperand(i)));

	impl.add(op);
	return true;
}

bool emit_get_mesh_payload_instruction(Converter::Impl &impl, const llvm::CallInst *instruction)
{
	// GetMeshPayload can only be called once per shader
	spv::Id type_id = impl.get_type_id(instruction->getType()->getPointerElementType());
	spv::Id var_id = impl.create_variable(spv::StorageClassTaskPayloadWorkgroupEXT, type_id);

	impl.rewrite_value(instruction, var_id);
	return true;
}

} // namespace dxil_spv
