From 12ba05385b30fb138f8df4b77077533a369684f5 Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Thu, 15 Feb 2024 21:54:41 +0100
Subject: [PATCH] fix: Saving not removing red highlighting

---
 .../include/content/views/view_patches.hpp    |  1 +
 .../content/providers/file_provider.cpp       |  7 +++--
 .../source/content/views/view_patches.cpp     | 30 ++++++++++++++-----
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/plugins/builtin/include/content/views/view_patches.hpp b/plugins/builtin/include/content/views/view_patches.hpp
index de5a00d87..8f8bc7862 100644
--- a/plugins/builtin/include/content/views/view_patches.hpp
+++ b/plugins/builtin/include/content/views/view_patches.hpp
@@ -17,6 +17,7 @@ namespace hex::plugin::builtin {
     private:
         u64 m_selectedPatch = 0x00;
         PerProvider<u32> m_numOperations;
+        PerProvider<u32> m_savedOperations;
     };
 
 }
\ No newline at end of file
diff --git a/plugins/builtin/source/content/providers/file_provider.cpp b/plugins/builtin/source/content/providers/file_provider.cpp
index 507392c06..4d764ab70 100644
--- a/plugins/builtin/source/content/providers/file_provider.cpp
+++ b/plugins/builtin/source/content/providers/file_provider.cpp
@@ -57,15 +57,16 @@ namespace hex::plugin::builtin {
     }
 
     void FileProvider::save() {
+        m_file.flush();
+
         #if defined(OS_WINDOWS)
             FILETIME ft;
             SYSTEMTIME st;
 
-            wolv::io::File file(m_path, wolv::io::File::Mode::Write);
-            if (file.isValid()) {
+            if (m_file.isValid()) {
                 GetSystemTime(&st);
                 if (SystemTimeToFileTime(&st, &ft)) {
-                    auto fileHandle = HANDLE(_get_osfhandle(_fileno(file.getHandle())));
+                    auto fileHandle = HANDLE(_get_osfhandle(_fileno(m_file.getHandle())));
                     SetFileTime(fileHandle, nullptr, nullptr, &ft);
                 }
             }
diff --git a/plugins/builtin/source/content/views/view_patches.cpp b/plugins/builtin/source/content/views/view_patches.cpp
index 53f7c4411..5f59afc74 100644
--- a/plugins/builtin/source/content/views/view_patches.cpp
+++ b/plugins/builtin/source/content/views/view_patches.cpp
@@ -38,7 +38,7 @@ namespace hex::plugin::builtin {
             }
         });
 
-        ImHexApi::HexEditor::addForegroundHighlightingProvider([](u64 offset, const u8* buffer, size_t, bool) -> std::optional<color_t> {
+        ImHexApi::HexEditor::addForegroundHighlightingProvider([this](u64 offset, const u8* buffer, size_t, bool) -> std::optional<color_t> {
             hex::unused(buffer);
 
             if (!ImHexApi::Provider::isValid())
@@ -49,18 +49,34 @@ namespace hex::plugin::builtin {
             offset -= provider->getBaseAddress();
 
             const auto &undoStack = provider->getUndoStack();
-            for (const auto &operation : undoStack.getAppliedOperations()) {
-                if (!operation->shouldHighlight())
-                    continue;
+            const auto stackSize = undoStack.getAppliedOperations().size();
+            const auto savedStackSize = m_savedOperations.get(provider);
 
-                if (operation->getRegion().overlaps(Region { offset, 1}))
-                    return ImGuiExt::GetCustomColorU32(ImGuiCustomCol_Patches);
+            if (stackSize == savedStackSize) {
+                // Do nothing
+            } else if (stackSize > savedStackSize) {
+                for (const auto &operation : undoStack.getAppliedOperations() | std::views::drop(savedStackSize)) {
+                    if (!operation->shouldHighlight())
+                        continue;
+
+                    if (operation->getRegion().overlaps(Region { offset, 1}))
+                        return ImGuiExt::GetCustomColorU32(ImGuiCustomCol_Patches);
+                }
+            } else {
+                for (const auto &operation : undoStack.getUndoneOperations() | std::views::reverse | std::views::take(savedStackSize - stackSize)) {
+                    if (!operation->shouldHighlight())
+                        continue;
+
+                    if (operation->getRegion().overlaps(Region { offset, 1}))
+                        return ImGuiExt::GetCustomColorU32(ImGuiCustomCol_Patches);
+                }
             }
 
             return std::nullopt;
         });
 
-        EventProviderSaved::subscribe([](auto *) {
+        EventProviderSaved::subscribe([this](prv::Provider *provider) {
+            m_savedOperations.get(provider) = provider->getUndoStack().getAppliedOperations().size();
             EventHighlightingChanged::post();
         });