From 206be8b110e2c58b626bd9104c949b68567b5261 Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Sun, 3 Jan 2021 02:37:37 +0100
Subject: [PATCH] Greatly improved hex editor byte highlighting performance

---
 include/lang/pattern_data.hpp    | 53 +++++++++++++++++++++++++++++++-
 include/views/view_hexeditor.hpp |  1 +
 source/views/view_hexeditor.cpp  | 27 +++++++++++-----
 3 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/include/lang/pattern_data.hpp b/include/lang/pattern_data.hpp
index b674b8908..791ac025c 100644
--- a/include/lang/pattern_data.hpp
+++ b/include/lang/pattern_data.hpp
@@ -69,12 +69,22 @@ namespace hex::lang {
         [[nodiscard]] virtual std::string getFormattedName() const = 0;
 
         virtual std::optional<u32> highlightBytes(size_t offset) {
-            if (offset >= this->getOffset() && offset < (this->getOffset() + this->getSize()))
+            auto currOffset = this->getOffset();
+            if (offset >= currOffset && offset < (currOffset + this->getSize()))
                 return this->getColor();
             else
                 return { };
         }
 
+        virtual std::map<u64, u32> getHighlightedAddresses() {
+            if (this->m_highlightedAddresses.empty()) {
+                for (u64 i = 0; i < this->getSize(); i++)
+                    this->m_highlightedAddresses.insert({ this->getOffset() + i, this->getColor() });
+            }
+
+            return this->m_highlightedAddresses;
+        }
+
         virtual void sort(ImGuiTableSortSpecs *sortSpecs, prv::Provider *provider) { }
 
         static bool sortPatternDataTable(ImGuiTableSortSpecs *sortSpecs, prv::Provider *provider, lang::PatternData* left, lang::PatternData* right) {
@@ -156,6 +166,7 @@ namespace hex::lang {
 
     protected:
         std::endian m_endian = std::endian::native;
+        std::map<u64, u32> m_highlightedAddresses;
 
     private:
         u64 m_offset;
@@ -231,6 +242,17 @@ namespace hex::lang {
                 return { };
         }
 
+        std::map<u64, u32> getHighlightedAddresses() override {
+            if (this->m_highlightedAddresses.empty()) {
+                auto ownAddresses = PatternData::getHighlightedAddresses();
+                auto pointedToAddresses = this->m_pointedAt->getHighlightedAddresses();
+
+                ownAddresses.merge(pointedToAddresses);
+                this->m_highlightedAddresses = ownAddresses;
+            }
+
+            return this->m_highlightedAddresses;
+        }
         [[nodiscard]] std::string getFormattedName() const override {
             return "Pointer";
         }
@@ -436,6 +458,15 @@ namespace hex::lang {
             return { };
         }
 
+        std::map<u64, u32> getHighlightedAddresses() override {
+            if (this->m_highlightedAddresses.empty()) {
+                for (auto &entry : this->m_entries) {
+                    this->m_highlightedAddresses.merge(entry->getHighlightedAddresses());
+                }
+            }
+
+            return this->m_highlightedAddresses;
+        }
         [[nodiscard]] std::string getFormattedName() const override {
             return this->m_entries[0]->getTypeName() + "[" + std::to_string(this->m_entries.size()) + "]";
         }
@@ -490,6 +521,16 @@ namespace hex::lang {
             return { };
         }
 
+        std::map<u64, u32> getHighlightedAddresses() override {
+            if (this->m_highlightedAddresses.empty()) {
+                for (auto &member : this->m_members) {
+                    this->m_highlightedAddresses.merge(member->getHighlightedAddresses());
+                }
+            }
+
+            return this->m_highlightedAddresses;
+        }
+
         void sort(ImGuiTableSortSpecs *sortSpecs, prv::Provider *provider) override {
             this->m_sortedMembers = this->m_members;
 
@@ -561,6 +602,16 @@ namespace hex::lang {
             return { };
         }
 
+        std::map<u64, u32> getHighlightedAddresses() override {
+            if (this->m_highlightedAddresses.empty()) {
+                for (auto &member : this->m_members) {
+                    this->m_highlightedAddresses.merge(member->getHighlightedAddresses());
+                }
+            }
+
+            return this->m_highlightedAddresses;
+        }
+
         void sort(ImGuiTableSortSpecs *sortSpecs, prv::Provider *provider) override {
             this->m_sortedMembers = this->m_members;
 
diff --git a/include/views/view_hexeditor.hpp b/include/views/view_hexeditor.hpp
index dbd5f5a14..59e3f65c1 100644
--- a/include/views/view_hexeditor.hpp
+++ b/include/views/view_hexeditor.hpp
@@ -32,6 +32,7 @@ namespace hex {
         imgui_addons::ImGuiFileBrowser m_fileBrowser;
 
         std::vector<lang::PatternData*> &m_patternData;
+        std::map<u64, u32> m_highlightedBytes;
 
         char m_searchStringBuffer[0xFFFF] = { 0 };
         char m_searchHexBuffer[0xFFFF] = { 0 };
diff --git a/source/views/view_hexeditor.cpp b/source/views/view_hexeditor.cpp
index 878c7dbf0..75ecf05e6 100644
--- a/source/views/view_hexeditor.cpp
+++ b/source/views/view_hexeditor.cpp
@@ -42,15 +42,19 @@ namespace hex {
         this->m_memoryEditor.HighlightFn = [](const ImU8 *data, size_t off, bool next) -> bool {
             ViewHexEditor *_this = (ViewHexEditor *) data;
 
-            for (auto& pattern : _this->m_patternData) {
-                if (next && pattern->highlightBytes(off - 1) != pattern->highlightBytes(off)) {
-                    return false;
-                }
+            std::optional<u32> currColor, prevColor;
+            if (_this->m_highlightedBytes.contains(off))
+                currColor = _this->m_highlightedBytes[off];
+            if (_this->m_highlightedBytes.contains(off - 1))
+                prevColor = _this->m_highlightedBytes[off - 1];
 
-                if (auto color = pattern->highlightBytes(off); color.has_value()) {
-                    _this->m_memoryEditor.HighlightColor = color.value();
-                    return true;
-                }
+            if (next && prevColor != currColor) {
+                return false;
+            }
+
+            if (currColor.has_value()) {
+                _this->m_memoryEditor.HighlightColor = currColor.value();
+                return true;
             }
 
             _this->m_memoryEditor.HighlightColor = 0x60C08080;
@@ -92,6 +96,13 @@ namespace hex {
                 View::doLater([] { ImGui::OpenPopup("Save Changes"); });
             }
         });
+
+        View::subscribeEvent(Events::PatternChanged, [this](const void *userData) {
+           this->m_highlightedBytes.clear();
+
+           for (const auto &pattern : this->m_patternData)
+               this->m_highlightedBytes.merge(pattern->getHighlightedAddresses());
+        });
     }
 
     ViewHexEditor::~ViewHexEditor() {