// Copyright Epic Games, Inc. All Rights Reserved.

/*=============================================================================
	D3D12RootSignature.h: D3D12 Root Signatures
=============================================================================*/

#pragma once

#include "D3D12RootSignatureDefinitions.h"
#include "D3D12Util.h"

// Root parameter keys grouped by visibility.
enum ERootParameterKeys
{
	PS_SRVs,
	PS_CBVs,
	PS_RootCBVs,
	PS_Samplers,
	VS_SRVs,
	VS_CBVs,
	VS_RootCBVs,
	VS_Samplers,
	VS_UAVs,
	GS_SRVs,
	GS_CBVs,
	GS_RootCBVs,
	GS_Samplers,
	MS_SRVs,
	MS_CBVs,
	MS_RootCBVs,
	MS_Samplers,
	AS_SRVs,
	AS_CBVs,
	AS_RootCBVs,
	AS_Samplers,
	ALL_SRVs,
	ALL_CBVs,
	ALL_RootCBVs,
	ALL_Samplers,
	ALL_UAVs, // non-VS stages (PS, CS, etc.)
	RPK_RootParameterKeyCount,
};

class FD3D12RootSignature;

class FD3D12RootSignatureDesc
{
public:
	explicit FD3D12RootSignatureDesc(const FD3D12QuantizedBoundShaderState& QBSS, const D3D12_RESOURCE_BINDING_TIER ResourceBindingTier);

	inline const D3D12_VERSIONED_ROOT_SIGNATURE_DESC& GetDesc() const { return RootDesc; }

	static constexpr uint32 MaxRootParameters = 32;	// Arbitrary max, increase as needed.

	inline int8 GetRootConstantsSlot() const { return RootConstantsSlot; }
	inline int8 GetDiagnosticBufferSlot() const { return DiagnosticBufferSlot; }
	inline int8 GetStaticShaderBindingSlot() const { return StaticShaderBindingSlot; }
	inline int8 GetStaticShaderBindingCount() const { return StaticShaderBindingCount; }

private:

	uint32 RootParametersSize;	// The size of all root parameters in the root signature. Size in DWORDs, the limit is 64.
	int8 RootConstantsSlot = -1;
	int8 DiagnosticBufferSlot = -1;
	int8 StaticShaderBindingSlot = -1;
	int8 StaticShaderBindingCount = -1;
	CD3DX12_ROOT_PARAMETER1 TableSlots[MaxRootParameters];
	CD3DX12_DESCRIPTOR_RANGE1 DescriptorRanges[MaxRootParameters];
	CD3DX12_VERSIONED_ROOT_SIGNATURE_DESC RootDesc;
};

class FD3D12RootSignature : public FD3D12AdapterChild
{
private:
	// Struct for all the useful info we want per shader stage.
	struct ShaderStage
	{
		// TODO: Make these arrays and index into them by type instead of individual variables.
		uint8 MaxCBVCount = 0;
		uint8 MaxSRVCount = 0;
		uint8 MaxSamplerCount = 0;
		uint8 MaxUAVCount = 0;
		CBVSlotMask CBVRegisterMask = 0;
		bool bVisible = false;
	};

public:
	explicit FD3D12RootSignature(FD3D12Adapter* InParent)
		: FD3D12AdapterChild(InParent)
	{}
	explicit FD3D12RootSignature(FD3D12Adapter* InParent, const FD3D12QuantizedBoundShaderState& InQBSS)
		: FD3D12AdapterChild(InParent)
	{
		Init(InQBSS);
	}

	void Init(const FD3D12QuantizedBoundShaderState& InQBSS);
	void Init(const D3D12_VERSIONED_ROOT_SIGNATURE_DESC& InDesc, ERootSignatureType InRootSignatureType);

	void InitStaticGraphicsRootSignature(EShaderBindingLayoutFlags InFlags);
	void InitStaticComputeRootSignatureDesc(EShaderBindingLayoutFlags InFlags);
#if D3D12_RHI_RAYTRACING
	void InitStaticRayTracingGlobalRootSignatureDesc(EShaderBindingLayoutFlags InFlags);
	void InitStaticRayTracingLocalRootSignatureDesc(EShaderBindingLayoutFlags InFlags);
#endif

	ID3D12RootSignature* GetRootSignature() const { return RootSignature.GetReference(); }
	ID3DBlob* GetRootSignatureBlob() const { return RootSignatureBlob.GetReference(); }

