AK: Make IntrusiveRedBlackTree capable of holding non-raw pointers

This is completely based on e4412f1f59
and will allow us to convert some AK::HashMap users in the kernel.

AK/IntrusiveDetails.h

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+namespace AK::Detail {
+
+template<typename T, typename Container>
+struct SubstituteIntrusiveContainerType {
+    using Type = Container;
+};
+
+template<typename T>
+struct SubstituteIntrusiveContainerType<T, NonnullRefPtr<T>> {
+    using Type = RefPtr<T>;
+};
+
+template<typename Container, bool _IsRaw>
+struct SelfReferenceIfNeeded {
+    Container reference = nullptr;
+};
+template<typename Container>
+struct SelfReferenceIfNeeded<Container, true> {
+};
+
+}

AK/IntrusiveList.h

@@ -9,6 +9,7 @@
 #include <AK/Assertions.h>
 #include <AK/BitCast.h>
 #include <AK/Forward.h>
+#include <AK/IntrusiveDetails.h>
 #include <AK/Noncopyable.h>
 #include <AK/StdLibExtras.h>
 
@@ -17,20 +18,10 @@ namespace AK {
 namespace Detail {
 template<typename T, typename Container = RawPtr<T>>
 class IntrusiveListNode;
-
-template<typename T, typename Container>
-struct SubstituteIntrusiveListNodeContainerType {
-    using Type = Container;
-};
-
-template<typename T>
-struct SubstituteIntrusiveListNodeContainerType<T, NonnullRefPtr<T>> {
-    using Type = RefPtr<T>;
-};
 }
 
 template<typename T, typename Container = RawPtr<T>>
-using IntrusiveListNode = Detail::IntrusiveListNode<T, typename Detail::SubstituteIntrusiveListNodeContainerType<T, Container>::Type>;
+using IntrusiveListNode = Detail::IntrusiveListNode<T, typename Detail::SubstituteIntrusiveContainerType<T, Container>::Type>;
 
 template<typename T, typename Container>
 class IntrusiveListStorage {
@@ -157,14 +148,6 @@ private:
     IntrusiveListStorage<T, Container> m_storage;
 };
 
-template<typename Contained, bool _IsRaw>
-struct SelfReferenceIfNeeded {
-    Contained reference = nullptr;
-};
-template<typename Contained>
-struct SelfReferenceIfNeeded<Contained, true> {
-};
-
 namespace Detail {
 
 template<typename T, typename Container>

AK/IntrusiveRedBlackTree.h

@@ -6,15 +6,21 @@
 
 #pragma once
 
+#include <AK/IntrusiveDetails.h>
 #include <AK/RedBlackTree.h>
 
 namespace AK {
 
-template<Integral K>
+namespace Detail {
+template<Integral K, typename V, typename Container = RawPtr<V>>
 class IntrusiveRedBlackTreeNode;
+}
 
-template<Integral K, typename V, IntrusiveRedBlackTreeNode<K> V::*member>
+template<Integral K, typename V, typename Container = RawPtr<V>>
-class IntrusiveRedBlackTree final : public BaseRedBlackTree<K> {
+using IntrusiveRedBlackTreeNode = Detail::IntrusiveRedBlackTreeNode<K, V, typename Detail::SubstituteIntrusiveContainerType<V, Container>::Type>;
+
+template<Integral K, typename V, typename Container, IntrusiveRedBlackTreeNode<K, V, Container> V::*member>
+class IntrusiveRedBlackTree : public BaseRedBlackTree<K> {
 
 public:
     IntrusiveRedBlackTree() = default;
@@ -24,9 +30,9 @@ public:
     }
 
     using BaseTree = BaseRedBlackTree<K>;
-    using TreeNode = IntrusiveRedBlackTreeNode<K>;
+    using TreeNode = IntrusiveRedBlackTreeNode<K, V, Container>;
 
-    V* find(K key)
+    Container find(K key)
     {
         auto* node = static_cast<TreeNode*>(BaseTree::find(this->m_root, key));
         if (!node)
@@ -34,7 +40,7 @@ public:
         return node_to_value(*node);
     }
 
-    V* find_largest_not_above(K key)
+    Container find_largest_not_above(K key)
     {
         auto* node = static_cast<TreeNode*>(BaseTree::find_largest_not_above(this->m_root, key));
         if (!node)
@@ -45,7 +51,10 @@ public:
     void insert(V& value)
     {
         auto& node = value.*member;
+        VERIFY(!node.m_in_tree);
         BaseTree::insert(&node);
+        if constexpr (!TreeNode::IsRaw)
+            node.m_self.reference = &value; // Note: Self-reference ensures that the object will keep a ref to itself when the Container is a smart pointer.
         node.m_in_tree = true;
     }
 
@@ -76,7 +85,7 @@ public:
             VERIFY(m_node);
             return *node_to_value(*m_node);
         }
-        ElementType* operator->()
+        auto operator->()
         {
             VERIFY(m_node);
             return node_to_value(*m_node);
@@ -117,6 +126,8 @@ public:
         node->right_child = nullptr;
         node->left_child = nullptr;
         node->m_in_tree = false;
+        if constexpr (!TreeNode::IsRaw)
+            node->m_self.reference = nullptr;
 
         return true;
     }
@@ -139,6 +150,8 @@ private:
         clear_nodes(static_cast<TreeNode*>(node->left_child));
         node->left_child = nullptr;
         node->m_in_tree = false;
+        if constexpr (!TreeNode::IsRaw)
+            node->m_self.reference = nullptr;
     }
 
     static V* node_to_value(TreeNode& node)
@@ -147,7 +160,9 @@ private:
     }
 };
 
-template<Integral K>
+namespace Detail {
+
+template<Integral K, typename V, typename Container>
 class IntrusiveRedBlackTreeNode : public BaseRedBlackTree<K>::Node {
 public:
     IntrusiveRedBlackTreeNode(K key)
@@ -165,10 +180,28 @@ public:
         return m_in_tree;
     }
 
+    static constexpr bool IsRaw = IsPointer<Container>;
+
+#ifndef __clang__
 private:
-    template<Integral TK, typename V, IntrusiveRedBlackTreeNode<TK> V::*member>
+    template<Integral TK, typename TV, typename TContainer, IntrusiveRedBlackTreeNode<TK, TV, TContainer> TV::*member>
-    friend class IntrusiveRedBlackTree;
+    friend class ::AK::IntrusiveRedBlackTree;
+#endif
+
     bool m_in_tree { false };
+    [[no_unique_address]] SelfReferenceIfNeeded<Container, IsRaw> m_self;
+};
+
+}
+
+// Specialise IntrusiveRedBlackTree for NonnullRefPtr
+// By default, red black trees cannot contain null entries anyway, so switch to RefPtr
+// and just make the user-facing functions deref the pointers.
+template<Integral K, typename V, IntrusiveRedBlackTreeNode<K, V, NonnullRefPtr<V>> V::*member>
+class IntrusiveRedBlackTree<K, V, NonnullRefPtr<V>, member> : public IntrusiveRedBlackTree<K, V, RefPtr<V>, member> {
+public:
+    [[nodiscard]] NonnullRefPtr<V> find(K key) const { return IntrusiveRedBlackTree<K, V, RefPtr<V>, member>::find(key).release_nonnull(); }
+    [[nodiscard]] NonnullRefPtr<V> find_largest_not_above(K key) const { return IntrusiveRedBlackTree<K, V, RefPtr<V>, member>::find_largest_not_above(key).release_nonnull(); }
 };
 
 }

