Utils: Add TreeModel iteration facility

Change-Id: Iaf115377de0e5fc0b004d5ea8ddc5c6eb31b5b6f
Reviewed-by: Christian Stenger <christian.stenger@theqtcompany.com>

src/libs/utils/treemodel.cpp

@@ -65,8 +65,7 @@ TreeItem *TreeItem::child(int pos) const
 {
     ensurePopulated();
     QTC_ASSERT(pos >= 0, return 0);
-    QTC_ASSERT(pos < m_children.size(), return 0);
-    return m_children.at(pos);
+    return pos < m_children.size() ? m_children.at(pos) : 0;
 }
 
 bool TreeItem::isLazy() const
@@ -268,7 +267,7 @@ TreeItem *TreeModel::itemFromIndex(const QModelIndex &idx) const
 QModelIndex TreeModel::indexFromItem(const TreeItem *item) const
 {
 //    CHECK(checkItem(item));
-    return indexFromItemHelper(item, m_root, QModelIndex());
+    return item ? indexFromItemHelper(item, m_root, QModelIndex()) : QModelIndex();
 }
 
 void TreeModel::appendItem(TreeItem *parent, TreeItem *item)
@@ -289,6 +288,18 @@ void TreeModel::updateItem(TreeItem *item)
     dataChanged(idx.sibling(idx.row(), 0), idx.sibling(idx.row(), item->columnCount() - 1));
 }
 
+UntypedTreeLevelItems TreeModel::untypedLevelItems(int level, TreeItem *start) const
+{
+    if (start == 0)
+        start = m_root;
+    return UntypedTreeLevelItems(start, level);
+}
+
+UntypedTreeLevelItems TreeModel::untypedLevelItems(TreeItem *start) const
+{
+    return UntypedTreeLevelItems(start, 1);
+}
+
 void TreeModel::removeItem(TreeItem *item)
 {
     QTC_ASSERT(item, return);
@@ -330,4 +341,98 @@ void TreeModel::checkIndex(const QModelIndex &index) const
     }
 }
 
+//
+// TreeLevelItems
+//
+
+UntypedTreeLevelItems::UntypedTreeLevelItems(TreeItem *item, int level)
+    : m_item(item), m_level(level)
+{}
+
+UntypedTreeLevelItems::const_iterator::const_iterator(TreeItem *base, int level)
+    : m_level(level)
+{
+    QTC_ASSERT(level > 0, return);
+    if (base) {
+        // "begin()"
+        m_depth = 0;
+        // Level x: The item m_item[x] is the m_pos[x]'th child of its
+        // parent, out of m_size[x].
+        m_pos[0] = 0;
+        m_size[0] = 1;
+        m_item[0] = base;
+        goDown();
+    } else {
+        // "end()"
+        m_depth = -1;
+    }
+}
+
+
+// Result is either an item of the target level, or 'end'.
+void UntypedTreeLevelItems::const_iterator::goDown()
+{
+    QTC_ASSERT(m_depth != -1, return);
+    QTC_ASSERT(m_depth < m_level, return);
+    do {
+        TreeItem *curr = m_item[m_depth];
+        int size = curr->rowCount();
+        if (size == 0) {
+            // This is a dead end not reaching to the desired level.
+            goUpNextDown();
+            return;
+        }
+        ++m_depth;
+        m_size[m_depth] = size;
+        m_pos[m_depth] = 0;
+        m_item[m_depth] = curr->child(0);
+    } while (m_depth < m_level);
+    // Did not reach the required level? Set to end().
+    if (m_depth != m_level)
+        m_depth = -1;
+}
+
+void UntypedTreeLevelItems::const_iterator::goUpNextDown()
+{
+    // Go up until we can move sidewards.
+    do {
+        --m_depth;
+        if (m_depth < 0)
+            return; // Solid end.
+    } while (++m_pos[m_depth] >= m_size[m_depth]);
+    m_item[m_depth] = m_item[m_depth - 1]->child(m_pos[m_depth]);
+    goDown();
+}
+
+bool UntypedTreeLevelItems::const_iterator::operator==(UntypedTreeLevelItems::const_iterator other) const
+{
+    if (m_depth != other.m_depth)
+        return false;
+    for (int i = 0; i <= m_depth; ++i)
+        if (m_item[i] != other.m_item[i])
+            return false;
+    return true;
+}
+
+void UntypedTreeLevelItems::const_iterator::operator++()
+{
+    QTC_ASSERT(m_depth == m_level, return);
+
+    int pos = ++m_pos[m_depth];
+    if (pos < m_size[m_depth])
+        m_item[m_depth] = m_item[m_depth - 1]->child(pos);
+    else
+        goUpNextDown();
+}
+
+UntypedTreeLevelItems::const_iterator UntypedTreeLevelItems::begin() const
+{
+    return const_iterator(m_item, m_level);
+}
+
+UntypedTreeLevelItems::const_iterator UntypedTreeLevelItems::end() const
+{
+    return const_iterator(0, m_level);
+}
+
 } // namespace Utils

