Reimplement search for GridImageModel

2 files changed, 148 insertions, 25 deletions

diff --git a/src/GridImagePackModel.cpp b/src/GridImagePackModel.cpp
index 4fee086a..59bfae80 100644
--- a/src/GridImagePackModel.cpp
+++ b/src/GridImagePackModel.cpp
@@ -4,11 +4,12 @@
 
 #include "GridImagePackModel.h"
 
-#include "Cache_p.h"
-#include "CompletionModelRoles.h"
+#include <QTextBoundaryFinder>
 
 #include <algorithm>
 
+#include "Cache_p.h"
+
 Q_DECLARE_METATYPE(StickerImage)
 
 GridImagePackModel::GridImagePackModel(const std::string &roomId, bool stickers, QObject *parent)
@@ -38,12 +39,52 @@ GridImagePackModel::GridImagePackModel(const std::string &roomId, bool stickers,
             rowToPack.push_back(packs.size());
         packs.push_back(std::move(newPack));
     }
+
+    // prepare search index
+
+    auto insertParts = [this](const QString &str, std::pair<std::uint32_t, std::uint32_t> id) {
+        QTextBoundaryFinder finder(QTextBoundaryFinder::BoundaryType::Word, str);
+        finder.toStart();
+        do {
+            auto start = finder.position();
+            finder.toNextBoundary();
+            auto end = finder.position();
+
+            auto ref = str.midRef(start, end - start).trimmed();
+            if (!ref.isEmpty())
+                trie_.insert<ElementRank::second>(ref.toUcs4(), id);
+        } while (finder.position() < str.size());
+    };
+
+    std::uint32_t packIndex = 0;
+    for (const auto &pack : packs) {
+        std::uint32_t imgIndex = 0;
+        for (const auto &img : pack.images) {
+            std::pair<std::uint32_t, std::uint32_t> key{packIndex, imgIndex};
+
+            QString string1 = img.second.toCaseFolded();
+            QString string2 = QString::fromStdString(img.first.body).toCaseFolded();
+
+            if (!string1.isEmpty()) {
+                trie_.insert<ElementRank::first>(string1.toUcs4(), key);
+                insertParts(string1, key);
+            }
+            if (!string2.isEmpty()) {
+                trie_.insert<ElementRank::first>(string2.toUcs4(), key);
+                insertParts(string2, key);
+            }
+
+            imgIndex++;
+        }
+        packIndex++;
+    }
 }
 
 int
 GridImagePackModel::rowCount(const QModelIndex &) const
 {
-    return (int)rowToPack.size();
+    return static_cast<int>(searchString_.isEmpty() ? rowToPack.size()
+                                                    : rowToFirstRowEntryFromSearch.size());
 }
 
 QHash<int, QByteArray>
