From 214e542da429b28da9fadb4cb3bcec7e2d70c648 Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Fri, 23 Feb 2024 17:49:20 +0100
Subject: [PATCH] impr: Make sure lots of yara matches doesn't lag out the
 editor

---
 .../include/content/views/view_yara.hpp       |  6 +-
 .../content/data_information_sections.cpp     |  2 +-
 .../source/content/views/view_yara.cpp        | 62 +++++++++++++++----
 3 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/plugins/yara_rules/include/content/views/view_yara.hpp b/plugins/yara_rules/include/content/views/view_yara.hpp
index e2dcd7689..3ff829d8d 100644
--- a/plugins/yara_rules/include/content/views/view_yara.hpp
+++ b/plugins/yara_rules/include/content/views/view_yara.hpp
@@ -7,6 +7,7 @@
 #include <hex/api/task_manager.hpp>
 
 #include <content/yara_rule.hpp>
+#include <wolv/container/interval_tree.hpp>
 
 namespace hex::plugin::yara {
 
@@ -22,8 +23,7 @@ namespace hex::plugin::yara {
         PerProvider<std::vector<YaraRule::Rule>> m_matchedRules;
         PerProvider<std::vector<std::string>> m_consoleMessages;
         PerProvider<u32> m_selectedRule;
-
-        PerProvider<std::vector<u32>> m_tooltipIds, m_highlightingIds;
+        PerProvider<wolv::container::IntervalTree<std::string>> m_highlights;
 
         TaskHolder m_matcherTask;
 
@@ -31,4 +31,4 @@ namespace hex::plugin::yara {
         void clearResult();
     };
 
-}
\ No newline at end of file
+}
diff --git a/plugins/yara_rules/source/content/data_information_sections.cpp b/plugins/yara_rules/source/content/data_information_sections.cpp
index 9827f3359..4943de2f4 100644
--- a/plugins/yara_rules/source/content/data_information_sections.cpp
+++ b/plugins/yara_rules/source/content/data_information_sections.cpp
@@ -13,7 +13,7 @@ namespace hex::plugin::yara {
 
     class InformationAdvancedFileInformation : public ContentRegistry::DataInformation::InformationSection {
     public:
-        InformationAdvancedFileInformation() : InformationSection("hex.yara.information_section.advanced_file_info") { }
+        InformationAdvancedFileInformation() : InformationSection("hex.yara.information_section.advanced_data_info") { }
         ~InformationAdvancedFileInformation() override = default;
 
         struct Category {
diff --git a/plugins/yara_rules/source/content/views/view_yara.cpp b/plugins/yara_rules/source/content/views/view_yara.cpp
index c89fc5be2..6de536fd6 100644
--- a/plugins/yara_rules/source/content/views/view_yara.cpp
+++ b/plugins/yara_rules/source/content/views/view_yara.cpp
@@ -82,6 +82,47 @@ namespace hex::plugin::yara {
                 return true;
             }
         });
+
+        ImHexApi::HexEditor::addBackgroundHighlightingProvider([this](u64 address, const u8 *, size_t size, bool) -> std::optional<color_t> {
+            auto &highlights = m_highlights.get();
+            const auto regions = highlights.overlapping({ address, address + (size - 1) });
+
+            constexpr static color_t YaraColor = 0x70B4771F;
+            if (regions.empty())
+                return std::nullopt;
+            else
+                return YaraColor;
+        });
+
+        ImHexApi::HexEditor::addTooltipProvider([this](u64 address, const u8 *, size_t size) {
+            if (m_matcherTask.isRunning())
+                return;
+
+            auto occurrences = m_highlights->overlapping({ address, (address + size - 1) });
+            if (occurrences.empty())
+                return;
+
+            ImGui::BeginTooltip();
+
+            for (const auto &occurrence : occurrences) {
+                ImGui::PushID(&occurrence);
+                if (ImGui::BeginTable("##tooltips", 1, ImGuiTableFlags_RowBg | ImGuiTableFlags_NoClip)) {
+                    ImGui::TableNextRow();
+                    ImGui::TableNextColumn();
+
+                    {
+                        const auto &tooltipValue = *occurrence.value;
+
+                        ImGuiExt::TextFormatted("{}", tooltipValue);
+                    }
+
+                    ImGui::EndTable();
+                }
+                ImGui::PopID();
+            }
+
+            ImGui::EndTooltip();
+        });
     }
 
     ViewYara::~ViewYara() {
@@ -194,13 +235,6 @@ namespace hex::plugin::yara {
     }
 
     void ViewYara::clearResult() {
-        for (const auto &id : *m_highlightingIds)
-            ImHexApi::HexEditor::removeBackgroundHighlight(id);
-        for (const auto &id : *m_tooltipIds)
-            ImHexApi::HexEditor::removeTooltip(id);
-
-        m_highlightingIds->clear();
-        m_tooltipIds->clear();
         m_matchedRules->clear();
         m_consoleMessages->clear();
     }
@@ -229,7 +263,7 @@ namespace hex::plugin::yara {
                         m_consoleMessages->emplace_back(error.message);
                     });
 
-                    break;
+                    return;
                 }
 
                 results.emplace_back(result.value());
@@ -246,11 +280,17 @@ namespace hex::plugin::yara {
                     m_consoleMessages->insert(m_consoleMessages->end(), result.consoleMessages.begin(), result.consoleMessages.end());
                 }
 
-                constexpr static color_t YaraColor = 0x70B4771F;
                 for (YaraRule::Rule &rule : *m_matchedRules) {
                     for (auto &match : rule.matches) {
-                        m_highlightingIds->push_back(ImHexApi::HexEditor::addBackgroundHighlight(match.region, YaraColor));
-                        m_tooltipIds->push_back(ImHexApi::HexEditor::addTooltip(match.region, hex::format("{0} : {1}   [{2}]", rule.identifier, fmt::join(rule.tags, ", "), match.variable), YaraColor));
+                        auto tags = hex::format("{}", fmt::join(rule.tags, ", "));
+                        m_highlights->insert(
+                            { match.region.getStartAddress(), match.region.getEndAddress() },
+                            hex::format("rule {0}{1} {{ {2} }}",
+                                rule.identifier,
+                                tags.empty() ? "" : hex::format(" : {}", tags),
+                                match.variable
+                            )
+                        );
                     }
                 }
             });