From caa0f015d4e30660ae7f7b34b175b404dcb38475 Mon Sep 17 00:00:00 2001
From: Arnaud Jamin <arnaud.jamin@gmail.com>
Date: Sun, 3 Dec 2023 23:56:15 -0500
Subject: [PATCH] CogImgui: Fix imgui windows not correctly placed at startup
 when using imgui viewports

---
 .../CogImgui/Private/CogImguiContext.cpp      | 29 +++++++++++++++----
 .../Source/CogImgui/Public/CogImguiContext.h  | 15 +++++++---
 .../CogWindow/Private/CogWindowManager.cpp    |  8 +++--
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/Plugins/Cog/Source/CogImgui/Private/CogImguiContext.cpp b/Plugins/Cog/Source/CogImgui/Private/CogImguiContext.cpp
index 6e7ce47..7a2125d 100644
--- a/Plugins/Cog/Source/CogImgui/Private/CogImguiContext.cpp
+++ b/Plugins/Cog/Source/CogImgui/Private/CogImguiContext.cpp
@@ -60,6 +60,7 @@ void FCogImguiContext::Initialize()
     IO.ConfigFlags |= ImGuiConfigFlags_NavNoCaptureKeyboard;
     IO.ConfigFlags |= ImGuiConfigFlags_NavEnableSetMousePos;
     IO.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
+    IO.ConfigFlags |= ImGuiConfigFlags_ViewportsEnable;
 
     IO.BackendFlags |= ImGuiBackendFlags_HasMouseCursors;
     IO.BackendFlags |= ImGuiBackendFlags_HasSetMousePos;
@@ -174,8 +175,19 @@ void FCogImguiContext::OnDisplayMetricsChanged(const FDisplayMetrics& DisplayMet
 }
 
 //--------------------------------------------------------------------------------------------------------------------------
-void FCogImguiContext::BeginFrame(float InDeltaTime)
+bool FCogImguiContext::BeginFrame(float InDeltaTime)
 {
+    //-------------------------------------------------------------------------------------------------------
+    // Skip the first frame, to let the main widget update its TickSpaceGeometry which is returned by the 
+    // plateform callback ImGui_GetWindowPos. When using viewports Imgui needs to know the main viewport 
+    // absolute position to correctly place the initial imgui windows. 
+    //-------------------------------------------------------------------------------------------------------
+    if (bIsFirstFrame)
+    {
+        bIsFirstFrame = false;
+        return false;
+    }
+
     ImGui::SetCurrentContext(ImGuiContext);
     ImPlot::SetImGuiContext(ImGuiContext);
     ImPlot::SetCurrentContext(PlotContext);
@@ -275,10 +287,11 @@ void FCogImguiContext::BeginFrame(float InDeltaTime)
         NewStyle.ScaleAllSizes(DpiScale);
     }
 
-
     ImGui::NewFrame();
 
     //DrawDebug();
+
+    return true;
 }
 
 //--------------------------------------------------------------------------------------------------------------------------
@@ -581,7 +594,7 @@ void FCogImguiContext::SetEnableInput(bool Value)
     if (bEnableInput)
     {
         FSlateThrottleManager::Get().DisableThrottle(true);
-        IsThrottleDisabled = true;
+        bIsThrottleDisabled = true;
 
         FSlateApplication& SlateApp = FSlateApplication::Get();
         
@@ -594,7 +607,7 @@ void FCogImguiContext::SetEnableInput(bool Value)
     }
     else
     {
-        if (IsThrottleDisabled)
+        if (bIsThrottleDisabled)
         {
             FSlateThrottleManager::Get().DisableThrottle(false);
         }
@@ -628,6 +641,12 @@ void FCogImguiContext::SetDPIScale(float Value)
 //--------------------------------------------------------------------------------------------------------------------------
 void FCogImguiContext::BuildFont()
 {
+    if (FontAtlasTexture != nullptr)
+    {
+        FontAtlasTexture->RemoveFromRoot();
+        FontAtlasTexture->ConditionalBeginDestroy();
+    }
+
     ImGuiIO& IO = ImGui::GetIO();
     IO.Fonts->Clear();
 
@@ -639,7 +658,7 @@ void FCogImguiContext::BuildFont()
     int32 TextureWidth, TextureHeight, BytesPerPixel;
     IO.Fonts->GetTexDataAsRGBA32(&TextureDataRaw, &TextureWidth, &TextureHeight, &BytesPerPixel);
 
-    UTexture2D* FontAtlasTexture = UTexture2D::CreateTransient(TextureWidth, TextureHeight, PF_R8G8B8A8, TEXT("ImGuiFontAtlas"));
+    FontAtlasTexture = UTexture2D::CreateTransient(TextureWidth, TextureHeight, PF_R8G8B8A8, TEXT("ImGuiFontAtlas"));
     FontAtlasTexture->Filter = TF_Bilinear;
     FontAtlasTexture->AddressX = TA_Wrap;
     FontAtlasTexture->AddressY = TA_Wrap;
diff --git a/Plugins/Cog/Source/CogImgui/Public/CogImguiContext.h b/Plugins/Cog/Source/CogImgui/Public/CogImguiContext.h
index e57018e..88d4b87 100644
--- a/Plugins/Cog/Source/CogImgui/Public/CogImguiContext.h
+++ b/Plugins/Cog/Source/CogImgui/Public/CogImguiContext.h
@@ -25,6 +25,7 @@ struct COGIMGUI_API FImGuiViewportData
 class COGIMGUI_API FCogImguiContext : public TSharedFromThis<FCogImguiContext>
 {
 public:
+
 	void Initialize();
 
 	void Shutdown();
@@ -37,7 +38,7 @@ public:
 
 	void SetShareMouse(bool Value);
 
-	void BeginFrame(float InDeltaTime);
+	bool BeginFrame(float InDeltaTime);
 
 	void EndFrame();
 
@@ -87,6 +88,9 @@ private:
 
 	static void ImGui_RenderWindow(ImGuiViewport* Viewport, void* Data);
 
+	UPROPERTY()
+	UTexture2D* FontAtlasTexture = nullptr;
+
 	TMap<TWeakPtr<SWindow>, ImGuiID> WindowToViewportMap;
 
 	TSharedPtr<IInputProcessor> InputProcessor = nullptr;
@@ -109,9 +113,12 @@ private:
 
 	bool bShareMouse = true;
 
-	float DpiScale = 1.f;
-
 	bool bRefreshDPIScale = false;
 
-	bool IsThrottleDisabled = false;
+	bool bIsThrottleDisabled = false;
+
+	bool bIsFirstFrame = true;
+
+	float DpiScale = 1.f;
+
 };
diff --git a/Plugins/Cog/Source/CogWindow/Private/CogWindowManager.cpp b/Plugins/Cog/Source/CogWindow/Private/CogWindowManager.cpp
index 2b3da9b..adfc930 100644
--- a/Plugins/Cog/Source/CogWindow/Private/CogWindowManager.cpp
+++ b/Plugins/Cog/Source/CogWindow/Private/CogWindowManager.cpp
@@ -152,9 +152,11 @@ void UCogWindowManager::Tick(float DeltaTime)
         Window->GameTick(DeltaTime);
     }
 
-    Context.BeginFrame(DeltaTime);
-    Render(DeltaTime);
-    Context.EndFrame();
+    if (Context.BeginFrame(DeltaTime))
+    {
+        Render(DeltaTime);
+        Context.EndFrame();
+    }
 }
 
 //--------------------------------------------------------------------------------------------------------------------------