@@ -59,30 +100,96 @@ QVariant
 GridImagePackModel::data(const QModelIndex &index, int role) const
 {
     if (index.row() < rowCount() && index.row() >= 0) {
-        const auto &pack = packs[rowToPack[index.row()]];
-        switch (role) {
-        case Roles::PackName:
-            return pack.packname;
-        case Roles::Row: {
-            std::size_t offset = static_cast<std::size_t>(index.row()) - pack.firstRow;
-            QList<StickerImage> imgs;
-            auto endOffset = std::min((offset + 1) * 3, pack.images.size());
-            for (std::size_t img = offset * 3; img < endOffset; img++) {
-                const auto &data = pack.images.at(img);
-                imgs.push_back({.url         = QString::fromStdString(data.first.url),
-                                .shortcode   = data.second,
-                                .body        = QString::fromStdString(data.first.body),
-                                .descriptor_ = std::vector{
-                                  pack.room_id,
-                                  pack.state_key,
-                                  data.second.toStdString(),
-                                }});
+        if (searchString_.isEmpty()) {
+            const auto &pack = packs[rowToPack[index.row()]];
+            switch (role) {
+            case Roles::PackName:
+                return pack.packname;
+            case Roles::Row: {
+                std::size_t offset = static_cast<std::size_t>(index.row()) - pack.firstRow;
+                QList<StickerImage> imgs;
+                auto endOffset = std::min((offset + 1) * columns, pack.images.size());
+                for (std::size_t img = offset * columns; img < endOffset; img++) {
+                    const auto &data = pack.images.at(img);
+                    imgs.push_back({.url         = QString::fromStdString(data.first.url),
+                                    .shortcode   = data.second,
+                                    .body        = QString::fromStdString(data.first.body),
+                                    .descriptor_ = std::vector{
+                                      pack.room_id,
+                                      pack.state_key,
+                                      data.second.toStdString(),
+                                    }});
+                }
+                return QVariant::fromValue(imgs);
+            }
+            default:
+                return {};
+            }
+        } else {
+            if (static_cast<size_t>(index.row()) >= rowToFirstRowEntryFromSearch.size())
+                return {};
+
+            const auto firstIndex = rowToFirstRowEntryFromSearch[index.row()];
+            const auto firstEntry = currentSearchResult[firstIndex];
+            const auto &pack      = packs[firstEntry.first];
+
+            switch (role) {
+            case Roles::PackName:
+                return pack.packname;
+            case Roles::Row: {
+                QList<StickerImage> imgs;
+                for (auto img = firstIndex;
+                     imgs.size() < columns && img < currentSearchResult.size() &&
+                     currentSearchResult[img].first == firstEntry.first;
+                     img++) {
+                    const auto &data = pack.images.at(currentSearchResult[img].second);
+                    imgs.push_back({.url         = QString::fromStdString(data.first.url),
+                                    .shortcode   = data.second,
+                                    .body        = QString::fromStdString(data.first.body),
+                                    .descriptor_ = std::vector{
+                                      pack.room_id,
+                                      pack.state_key,
+                                      data.second.toStdString(),
+                                    }});
+                }
+                return QVariant::fromValue(imgs);
+            }
+            default:
+                return {};
             }
-            return QVariant::fromValue(imgs);
-        }
-        default:
-            return {};
         }
     }
     return {};
 }
+
+void
+GridImagePackModel::setSearchString(QString key)
+{
+    beginResetModel();
+    currentSearchResult.clear();
+    rowToFirstRowEntryFromSearch.clear();
+    searchString_ = key;
+
+    if (!key.isEmpty()) {
+        auto searchParts = key.toCaseFolded().toUcs4();
+        auto tempResults =
+          trie_.search(searchParts, static_cast<std::size_t>(columns * columns * 4));
+        std::ranges::sort(tempResults);
+        currentSearchResult = std::move(tempResults);
+
+        std::size_t lastPack = -1;
+        int columnIndex      = 0;
+        for (std::size_t i = 0; i < currentSearchResult.size(); i++) {
+            auto elem = currentSearchResult[i];
+            if (elem.first != lastPack || columnIndex == columns) {
+                columnIndex = 0;
+                lastPack    = elem.first;
+                rowToFirstRowEntryFromSearch.push_back(i);
+            }
+            columnIndex++;
+        }
+    }
+
+    endResetModel();
+    emit newSearchString();
+}
diff --git a/src/GridImagePackModel.h b/src/GridImagePackModel.h
index 1345b103..06dfe734 100644
--- a/src/GridImagePackModel.h
+++ b/src/GridImagePackModel.h
@@ -5,11 +5,14 @@
 #pragma once
 
 #include <QAbstractListModel>
+#include <QMultiMap>
 #include <QObject>
 #include <QString>
 
 #include <mtx/events/mscs/image_packs.hpp>
 
+#include "CompletionProxyModel.h"
+
 struct StickerImage
 {
     Q_GADGET
@@ -41,6 +44,8 @@ public:
 class GridImagePackModel final : public QAbstractListModel
 {
     Q_OBJECT
+    Q_PROPERTY(QString searchString READ searchString WRITE setSearchString NOTIFY newSearchString)
+
 public:
     enum Roles
     {
@@ -53,6 +58,12 @@ public:
     int rowCount(const QModelIndex &parent = QModelIndex()) const override;
     QVariant data(const QModelIndex &index, int role) const override;
 
+    QString searchString() const { return searchString_; }
+    void setSearchString(QString newValue);
+
+signals:
+    void newSearchString();
+
 private:
     std::string room_id;
 
@@ -69,4 +80,9 @@ private:
     std::vector<PackDesc> packs;
     std::vector<size_t> rowToPack;
     int columns = 3;
+
+    QString searchString_;
+    trie<uint, std::pair<std::uint32_t, std::uint32_t>> trie_;
+    std::vector<std::pair<std::uint32_t, std::uint32_t>> currentSearchResult;
+    std::vector<std::size_t> rowToFirstRowEntryFromSearch;
 };