Fixed texture stage count determination

DDrawCompat/D3dDdi/DeviceState.cpp

@@ -47,11 +47,13 @@ namespace D3dDdi
 		, m_vertexShaderConst{}
 		, m_vertexShaderConstB{}
 		, m_vertexShaderConstI{}
+		, m_vertexDecl(nullptr)
 		, m_changedStates(0)
 		, m_maxChangedTextureStage(0)
 		, m_changedTextureStageStates{}
 		, m_vsVertexFixup(createVertexShader(g_vsVertexFixup))
 		, m_textureResource{}
+		, m_pixelShader(nullptr)
 		, m_isLocked(false)
 		, m_spriteMode(false)
 	{
@@ -275,11 +277,15 @@ namespace D3dDdi
 		return m_textureResource[stage];
 	}
 
+	UINT DeviceState::getTextureStageCount() const
+	{
+		return m_pixelShader ? m_pixelShader->textureStageCount : 0;
+	}
+
 	const DeviceState::VertexDecl& DeviceState::getVertexDecl() const
 	{
 		static const VertexDecl emptyDecl = {};
-		auto it = m_vertexShaderDecls.find(m_app.vertexShaderDecl);
+		return m_vertexDecl ? *m_vertexDecl : emptyDecl;
-		return it != m_vertexShaderDecls.end() ? it->second : emptyDecl;
 	}
 
 	bool DeviceState::isColorKeyUsed()
@@ -290,7 +296,8 @@ namespace D3dDdi
 		}
 
 		bool used = false;
-		for (UINT i = 0; i < getVertexDecl().textureStageCount && !used; ++i)
+		UINT textureStageCount = getTextureStageCount();
+		for (UINT i = 0; i < textureStageCount && !used; ++i)
 		{
 			used = !m_app.textureStageState[i][D3DDDITSS_DISABLETEXTURECOLORKEY];
 		}
@@ -426,15 +433,17 @@ namespace D3dDdi
 
 	HRESULT DeviceState::pfnCreatePixelShader(D3DDDIARG_CREATEPIXELSHADER* data, const UINT* code)
 	{
+		ShaderAssembler shaderAssembler(code, data->CodeSize);
 		LOG_DEBUG << "Pixel shader bytecode: " << Compat::hexDump(code, data->CodeSize);
-		LOG_DEBUG << ShaderAssembler(code, data->CodeSize).disassemble();
+		LOG_DEBUG << shaderAssembler.disassemble();
 		HRESULT result = m_device.getOrigVtable().pfnCreatePixelShader(m_device, data, code);
 		if (SUCCEEDED(result))
 		{
 			m_pixelShaders.emplace(data->ShaderHandle,
 				PixelShader{ std::vector<UINT>(code, code + data->CodeSize / 4),
 				std::unique_ptr<void, ResourceDeleter>(
-					nullptr, ResourceDeleter(m_device, m_device.getOrigVtable().pfnDeleteVertexShaderFunc)) });
+					nullptr, ResourceDeleter(m_device, m_device.getOrigVtable().pfnDeleteVertexShaderFunc)),
+				    shaderAssembler.getTextureStageCount(), false });
 		}
 		return result;
 	}
