CogImGui: Fix crash when restating level

Plugins/Cog/Source/CogImgui/Private/CogImguiModule.cpp

@@ -37,7 +37,7 @@ void FCogImguiModule::Initialize()
 }
 
 //--------------------------------------------------------------------------------------------------------------------------
-TSharedPtr<SCogImguiWidget> FCogImguiModule::CreateImGuiViewport(UGameViewportClient* GameViewport, FCogImguiRenderFunction Render, ImFontAtlas* FontAtlas /*= nullptr*/)
+TSharedPtr<SCogImguiWidget> FCogImguiModule::CreateImGuiWidget(UGameViewportClient* GameViewport, FCogImguiRenderFunction Render, ImFontAtlas* FontAtlas /*= nullptr*/)
 {
     if (bIsInitialized == false)
     {
@@ -72,6 +72,24 @@ TSharedPtr<SCogImguiWidget> FCogImguiModule::CreateImGuiViewport(UGameViewportCl
     return ImguiWidget;
 }
 
+//--------------------------------------------------------------------------------------------------------------------------
+void FCogImguiModule::DestroyImGuiWidget(TSharedPtr<SCogImguiWidget> ImGuiWidget)
+{
+    UGameViewportClient* Viewport = ImGuiWidget->GetGameViewport().Get();
+    if (Viewport == nullptr)
+    {
+        return;
+    }
+
+    TSharedPtr<SWidget> ParentWidget = ImGuiWidget->GetParentWidget();
+    if (ParentWidget.IsValid() == false)
+    {
+        return;
+    }
+
+    Viewport->RemoveViewportWidgetContent(ParentWidget.ToSharedRef());
+}
+
 //--------------------------------------------------------------------------------------------------------------------------
 #undef LOCTEXT_NAMESPACE
 

Plugins/Cog/Source/CogImgui/Public/CogImguiModule.h

@@ -23,7 +23,9 @@ public:
 	virtual void ShutdownModule() override;
 
     //----------------------------------------------------------------------------------------------------------------------
-    TSharedPtr<SCogImguiWidget> CreateImGuiViewport(UGameViewportClient* GameViewport,  FCogImguiRenderFunction Render, ImFontAtlas* FontAtlas = nullptr);
+    TSharedPtr<SCogImguiWidget> CreateImGuiWidget(UGameViewportClient* GameViewport,  FCogImguiRenderFunction Render, ImFontAtlas* FontAtlas = nullptr);
+
+    void DestroyImGuiWidget(TSharedPtr<SCogImguiWidget> ImGuiWidget);
 
     FCogImguiTextureManager& GetTextureManager() { return TextureManager; }
     

Plugins/Cog/Source/CogImgui/Public/CogImguiWidget.h

@@ -65,6 +65,8 @@ public:
 
     void SetAsCurrentContext();
 
+    TWeakObjectPtr<UGameViewportClient> GetGameViewport() const { return GameViewport; }
+
 protected:
 
     FVector2D TransformScreenPointToImGui(const FGeometry& MyGeometry, const FVector2D& Point) const;

Plugins/Cog/Source/CogWindow/Private/CogWindowManager.cpp

@@ -38,7 +38,7 @@ void UCogWindowManager::PostInitProperties()
 //--------------------------------------------------------------------------------------------------------------------------
 void UCogWindowManager::InitializeInternal()
 {
-    ImGuiWidget = FCogImguiModule::Get().CreateImGuiViewport(GEngine->GameViewport, [this](float DeltaTime) { Render(DeltaTime); });
+    ImGuiWidget = FCogImguiModule::Get().CreateImGuiWidget(GEngine->GameViewport, [this](float DeltaTime) { Render(DeltaTime); });
 
     ImGuiSettingsHandler IniHandler;
     IniHandler.TypeName = "Cog";
@@ -98,6 +98,8 @@ void UCogWindowManager::Shutdown()
         IConsoleManager::Get().UnregisterConsoleObject(ConsoleCommand);
     }
 
+    FCogImguiModule::Get().DestroyImGuiWidget(ImGuiWidget);
+
     SaveConfig();
 }