RBTree.h

Go to the documentation of this file.

/*
 * /phoenix/include/common/RBTree.h
 *
 * This file is a part of Phoenix operating system.
 * Copyright (c) 2011-2012, Artyom Lebedev <artyom.lebedev@gmail.com>
 * All rights reserved.
 * See COPYING file for copyright details.
 */

#ifndef RBTREE_H_
#define RBTREE_H_

class RBTreeBase {
public:
    bool Validate();

protected:
    class EntryBase {
    protected:
        friend class RBTreeBase;

        u8 isRed:1, 
           isWired:1; 
        EntryBase *parent, *child[2];

        EntryBase() { isWired = false; }
    };

    RBTreeBase();

    virtual int Compare(EntryBase *e1, EntryBase *e2) = 0;

    virtual int Compare(EntryBase *e, void *key) = 0;

    EntryBase *InsertNode(EntryBase *node);

    EntryBase *GetNextNode(EntryBase *node = 0);

    EntryBase *Lookup(void *key);

    void Delete(EntryBase *entry);

    EntryBase *Lowest();

    EntryBase *Highest();

private:
    EntryBase *_root;
    size_t _nodesCount;
    unsigned _generation;

    void _RebalanceInsertion(EntryBase *node);

    void _RebalanceDeletion(EntryBase *node);

    inline void _Rotate(EntryBase *node, int dir) {
        ASSERT(dir == 0 || dir == 1);
        EntryBase *x = node->child[dir];
        ASSERT(x);
        x->parent = node->parent;
        if (node->parent) {
            if (node->parent->child[0] == node) {
                node->parent->child[0] = x;
            } else {
                ASSERT(node->parent->child[1] == node);
                node->parent->child[1] = x;
            }
        } else {
            ASSERT(_root == node);
            _root = x;
        }

        node->child[dir] = x->child[!dir];
        if (node->child[dir]) {
            node->child[dir]->parent = node;
        }

        x->child[!dir] = node;
        node->parent = x;
    }

    inline void _CheckRebalanceInsertion(EntryBase *node) {
        ASSERT(node->isRed);
        if (node->parent && node->parent->isRed && node->parent->parent) {
            _RebalanceInsertion(node);
        }
    }
};

template <class T, int (T::*Comparator)(T &obj),
          typename key_t, int (T::*KeyComparator)(key_t &key)>
class RBTree : public RBTreeBase {
public:
    class Entry : public EntryBase {
    protected:
        friend class RBTree;

        T *obj;
    };

    inline RBTree() : RBTreeBase() { }

    virtual int Compare(EntryBase *e1, EntryBase *e2)
    {
        return (static_cast<Entry *>(e1)->obj->*Comparator)(*static_cast<Entry *>(e2)->obj);
    }

    virtual int Compare(EntryBase *e, void *key)
    {
        if (!KeyComparator) {
            FAULT("Key comparator not provided");
        }
        return (static_cast<Entry *>(e)->obj->*KeyComparator)(*static_cast<key_t *>(key));
    }

    inline T *InsertProbe(T *obj, Entry *e)
    {
        e->obj = obj;
        return static_cast<Entry *>(InsertNode(e))->obj;
    }

    inline T *Insert(T *obj, Entry *e)
    {
        e->obj = obj;
        if (static_cast<Entry *>(InsertNode(e)) != e) {
            return 0;
        }
        return obj;
    }

    inline T *Lookup(key_t &key)
    {
        EntryBase *e = RBTreeBase::Lookup(&key);
        if (!e) {
            return 0;
        }
        return static_cast<Entry *>(e)->obj;
    }

    inline void Delete(Entry *e)
    {
        RBTreeBase::Delete(static_cast<EntryBase *>(e));
    }

    inline T *Delete(key_t &key)
    {
        EntryBase *e = RBTreeBase::Lookup(&key);
        if (!e) {
            return 0;
        }
        RBTreeBase::Delete(e);
        return static_cast<Entry *>(e)->obj;
    }

    inline T *Lowest() {
        EntryBase *e = RBTreeBase::Lowest();
        if (!e) {
            return 0;
        }
        return static_cast<Entry *>(e)->obj;
    }

    inline T *Highest() {
        EntryBase *e = RBTreeBase::Highest();
        if (!e) {
            return 0;
        }
        return static_cast<Entry *>(e)->obj;
    }

    /* Iteration interface. */

    class Iterator {
    public:
        inline Iterator(RBTree &tree, bool isStart) :
            _tree(tree)
        {
            if (isStart) {
                e = static_cast<Entry *>(_tree.GetNextNode());
            } else {
                e = 0;
            }
        }

        inline bool operator !=(Iterator &iter) { return e != iter.e; }

        inline void operator ++() {
            if (e) {
                e = static_cast<Entry *>(_tree.GetNextNode(e));
            }
        }

        inline T &operator *() { return *e->obj; }

    private:
        RBTree &_tree;
        Entry *e;
    };

    inline Iterator begin() { return Iterator(*this, true); }
    inline Iterator end() { return Iterator(*this, false); }
};

template <class T, int (T::*Comparator)(T &obj),
          typename key_t, int (T::*KeyComparator)(key_t &key)>
static inline typename RBTree<T, Comparator, key_t, KeyComparator>::Iterator
begin(RBTree<T, Comparator, key_t, KeyComparator> &tree)
{
    return tree.begin();
}

template <class T, int (T::*Comparator)(T &obj),
          typename key_t, int (T::*KeyComparator)(key_t &key)>
static inline typename RBTree<T, Comparator, key_t, KeyComparator>::Iterator
end(RBTree<T, Comparator, key_t, KeyComparator> &tree)
{
    return tree.end();
}

#endif /* RBTREE_H_ */