rbtree.h

/*
 * Copyright (c) 2018-2026 Heng Yuan
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef COOK_MEM_RBTREE_H
#define COOK_MEM_RBTREE_H
 
#include <cstddef>
 
#include "exception.h"
 
// Marks a private field that is intentionally never used, so that compilers
// do not warn about it.
#if defined(__GNUC__)
#define COOKMEM_UNUSED_FIELD __attribute__ ((unused))
#else
#define COOKMEM_UNUSED_FIELD
#endif
 
namespace cookmem
{
 
class PtrRBNode
{
public:
    typedef PtrRBNode       SelfType;
    typedef std::size_t     ValueType;
 
    void
    init (ValueType value, SelfType* parentNode)
    {
        m_size = value;
        m_parent = parentNode;
        m_children[0] = nullptr;
        m_children[1] = nullptr;
        m_next = nullptr;
        m_red = true;
        m_inList = false;
    }
 
    ValueType
    getValue () const
    {
        return m_size;
    }
 
    SelfType**
    getChildren ()
    {
        return m_children;
    }
 
    SelfType* const *
    getChildren () const
    {
        return m_children;
    }
 
    SelfType*
    getLeft ()
    {
        return m_children[0];
    }
 
    const SelfType*
    getLeft () const
    {
        return m_children[0];
    }
 
    void
    setLeft (SelfType* left)
    {
        m_children[0] = left;
    }
 
    SelfType*
    getRight ()
    {
        return m_children[1];
    }
 
    const SelfType*
    getRight () const
    {
        return m_children[1];
    }
 
    void
    setRight (SelfType* right)
    {
        m_children[1] = right;
    }
 
    SelfType*
    getParent ()
    {
        return m_parent;
    }
 
    const SelfType*
    getParent () const
    {
        return m_parent;
    }
 
    void
    setParent (SelfType* parentNode)
    {
        m_parent = parentNode;
    }
 
    void
    addList (SelfType* other)
    {
        other->m_next = m_next;
        other->m_children[0] = this;
        other->m_inList = true;
        if (m_next != nullptr)
        {
            m_next->m_children[0] = other;
        }
        m_next = other;
    }
 
    bool
    hasList () const
    {
        return m_next != nullptr;
    }
 
    bool
    isInList () const
    {
        return m_inList;
    }
 
    SelfType*
    getNext ()
    {
        return m_next;
    }
 
    const SelfType*
    getNext () const
    {
        return m_next;
    }
 
    SelfType*
    removeFromList ()
    {
        SelfType*   retNode = m_next;
        SelfType*   nextNode = retNode->m_next;
        m_next = nextNode;
        if (nextNode != nullptr)
        {
            nextNode->m_children[0] = this;
        }
        retNode->m_next = nullptr;
        retNode->m_inList = false;
 
        return retNode;
    }
 
    void
    unlinkFromList ()
    {
        COOKMEM_ASSERT (m_inList);
 
        SelfType*   prevNode = m_children[0];
        SelfType*   nextNode = m_next;
        prevNode->m_next = nextNode;
        if (nextNode != nullptr)
        {
            nextNode->m_children[0] = prevNode;
        }
        m_next = nullptr;
        m_inList = false;
    }
 
    SelfType*
    promoteFirstListNode ()
    {
        SelfType*   replacementNode = m_next;
 
        // The rest of the list stays attached to replacementNode, and the
        // first node of the rest already points back to replacementNode.
        replacementNode->m_parent = m_parent;
        replacementNode->m_children[0] = m_children[0];
        replacementNode->m_children[1] = m_children[1];
        replacementNode->m_red = m_red;
        replacementNode->m_inList = false;
 
        m_next = nullptr;
        return replacementNode;
    }
 
    bool
    isRed () const
    {
        return m_red;
    }
 
    void
    setRed (bool red)
    {
        m_red = red;
    }
 
private:
    std::size_t     m_reserved COOKMEM_UNUSED_FIELD;
    std::size_t     m_size;
    SelfType*       m_children[2];
    SelfType*       m_parent;
    SelfType*       m_next;
    bool            m_red;
    bool            m_inList;
};
 
template<class Node>
class RedBlackTree
{
public:
    typedef Node                            NodeType;
    typedef typename NodeType::ValueType    ValueType;
 
public:
    RedBlackTree ()
    : m_root (nullptr)
    {
    }
 
    void
    add (void* ptr, ValueType value)
    {
        NodeType* parentNode = searchClosest (value);
        NodeType* newNode = reinterpret_cast<NodeType*>(ptr);
        newNode->init (value, parentNode);
 
        if (parentNode == nullptr)
        {
            m_root = newNode;
        }
        else if (value < parentNode->getValue ())
        {
            parentNode->setLeft (newNode);
        }
        else if (value > parentNode->getValue ())
        {
            parentNode->setRight (newNode);
        }
        else    // value == parentNode->getValue ()
        {
            // Equal values are kept in a list and do not affect tree balance.
            parentNode->addList (newNode);
            return;
        }
 
        if (newNode->getParent () == nullptr)
        {
            newNode->setRed (false);
            return;
        }
        if (newNode->getParent ()->getParent () == nullptr)
        {
            return;
        }
 
        balanceInsert (newNode);
    }
    NodeType* getRoot()
    {
        return this->m_root;
    }
 
    NodeType*
    remove (ValueType& size, ValueType limit = ~static_cast<ValueType> (0))
    {
        NodeType*   lastBigNode = nullptr;
        NodeType*   node = m_root;
 
        while (node != nullptr)
        {
            if (node->getValue () < size)
            {
                node = node->getRight ();
            }
            else if (node->getValue () > size)
            {
                lastBigNode = node;
                node = node->getLeft ();
            }
            else
            {
                break;
            }
        }
 
        if (node == nullptr)
        {
            if (lastBigNode == nullptr)
            {
                return nullptr;
            }
            else
            {
                node = lastBigNode;
            }
        }
 
        if (node->getValue () >= limit)
        {
            return nullptr;
        }
 
        if (node->hasList ())
        {
            NodeType* listNode = node->removeFromList ();
            size = node->getValue ();
            return listNode;
        }
 
        size = node->getValue ();
        removeNode (node);
 
        return node;
    }
 
    void
    remove (void* ptr)
    {
        NodeType*   targetNode = reinterpret_cast<NodeType*>(ptr);
 
        if (targetNode->isInList ())
        {
            // The node is in the DLL of a tree node.  Just unlink it.
            targetNode->unlinkFromList ();
            return;
        }
 
        // The node is in the tree.  Since nodes have parent pointers, there
        // is no need to search from the root.
        removeNode (targetNode);
    }
 
    bool
    contains (void* ptr)
    {
        if (m_root == nullptr)
        {
            return false;
        }
 
        NodeType* node2 = reinterpret_cast<NodeType*>(ptr);
        const ValueType value = node2->getValue ();
 
        NodeType* node = m_root;
        while (node != nullptr)
        {
            if (value < node->getValue ())
            {
                node = node->getLeft ();
            }
            else if (value > node->getValue ())
            {
                node = node->getRight ();
            }
            else
            {
                do
                {
                    if (node == node2)
                    {
                        return true;
                    }
                    node = node->getNext ();
                }
                while (node != nullptr);
            }
        }
        return false;
    }
 
    bool
    canGrowInPlace (void* ptr, ValueType newValue)
    {
        NodeType* node = reinterpret_cast<NodeType*>(ptr);
 
        // All nodes in a list must have the same value.
        if (node->isInList () || node->hasList ())
        {
            return false;
        }
 
        // The new value must stay below the value of the next larger node.
        if (node->getRight () != nullptr)
        {
            return newValue < getLeftMostNode (node->getRight ())->getValue ();
        }
 
        NodeType* parentNode = node->getParent ();
        while (parentNode != nullptr && node == parentNode->getRight ())
        {
            node = parentNode;
            parentNode = node->getParent ();
        }
        return parentNode == nullptr || newValue < parentNode->getValue ();
    }
 
    bool
    isEmpty () const
    {
        return m_root == nullptr;
    }
 
private:
 
    void
    updateParent (NodeType* newChildNode, NodeType* oldChildNode)
    {
        NodeType* parentNode = oldChildNode->getParent ();
        if (parentNode == nullptr)
        {
            m_root = newChildNode;
        }
        else if (oldChildNode == parentNode->getLeft ())
        {
            parentNode->setLeft (newChildNode);
        }
        else if (oldChildNode == parentNode->getRight ())
        {
            parentNode->setRight (newChildNode);
        }
        else
        {
            COOKMEM_ASSERT (false);
        }
        if (newChildNode != nullptr)
        {
            newChildNode->setParent (parentNode);
        }
    }
 
    NodeType*
    getLeftMostNode (NodeType* node)
    {
        while (node->getLeft () != nullptr)
        {
            node = node->getLeft ();
        }
        return node;
    }
 
    NodeType*
    searchClosest (ValueType key)
    {
        if (m_root == nullptr)
        {
            return nullptr;
        }
 
        NodeType* node = m_root;
 
        for (;;)
        {
            if (key < node->getValue ())
            {
                if (node->getLeft () == nullptr)
                {
                    return node;
                }
                else
                {
                    node = node->getLeft ();
                }
            }
            else if (key > node->getValue ())
            {
                if (node->getRight () == nullptr)
                {
                    return node;
                }
                else
                {
                    node = node->getRight ();
                }
            }
            else
            {
                return node;
            }
        }
    }
 
    void
    balanceInsert (NodeType* node)
    {
        while (node != m_root &&
               node->getParent ()->isRed ())
        {
            NodeType* grandParentNode = node->getParent ()->getParent ();
            if (node->getParent () == grandParentNode->getLeft ())
            {
                NodeType* uncleNode = grandParentNode->getRight ();
                if (uncleNode != nullptr &&
                    uncleNode->isRed ())
                {
                    uncleNode->setRed (false);
                    node->getParent ()->setRed (false);
                    grandParentNode->setRed (true);
                    node = node->getParent ()->getParent ();
                }
                else
                {
                    if (node == node->getParent ()->getRight ())
                    {
                        node = node->getParent ();
                        doLeftRotation(node);
                    }
                    node->getParent ()->setRed (false);
                    node->getParent ()->getParent ()->setRed (true);
                    doRightRotation (node->getParent ()->getParent ());
                }
            }
            else
            {
                NodeType* uncleNode = grandParentNode->getLeft ();
                if (uncleNode != nullptr &&
                    uncleNode->isRed ())
                {
                    uncleNode->setRed (false);
                    node->getParent ()->setRed (false);
                    grandParentNode->setRed (true);
                    node = node->getParent ()->getParent ();
                }
                else
                {
                    if (node == node->getParent ()->getLeft ())
                    {
                        node = node->getParent ();
                        doRightRotation (node);
                    }
                    node->getParent ()->setRed (false);
                    node->getParent ()->getParent ()->setRed (true);
                    doLeftRotation (node->getParent ()->getParent ());
                }
            }
 
            if (node == m_root)
            {
                break;
            }
        }
 
        m_root->setRed (false);
    }
 
    void
    removeNode (NodeType* node)
    {
        if (node->hasList ())
        {
            // Replace the node with the first node of its list.  The tree
            // structure and colors do not change.
            NodeType* replacementNode = node->promoteFirstListNode ();
            if (replacementNode->getLeft () != nullptr)
            {
                replacementNode->getLeft ()->setParent (replacementNode);
            }
            if (replacementNode->getRight () != nullptr)
            {
                replacementNode->getRight ()->setParent (replacementNode);
            }
            updateParent (replacementNode, node);
            return;
        }
 
        NodeType* replacementNode = node;
        bool replacementWasRed = replacementNode->isRed ();
        NodeType* childNode;
        NodeType* childParentNode;
 
        if (node->getLeft () == nullptr)
        {
            childNode = node->getRight ();
            childParentNode = node->getParent ();
            updateParent (childNode, node);
        }
        else if (node->getRight () == nullptr)
        {
            childNode = node->getLeft ();
            childParentNode = node->getParent ();
            updateParent (childNode, node);
        }
        else
        {
            replacementNode = getLeftMostNode (node->getRight ());
            replacementWasRed = replacementNode->isRed ();
            childNode = replacementNode->getRight ();
 
            if (replacementNode->getParent () == node)
            {
                childParentNode = replacementNode;
                if (childNode != nullptr)
                {
                    childNode->setParent (replacementNode);
                }
            }
            else
            {
                childParentNode = replacementNode->getParent ();
                updateParent (childNode, replacementNode);
                replacementNode->setRight (node->getRight ());
                replacementNode->getRight ()->setParent (replacementNode);
            }
 
            updateParent (replacementNode, node);
            replacementNode->setLeft (node->getLeft ());
            replacementNode->getLeft ()->setParent (replacementNode);
            replacementNode->setRed (node->isRed ());
        }
 
        if (!replacementWasRed)
        {
            balanceRemove (childNode, childParentNode);
        }
        if (childNode != nullptr)
        {
            childNode->setRed (false);
        }
        if (m_root != nullptr)
        {
            m_root->setRed (false);
        }
    }
 
    static bool
    isBlack (const NodeType* node)
    {
        return node == nullptr || !node->isRed ();
    }
 
    void
    balanceRemove (NodeType* node, NodeType* parentNode)
    {
        while (node != m_root && isBlack (node))
        {
            if (parentNode == nullptr)
            {
                break;
            }
 
            if (node == parentNode->getLeft ())
            {
                NodeType* siblingNode = parentNode->getRight ();
                if (siblingNode != nullptr && siblingNode->isRed ())
                {
                    siblingNode->setRed (false);
                    parentNode->setRed (true);
                    doLeftRotation (parentNode);
                    siblingNode = parentNode->getRight ();
                }
 
                if (siblingNode == nullptr ||
                    (isBlack (siblingNode->getLeft ()) &&
                     isBlack (siblingNode->getRight ())))
                {
                    if (siblingNode != nullptr)
                    {
                        siblingNode->setRed (true);
                    }
                    node = parentNode;
                    parentNode = node->getParent ();
                }
                else
                {
                    if (isBlack (siblingNode->getRight ()))
                    {
                        if (siblingNode->getLeft () != nullptr)
                        {
                            siblingNode->getLeft ()->setRed (false);
                        }
                        siblingNode->setRed (true);
                        doRightRotation (siblingNode);
                        siblingNode = parentNode->getRight ();
                    }
                    siblingNode->setRed (parentNode->isRed ());
                    parentNode->setRed (false);
                    if (siblingNode->getRight () != nullptr)
                    {
                        siblingNode->getRight ()->setRed (false);
                    }
                    doLeftRotation (parentNode);
                    node = m_root;
                    parentNode = nullptr;
                }
            }
            else
            {
                NodeType* siblingNode = parentNode->getLeft ();
                if (siblingNode != nullptr && siblingNode->isRed ())
                {
                    siblingNode->setRed (false);
                    parentNode->setRed (true);
                    doRightRotation (parentNode);
                    siblingNode = parentNode->getLeft ();
                }
 
                if (siblingNode == nullptr ||
                    (isBlack (siblingNode->getLeft ()) &&
                     isBlack (siblingNode->getRight ())))
                {
                    if (siblingNode != nullptr)
                    {
                        siblingNode->setRed (true);
                    }
                    node = parentNode;
                    parentNode = node->getParent ();
                }
                else
                {
                    if (isBlack (siblingNode->getLeft ()))
                    {
                        if (siblingNode->getRight () != nullptr)
                        {
                            siblingNode->getRight ()->setRed (false);
                        }
                        siblingNode->setRed (true);
                        doLeftRotation (siblingNode);
                        siblingNode = parentNode->getLeft ();
                    }
                    siblingNode->setRed (parentNode->isRed ());
                    parentNode->setRed (false);
                    if (siblingNode->getLeft () != nullptr)
                    {
                        siblingNode->getLeft ()->setRed (false);
                    }
                    doRightRotation (parentNode);
                    node = m_root;
                    parentNode = nullptr;
                }
            }
        }
 
        if (node != nullptr)
        {
            node->setRed (false);
        }
    }
 
    void
    doLeftRotation (NodeType* nodeA)
    {
        NodeType* nodeC = nodeA->getRight ();
 
        NodeType* nodeD = nodeC->getLeft ();
        nodeA->setRight (nodeD);
        if (nodeD != nullptr)
        {
            nodeD->setParent (nodeA);
        }
 
        updateParent (nodeC, nodeA);
 
        nodeC->setLeft (nodeA);
        nodeA->setParent (nodeC);
    }
 
    void
    doRightRotation (NodeType* nodeA)
    {
        NodeType* nodeB = nodeA->getLeft ();
 
        NodeType* nodeE = nodeB->getRight ();
        nodeA->setLeft (nodeE);
        if (nodeE != nullptr)
        {
            nodeE->setParent (nodeA);
        }
 
        updateParent (nodeB, nodeA);
 
        nodeB->setRight (nodeA);
        nodeA->setParent (nodeB);
    }
 
private:
    NodeType *m_root;
};
 
typedef RedBlackTree<PtrRBNode> PtrRBTree;
 
} // namespace cookmem
 
#endif // COOK_MEM_RBTREE_H