@@ -459,10 +468,6 @@ namespace D3dDdi
 			{
 				decl.texCoordOffset[vertexElements[i].UsageIndex] = vertexElements[i].Offset;
 				decl.texCoordType[vertexElements[i].UsageIndex] = vertexElements[i].Type;
-				if (vertexElements[i].UsageIndex >= decl.textureStageCount)
-				{
-					decl.textureStageCount = vertexElements[i].UsageIndex + 1;
-				}
 			}
 			else if (D3DDECLUSAGE_POSITIONT == vertexElements[i].Usage)
 			{
@@ -487,6 +492,10 @@ namespace D3dDdi
 			auto it = m_pixelShaders.find(shader);
 			if (it != m_pixelShaders.end())
 			{
+				if (m_pixelShader == &it->second)
+				{
+					m_pixelShader = nullptr;
+				}
 				if (it->second.modifiedPixelShader)
 				{
 					deleteShader(it->second.modifiedPixelShader.release(), &State::pixelShader,
@@ -500,9 +509,14 @@ namespace D3dDdi
 
 	HRESULT DeviceState::pfnDeleteVertexShaderDecl(HANDLE shader)
 	{
+		const bool isCurrent = m_app.vertexShaderDecl == shader;
 		HRESULT result = deleteShader(shader, &State::vertexShaderDecl, m_device.getOrigVtable().pfnDeleteVertexShaderDecl);
 		if (SUCCEEDED(result))
 		{
+			if (isCurrent)
+			{
+				m_vertexDecl = nullptr;
+			}
 			m_vertexShaderDecls.erase(shader);
 		}
 		return result;
@@ -525,6 +539,8 @@ namespace D3dDdi
 	{
 		m_app.pixelShader = shader;
 		m_changedStates |= CS_SHADER;
+		auto it = m_pixelShaders.find(shader);
+		m_pixelShader = it != m_pixelShaders.end() ? &it->second : nullptr;
 		return S_OK;
 	}
 
@@ -636,6 +652,8 @@ namespace D3dDdi
 	{
 		m_app.vertexShaderDecl = shader;
 		m_changedStates |= CS_SHADER;
+		auto it = m_vertexShaderDecls.find(shader);
+		m_vertexDecl = it != m_vertexShaderDecls.end() ? &it->second : nullptr;
 		return S_OK;
 	}
 
@@ -669,7 +687,8 @@ namespace D3dDdi
 
 	void DeviceState::prepareTextures()
 	{
-		for (UINT stage = 0; stage < getVertexDecl().textureStageCount; ++stage)
+		UINT textureStageCount = getTextureStageCount();
+		for (UINT stage = 0; stage < textureStageCount; ++stage)
 		{
 			auto resource = getTextureResource(stage);
 			if (resource)

DDrawCompat/D3dDdi/DeviceState.h

@@ -80,7 +80,6 @@ namespace D3dDdi
 			std::vector<D3DDDIVERTEXELEMENT> elements;
 			std::array<UINT, 8> texCoordOffset;
 			std::array<UINT, 8> texCoordType;
-			UINT textureStageCount;
 			bool isTransformed;
 		};
 
@@ -130,6 +129,7 @@ namespace D3dDdi
 		const State& getAppState() const { return m_app; }
 		const State& getCurrentState() const { return m_current; }
 		Resource* getTextureResource(UINT stage);
+		UINT getTextureStageCount() const;
 		const VertexDecl& getVertexDecl() const;
 		HANDLE getVertexFixupDecl() const { return m_vsVertexFixup.get(); }
 		bool isLocked() const { return m_isLocked; }
@@ -153,6 +153,7 @@ namespace D3dDdi
 		{
 			std::vector<UINT> tokens;
 			std::unique_ptr<void, ResourceDeleter> modifiedPixelShader;
+			UINT textureStageCount;
 			bool isModified;
 		};
 
@@ -217,6 +218,7 @@ namespace D3dDdi
 		std::array<ShaderConstB, 16> m_vertexShaderConstB;
 		std::array<ShaderConstI, 16> m_vertexShaderConstI;
 		std::map<HANDLE, VertexDecl> m_vertexShaderDecls;
+		VertexDecl* m_vertexDecl;
 		UINT m_changedStates;
 		UINT m_maxChangedTextureStage;
 		UINT m_usedTextureStages;
@@ -225,6 +227,7 @@ namespace D3dDdi
 		std::unique_ptr<void, ResourceDeleter> m_vsVertexFixup;
 		std::array<Resource*, 8> m_textureResource;
 		std::map<HANDLE, PixelShader> m_pixelShaders;
+		PixelShader* m_pixelShader;
 		bool m_isLocked;
 		bool m_spriteMode;
 	};

DDrawCompat/D3dDdi/DrawPrimitive.cpp

@@ -829,9 +829,10 @@ namespace D3dDdi
 		auto& state = m_device.getState();
 		auto& appState = state.getAppState();
 		auto& decl = state.getVertexDecl();
+		auto textureStageCount = state.getTextureStageCount();
 		auto vertices = m_streamSource.vertices + baseVertexIndex * m_streamSource.stride;
 
-		for (UINT stage = 0; stage < decl.textureStageCount; ++stage)
+		for (UINT stage = 0; stage < textureStageCount; ++stage)
 		{
 			const UINT D3DDECLTYPE_FLOAT2 = 1;
 			if (!appState.textures[stage] ||

DDrawCompat/D3dDdi/ShaderAssembler.cpp

@@ -696,6 +696,12 @@ namespace D3dDdi
 		return static_cast<ShaderType>(m_tokens.front() >> 16);
 	}
 
+	UINT ShaderAssembler::getTextureStageCount()
+	{
+		auto usedSamplers = getUsedRegisterNumbers(D3DSPR_SAMPLER);
+		return usedSamplers.empty() ? 0 : (*usedSamplers.rbegin() + 1);
+	}
+
 	template <typename Token>
 	Token ShaderAssembler::getToken(UINT offset) const
 	{

DDrawCompat/D3dDdi/ShaderAssembler.h

@@ -16,6 +16,7 @@ namespace D3dDdi
 
 		bool addAlphaTest(UINT alphaRef);
 		std::string disassemble();
+		UINT getTextureStageCount();
 		const std::vector<UINT>& getTokens() const { return m_tokens; }
 
 	private: