avltree.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_PTR_AVLTREE_H
#define COOK_PTR_AVLTREE_H
 
#include <cstddef>
#include <cstdint>
#include <cstdio>
 
#include "exception.h"
 
namespace cookmem
{
 
class PtrAVLTree
{
private:
    struct Node
    {
        std::size_t     reserved;
        std::size_t     size;
        Node*           left;
        Node*           right;
        Node*           next;
        std::int16_t    height;
    };
 
    struct RemoveStruct
    {
        std::size_t     size;
        Node*           node;
 
        RemoveStruct (std::size_t s)
        : size(s), node(nullptr)
        {
        }
    };
 
public:
    PtrAVLTree()
    : m_root(nullptr)
    {
    }
 
    ~PtrAVLTree() {}
 
    void
    add (void* ptr, std::size_t size)
    {
        Node*   node = reinterpret_cast<Node*>(ptr);
        node->size = size;
 
        if (m_root == nullptr)
        {
            node->left = nullptr;
            node->right = nullptr;
            node->next = nullptr;
            node->height = 1;
            m_root = node;
 
            return;
        }
 
        Node* root;
        Node* stack[sizeof(Node*) * 8 + 1];
        std::int16_t depth;
 
        root = stack[0] = m_root;
 
        for (depth = 0; ; ++depth)
        {
            if (root == nullptr)
            {
                setParent (node, stack[depth - 1]);
                node->left = nullptr;
                node->right = nullptr;
                node->next = nullptr;
                node->height = 1;
 
                root = balance (stack, depth);
                if (root != nullptr)
                {
                    m_root = root;
                }
                return;
            }
 
            if (size < root->size)
            {
                root = stack[depth + 1] = root->left;
                continue;
            }
            else if (size > root->size)
            {
                root = stack[depth + 1] = root->right;
                continue;
            }
            else
            {
                // The current root matches the node's size.  In this case we
                // just add node to the next pointer of the root.
                //
                Node* next = root->next;
                root->next = node;
                node->next = next;
                node->left = root;
                node->height = -1;      // indicating that this node is in a DLL
                if (next)
                {
                    next->left = node;
                }
 
                // We can simply return since the tree structure is not
                // changed.
                return;
            }
        }
        throw Exception (MEM_ERROR_GENERAL, "Invalid PtrAVLTree.", __FILE__, __LINE__);
    }
 
    void*
    remove (std::size_t& size, std::size_t limit = ~static_cast<std::size_t> (0))
    {
        Node* stack[sizeof(Node*) * 8 + 1];
        Node* root = m_root;
        Node* candidate = nullptr;
        std::int16_t depth = 0;
        std::int16_t candidateDepth = 0;
 
        while (root != nullptr)
        {
            stack[depth] = root;
 
            if (size <= root->size)
            {
                candidate = root;
                candidateDepth = depth;
                root = root->left;
            }
            else
            {
                root = root->right;
            }
            ++depth;
        }
 
        if (candidate == nullptr || candidate->size >= limit)
        {
            return nullptr;
        }
 
        // The candidate is the smallest node whose size is at least the key.
        if (candidate->next)
        {
            Node* next = candidate->next;
            candidate->next = next->next;
            if (candidate->next != nullptr)
            {
                candidate->next->left = candidate;
            }
 
            size = next->size;
            return next;
        }
 
        size = candidate->size;
        removeNode (candidate, stack, candidateDepth);
        return candidate;
    }
 
    void
    remove (void* ptr)
    {
        Node* node = reinterpret_cast<Node*>(ptr);
 
        if (node->height == -1)
        {
            // this is a simple case that node is in a DLL
            Node* prev = node->left;
            Node* next = node->next;
            prev->next = next;
            if (next)
            {
                next->left = prev;
            }
            return;
        }
 
        // The node is not in a DLL.  So this is a more complicate case that
        // need to do the removal from the root (since there is a potential
        // need to balance the tree from root.
 
        Node* root;
        Node* stack[sizeof(Node*) * 8 + 1];
        std::int16_t depth;
 
        std::size_t size = node->size;
 
        COOKMEM_ASSERT(m_root != nullptr);
 
        root = stack[0] = m_root;
 
        for (depth = 0; ; ++depth)
        {
            if (root == nullptr)
            {
                throw Exception (MEM_ERROR_GENERAL, "pointer not found.", __FILE__, __LINE__);
            }
 
            if (size < root->size)
            {
                root = stack[depth + 1] = root->left;
                continue;
            }
            else if (size > root->size)
            {
                root = stack[depth + 1] = root->right;
                continue;
            }
 
            break;
        }
 
        COOKMEM_ASSERT(root == node);
 
        if (root->next)
        {
            // Just use the next node to replace the current node.
 
            Node* next = root->next;
            next->left = root->left;
            next->right = root->right;
            next->height = root->height;
 
            if (depth == 0)
            {
                m_root = next;
            }
            else
            {
                setParent (next, stack[depth - 1]);
            }
            return;
        }
 
        removeNode (root, stack, depth);
    }
 
    bool
    contains (void* ptr)
    {
        if (m_root == nullptr)
        {
            return false;
        }
 
        Node* node = reinterpret_cast<Node*>(ptr);
        const std::size_t size = node->size;
 
        Node* root = m_root;
        while (root != nullptr)
        {
            if (size < root->size)
            {
                root = root->left;
            }
            else if (size > root->size)
            {
                root = root->right;
            }
            else
            {
                do
                {
                    if (root == node)
                    {
                        return true;
                    }
                    root = root->next;
                }
                while (root != nullptr);
            }
        }
        return false;
    }
 
    bool
    canGrowInPlace (void* ptr, std::size_t newSize)
    {
        Node* node = reinterpret_cast<Node*>(ptr);
 
        // All nodes in a DLL must have the same size.
        if (node->height == -1 || node->next != nullptr)
        {
            return false;
        }
 
        // Find the next larger node.  There are no parent pointers, so
        // search from the root and remember the last node where the search
        // went left.
        std::size_t size = node->size;
        Node* nextLarger = nullptr;
        Node* root = m_root;
        while (root != node)
        {
            if (root == nullptr)
            {
                return false;
            }
            if (size < root->size)
            {
                nextLarger = root;
                root = root->left;
            }
            else
            {
                root = root->right;
            }
        }
 
        if (node->right != nullptr)
        {
            nextLarger = node->right;
            while (nextLarger->left != nullptr)
            {
                nextLarger = nextLarger->left;
            }
        }
 
        return nextLarger == nullptr || newSize < nextLarger->size;
    }
 
    bool
    isEmpty ()
    {
        return m_root == nullptr;
    }
 
    void
    printGraph ()
    {
        printf ("graph G {\n");
        printNode (m_root);
        printf ("}\n");
    }
private:
    inline static std::int16_t
    getMax (std::int16_t a, std::int16_t b)
    {
        return (a < b) ? b : a;
    }
 
    inline static std::int16_t
    getHeight (const Node* node)
    {
        return node ? node->height : 0;
    }
 
    inline static void
    updateHeight (Node* node)
    {
        node->height = getMax (getHeight (node->left), getHeight (node->right)) + 1;
    }
 
    inline static void
    setParent (Node* node, Node* parent)
    {
        if (node->size < parent->size)
        {
            parent->left = node;
        }
        else
        {
            parent->right = node;
        }
    }
 
    inline static void
    setLeftChild (Node* root, Node* left)
    {
        root->left = left;
    }
 
    inline static void
    setRightChild (Node* root, Node* right)
    {
        root->right = right;
    }
 
    inline static Node*
    balance (Node** stack, std::int16_t depth)
    {
        for (;;)
        {
            --depth;
            Node* root = stack[depth];
            Node* left = root->left;
            Node* right = root->right;
 
            std::int16_t oldHeight = root->height;
            int16_t diff = getHeight (left) - getHeight (right);
            if (diff > 1)
            {
                if (getHeight (left->left) >= getHeight (left->right))
                {
                    root = rotateWithLeftChild (root);
                }
                else
                {
                    root = doubleWithLeftChild (root);
                }
                // update the parent
                if (depth == 0)
                {
                    return root;
                }
                setParent (root, stack[depth - 1]);
                continue;
            }
            else if (diff < -1)
            {
                if (getHeight (right->right) >= getHeight (right->left))
                {
                    root = rotateWithRightChild (root);
                }
                else
                {
                    root = doubleWithRightChild (root);
                }
                // update the parent
                if (depth == 0)
                {
                    return root;
                }
                setParent (root, stack[depth - 1]);
                continue;
            }
            else
            {
                updateHeight (root);
            }
 
            if (depth == 0)
            {
                return root;
            }
 
            if (oldHeight == root->height)
                break;
        }
        return nullptr;
    }
 
    inline static Node*
    rotateWithLeftChild (Node* root)
    {
        Node* left = root->left;
        setLeftChild (root, left->right);
        setRightChild (left, root);
        updateHeight (root);
        updateHeight (left);
        return left;
    }
 
    inline static Node*
    rotateWithRightChild (Node* root)
    {
        Node* right = root->right;
        setRightChild (root, right->left);
        setLeftChild (right, root);
        updateHeight (root);
        updateHeight (right);
        return right;
    }
 
    inline static Node*
    doubleWithLeftChild (Node* root)
    {
        setLeftChild (root, rotateWithRightChild (root->left));
        return rotateWithLeftChild (root);
    }
 
    inline static Node*
    doubleWithRightChild (Node* root)
    {
        setRightChild (root, rotateWithLeftChild (root->right));
        return rotateWithRightChild (root);
    }
 
    static void
    printNode (Node* root)
    {
        if (!root)
            return;
        printf ("%016zx -- {", (std::size_t)root);
        if (root->left)
        {
            printf (" %016zx", (std::size_t)root->left);
        }
        if (root->right)
        {
            printf (" %016zx", (std::size_t)root->right);
        }
        printf (" }\n");
        printf ("%016zx [ label = %zd ]\n",
                (std::size_t)root,
                root->size);
        printNode (root->left);
        printNode (root->right);
    }
 
    void
    removeNode (Node* root, Node** stack, int depth)
    {
        Node* left = root->left;
        Node* right = root->right;
 
        Node* newRoot;
        if (left == nullptr)
        {
            newRoot = right;
        }
        else if (right == nullptr)
        {
            newRoot = left;
        }
        else
        {
            // this is a more complicated case.
            // the solution is the move the smallest of the right branch
            // here.
            std::int16_t newDepth;
            Node* newStack[sizeof (Node*) * 8 + 1];
            newStack[0] = right;
            for (newDepth = 0; newStack[newDepth]->left != nullptr; ++newDepth)
            {
                newStack[newDepth + 1] = newStack[newDepth]->left;
            }
            Node* min = newStack[newDepth];
            if (min == right)
            {
                newRoot = right;
                setLeftChild (newRoot, left);
                updateHeight (newRoot);
            }
            else
            {
                // first detach min from the parent
                Node* parent = newStack[newDepth - 1];
                parent->left = min->right;
 
                // rebalance the right
                Node* newRight = balance (newStack, newDepth);
                if (newRight != nullptr)
                {
                    right = newRight;
                }
 
                setLeftChild (min, left);
                setRightChild (min, right);
                updateHeight (min);
 
                newRoot = min;
            }
        }
 
        if (depth == 0)
        {
            m_root = newRoot;
            if (newRoot != nullptr)
            {
                stack[0] = newRoot;
                newRoot = balance (stack, 1);
                if (newRoot != nullptr)
                {
                    m_root = newRoot;
                }
            }
            return;
        }
        if (newRoot == nullptr)
        {
            // so root is the leaf.
            if (root->size < stack[depth - 1]->size)
            {
                stack[depth - 1]->left = nullptr;
            }
            else
            {
                stack[depth - 1]->right = nullptr;
            }
            newRoot = balance (stack, depth);
            if (newRoot != nullptr)
            {
                m_root = newRoot;
            }
            return;
        }
 
        setParent (newRoot, stack[depth - 1]);
 
        stack[depth] = newRoot;
        newRoot = balance (stack, depth + 1);
 
        if (newRoot != nullptr)
        {
            m_root = newRoot;
        }
    }
 
private:
    // node entries
    Node*       m_root;
};
 
}   // namespace cookmem
 
#endif  // COOK_PTR_AVLTREE_H