Cache shader color key in surface

Dllmain/BuildNo.rc

@@ -1 +1 @@
-#define BUILD_NUMBER 6966 
+#define BUILD_NUMBER 6967 

ddraw/IDirect3DDeviceX.cpp

@@ -1717,14 +1717,13 @@ HRESULT m_IDirect3DDeviceX::EndScene()
 
 			if (PrimarySurface->IsSurfaceDirty())
 			{
-				DrawStates.ddpfPixelFormat.dwSize = sizeof(DDPIXELFORMAT);
-				PrimarySurface->GetPixelFormat(&DrawStates.ddpfPixelFormat);
+				PrimarySurface->GetColorKeyForShader(DrawStates.lowColorKey, DrawStates.highColorKey, true);
 
-				SetDrawStates(0, D3DDP_DXW_DRAW2DSURFACE, 9);
+				SetDrawStates(0, D3DDP_DXW_DRAW2DSURFACE | D3DDP_DXW_COLORKEYENABLE, 9);
 
 				ddrawParent->Draw2DSurface(PrimarySurface);
 
-				RestoreDrawStates(0, D3DDP_DXW_DRAW2DSURFACE, 9);
+				RestoreDrawStates(0, D3DDP_DXW_DRAW2DSURFACE | D3DDP_DXW_COLORKEYENABLE, 9);
 			}
 		}
 
@@ -3346,10 +3345,7 @@ void m_IDirect3DDeviceX::ResetDevice()
 inline void m_IDirect3DDeviceX::UpdateDrawFlags(DWORD& dwFlags)
 {
 	// Check for color key
-	DrawStates.ddpfPixelFormat.dwSize = sizeof(DDPIXELFORMAT);
-	if (rsColorKeyEnabled && CurrentTextureSurfaceX &&
-		CurrentTextureSurfaceX->GetColorKey(DrawStates.dwColorSpaceLowValue, DrawStates.dwColorSpaceHighValue) &&
-		SUCCEEDED(CurrentTextureSurfaceX->GetPixelFormat(&DrawStates.ddpfPixelFormat)))
+	if (rsColorKeyEnabled && CurrentTextureSurfaceX && CurrentTextureSurfaceX->GetColorKeyForShader(DrawStates.lowColorKey, DrawStates.highColorKey, false))
 	{
 		dwFlags |= D3DDP_DXW_COLORKEYENABLE;
 	}
@@ -3387,12 +3383,8 @@ inline void m_IDirect3DDeviceX::SetDrawStates(DWORD dwVertexTypeDesc, DWORD dwFl
 		if (colorkeyPixelShader && *colorkeyPixelShader)
 		{
 			(*d3d9Device)->SetPixelShader(*colorkeyPixelShader);
-
-			// Set color key
-			float highColorKey[4], lowColorKey[4];
-			GetColorKeyArray(lowColorKey, highColorKey, DrawStates.dwColorSpaceLowValue, DrawStates.dwColorSpaceHighValue, DrawStates.ddpfPixelFormat);
-			(*d3d9Device)->SetPixelShaderConstantF(0, lowColorKey, 1);
-			(*d3d9Device)->SetPixelShaderConstantF(1, highColorKey, 1);
+			(*d3d9Device)->SetPixelShaderConstantF(0, DrawStates.lowColorKey, 1);
+			(*d3d9Device)->SetPixelShaderConstantF(1, DrawStates.highColorKey, 1);
 		}
 	}
 	if (dwFlags & D3DDP_DXW_DRAW2DSURFACE)
@@ -3411,19 +3403,6 @@ inline void m_IDirect3DDeviceX::SetDrawStates(DWORD dwVertexTypeDesc, DWORD dwFl
 		(*d3d9Device)->SetRenderState(D3DRS_ALPHATESTENABLE, FALSE);
 		(*d3d9Device)->SetRenderState(D3DRS_ALPHABLENDENABLE, FALSE);
 		(*d3d9Device)->SetRenderState(D3DRS_FOGENABLE, FALSE);
-
-		if (!colorkeyPixelShader || !*colorkeyPixelShader)
-		{
-			colorkeyPixelShader = ddrawParent->GetColorKeyShader();
-		}
-		if (colorkeyPixelShader && *colorkeyPixelShader)
-		{
-			(*d3d9Device)->SetPixelShader(*colorkeyPixelShader);
-			float highColorKey[4], lowColorKey[4];
-			GetColorKeyArray(lowColorKey, highColorKey, 0x00000000, 0x00000000, DrawStates.ddpfPixelFormat);
-			(*d3d9Device)->SetPixelShaderConstantF(0, lowColorKey, 1);
-			(*d3d9Device)->SetPixelShaderConstantF(1, highColorKey, 1);
-		}
 	}
 }
 