	inline uint32 SamplerRDTBindSlot(EShaderFrequency ShaderStage) const
	{
		switch (ShaderStage)
		{
		case SF_Vertex: return BindSlotMap[VS_Samplers];
		case SF_Mesh: return BindSlotMap[MS_Samplers];
		case SF_Amplification: return BindSlotMap[AS_Samplers];
		case SF_Pixel: return BindSlotMap[PS_Samplers];
		case SF_Geometry: return BindSlotMap[GS_Samplers];
		case SF_Compute: return BindSlotMap[ALL_Samplers];

		default: check(false);
			return UINT_MAX;
		}
	}

	inline uint32 SRVRDTBindSlot(EShaderFrequency ShaderStage) const
	{
		switch (ShaderStage)
		{
		case SF_Vertex: return BindSlotMap[VS_SRVs];
		case SF_Mesh: return BindSlotMap[MS_SRVs];
		case SF_Amplification: return BindSlotMap[AS_SRVs];
		case SF_Pixel: return BindSlotMap[PS_SRVs];
		case SF_Geometry: return BindSlotMap[GS_SRVs];
		case SF_Compute: return BindSlotMap[ALL_SRVs];

		default: check(false);
			return UINT_MAX;
		}
	}

	inline uint32 CBVRDTBindSlot(EShaderFrequency ShaderStage) const
	{
		switch (ShaderStage)
		{
		case SF_Vertex: return BindSlotMap[VS_CBVs];
		case SF_Mesh: return BindSlotMap[MS_CBVs];
		case SF_Amplification: return BindSlotMap[AS_CBVs];
		case SF_Pixel: return BindSlotMap[PS_CBVs];
		case SF_Geometry: return BindSlotMap[GS_CBVs];
		case SF_Compute: return BindSlotMap[ALL_CBVs];

		default: check(false);
			return UINT_MAX;
		}
	}

	inline uint32 CBVRDBaseBindSlot(EShaderFrequency ShaderStage) const
	{
		switch (ShaderStage)
		{
		case SF_Vertex: return BindSlotMap[VS_RootCBVs];
		case SF_Mesh: return BindSlotMap[MS_RootCBVs];
		case SF_Amplification: return BindSlotMap[AS_RootCBVs];
		case SF_Pixel: return BindSlotMap[PS_RootCBVs];
		case SF_Geometry: return BindSlotMap[GS_RootCBVs];

		case SF_NumFrequencies:
		case SF_Compute: return BindSlotMap[ALL_RootCBVs];

		default: check(false);
			return UINT_MAX;
		}
	}

	inline uint32 CBVRDBindSlot(EShaderFrequency ShaderStage, uint32 BufferIndex) const
	{
		// This code assumes that all Root CBVs for a particular stage are contiguous in the root signature (thus indexable by the buffer index).
		return CBVRDBaseBindSlot(ShaderStage) + BufferIndex;
	}

	inline uint32 UAVRDTBindSlot(EShaderFrequency ShaderStage) const
	{
		check(ShaderStage == SF_Pixel || ShaderStage == SF_Vertex || ShaderStage == SF_Compute);
		const uint32 MapSlotIndex = ShaderStage == SF_Vertex ? VS_UAVs : ALL_UAVs;
		return BindSlotMap[MapSlotIndex];
	}

	inline static bool IsValidBindSlot(uint32 BindSlotMapIndex)
	{
		return BindSlotMapIndex < InvalidBindSlotMapIndex;
	}

	inline bool HasUAVs() const { return bHasUAVs; }
	inline bool HasSRVs() const { return bHasSRVs; }
	inline bool HasCBVs() const { return bHasCBVs; }
	inline bool HasRootCBs() const { return bHasRootCBs; }
	inline bool HasTableResources() const { return bHasUAVs || bHasSRVs; }
	inline bool HasTableConstants() const { return bHasCBVs; }
	inline bool HasSamplers() const { return bHasSamplers; }

#if PLATFORM_SUPPORTS_BINDLESS_RENDERING
	bool UsesDynamicResources() const { return bUsesDynamicResources; }
	bool UsesDynamicSamplers() const { return bUsesDynamicSamplers; }
#else
	constexpr bool UsesDynamicResources() const { return false; }
	constexpr bool UsesDynamicSamplers() const { return false; }
#endif

