From 3f33545f8b63c30be8303b95dae69ffa9635048e Mon Sep 17 00:00:00 2001
From: Sebastian <sebastian.e.gross@gmail.com>
Date: Sat, 28 Oct 2017 22:25:44 +0100
Subject: [PATCH] Added automatic ImGui context switching when updating
 different editor and PIE worlds. Allows to use ImGui API during world update
 with controls being drawn in correct viewport.

---
 Source/ImGui/Private/ImGuiContextManager.cpp  | 33 +++++++++++++++++++
 Source/ImGui/Private/ImGuiContextManager.h    |  6 +++-
 .../Private/Utilities/WorldContextIndex.h     | 17 +++++++++-
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/Source/ImGui/Private/ImGuiContextManager.cpp b/Source/ImGui/Private/ImGuiContextManager.cpp
index 6145c5e..075fff9 100644
--- a/Source/ImGui/Private/ImGuiContextManager.cpp
+++ b/Source/ImGui/Private/ImGuiContextManager.cpp
@@ -55,8 +55,19 @@ namespace
 #endif // WITH_EDITOR
 }
 
+FImGuiContextManager::FImGuiContextManager()
+{
+#if WITH_EDITOR
+	FWorldDelegates::OnWorldTickStart.AddRaw(this, &FImGuiContextManager::OnWorldTickStart);
+#endif
+}
+
 FImGuiContextManager::~FImGuiContextManager()
 {
+#if WITH_EDITOR
+	// Order matters because contexts can be created during World Tick Start events.
+	FWorldDelegates::OnWorldTickStart.RemoveAll(this);
+#endif
 	Contexts.Empty();
 	ImGui::Shutdown();
 }
@@ -77,6 +88,16 @@ void FImGuiContextManager::Tick(float DeltaSeconds)
 	}
 }
 
+#if WITH_EDITOR
+void FImGuiContextManager::OnWorldTickStart(ELevelTick TickType, float DeltaSeconds)
+{
+	if (GWorld)
+	{
+		GetWorldContextProxy(*GWorld).SetAsCurrent();
+	}
+}
+#endif // WITH_EDITOR
+
 #if WITH_EDITOR
 FImGuiContextManager::FContextData& FImGuiContextManager::GetEditorContextData()
 {
@@ -109,6 +130,18 @@ FImGuiContextManager::FContextData& FImGuiContextManager::GetWorldContextData(co
 {
 	using namespace Utilities;
 
+#if WITH_EDITOR
+	if (World.WorldType == EWorldType::Editor)
+	{
+		if (OutIndex)
+		{
+			*OutIndex = Utilities::EDITOR_CONTEXT_INDEX;
+		}
+
+		return GetEditorContextData();
+	}
+#endif
+
 	const FWorldContext* WorldContext = GetWorldContext(World);
 	const int32 Index = GetWorldContextIndex(*WorldContext);
 
diff --git a/Source/ImGui/Private/ImGuiContextManager.h b/Source/ImGui/Private/ImGuiContextManager.h
index 68c888f..ea94e40 100644
--- a/Source/ImGui/Private/ImGuiContextManager.h
+++ b/Source/ImGui/Private/ImGuiContextManager.h
@@ -11,7 +11,7 @@ class FImGuiContextManager
 {
 public:
 
-	FImGuiContextManager() = default;
+	FImGuiContextManager();
 
 	FImGuiContextManager(const FImGuiContextManager&) = delete;
 	FImGuiContextManager& operator=(const FImGuiContextManager&) = delete;
@@ -86,6 +86,10 @@ private:
 
 #endif // WITH_EDITOR
 
+#if WITH_EDITOR
+	void OnWorldTickStart(ELevelTick TickType, float DeltaSeconds);
+#endif
+
 #if WITH_EDITOR
 	FContextData& GetEditorContextData();
 #endif
diff --git a/Source/ImGui/Private/Utilities/WorldContextIndex.h b/Source/ImGui/Private/Utilities/WorldContextIndex.h
index 6b2a6fb..011ef8e 100644
--- a/Source/ImGui/Private/Utilities/WorldContextIndex.h
+++ b/Source/ImGui/Private/Utilities/WorldContextIndex.h
@@ -38,7 +38,22 @@ namespace Utilities
 		// In single-PIE with dedicated server or multi-PIE sessions worlds have PIEInstance starting from 1 for server
 		// and 2+ for clients, what maps directly to our index.
 
-		return WorldContext.WorldType == EWorldType::PIE ? FMath::Max(WorldContext.PIEInstance, 1) : STANDALONE_GAME_CONTEXT_INDEX;
+		switch (WorldContext.WorldType)
+		{
+		case EWorldType::PIE:
+			return FMath::Max(WorldContext.PIEInstance, 1);
+		case EWorldType::Game:
+			return STANDALONE_GAME_CONTEXT_INDEX;
+		case EWorldType::Editor:
+			return EDITOR_CONTEXT_INDEX;
+		default:
+			return INVALID_CONTEXT_INDEX;
+		}
+	}
+
+	int32 GetWorldContextIndex(const UWorld& World)
+	{
+		return (World.WorldType == EWorldType::Editor) ? EDITOR_CONTEXT_INDEX : GetWorldContextIndex(World.GetGameInstance());
 	}
 
 #else