@@ -3461,8 +3440,6 @@ inline void m_IDirect3DDeviceX::RestoreDrawStates(DWORD dwVertexTypeDesc, DWORD
 		(*d3d9Device)->SetRenderState(D3DRS_LIGHTING, DrawStates.rsLighting);
 		(*d3d9Device)->SetSamplerState(0, D3DSAMP_MAGFILTER, DrawStates.ssMagFilter);
 
-		(*d3d9Device)->SetPixelShader(nullptr);
-
 		SetTexture(0, AttachedTexture[0]);
 		SetTexture(1, AttachedTexture[1]);
 	}

ddraw/IDirect3DDeviceX.h

@@ -31,9 +31,8 @@ class m_IDirect3DDeviceX : public IUnknown, public AddressLookupTableDdrawObject
 		DWORD rsAlphaTestEnable = 0;
 		DWORD rsFogEnable = 0;
 		DWORD ssMagFilter = 0;
-		DWORD dwColorSpaceLowValue = 0;
-		DWORD dwColorSpaceHighValue = 0;
-		DDPIXELFORMAT ddpfPixelFormat = {};
+		float lowColorKey[4] = {};
+		float highColorKey[4] = {};
 	} DrawStates;
 
 	// Store d3d device version wrappers

ddraw/IDirectDrawSurfaceX.cpp

@@ -1681,6 +1681,9 @@ HRESULT m_IDirectDrawSurfaceX::GetColorKey(DWORD dwFlags, LPDDCOLORKEY lpDDColor
 			return DDERR_INVALIDPARAMS;
 		}
 
+		// Reset shader flag
+		ShaderColorKey.IsSet = false;
+
 		// Get color key index
 		DWORD dds = 0;
 		switch (dwFlags)
@@ -1731,6 +1734,58 @@ HRESULT m_IDirectDrawSurfaceX::GetColorKey(DWORD dwFlags, LPDDCOLORKEY lpDDColor
 	return ProxyInterface->GetColorKey(dwFlags, lpDDColorKey);
 }
 