	inline bool HasVS() const { return Stage[SF_Vertex].bVisible; }
	inline bool HasMS() const { return Stage[SF_Mesh].bVisible; }
	inline bool HasAS() const { return Stage[SF_Amplification].bVisible; }
	inline bool HasGS() const { return Stage[SF_Geometry].bVisible; }
	inline bool HasPS() const { return Stage[SF_Pixel].bVisible; }
	inline bool HasCS() const { return Stage[SF_Compute].bVisible; }	// Root signatures can be used for Graphics and/or Compute because they exist in separate bind spaces.
	inline uint32 MaxSamplerCount(uint32 ShaderStage) const { check(ShaderStage != SF_NumFrequencies); return Stage[ShaderStage].MaxSamplerCount; }
	inline uint32 MaxSRVCount(uint32 ShaderStage) const { check(ShaderStage != SF_NumFrequencies); return Stage[ShaderStage].MaxSRVCount; }
	inline uint32 MaxCBVCount(uint32 ShaderStage) const { check(ShaderStage != SF_NumFrequencies); return Stage[ShaderStage].MaxCBVCount; }
	inline uint32 MaxUAVCount(uint32 ShaderStage) const { check(ShaderStage != SF_NumFrequencies); return Stage[ShaderStage].MaxUAVCount; }
	inline CBVSlotMask CBVRegisterMask(uint32 ShaderStage) const { check(ShaderStage != SF_NumFrequencies); return Stage[ShaderStage].CBVRegisterMask; }

	uint32 GetBindSlotOffsetInBytes(uint8 BindSlotIndex) const { check(BindSlotIndex < UE_ARRAY_COUNT(BindSlotOffsetsInDWORDs)); return 4 * BindSlotOffsetsInDWORDs[BindSlotIndex]; }
	uint32 GetTotalRootSignatureSizeInBytes() const { return 4 * TotalRootSignatureSizeInDWORDs; }

	inline int8 GetRootConstantsSlot() const { return RootConstantsSlot; }
	inline int8 GetStaticShaderBindingSlot() const { return StaticShaderBindingSlot; }
	inline int8 GetStaticShaderBindingCount() const { return StaticShaderBindingCount; }

	// Returns root parameter slot for the internal shader diagnostic buffer (used for asserts, etc.) or -1 if not available.
	inline int8 GetDiagnosticBufferSlot() const { return DiagnosticBufferSlot; }

private:
	void AnalyzeSignature(const D3D12_VERSIONED_ROOT_SIGNATURE_DESC& Desc, ERootSignatureType InRootSignatureType);

	template<typename RootSignatureDescType>
	void InternalAnalyzeSignature(const RootSignatureDescType& Desc, ERootSignatureType InRootSignatureType);

	inline bool HasVisibility(const D3D12_SHADER_VISIBILITY& ParameterVisibility, const D3D12_SHADER_VISIBILITY& Visibility) const
	{
		return ParameterVisibility == D3D12_SHADER_VISIBILITY_ALL || ParameterVisibility == Visibility;
	}

	inline void SetSamplersRDTBindSlot(EShaderFrequency SF, uint8 RootParameterIndex)
	{
		uint8* pBindSlot = nullptr;
		switch (SF)
		{
		case SF_Vertex: pBindSlot = &BindSlotMap[VS_Samplers]; break;
		case SF_Mesh: pBindSlot = &BindSlotMap[MS_Samplers]; break;
		case SF_Amplification: pBindSlot = &BindSlotMap[AS_Samplers]; break;
		case SF_Pixel: pBindSlot = &BindSlotMap[PS_Samplers]; break;
		case SF_Geometry: pBindSlot = &BindSlotMap[GS_Samplers]; break;

		case SF_Compute:
		case SF_NumFrequencies: pBindSlot = &BindSlotMap[ALL_Samplers]; break;

		default: check(false);
			return;
		}

		check(*pBindSlot == InvalidBindSlotMapIndex);
		*pBindSlot = RootParameterIndex;

		bHasSamplers = true;
	}

	inline void SetSRVRDTBindSlot(EShaderFrequency SF, uint8 RootParameterIndex)
	{
		uint8* pBindSlot = nullptr;
		switch (SF)
		{
		case SF_Vertex: pBindSlot = &BindSlotMap[VS_SRVs]; break;
		case SF_Mesh: pBindSlot = &BindSlotMap[MS_SRVs]; break;
		case SF_Amplification: pBindSlot = &BindSlotMap[AS_SRVs]; break;
		case SF_Pixel: pBindSlot = &BindSlotMap[PS_SRVs]; break;
		case SF_Geometry: pBindSlot = &BindSlotMap[GS_SRVs]; break;

		case SF_Compute:
		case SF_NumFrequencies: pBindSlot = &BindSlotMap[ALL_SRVs]; break;

		default: check(false);
			return;
		}

		check(*pBindSlot == InvalidBindSlotMapIndex);
		*pBindSlot = RootParameterIndex;

		bHasSRVs = true;
	}

