From cacc63ef0faa33e3e3160d94e53e92b38dec8cae Mon Sep 17 00:00:00 2001
From: vuurvlieg <vuurvli3g@protonmail.com>
Date: Tue, 2 Apr 2024 16:12:26 +0200
Subject: [PATCH] Add walk methods to Group for convenient & efficient tree
 traversal

---
 src/core/Group.cpp  |  10 +++++
 src/core/Group.h    | 102 ++++++++++++++++++++++++++++++++++++++++++++
 tests/TestGroup.cpp |  99 ++++++++++++++++++++++++++++++++++++++++++
 tests/TestGroup.h   |   1 +
 4 files changed, 212 insertions(+)

diff --git a/src/core/Group.cpp b/src/core/Group.cpp
index 4f8434bab..96287ef15 100644
--- a/src/core/Group.cpp
+++ b/src/core/Group.cpp
@@ -558,6 +558,16 @@ bool Group::hasChildren() const
     return !children().isEmpty();
 }
 
+bool Group::isDescendantOf(const Group* group) const
+{
+    for(const Group* parent = m_parent; parent; parent = parent->m_parent) {
+        if (parent == group) {
+            return true;
+        }
+    }
+    return false;
+}
+
 Database* Group::database()
 {
     return m_db;
diff --git a/src/core/Group.h b/src/core/Group.h
index e10aafd87..3a0e570d1 100644
--- a/src/core/Group.h
+++ b/src/core/Group.h
@@ -20,11 +20,24 @@
 #define KEEPASSX_GROUP_H
 
 #include <QPointer>
+#include <QList>
+#include <utility>
 
 #include "core/CustomData.h"
 #include "core/Database.h"
 #include "core/Entry.h"
 
+
+class Entry;
+class Group;
+
+
+template <typename TCallable> concept CGroupVisitor = std::is_invocable_v<TCallable, Group*>;
+template <typename TCallable> concept CGroupConstVisitor = std::is_invocable_v<TCallable, const Group*>;
+template <typename TCallable> concept CEntryVisitor = std::is_invocable_v<TCallable, Entry*>;
+template <typename TCallable> concept CEntryConstVisitor = std::is_invocable_v<TCallable, const Entry*>;
+
+
 class Group : public ModifiableObject
 {
     Q_OBJECT
@@ -153,6 +166,7 @@ public:
     void setParent(Group* parent, int index = -1, bool trackPrevious = true);
     QStringList hierarchy(int height = -1) const;
     bool hasChildren() const;
+    bool isDescendantOf(const Group* group) const;
 
     Database* database();
     const Database* database() const;
@@ -165,6 +179,41 @@ public:
     QList<Entry*> entriesRecursive(bool includeHistoryItems = false) const;
     QList<const Group*> groupsRecursive(bool includeSelf) const;
     QList<Group*> groupsRecursive(bool includeSelf);
+
+    // walk methods for traversing the tree efficiently (depth-first search)
+    template <CGroupVisitor TGroupCallable, CEntryVisitor TEntryCallable>
+    bool walk(bool includeSelf, TGroupCallable&& groupVisitor, TEntryCallable&& entryVisitor)
+    {
+        return walk<TGroupCallable, TEntryCallable, false, true, true>(
+            includeSelf, std::forward<TGroupCallable>(groupVisitor), std::forward<TEntryCallable>(entryVisitor));
+    }
+    template <CGroupConstVisitor TGroupCallable, CEntryConstVisitor TEntryCallable>
+    bool walk(bool includeSelf, TGroupCallable&& groupVisitor, TEntryCallable&& entryVisitor) const
+    {
+        return walk<TGroupCallable, TEntryCallable, true, true, true>(
+            includeSelf, std::forward<TGroupCallable>(groupVisitor), std::forward<TEntryCallable>(entryVisitor));
+    }
+    template <CGroupConstVisitor TGroupCallable> bool walkGroups(bool includeSelf, TGroupCallable&& groupVisitor) const
+    {
+        return walk<TGroupCallable, void*, true, true, false>(
+            includeSelf, std::forward<TGroupCallable>(groupVisitor), nullptr);
+    }
+    template <CGroupVisitor TGroupCallable> bool walkGroups(bool includeSelf, TGroupCallable&& groupVisitor)
+    {
+        return walk<TGroupCallable, void*, false, true, false>(
+            includeSelf, std::forward<TGroupCallable>(groupVisitor), nullptr);
+    }
+    template <CEntryConstVisitor TEntryCallable> bool walkEntries(TEntryCallable&& entryVisitor) const
+    {
+        return walk<void*, TEntryCallable, true, false, true>(
+            true, nullptr, std::forward<TEntryCallable>(entryVisitor));
+    }
+    template <CEntryVisitor TEntryCallable> bool walkEntries(TEntryCallable&& entryVisitor)
+    {
+        return walk<void*, TEntryCallable, false, false, true>(
+            true, nullptr, std::forward<TEntryCallable>(entryVisitor));
+    }
+
     QSet<QUuid> customIconsRecursive() const;
     QList<QString> usernamesRecursive(int topN = -1) const;
 
@@ -210,6 +259,8 @@ private slots:
     void updateTimeinfo();
 
 private:
+    template <typename TGroupCallable, typename TEntryCallable, bool kIsConst, bool kVisitGroups, bool kVisitEntries>
+    bool walk(bool includeSelf, TGroupCallable&& groupVisitor, TEntryCallable&& entryVisitor) const;
     template <class P, class V> bool set(P& property, const V& value, bool preserveTimeinfo = false);
 
     void emitModifiedEx(bool preserveTimeinfo);
@@ -240,4 +291,55 @@ private:
 
 Q_DECLARE_OPERATORS_FOR_FLAGS(Group::CloneFlags)
 
+// helpers to support non-bool returning callables
+template <bool kDefaultRetVal, typename TCallable, typename... Args>
+bool visitorPredicateImpl(std::true_type, TCallable&& callable, Args&&... args)
+{
+    return callable(std::forward<Args>(args)...);
+}
+
+template <bool kDefaultRetVal, typename TCallable, typename... Args>
+bool visitorPredicateImpl(std::false_type, TCallable&& callable, Args&&... args)
+{
+    callable(std::forward<Args>(args)...);
+    return kDefaultRetVal;
+}
+
+template <bool kDefaultRetVal, typename TCallable, typename... Args>
+bool visitorPredicate(TCallable&& callable, Args&&... args)
+{
+    using RetType = decltype(callable(args...));
+    return visitorPredicateImpl<kDefaultRetVal>(
+        std::is_same<RetType, bool>{}, std::forward<TCallable>(callable), std::forward<Args>(args)...);
+}
+
+template<typename TGroupCallable, typename TEntryCallable, bool kIsConst, bool kVisitGroups, bool kVisitEntries>
+bool Group::walk(bool includeSelf, TGroupCallable&& groupVisitor, TEntryCallable&& entryVisitor) const
+{
+    using GroupType = typename std::conditional<kIsConst,const Group, Group>::type;
+    QList<Group*> groupsToVisit;
+    if (includeSelf) {
+        groupsToVisit.append(const_cast<Group*>(this));
+    } else {
+        groupsToVisit.append(m_children);
+    }
+    while (!groupsToVisit.isEmpty()) {
+        GroupType* group = groupsToVisit.takeLast(); // right-to-left
+        if constexpr (kVisitGroups) {
+            if (visitorPredicate<false>(groupVisitor, group)) {
+                return true;
+            }
+        }
+        if constexpr (kVisitEntries) {
+            for (auto* entry : group->m_entries) {
+                if (visitorPredicate<false>(entryVisitor, entry)) {
+                    return true;
+                }
+            }
+        }
+        groupsToVisit.append(group->m_children);
+    }
+    return false;
+}
+
 #endif // KEEPASSX_GROUP_H
diff --git a/tests/TestGroup.cpp b/tests/TestGroup.cpp
index 7b5daf803..467064709 100644
--- a/tests/TestGroup.cpp
+++ b/tests/TestGroup.cpp
@@ -20,6 +20,7 @@
 #include "mock/MockClock.h"
 
 #include <QSet>
+#include <QVector>
 #include <QSignalSpy>
 #include <QtTestGui>
 
@@ -1381,3 +1382,101 @@ void TestGroup::testTimeinfoChanges()
     QCOMPARE(root->timeInfo().lastModificationTime(), startTime);
     QCOMPARE(subgroup1->timeInfo().lastModificationTime(), startTime);
 }
+
+void TestGroup::testWalk()
+{
+    QScopedPointer<Group> root(new Group());
+    size_t totalGroups{1}, totalEntries{0};
+    for (int i = 0; i < 3; ++i) {
+        Group* subgroup = new Group();
+        subgroup->setParent(root.data());
+        ++totalGroups;
+        int rows = i + 1;
+        int columns = i;
+        QVector<Group*> groupsVec;
+        groupsVec.resize(rows * columns);
+        for (int r = 0; r < rows; ++r) {
+            for (int c = 0; c < columns; ++c) {
+                int index = r * columns + c;
+                Group* group = new Group();
+                groupsVec[index] = group;
+                group->setParent(c > 0 ? groupsVec[index - 1] : subgroup);
+                int entryCount = std::max<int>(1, c + 1 >= columns ? 20 - (i * 3) : c + r);
+                for (int e = 0; e < entryCount; ++e) {
+                    Entry* entry = new Entry();
+                    entry->setGroup(group);
+                }
+                totalEntries += entryCount;
+            }
+        }
+        totalGroups += groupsVec.size();
+    }
+    
+    
+    size_t groupCount{0}, entryCount{0};
+    auto groupCounter = [&](Group* group){ groupCount += 1; };
+    auto entryCounter = [&](const Entry* entry){ entryCount += 1; };
+    bool shouldHaveStopped = false;
+    bool calledAfterStopped = false;
+    auto groupStopHalfway = [&](Group* group) {
+        groupCounter(group);
+        if (groupCount >= totalGroups / 2) {
+            if (shouldHaveStopped) {
+                calledAfterStopped = true;
+            } else {
+                shouldHaveStopped = true;
+                calledAfterStopped = false;
+            }
+            return true;
+        }
+        return false;
+    };
+    auto entryStopHalfWay = [&](Entry* entry) {
+        entryCounter(entry);
+        if (entryCount >= totalEntries / 2) {
+            if (shouldHaveStopped) {
+                calledAfterStopped = true;
+            } else {
+                shouldHaveStopped = true;
+                calledAfterStopped = false;
+            }
+            return true;
+        }
+        return false;
+    };
+
+
+    bool result = root->walk(true, groupCounter, entryCounter);
+    // walk should not stopped
+    QCOMPARE(result, false);
+    // walk should have visited all groups & entries
+    QCOMPARE(groupCount, totalGroups);
+    QCOMPARE(entryCount, totalEntries);
+
+    groupCount = entryCount = 0;
+    result = root->walkGroups(true, groupCounter);
+    QCOMPARE(result, false);
+    QCOMPARE(groupCount, totalGroups);
+    result = const_cast<const Group*>(root.data())->walkEntries(entryCounter);
+    QCOMPARE(result, false);
+    QCOMPARE(entryCount, totalEntries);
+
+    groupCount = entryCount = 0;
+    result = root->walk(false, groupStopHalfway, entryStopHalfWay);
+    // should have stopped
+    QCOMPARE(result, true);
+    // should not have been called after stopped
+    QCOMPARE(calledAfterStopped, false);
+
+    groupCount = entryCount = 0;
+    shouldHaveStopped = false;
+    result = root->walkGroups(false, groupStopHalfway);
+    QCOMPARE(result, true);
+    QCOMPARE(calledAfterStopped, false);
+
+    groupCount = entryCount = 0;
+    shouldHaveStopped = false;
+    result = root->walkEntries(entryStopHalfWay);
+    QCOMPARE(result, true);
+    QCOMPARE(calledAfterStopped, false);
+}
diff --git a/tests/TestGroup.h b/tests/TestGroup.h
index 1f01d9491..8442e7d37 100644
--- a/tests/TestGroup.h
+++ b/tests/TestGroup.h
@@ -51,6 +51,7 @@ private slots:
     void testPreviousParentGroup();
     void testAutoTypeState();
     void testTimeinfoChanges();
+    void testWalk();
 };
 
 #endif // KEEPASSX_TESTGROUP_H