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=( 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()
        {
            Trie<T> *n = 0;
            for( int i = 0; i < 256; i++ )
            {
                if( data.map->z( i ) )
                    return TrieIterator( TrieIteratorData<T>( data.map->z( i )->map, data ) );
            }
            if( !data.parent || !data.parent->map )
                return TrieIterator( TrieIteratorData<T>( 0, TrieIteratorData<T>() ) );
            TrieIteratorData<T> d = data;
            do
            {
                for( int i = 0; i < 256; i++ )
                {
                    if( d.parent->map->z( i ) && d.parent->map->z( i )->map == d.map )
                    {
                        for( int j = i + 1; j < 256; j++ )
                        {
                            if( d.parent->map->z( j ) )
                                return TrieIterator( TrieIteratorData<T>( d.parent->map->z( j )->map, *d.parent ) );
                        }
                        d = *d.parent;
                        break;
                    }
                }
            } while( d.parent && d.parent->map );
            return TrieIterator( TrieIteratorData<T>( 0, TrieIteratorData<T>() ) );
        }

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

        TrieIterator< T > &operator++( int ) // postfix
        {
            *this = next();
            return *this;
        }

        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
    {
    private:
        CharMap< Trie< T >, T > *map;
        int ref;

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

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

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

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

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

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

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

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

        Trie *getThis()
        {
            ref++;
            return this;
        }

        Trie *release()
        {
            if( !--ref )
                delete this;
            return 0;
        }
        friend TrieIterator<T>;
    };
}