Tests/AK/TestIntrusiveRedBlackTree.cpp

@@ -18,20 +18,21 @@ public:
     {
     }
 
-    IntrusiveRedBlackTreeNode<int> m_tree_node;
+    IntrusiveRedBlackTreeNode<int, IntrusiveTest, RawPtr<IntrusiveTest>> m_tree_node;
     int m_some_value;
 };
+using IntrusiveRBTree = IntrusiveRedBlackTree<int, IntrusiveTest, RawPtr<IntrusiveTest>, &IntrusiveTest::m_tree_node>;
 
 TEST_CASE(construct)
 {
-    IntrusiveRedBlackTree<int, IntrusiveTest, &IntrusiveTest::m_tree_node> empty;
+    IntrusiveRBTree empty;
     EXPECT(empty.is_empty());
     EXPECT(empty.size() == 0);
 }
 
 TEST_CASE(ints)
 {
-    IntrusiveRedBlackTree<int, IntrusiveTest, &IntrusiveTest::m_tree_node> test;
+    IntrusiveRBTree test;
     IntrusiveTest first { 1, 10 };
     test.insert(first);
     IntrusiveTest second { 3, 20 };
@@ -51,7 +52,7 @@ TEST_CASE(ints)
 
 TEST_CASE(largest_smaller_than)
 {
-    IntrusiveRedBlackTree<int, IntrusiveTest, &IntrusiveTest::m_tree_node> test;
+    IntrusiveRBTree test;
     IntrusiveTest first { 1, 10 };
     test.insert(first);
     IntrusiveTest second { 11, 20 };
@@ -71,7 +72,7 @@ TEST_CASE(largest_smaller_than)
 TEST_CASE(key_ordered_iteration)
 {
     constexpr auto amount = 10000;
-    IntrusiveRedBlackTree<int, IntrusiveTest, &IntrusiveTest::m_tree_node> test;
+    IntrusiveRBTree test;
     NonnullOwnPtrVector<IntrusiveTest> m_entries;
     Array<int, amount> keys {};
 
@@ -104,7 +105,7 @@ TEST_CASE(key_ordered_iteration)
 
 TEST_CASE(clear)
 {
-    IntrusiveRedBlackTree<int, IntrusiveTest, &IntrusiveTest::m_tree_node> test;
+    IntrusiveRBTree test;
     NonnullOwnPtrVector<IntrusiveTest> m_entries;
     for (size_t i = 0; i < 1000; i++) {
         auto entry = make<IntrusiveTest>(i, i);
@@ -114,3 +115,81 @@ TEST_CASE(clear)
     test.clear();
     EXPECT_EQ(test.size(), 0u);
 }
+
+class IntrusiveRefPtrTest : public RefCounted<IntrusiveRefPtrTest> {
+public:
+    IntrusiveRefPtrTest(int key)
+        : m_tree_node(key)
+    {
+    }
+
+    IntrusiveRedBlackTreeNode<int, IntrusiveRefPtrTest, RefPtr<IntrusiveRefPtrTest>> m_tree_node;
+};
+using IntrusiveRefPtrRBTree = IntrusiveRedBlackTree<int, IntrusiveRefPtrTest, RefPtr<IntrusiveRefPtrTest>, &IntrusiveRefPtrTest::m_tree_node>;
+
+TEST_CASE(intrusive_ref_ptr_no_ref_leaks)
+{
+    auto item = adopt_ref(*new IntrusiveRefPtrTest(0));
+    EXPECT_EQ(1u, item->ref_count());
+    IntrusiveRefPtrRBTree ref_tree;
+
+    ref_tree.insert(*item);
+    EXPECT_EQ(2u, item->ref_count());
+
+    ref_tree.remove(0);
+    EXPECT_EQ(1u, item->ref_count());
+}
+
+TEST_CASE(intrusive_ref_ptr_clear)
+{
+    auto item = adopt_ref(*new IntrusiveRefPtrTest(0));
+    EXPECT_EQ(1u, item->ref_count());
+    IntrusiveRefPtrRBTree ref_tree;
+
+    ref_tree.insert(*item);
+    EXPECT_EQ(2u, item->ref_count());
+
+    ref_tree.clear();
+    EXPECT_EQ(1u, item->ref_count());
+}
+
+TEST_CASE(intrusive_ref_ptr_destructor)
+{
+    auto item = adopt_ref(*new IntrusiveRefPtrTest(0));
+    EXPECT_EQ(1u, item->ref_count());
+
+    {
+        IntrusiveRefPtrRBTree ref_tree;
+        ref_tree.insert(*item);
+        EXPECT_EQ(2u, item->ref_count());
+    }
+
+    EXPECT_EQ(1u, item->ref_count());
+}
+
+class IntrusiveNonnullRefPtrTest : public RefCounted<IntrusiveNonnullRefPtrTest> {
+public:
+    IntrusiveNonnullRefPtrTest(int key)
+        : m_tree_node(key)
+    {
+    }
+
+    IntrusiveRedBlackTreeNode<int, IntrusiveNonnullRefPtrTest, NonnullRefPtr<IntrusiveNonnullRefPtrTest>> m_tree_node;
+};
+using IntrusiveNonnullRefPtrRBTree = IntrusiveRedBlackTree<int, IntrusiveNonnullRefPtrTest, NonnullRefPtr<IntrusiveNonnullRefPtrTest>, &IntrusiveNonnullRefPtrTest::m_tree_node>;
+
+TEST_CASE(intrusive_nonnull_ref_ptr_intrusive)
+{
+    auto item = adopt_ref(*new IntrusiveNonnullRefPtrTest(0));
+    EXPECT_EQ(1u, item->ref_count());
+    IntrusiveNonnullRefPtrRBTree nonnull_ref_tree;
+
+    nonnull_ref_tree.insert(*item);
+    EXPECT_EQ(2u, item->ref_count());
+    EXPECT(!nonnull_ref_tree.is_empty());
+
+    nonnull_ref_tree.remove(0);
+    EXPECT_EQ(1u, item->ref_count());
+
+    EXPECT(nonnull_ref_tree.is_empty());
+}