	inline void SetCBVRDTBindSlot(EShaderFrequency SF, uint8 RootParameterIndex)
	{
		uint8* pBindSlot = nullptr;
		switch (SF)
		{
		case SF_Vertex: pBindSlot = &BindSlotMap[VS_CBVs]; break;
		case SF_Mesh: pBindSlot = &BindSlotMap[MS_CBVs]; break;
		case SF_Amplification: pBindSlot = &BindSlotMap[AS_CBVs]; break;
		case SF_Pixel: pBindSlot = &BindSlotMap[PS_CBVs]; break;
		case SF_Geometry: pBindSlot = &BindSlotMap[GS_CBVs]; break;

		case SF_Compute:
		case SF_NumFrequencies: pBindSlot = &BindSlotMap[ALL_CBVs]; break;

		default: check(false);
			return;
		}

		check(*pBindSlot == InvalidBindSlotMapIndex);
		*pBindSlot = RootParameterIndex;

		bHasCBVs = true;
	}

	inline void SetCBVRDBindSlot(EShaderFrequency SF, uint8 RootParameterIndex)
	{
		uint8* pBindSlot = nullptr;
		switch (SF)
		{
		case SF_Vertex: pBindSlot = &BindSlotMap[VS_RootCBVs]; break;
		case SF_Mesh: pBindSlot = &BindSlotMap[MS_RootCBVs]; break;
		case SF_Amplification: pBindSlot = &BindSlotMap[AS_RootCBVs]; break;
		case SF_Pixel: pBindSlot = &BindSlotMap[PS_RootCBVs]; break;
		case SF_Geometry: pBindSlot = &BindSlotMap[GS_RootCBVs]; break;

		case SF_Compute:
		case SF_NumFrequencies: pBindSlot = &BindSlotMap[ALL_RootCBVs]; break;

		default: check(false);
			return;
		}

		check(*pBindSlot == InvalidBindSlotMapIndex);
		*pBindSlot = RootParameterIndex;

		bHasRootCBs = true;
	}

	inline void SetUAVRDTBindSlot(EShaderFrequency SF, uint8 RootParameterIndex)
	{
		check(SF == SF_Pixel || SF == SF_Vertex || SF == SF_Compute || SF == SF_NumFrequencies);

		const uint32 MapSlotIndex = SF == SF_Vertex ? VS_UAVs : ALL_UAVs;
		uint8* pBindSlot = &BindSlotMap[MapSlotIndex];

		check(*pBindSlot == InvalidBindSlotMapIndex);
		*pBindSlot = RootParameterIndex;

		bHasUAVs = true;
	}

	inline void SetMaxSamplerCount(EShaderFrequency SF, uint8 Count)
	{
		if (SF == SF_NumFrequencies)
		{
			// Update all counts for all stages.
			for (uint32 s = SF_Vertex; s <= SF_Compute; s++)
			{
				Stage[s].MaxSamplerCount = Count;
			}
		}
		else
		{
			Stage[SF].MaxSamplerCount = Count;
		}
	}

	inline void SetMaxSRVCount(EShaderFrequency SF, uint8 Count)
	{
		if (SF == SF_NumFrequencies)
		{
			// Update all counts for all stages.
			for (uint32 s = SF_Vertex; s <= SF_Compute; s++)
			{
				Stage[s].MaxSRVCount = Count;
			}
		}
		else
		{
			Stage[SF].MaxSRVCount = Count;
		}
	}

	// Update the mask that indicates what shader registers are used in the descriptor table.
	template<typename DescriptorRangeType>
	inline void UpdateCBVRegisterMaskWithDescriptorRange(EShaderFrequency SF, const DescriptorRangeType& Range)
	{
		const uint32 StartRegister = Range.BaseShaderRegister;
		const uint32 EndRegister = StartRegister + Range.NumDescriptors;
		const uint32 StartStage = (SF == SF_NumFrequencies) ? SF_Vertex : SF;
		const uint32 EndStage = (SF == SF_NumFrequencies) ? SF_Compute : SF;
		for (uint32 CurrentStage = StartStage; CurrentStage <= EndStage; CurrentStage++)
		{
			for (uint32 Register = StartRegister; Register < EndRegister; Register++)
			{
				// The bit shouldn't already be set for the current register.
				check((Stage[CurrentStage].CBVRegisterMask & (1 << Register)) == 0); 
				Stage[CurrentStage].CBVRegisterMask |= (1 << Register);
			}
		}
	}