src/libs/utils/treemodel.h

@@ -32,9 +32,13 @@
 #define UTILS_TREEMODEL_H
 
 #include "utils_global.h"
+#include "algorithm.h"
 
 #include <QAbstractItemModel>
 
+#include <functional>
+#include <iterator>
+
 namespace Utils {
 
 class QTCREATOR_UTILS_EXPORT TreeItem
@@ -81,6 +85,76 @@ private:
     friend class TreeModel;
 };
 
+class QTCREATOR_UTILS_EXPORT UntypedTreeLevelItems
+{
+public:
+    explicit UntypedTreeLevelItems(TreeItem *item, int level = 1);
+
+    typedef TreeItem *value_type;
+
+    class const_iterator
+    {
+    public:
+        typedef std::forward_iterator_tag iterator_category;
+        typedef TreeItem *value_type;
+        typedef std::ptrdiff_t difference_type;
+        typedef const value_type *pointer;
+        typedef const value_type &reference;
+
+        const_iterator(TreeItem *base, int level);
+
+        TreeItem *operator*() { return m_item[m_depth]; }
+        void operator++();
+
+        bool operator==(const_iterator other) const;
+        bool operator!=(const_iterator other) const { return !operator==(other); }
+
+    private:
+        void goDown();
+        void goUpNextDown();
+
+        int m_level;
+        int m_depth;
+        TreeItem *m_item[8];
+        int m_pos[8];
+        int m_size[8];
+    };
+
+    const_iterator begin() const;
+    const_iterator end() const;
+
+private:
+    TreeItem *m_item;
+    int m_level;
+};
+
+template <class T>
+class QTCREATOR_UTILS_EXPORT TreeLevelItems
+{
+public:
+    typedef T value_type;
+
+    explicit TreeLevelItems(const UntypedTreeLevelItems &items) : m_items(items) {}
+
+    struct const_iterator : public UntypedTreeLevelItems::const_iterator
+    {
+        typedef std::forward_iterator_tag iterator_category;
+        typedef T value_type;
+        typedef std::ptrdiff_t difference_type;
+        typedef const value_type *pointer;
+        typedef const value_type &reference;
+
+        const_iterator(UntypedTreeLevelItems::const_iterator it) : UntypedTreeLevelItems::const_iterator(it) {}
+        T operator*() { return static_cast<T>(UntypedTreeLevelItems::const_iterator::operator*()); }
+    };
+
+    const_iterator begin() const { return const_iterator(m_items.begin()); }
+    const_iterator end() const { return const_iterator(m_items.end()); }
+
+private:
+    UntypedTreeLevelItems m_items;
+};
+
 class QTCREATOR_UTILS_EXPORT TreeModel : public QAbstractItemModel
 {
 public:
@@ -105,6 +179,27 @@ public:
     void removeItem(TreeItem *item); // item is not destroyed.
     void updateItem(TreeItem *item); // call to trigger dataChanged
 
+    UntypedTreeLevelItems untypedLevelItems(int level = 0, TreeItem *start = 0) const;
+    UntypedTreeLevelItems untypedLevelItems(TreeItem *start) const;
+
+    template <class T>
+    TreeLevelItems<T> treeLevelItems(int level, TreeItem *start = 0) const
+    {
+        return TreeLevelItems<T>(untypedLevelItems(level, start));
+    }
+
+    template <class T>
+    TreeLevelItems<T> treeLevelItems(TreeItem *start) const
+    {
+        return TreeLevelItems<T>(untypedLevelItems(start));
+    }
+
+    template <class T>
+    T findItemAtLevel(int level, std::function<bool(T)> f, TreeItem *start = 0) const
+    {
+        return Utils::findOrDefault(treeLevelItems<T>(level, start), f);
+    }
+
 private:
     QModelIndex indexFromItemHelper(const TreeItem *needle,
         TreeItem *parentItem, const QModelIndex &parentIndex) const;