intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 298 insertions, 0 deletions

diff --git a/library/cpp/containers/intrusive_rb_tree/rb_tree_ut.cpp b/library/cpp/containers/intrusive_rb_tree/rb_tree_ut.cpp
new file mode 100644
index 0000000000..c34ed1fd9b
--- /dev/null
+++ b/library/cpp/containers/intrusive_rb_tree/rb_tree_ut.cpp
@@ -0,0 +1,298 @@
+#include "rb_tree.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+
+#include <util/random/fast.h>
+#include <util/random/easy.h>
+#include <util/random/shuffle.h>
+
+class TRedBlackTreeTest: public TTestBase {
+    struct TCmp {
+        template <class T>
+        static inline bool Compare(const T& l, const T& r) {
+            return l.N < r.N;
+        }
+
+        template <class T>
+        static inline bool Compare(const T& l, int r) {
+            return l.N < r;
+        }
+
+        template <class T>
+        static inline bool Compare(int l, const T& r) {
+            return l < r.N;
+        }
+    };
+
+    class TNode: public TRbTreeItem<TNode, TCmp> {
+    public:
+        inline TNode(int n) noexcept
+            : N(n)
+        {
+        }
+
+        int N;
+    };
+
+    using TTree = TRbTree<TNode, TCmp>;
+
+    UNIT_TEST_SUITE(TRedBlackTreeTest);
+    UNIT_TEST(TestEmpty)
+    UNIT_TEST(TestInsert)
+    UNIT_TEST(TestErase)
+    UNIT_TEST(TestFind)
+    UNIT_TEST(TestStress)
+    UNIT_TEST(TestGettingIndexWithDifferentValues)
+    UNIT_TEST(TestCheckChildrenAfterErase)
+    UNIT_TEST(TestGettingIndexWithDifferentValuesAfterErase)
+    UNIT_TEST(TestGettingIndexWithEqualValues)
+    UNIT_TEST(TestLessCountOnEmptyTree)
+    UNIT_TEST_SUITE_END();
+
+private:
+    inline void TestStress() {
+        TVector<TSimpleSharedPtr<TNode>> nodes;
+
+        for (int i = 0; i < 1000; ++i) {
+            nodes.push_back(new TNode(i));
+        }
+
+        TTree tree;
+        TReallyFastRng32 rnd(Random());
+
+        for (size_t i = 0; i < 1000000; ++i) {
+            tree.Insert(nodes[rnd.Uniform(nodes.size())].Get());
+        }
+
+        for (TTree::TConstIterator it = tree.Begin(); it != tree.End();) {
+            const int v1 = it->N;
+
+            if (++it == tree.End()) {
+                break;
+            }
+
+            const int v2 = it->N;
+
+            UNIT_ASSERT(v1 < v2);
+        }
+    }
+
+    inline void TestGettingIndexWithDifferentValues() {
+        TVector<TSimpleSharedPtr<TNode>> nodes;
+        size_t N = 1000;
+
+        for (size_t i = 0; i < N; ++i) {
+            nodes.push_back(new TNode(int(i)));
+        }
+
+        TTree tree;
+        Shuffle(nodes.begin(), nodes.end());
+
+        for (size_t i = 0; i < N; ++i) {
+            tree.Insert(nodes[i].Get());
+        }
+
+        for (size_t i = 0; i < N; ++i) {
+            UNIT_ASSERT_EQUAL(tree.LessCount(i), i);
+            UNIT_ASSERT_EQUAL(tree.NotGreaterCount(i), i + 1);
+            UNIT_ASSERT_EQUAL(tree.GreaterCount(i), N - i - 1);
+            UNIT_ASSERT_EQUAL(tree.NotLessCount(i), N - i);
+
+            auto nodePtr = tree.Find(i);
+            UNIT_ASSERT_EQUAL(tree.GetIndex(nodePtr), i);
+            UNIT_ASSERT_EQUAL(tree.GetIndex(nodes[i].Get()), static_cast<size_t>(nodes[i]->N));
+        }
+    }
+
+    inline void TestCheckChildrenAfterErase() {
+        TVector<TSimpleSharedPtr<TNode>> nodes;
+        size_t N = 1000;
+
+        for (size_t i = 0; i < N; ++i) {
+            nodes.push_back(new TNode(int(i)));
+        }
+
+        TTree tree;
+        Shuffle(nodes.begin(), nodes.end());
+
+        for (size_t i = 0; i < N; ++i) {
+            tree.Insert(nodes[i].Get());
+        }
+        auto checker = [](const TTree& tree) {
+            for (auto node = tree.Begin(); node != tree.End(); ++node) {
+                size_t childrens = 1;
+                if (node->Left_) {
+                    childrens += node->Left_->Children_;
+                }
+                if (node->Right_) {
+                    childrens += node->Right_->Children_;
+                }
+                UNIT_ASSERT_VALUES_EQUAL(childrens, node->Children_);
+            }
+        };
+
+        for (auto node : nodes) {
+            tree.Erase(node.Get());
+            checker(tree);
+        }
+    }
+
+    inline void TestGettingIndexWithDifferentValuesAfterErase() {
+        TVector<TSimpleSharedPtr<TNode>> nodes;
+        size_t N = 1000;
+
+        for (size_t i = 0; i < N; ++i) {
+            nodes.push_back(new TNode(int(i)));
+        }
+
+        TTree tree;
+        Shuffle(nodes.begin(), nodes.end());
+
+        for (size_t i = 0; i < N; ++i) {
+            tree.Insert(nodes[i].Get());
+        }
+        {
+            size_t index = 0;
+            for (auto node = tree.Begin(); node != tree.End(); ++node, ++index) {
+                UNIT_ASSERT_VALUES_EQUAL(tree.GetIndex(&*node), index);
+                UNIT_ASSERT_VALUES_EQUAL(tree.ByIndex(index)->N, node->N);
+                UNIT_ASSERT_VALUES_EQUAL(node->N, index);
+            }
+        }
+
+        for (size_t i = 1; i < N; i += 2) {
+            auto* node = tree.Find(i);
+            UNIT_ASSERT_VALUES_EQUAL(node->N, i);
+            tree.Erase(node);
+        }
+        {
+            size_t index = 0;
+            for (auto node = tree.Begin(); node != tree.End(); ++node, ++index) {
+                UNIT_ASSERT_VALUES_EQUAL(tree.GetIndex(&*node), index);
+                UNIT_ASSERT_VALUES_EQUAL(tree.ByIndex(index)->N, node->N);
+                UNIT_ASSERT_VALUES_EQUAL(node->N, 2 * index);
+            }
+        }
+    }
+
+    inline void TestGettingIndexWithEqualValues() {
+        TVector<TSimpleSharedPtr<TNode>> nodes;
+        size_t N = 1000;
+
+        for (size_t i = 0; i < N; ++i) {
+            nodes.push_back(new TNode(0));
+        }
+
+        TTree tree;
+
+        for (size_t i = 0; i < N; ++i) {
+            tree.Insert(nodes[i].Get());
+        }
+
+        for (size_t i = 0; i < N; ++i) {
+            UNIT_ASSERT_EQUAL(tree.LessCount(nodes[i]->N), 0);
+            UNIT_ASSERT_EQUAL(tree.NotGreaterCount(nodes[i]->N), N);
+            UNIT_ASSERT_EQUAL(tree.GreaterCount(nodes[i]->N), 0);
+            UNIT_ASSERT_EQUAL(tree.NotLessCount(nodes[i]->N), N);
+
+            UNIT_ASSERT_EQUAL(tree.LessCount(*nodes[i].Get()), 0);
+            UNIT_ASSERT_EQUAL(tree.NotGreaterCount(*nodes[i].Get()), N);
+            UNIT_ASSERT_EQUAL(tree.GreaterCount(*nodes[i].Get()), 0);
+            UNIT_ASSERT_EQUAL(tree.NotLessCount(*nodes[i].Get()), N);
+        }
+    }
+
+    inline void TestFind() {
+        TTree tree;
+
+        {
+            TNode n1(1);
+            TNode n2(2);
+            TNode n3(3);
+
+            tree.Insert(n1);
+            tree.Insert(n2);
+            tree.Insert(n3);
+
+            UNIT_ASSERT_EQUAL(tree.Find(1)->N, 1);
+            UNIT_ASSERT_EQUAL(tree.Find(2)->N, 2);
+            UNIT_ASSERT_EQUAL(tree.Find(3)->N, 3);
+
+            UNIT_ASSERT(!tree.Find(0));
+            UNIT_ASSERT(!tree.Find(4));
+            UNIT_ASSERT(!tree.Find(1234567));
+        }
+
+        UNIT_ASSERT(tree.Empty());
+    }
+
+    inline void TestEmpty() {
+        TTree tree;
+
+        UNIT_ASSERT(tree.Empty());
+        UNIT_ASSERT_EQUAL(tree.Begin(), tree.End());
+    }
+
+    inline void TestInsert() {
+        TTree tree;
+
+        {
+            TNode n1(1);
+            TNode n2(2);
+            TNode n3(3);
+
+            tree.Insert(n1);
+            tree.Insert(n2);
+            tree.Insert(n3);
+
+            TTree::TConstIterator it = tree.Begin();
+
+            UNIT_ASSERT_EQUAL((it++)->N, 1);
+            UNIT_ASSERT_EQUAL((it++)->N, 2);
+            UNIT_ASSERT_EQUAL((it++)->N, 3);
+            UNIT_ASSERT_EQUAL(it, tree.End());
+        }
+
+        UNIT_ASSERT(tree.Empty());
+    }
+
+    inline void TestErase() {
+        TTree tree;
+
+        {
+            TNode n1(1);
+            TNode n2(2);
+            TNode n3(3);
+
+            tree.Insert(n1);
+            tree.Insert(n2);
+            tree.Insert(n3);
+
+            TTree::TIterator it = tree.Begin();
+
+            tree.Erase(it++);
+
+            UNIT_ASSERT_EQUAL(it, tree.Begin());
+            UNIT_ASSERT_EQUAL(it->N, 2);
+
+            tree.Erase(it++);
+
+            UNIT_ASSERT_EQUAL(it, tree.Begin());
+            UNIT_ASSERT_EQUAL(it->N, 3);
+
+            tree.Erase(it++);
+
+            UNIT_ASSERT_EQUAL(it, tree.Begin());
+            UNIT_ASSERT_EQUAL(it, tree.End());
+        }
+
+        UNIT_ASSERT(tree.Empty());
+    }
+
+    inline void TestLessCountOnEmptyTree() {
+        TTree tree;
+        UNIT_ASSERT_VALUES_EQUAL(0, tree.LessCount(TNode(1)));
+    }
+};
+
+UNIT_TEST_SUITE_REGISTRATION(TRedBlackTreeTest);