	// Update the mask that indicates what shader registers are used in the root descriptor.
	template<typename DescriptorType>
	inline void UpdateCBVRegisterMaskWithDescriptor(EShaderFrequency SF, const DescriptorType& Descriptor)
	{
		const uint32 StartStage = (SF == SF_NumFrequencies) ? SF_Vertex : SF;
		const uint32 EndStage = (SF == SF_NumFrequencies) ? SF_Compute : SF;
		const uint32& Register = Descriptor.ShaderRegister;
		for (uint32 CurrentStage = StartStage; CurrentStage <= EndStage; CurrentStage++)
		{
			// The bit shouldn't already be set for the current register.
			check((Stage[CurrentStage].CBVRegisterMask & (1 << Register)) == 0); 
			Stage[CurrentStage].CBVRegisterMask |= (1 << Register);
		}
	}

	inline void SetMaxCBVCount(EShaderFrequency SF, uint8 Count)
	{
		if (SF == SF_NumFrequencies)
		{
			// Update all counts for all stages.
			for (uint32 s = SF_Vertex; s <= SF_Compute; s++)
			{
				Stage[s].MaxCBVCount = Count;
			}
		}
		else
		{
			Stage[SF].MaxCBVCount = Count;
		}
	}

	inline void IncrementMaxCBVCount(EShaderFrequency SF, uint8 Count)
	{
		if (SF == SF_NumFrequencies)
		{
			// Update all counts for all stages.
			for (uint32 s = SF_Vertex; s <= SF_Compute; s++)
			{
				Stage[s].MaxCBVCount += Count;
			}
		}
		else
		{
			Stage[SF].MaxCBVCount += Count;
		}
	}

	inline void SetMaxUAVCount(EShaderFrequency SF, uint8 Count)
	{
		if (SF == SF_NumFrequencies)
		{
			// Update all counts for all stages.
			for (uint32 s = SF_Vertex; s <= SF_Compute; s++)
			{
				Stage[s].MaxUAVCount = Count;
			}
		}
		else
		{
			Stage[SF].MaxUAVCount = Count;
		}
	}

	TRefCountPtr<ID3D12RootSignature> RootSignature;
	uint8 BindSlotMap[RPK_RootParameterKeyCount];	// This map uses an enum as a key to lookup the root parameter index
	static constexpr uint8 InvalidBindSlotMapIndex = 0xFF;
	ShaderStage Stage[SF_NumFrequencies];
	TRefCountPtr<ID3DBlob> RootSignatureBlob;

	uint8 BindSlotOffsetsInDWORDs[FD3D12RootSignatureDesc::MaxRootParameters] = {};
	uint8 TotalRootSignatureSizeInDWORDs = 0;
	int8 RootConstantsSlot = -1;
	int8 DiagnosticBufferSlot = -1;
	int8 StaticShaderBindingSlot = -1;
	int8 StaticShaderBindingCount = -1;

	uint8 bHasUAVs : 1;
	uint8 bHasSRVs : 1;
	uint8 bHasCBVs : 1;
	uint8 bHasRootCBs : 1;
	uint8 bHasSamplers : 1;
#if PLATFORM_SUPPORTS_BINDLESS_RENDERING
	uint8 bUsesDynamicResources : 1;
	uint8 bUsesDynamicSamplers : 1;
#endif
};

class FD3D12RootSignatureManager : public FD3D12AdapterChild
{
public:
	explicit FD3D12RootSignatureManager(FD3D12Adapter* InParent)
		: FD3D12AdapterChild(InParent)
	{
	}
	~FD3D12RootSignatureManager()
	{
		check(RootSignatureMap.Num() == 0);
	}
	void Destroy();

	FD3D12RootSignature* GetRootSignature(const FD3D12QuantizedBoundShaderState& QBSS);
	FD3D12QuantizedBoundShaderState GetQuantizedBoundShaderState(const FD3D12RootSignature* const RootSignature);

private:
	FCriticalSection CS;
	FD3D12RootSignature* CreateRootSignature(const FD3D12QuantizedBoundShaderState& QBSS);

	TMap<FD3D12QuantizedBoundShaderState, FD3D12RootSignature*> RootSignatureMap;
};