KSG-Framework/Trie.h

#pragma once

#include "CharMap.h"

namespace Framework
{
    template<class T> class Trie;

    template<class T> struct TrieIteratorData
    {
        CharMap<Trie<T>, T>* map;
        TrieIteratorData<T>* parent;

        TrieIteratorData()
        {
            map = 0;
            parent = 0;
        }

        TrieIteratorData(CharMap<Trie<T>, T>* zMap, TrieIteratorData<T> parent)
        {
            map = zMap;
            this->parent = new TrieIteratorData<T>(parent);
        }

        TrieIteratorData(const TrieIteratorData& data)
        {
            map = data.map;
            if (data.parent)
                parent = new TrieIteratorData(*data.parent);
            else
                parent = 0;
        }

        ~TrieIteratorData()
        {
            delete parent;
        }

        TrieIteratorData& operator=(const TrieIteratorData& data)
        {
            map = data.map;
            TrieIteratorData* tmp = parent;
            if (data.parent)
                parent = new TrieIteratorData(*data.parent);
            else
                parent = 0;
            delete tmp;
            return *this;
        }

        void set(CharMap<Trie<T>, T>* zMap, TrieIteratorData parent)
        {
            map = zMap;
            parent = new TrieIteratorData(parent);
        }
    };

    template<class T> class TrieIterator
    {
    private:
        TrieIteratorData<T> data;

    public:
        TrieIterator(TrieIteratorData<T> data)
        {
            this->data = data;
        }

        TrieIterator next()
        {
            for (unsigned char i = 0; true; i++)
            {
                if (data.map->z(i))
                    return TrieIterator(
                        TrieIteratorData<T>(data.map->z(i)->map, data));
                if (i == 255) break;
            }
            if (!data.parent || !data.parent->map)
                return TrieIterator(
                    TrieIteratorData<T>(0, TrieIteratorData<T>()));
            TrieIteratorData<T> d = data;
            do
            {
                for (unsigned char i = 0; true; i++)
                {
                    if (d.parent->map->z(i)
                        && d.parent->map->z(i)->map == d.map)
                    {
                        if (i < 256)
                        {
                            for (unsigned char j = (unsigned char)(i + 1); true;
                                 j++)
                            {
                                if (d.parent->map->z(j))
                                    return TrieIterator(TrieIteratorData<T>(
                                        d.parent->map->z(j)->map, *d.parent));
                                if (j == 255) break;
                            }
                        }
                        d = *d.parent;
                        break;
                    }
                    if (i == 255) break;
                }
            } while (d.parent && d.parent->map);
            return TrieIterator(TrieIteratorData<T>(0, TrieIteratorData<T>()));
        }

        TrieIterator<T>& operator++() //! prefix
        {
            data = next().data;
            return *this;
        }

        TrieIterator<T> operator++(int) //! postfix
        {
            TrieIterator<T> temp(*this);
            data = next().data;
            return temp;
        }

        operator bool()
        {
            return data.map != 0;
        }

        operator T*()
        {
            return data.map->zValue();
        }

        T* operator->()
        {
            return data.map->zValue();
        }

        T* val()
        {
            return data.map->zValue();
        }
    };

    template<class T> class Trie : public virtual ReferenceCounter
    {
    private:
        CharMap<Trie<T>, T>* map;

    public:
        Trie()
            : ReferenceCounter()
        {
            map = new CharMap<Trie<T>, T>();
        }

        ~Trie()
        {
            map->release();
        }

        void set(const char* addr, int addrLen, T* value)
        {
            if (!addrLen)
                map->setValue(value);
            else
            {
                if (!map->z(*addr)) map->set(*addr, new Trie<T>());
                map->z(*addr)->set(addr + 1, addrLen - 1, value);
            }
        }

        void remove(const char* addr, int addrLen)
        {
            if (!addrLen)
                map->setValue(0);
            else
            {
                if (!map->z(*addr)) return;
                map->z(*addr)->remove(addr + 1, addrLen - 1);
            }
        }

        T* get(const char* addr, int addrLen)
        {
            if (!addrLen)
                return map->getValue();
            else
            {
                if (!map->z(*addr)) return 0;
                return map->z(*addr)->get(addr + 1, addrLen - 1);
            }
        }

        T* z(const char* addr, int addrLen)
        {
            if (!addrLen)
                return map->zValue();
            else
            {
                if (!map->z(*addr)) return 0;
                return map->z(*addr)->z(addr + 1, addrLen - 1);
            }
        }

        void leeren()
        {
            map->release();
            map = new CharMap<Trie<T>, T>();
        }

        TrieIterator<T> getIterator()
        {
            return TrieIterator<T>(
                TrieIteratorData<T>(map, TrieIteratorData<T>()));
        }

        friend TrieIterator<T>;
    };
} // namespace Framework