+bool m_IDirectDrawSurfaceX::GetColorKeyForShader(float(&lowColorKey)[4], float(&highColorKey)[4], bool PrimarySurface)
+{
+	// Primary 2D surface background color
+	if (PrimarySurface && IsPrimarySurface())
+	{
+		if (!primary.ShaderColorKey.IsSet)
+		{
+			GetColorKeyArray(primary.ShaderColorKey.lowColorKey, primary.ShaderColorKey.highColorKey, 0x00000000, 0x00000000, surfaceDesc2.ddpfPixelFormat);
+			primary.ShaderColorKey.IsSet = true;
+		}
+		lowColorKey[0] = primary.ShaderColorKey.lowColorKey[0];
+		lowColorKey[1] = primary.ShaderColorKey.lowColorKey[1];
+		lowColorKey[2] = primary.ShaderColorKey.lowColorKey[2];
+		lowColorKey[3] = primary.ShaderColorKey.lowColorKey[3];
+		highColorKey[0] = primary.ShaderColorKey.highColorKey[0];
+		highColorKey[1] = primary.ShaderColorKey.highColorKey[1];
+		highColorKey[2] = primary.ShaderColorKey.highColorKey[2];
+		highColorKey[3] = primary.ShaderColorKey.highColorKey[3];
+		return true;
+	}
+
+	// Surface low and high color space
+	if (!ShaderColorKey.IsSet)
+	{
+		if (surfaceDesc2.ddsCaps.dwCaps & DDSD_CKSRCBLT)
+		{
+			GetColorKeyArray(ShaderColorKey.lowColorKey, ShaderColorKey.highColorKey,
+				surfaceDesc2.ddckCKSrcBlt.dwColorSpaceLowValue, surfaceDesc2.ddckCKSrcBlt.dwColorSpaceHighValue, surfaceDesc2.ddpfPixelFormat);
+			ShaderColorKey.IsSet = true;
+		}
+		else if (surfaceDesc2.ddsCaps.dwCaps & DDCKEY_SRCOVERLAY)
+		{
+			GetColorKeyArray(ShaderColorKey.lowColorKey, ShaderColorKey.highColorKey,
+				surfaceDesc2.ddckCKSrcOverlay.dwColorSpaceLowValue, surfaceDesc2.ddckCKSrcOverlay.dwColorSpaceHighValue, surfaceDesc2.ddpfPixelFormat);
+			ShaderColorKey.IsSet = true;
+		}
+		else
+		{
+			return false;
+		}
+	}
+	lowColorKey[0] = ShaderColorKey.lowColorKey[0];
+	lowColorKey[1] = ShaderColorKey.lowColorKey[1];
+	lowColorKey[2] = ShaderColorKey.lowColorKey[2];
+	lowColorKey[3] = ShaderColorKey.lowColorKey[3];
+	highColorKey[0] = ShaderColorKey.highColorKey[0];
+	highColorKey[1] = ShaderColorKey.highColorKey[1];
+	highColorKey[2] = ShaderColorKey.highColorKey[2];
+	highColorKey[3] = ShaderColorKey.highColorKey[3];
+	return true;
+}
+
 HRESULT m_IDirectDrawSurfaceX::GetDC(HDC FAR* lphDC)
 {
 	Logging::LogDebug() << __FUNCTION__ << " (" << this << ")";
@@ -4256,6 +4311,10 @@ void m_IDirectDrawSurfaceX::ReleaseD9Surface(bool BackupData, bool DeviceLost)
 	LastLock.bEvenScanlines = false;
 	LastLock.bOddScanlines = false;
 
+	// Reset shader flag
+	ShaderColorKey.IsSet = false;
+	primary.ShaderColorKey.IsSet = false;
+
 	// Reset display flags
 	if (ResetDisplayFlags)
 	{

ddraw/IDirectDrawSurfaceX.h

@@ -66,10 +66,18 @@ class m_IDirectDrawSurfaceX : public IUnknown, public AddressLookupTableDdrawObj
 		bool isAttachedSurfaceAdded = false;
 	};
 
+	struct COLORKEY
+	{
+		bool IsSet = false;
+		float lowColorKey[4] = {};
+		float highColorKey[4] = {};
+	};
+
 	// Extra Direct3D9 devices used in the primary surface
 	struct D9PRIMARY
 	{
 		DWORD LastPaletteUSN = 0;								// The USN that was used last time the palette was updated
+		COLORKEY ShaderColorKey;								// Used to store color key array for shader
 		LPDIRECT3DSURFACE9 BlankSurface = nullptr;				// Blank surface used for clearing main surface
 		LPDIRECT3DTEXTURE9 PaletteTexture = nullptr;			// Extra surface texture used for storing palette entries for the pixel shader
 	};
@@ -134,6 +142,7 @@ class m_IDirectDrawSurfaceX : public IUnknown, public AddressLookupTableDdrawObj
 	DDRAWEMULATELOCK EmuLock;							// For aligning bits after a lock for games that hard code the pitch
 	std::vector<byte> ByteArray;						// Memory used for coping from one surface to the same surface
 	std::vector<byte> Backup;							// Memory used for backing up the surfaceTexture
+	COLORKEY ShaderColorKey;							// Used to store color key array for shader
 	SURFACECREATE ShouldEmulate = SC_NOT_CREATED;		// Used to help determine if surface should be emulated
 
 	// Extra Direct3D9 devices used in the primary surface
@@ -406,22 +415,7 @@ class m_IDirectDrawSurfaceX : public IUnknown, public AddressLookupTableDdrawObj
 	inline bool IsEmulationDCReady() { return (IsUsingEmulation() && !surface.emu->UsingGameDC); }
 	inline bool IsSurface3DDevice() { return Is3DRenderingTarget; }
 	inline bool IsSurfaceDirty() { return surface.IsDirtyFlag; }
-	inline bool GetColorKey(DWORD& ColorSpaceLowValue, DWORD& ColorSpaceHighValue)
-	{
-		if (surfaceDesc2.ddsCaps.dwCaps & DDSD_CKSRCBLT)
-		{
-			ColorSpaceLowValue = surfaceDesc2.ddckCKSrcBlt.dwColorSpaceLowValue;
-			ColorSpaceHighValue = surfaceDesc2.ddckCKSrcBlt.dwColorSpaceHighValue;
-			return true;
-		}
-		else if (surfaceDesc2.ddsCaps.dwCaps & DDCKEY_SRCOVERLAY)
-		{
-			ColorSpaceLowValue = surfaceDesc2.ddckCKSrcOverlay.dwColorSpaceLowValue;
-			ColorSpaceHighValue = surfaceDesc2.ddckCKSrcOverlay.dwColorSpaceHighValue;
-			return true;
-		}
-		return false;
-	}
+	bool GetColorKeyForShader(float(&lowColorKey)[4], float(&highColorKey)[4], bool PrimarySurface);
 	inline bool GetSurfaceSetSize(DWORD& Width, DWORD& Height)
 	{
 		if ((surfaceDesc2.dwFlags & (DDSD_WIDTH | DDSD_HEIGHT)) == (DDSD_WIDTH | DDSD_HEIGHT) &&