From 7220b809991c0a6bd1442eea64ec58aed1bd458f Mon Sep 17 00:00:00 2001
From: narzoul <narzoul@users.noreply.github.com>
Date: Thu, 22 Sep 2016 23:10:00 +0200
Subject: [PATCH] Tagged all DirectDraw objects for tracking lifetime

---
 DDrawCompat/DDraw/ActivateAppHandler.cpp      |  3 +-
 DDrawCompat/DDraw/DirectDraw.cpp              | 50 ++++++++++++---
 DDrawCompat/DDraw/DirectDraw.h                |  9 +++
 .../DDraw/Surfaces/FullScreenTagSurface.cpp   | 62 -------------------
 .../DDraw/Surfaces/FullScreenTagSurface.h     | 17 -----
 DDrawCompat/DDraw/Surfaces/Surface.cpp        | 21 ++++++-
 DDrawCompat/DDraw/Surfaces/Surface.h          |  4 ++
 DDrawCompat/DDraw/Surfaces/TagSurface.cpp     | 50 +++++++++++++++
 DDrawCompat/DDraw/Surfaces/TagSurface.h       | 15 +++++
 DDrawCompat/DDrawCompat.vcxproj               |  4 +-
 DDrawCompat/DDrawCompat.vcxproj.filters       |  4 +-
 11 files changed, 145 insertions(+), 94 deletions(-)
 delete mode 100644 DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.cpp
 delete mode 100644 DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.h
 create mode 100644 DDrawCompat/DDraw/Surfaces/TagSurface.cpp
 create mode 100644 DDrawCompat/DDraw/Surfaces/TagSurface.h

diff --git a/DDrawCompat/DDraw/ActivateAppHandler.cpp b/DDrawCompat/DDraw/ActivateAppHandler.cpp
index ec14d24..250a80b 100644
--- a/DDrawCompat/DDraw/ActivateAppHandler.cpp
+++ b/DDrawCompat/DDraw/ActivateAppHandler.cpp
@@ -4,7 +4,6 @@
 #include "DDraw/ActivateAppHandler.h"
 #include "DDraw/DirectDraw.h"
 #include "DDraw/DisplayMode.h"
-#include "DDraw/Surfaces/FullScreenTagSurface.h"
 #include "DDraw/Surfaces/PrimarySurface.h"
 #include "DDraw/Surfaces/SurfaceImpl.h"
 #include "Gdi/Gdi.h"
@@ -93,7 +92,7 @@ namespace
 			Gdi::disableEmulation();
 		}
 
-		auto dd(DDraw::FullScreenTagSurface::getFullScreenDirectDraw());
+		auto dd(DDraw::getFullScreenDirectDraw());
 		if (dd)
 		{
 			if (isActivated)
diff --git a/DDrawCompat/DDraw/DirectDraw.cpp b/DDrawCompat/DDraw/DirectDraw.cpp
index aa7169a..aa602c2 100644
--- a/DDrawCompat/DDraw/DirectDraw.cpp
+++ b/DDrawCompat/DDraw/DirectDraw.cpp
@@ -4,12 +4,38 @@
 #include "DDraw/DirectDraw.h"
 #include "DDraw/DirectDrawSurface.h"
 #include "DDraw/DisplayMode.h"
-#include "DDraw/Surfaces/FullScreenTagSurface.h"
+#include "DDraw/Surfaces/TagSurface.h"
 #include "DDraw/Surfaces/PrimarySurface.h"
 #include "DDraw/Surfaces/Surface.h"
 
 namespace DDraw
 {
+	TagSurface* g_fullScreenTagSurface = nullptr;
+
+	template <typename TDirectDraw>
+	void* getDdObject(TDirectDraw& dd)
+	{
+		return reinterpret_cast<void**>(&dd)[1];
+	}
+
+	template void* getDdObject(IDirectDraw&);
+	template void* getDdObject(IDirectDraw2&);
+	template void* getDdObject(IDirectDraw4&);
+	template void* getDdObject(IDirectDraw7&);
+
+	CompatPtr<IDirectDraw7> getFullScreenDirectDraw()
+	{
+		return g_fullScreenTagSurface ? g_fullScreenTagSurface->getDirectDraw() : nullptr;
+	}
+
+	void onRelease(TagSurface& dd)
+	{
+		if (&dd == g_fullScreenTagSurface)
+		{
+			g_fullScreenTagSurface = nullptr;
+		}
+	}
+
 	template <typename TDirectDraw>
 	void DirectDraw<TDirectDraw>::setCompatVtable(Vtable<TDirectDraw>& vtable)
 	{
@@ -85,16 +111,26 @@ namespace DDraw
 		HRESULT result = s_origVtable.SetCooperativeLevel(This, hWnd, dwFlags);
 		if (SUCCEEDED(result))
 		{
-			if (dwFlags & DDSCL_FULLSCREEN)
+			void* ddObject = getDdObject(*This);
+			TagSurface* tagSurface = TagSurface::get(ddObject);
+			if (!tagSurface)
 			{
 				CompatPtr<IDirectDraw> dd(Compat::queryInterface<IDirectDraw>(This));
-				DDraw::FullScreenTagSurface::create(*dd);
-				ActivateAppHandler::setFullScreenCooperativeLevel(hWnd, dwFlags);
+				TagSurface::create(*dd);
+				tagSurface = TagSurface::get(ddObject);
 			}
-			else if (CompatPtr<IDirectDraw7>(Compat::queryInterface<IDirectDraw7>(This)).get() ==
-				DDraw::FullScreenTagSurface::getFullScreenDirectDraw().get())
+
+			if (tagSurface)
 			{
-				DDraw::FullScreenTagSurface::destroy();
+				if (dwFlags & DDSCL_FULLSCREEN)
+				{
+					g_fullScreenTagSurface = tagSurface;
+					ActivateAppHandler::setFullScreenCooperativeLevel(hWnd, dwFlags);
+				}
+				else if ((dwFlags & DDSCL_NORMAL) && tagSurface == g_fullScreenTagSurface)
+				{
+					g_fullScreenTagSurface = nullptr;
+				}
 			}
 		}
 		return result;
diff --git a/DDrawCompat/DDraw/DirectDraw.h b/DDrawCompat/DDraw/DirectDraw.h
index 1f73aab..7953912 100644
--- a/DDrawCompat/DDraw/DirectDraw.h
+++ b/DDrawCompat/DDraw/DirectDraw.h
@@ -1,11 +1,20 @@
 #pragma once
 
+#include "Common/CompatPtr.h"
 #include "Common/CompatVtable.h"
 #include "DDraw/Visitors/DirectDrawVtblVisitor.h"
 #include "DDraw/Types.h"
 
 namespace DDraw
 {
+	class TagSurface;
+
+	template <typename TDirectDraw>
+	void* getDdObject(TDirectDraw& dd);
+
+	CompatPtr<IDirectDraw7> getFullScreenDirectDraw();
+	void onRelease(TagSurface& dd);
+
 	template <typename TDirectDraw>
 	class DirectDraw: public CompatVtable<Vtable<TDirectDraw>>
 	{
diff --git a/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.cpp b/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.cpp
deleted file mode 100644
index f9454f4..0000000
--- a/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.cpp
+++ /dev/null
@@ -1,62 +0,0 @@
-#include "DDraw/Surfaces/SurfaceImpl.h"
-#include "DDraw/Surfaces/FullScreenTagSurface.h"
-
-namespace
-{
-	CompatWeakPtr<IDirectDrawSurface> g_surface = nullptr;
-}
-
-namespace DDraw
-{
-	FullScreenTagSurface::~FullScreenTagSurface()
-	{
-		g_surface = nullptr;
-	}
-
-	HRESULT FullScreenTagSurface::create(CompatRef<IDirectDraw> dd)
-	{
-		destroy();
-
-		DDSURFACEDESC desc = {};
-		desc.dwSize = sizeof(desc);
-		desc.dwFlags = DDSD_WIDTH | DDSD_HEIGHT | DDSD_CAPS;
-		desc.dwWidth = 1;
-		desc.dwHeight = 1;
-		desc.ddsCaps.dwCaps = DDSCAPS_OFFSCREENPLAIN | DDSCAPS_SYSTEMMEMORY;
-
-		IDirectDrawSurface* surface = nullptr;
-		HRESULT result = Surface::create(dd, desc, surface);
-		if (SUCCEEDED(result))
-		{
-			CompatPtr<IDirectDrawSurface7> surface7(Compat::queryInterface<IDirectDrawSurface7>(surface));
-			std::unique_ptr<Surface> privateData(new FullScreenTagSurface());
-			attach(*surface7, privateData);
-			g_surface = surface;
-		}
-		return result;
-	}
-
-	void FullScreenTagSurface::destroy()
-	{
-		g_surface.release();
-	}
-
-	CompatPtr<IDirectDraw7> FullScreenTagSurface::getFullScreenDirectDraw()
-	{
-		if (!g_surface)
-		{
-			return nullptr;
-		}
-
-		CompatPtr<IUnknown> dd = nullptr;
-		auto tagSurface(getFullScreenTagSurface());
-		tagSurface.get()->lpVtbl->GetDDInterface(tagSurface, reinterpret_cast<void**>(&dd.getRef()));
-		return CompatPtr<IDirectDraw7>(Compat::queryInterface<IDirectDraw7>(dd.get()));
-	}
-
-	CompatPtr<IDirectDrawSurface7> FullScreenTagSurface::getFullScreenTagSurface()
-	{
-		return CompatPtr<IDirectDrawSurface7>(
-			Compat::queryInterface<IDirectDrawSurface7>(g_surface.get()));
-	}
-}
diff --git a/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.h b/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.h
deleted file mode 100644
index f699c0c..0000000
--- a/DDrawCompat/DDraw/Surfaces/FullScreenTagSurface.h
+++ /dev/null
@@ -1,17 +0,0 @@
-#pragma once
-
-#include "DDraw/Surfaces/Surface.h"
-
-namespace DDraw
-{
-	class FullScreenTagSurface : public Surface
-	{
-	public:
-		virtual ~FullScreenTagSurface();
-
-		static HRESULT create(CompatRef<IDirectDraw> dd);
-		static void destroy();
-		static CompatPtr<IDirectDraw7> getFullScreenDirectDraw();
-		static CompatPtr<IDirectDrawSurface7> getFullScreenTagSurface();
-	};
-}
diff --git a/DDrawCompat/DDraw/Surfaces/Surface.cpp b/DDrawCompat/DDraw/Surfaces/Surface.cpp
index b766e3a..357b24d 100644
--- a/DDrawCompat/DDraw/Surfaces/Surface.cpp
+++ b/DDrawCompat/DDraw/Surfaces/Surface.cpp
@@ -1,6 +1,7 @@
 #include <initguid.h>
 
 #include "Common/CompatPtr.h"
+#include "DDraw/DirectDraw.h"
 #include "DDraw/DisplayMode.h"
 #include "DDraw/Surfaces/Surface.h"
 #include "DDraw/Surfaces/SurfaceImpl.h"
@@ -78,7 +79,8 @@ namespace DDraw
 	}
 
 	Surface::Surface()
-		: m_ddId()
+		: m_dds(nullptr)
+		, m_ddId()
 		, m_ddObject(nullptr)
 		, m_refCount(0)
 	{
@@ -104,8 +106,10 @@ namespace DDraw
 			privateData->m_impl4->m_data = privateData.get();
 			privateData->m_impl7->m_data = privateData.get();
 
+			privateData->m_dds = CompatPtr<IDirectDrawSurface>(
+				Compat::queryInterface<IDirectDrawSurface>(&dds));
 			privateData->m_ddId = getDdIidFromVtablePtr(reinterpret_cast<void**>(dd.get())[0]);
-			privateData->m_ddObject = reinterpret_cast<void**>(dd.get())[1];
+			privateData->m_ddObject = DDraw::getDdObject(*CompatPtr<IDirectDraw>(dd));
 
 			privateData.release();
 		}
@@ -181,6 +185,19 @@ namespace DDraw
 	template <>
 	SurfaceImpl<IDirectDrawSurface7>* Surface::getImpl<IDirectDrawSurface7>() const { return m_impl7.get(); }
 
+	CompatPtr<IDirectDraw7> Surface::getDirectDraw() const
+	{
+		auto dds(getDirectDrawSurface());
+		CompatPtr<IUnknown> dd;
+		m_impl7->GetDDInterface(dds, reinterpret_cast<void**>(&dd.getRef()));
+		return CompatPtr<IDirectDraw7>(dd);
+	}
+
+	CompatPtr<IDirectDrawSurface7> Surface::getDirectDrawSurface() const
+	{
+		return CompatPtr<IDirectDrawSurface7>(Compat::queryInterface<IDirectDrawSurface7>(m_dds));
+	}
+
 	template <typename TSurface>
 	Surface* Surface::getSurface(TSurface& dds)
 	{
diff --git a/DDrawCompat/DDraw/Surfaces/Surface.h b/DDrawCompat/DDraw/Surfaces/Surface.h
index 41343e0..8a201e6 100644
--- a/DDrawCompat/DDraw/Surfaces/Surface.h
+++ b/DDrawCompat/DDraw/Surfaces/Surface.h
@@ -29,6 +29,9 @@ namespace DDraw
 		template <typename TSurface>
 		static Surface* getSurface(TSurface& dds);
 
+		CompatPtr<IDirectDraw7> getDirectDraw() const;
+		CompatPtr<IDirectDrawSurface7> getDirectDrawSurface() const;
+
 		template <typename TSurface>
 		SurfaceImpl<TSurface>* getImpl() const;
 
@@ -51,6 +54,7 @@ namespace DDraw
 			IDirectDrawSurface7* surface, DDSURFACEDESC2* desc, void* rootSurface);
 		virtual void createImpl();
 
+		IDirectDrawSurface* m_dds;
 		IID m_ddId;
 		void* m_ddObject;
 		DWORD m_refCount;
diff --git a/DDrawCompat/DDraw/Surfaces/TagSurface.cpp b/DDrawCompat/DDraw/Surfaces/TagSurface.cpp
new file mode 100644
index 0000000..743af80
--- /dev/null
+++ b/DDrawCompat/DDraw/Surfaces/TagSurface.cpp
@@ -0,0 +1,50 @@
+#include <algorithm>
+#include <map>
+
+#include "DDraw/DirectDraw.h"
+#include "DDraw/Surfaces/SurfaceImpl.h"
+#include "DDraw/Surfaces/TagSurface.h"
+
+namespace
+{
+	std::map<void*, DDraw::TagSurface*> g_tagSurfaces;
+}
+
+namespace DDraw
+{
+	TagSurface::~TagSurface()
+	{
+		std::find_if(g_tagSurfaces.begin(), g_tagSurfaces.end(),
+			[=](auto& i) { return i.second == this; })->second;
+		DDraw::onRelease(*this);
+		g_tagSurfaces.erase(std::find_if(g_tagSurfaces.begin(), g_tagSurfaces.end(),
+			[=](auto& i) { return i.second == this; }));
+	}
+
+	HRESULT TagSurface::create(CompatRef<IDirectDraw> dd)
+	{
+		DDSURFACEDESC desc = {};
+		desc.dwSize = sizeof(desc);
+		desc.dwFlags = DDSD_WIDTH | DDSD_HEIGHT | DDSD_CAPS;
+		desc.dwWidth = 1;
+		desc.dwHeight = 1;
+		desc.ddsCaps.dwCaps = DDSCAPS_OFFSCREENPLAIN | DDSCAPS_SYSTEMMEMORY;
+
+		IDirectDrawSurface* surface = nullptr;
+		HRESULT result = Surface::create(dd, desc, surface);
+		if (SUCCEEDED(result))
+		{
+			std::unique_ptr<Surface> privateData(new TagSurface());
+			g_tagSurfaces[getDdObject(dd.get())] = static_cast<TagSurface*>(privateData.get());
+			CompatPtr<IDirectDrawSurface7> surface7(Compat::queryInterface<IDirectDrawSurface7>(surface));
+			attach(*surface7, privateData);
+		}
+		return result;
+	}
+
+	TagSurface* TagSurface::get(void* ddObject)
+	{
+		auto it = g_tagSurfaces.find(ddObject);
+		return it != g_tagSurfaces.end() ? it->second : nullptr;
+	}
+}
diff --git a/DDrawCompat/DDraw/Surfaces/TagSurface.h b/DDrawCompat/DDraw/Surfaces/TagSurface.h
new file mode 100644
index 0000000..29d4194
--- /dev/null
+++ b/DDrawCompat/DDraw/Surfaces/TagSurface.h
@@ -0,0 +1,15 @@
+#pragma once
+
+#include "DDraw/Surfaces/Surface.h"
+
+namespace DDraw
+{
+	class TagSurface : public Surface
+	{
+	public:
+		virtual ~TagSurface();
+
+		static HRESULT create(CompatRef<IDirectDraw> dd);
+		static TagSurface* get(void* ddObject);
+	};
+}
diff --git a/DDrawCompat/DDrawCompat.vcxproj b/DDrawCompat/DDrawCompat.vcxproj
index 59a98ac..f6deb79 100644
--- a/DDrawCompat/DDrawCompat.vcxproj
+++ b/DDrawCompat/DDrawCompat.vcxproj
@@ -182,7 +182,7 @@
     <ClInclude Include="DDraw\Hooks.h" />
     <ClInclude Include="DDraw\Repository.h" />
     <ClInclude Include="DDraw\ScopedThreadLock.h" />
-    <ClInclude Include="DDraw\Surfaces\FullScreenTagSurface.h" />
+    <ClInclude Include="DDraw\Surfaces\TagSurface.h" />
     <ClInclude Include="DDraw\Surfaces\PrimarySurface.h" />
     <ClInclude Include="DDraw\Surfaces\PrimarySurfaceImpl.h" />
     <ClInclude Include="DDraw\Surfaces\Surface.h" />
@@ -234,7 +234,7 @@
     <ClCompile Include="DDraw\Repository.cpp" />
     <ClCompile Include="DDraw\IReleaseNotifier.cpp" />
     <ClCompile Include="DDraw\RealPrimarySurface.cpp" />
-    <ClCompile Include="DDraw\Surfaces\FullScreenTagSurface.cpp" />
+    <ClCompile Include="DDraw\Surfaces\TagSurface.cpp" />
     <ClCompile Include="DDraw\Surfaces\PrimarySurface.cpp" />
     <ClCompile Include="DDraw\Surfaces\PrimarySurfaceImpl.cpp" />
     <ClCompile Include="DDraw\Surfaces\Surface.cpp" />
diff --git a/DDrawCompat/DDrawCompat.vcxproj.filters b/DDrawCompat/DDrawCompat.vcxproj.filters
index 5c32915..8c089aa 100644
--- a/DDrawCompat/DDrawCompat.vcxproj.filters
+++ b/DDrawCompat/DDrawCompat.vcxproj.filters
@@ -255,7 +255,7 @@
     <ClInclude Include="DDraw\Surfaces\PrimarySurfaceImpl.h">
       <Filter>Header Files\DDraw\Surfaces</Filter>
     </ClInclude>
-    <ClInclude Include="DDraw\Surfaces\FullScreenTagSurface.h">
+    <ClInclude Include="DDraw\Surfaces\TagSurface.h">
       <Filter>Header Files\DDraw\Surfaces</Filter>
     </ClInclude>
   </ItemGroup>
@@ -386,7 +386,7 @@
     <ClCompile Include="DDraw\Surfaces\PrimarySurfaceImpl.cpp">
       <Filter>Source Files\DDraw\Surfaces</Filter>
     </ClCompile>
-    <ClCompile Include="DDraw\Surfaces\FullScreenTagSurface.cpp">
+    <ClCompile Include="DDraw\Surfaces\TagSurface.cpp">
       <Filter>Source Files\DDraw\Surfaces</Filter>
     </ClCompile>
   </ItemGroup>