wb-trie.h

Go to the documentation of this file.

// 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.
//
// Copyright 2014-2015 Tsinghua University
// Author: wb.th08@gmail.com (Bin Wang), ozj@tsinghua.edu.cn (Zhijian Ou)
//
// All h, cpp, cc, and script files (e.g. bat, sh, pl, py) should include the above
// license declaration. Different coding language may use different comment styles.




#ifndef _WB_TRIE_H_
#define _WB_TRIE_H_
#include "wb-lhash.h"

namespace wb
{

#define _wb_TRIE                wb::Trie<KeyT, DataT>
#define _wb_LHASH_FOR_TRIE      wb::LHash<KeyT, wb::Trie<KeyT, DataT>*>
#define _wb_LHASHITER_FOR_TRIE  wb::LHashIter<KeyT, wb::Trie<KeyT, DataT>*>
#define _wb_UNIT_FOR_TRIE       wb::LHash::Unit<KeyT, wb::Trie<KeyT, DataT>*>

    template <class KeyT, class DataT> class Trie;
    template <class KeyT, class DataT> class TrieIter;
    template <class KeyT, class DataT> class TrieIter2;

    template <class KeyT, class DataT>
    class Trie
    {
        friend class TrieIter<KeyT, DataT>;
        friend class TrieIter2<KeyT, DataT>;
    private:
        DataT m_value; 
        _wb_LHASH_FOR_TRIE *m_phash; 
    public:
        Trie() :m_phash(NULL){ Map_noKey(m_value); }
        ~Trie() { Release();  }
        void Release()
        {
            if (m_phash)
            {
                KeyT key;
                _wb_TRIE **ppsub;
                _wb_LHASHITER_FOR_TRIE iter(m_phash);
                while (ppsub = iter.Next(key)) {
                    (*ppsub)->Release();
                }
            }

            Map_noKey(m_value);
            if (m_phash) {
                delete m_phash;
                m_phash = NULL;
            }
        }
        size_t TotalMemCost()
        {
            size_t nSize = sizeof(*this);
            if (m_phash) {
                KeyT key;
                nSize += m_phash->TotalMemCost();
                _wb_TRIE **ppsub;
                _wb_LHASHITER_FOR_TRIE iter(m_phash);
                while (ppsub = iter.Next(key)) {
                    nSize += (*ppsub)->TotalMemCost();
                }
            }
            return nSize;
        }
        /* For Trie, clean is equal to release */
        void Clean() { Release();  }
        void Fill(DataT d)
        {
            KeyT key;
            _wb_TRIE **ppsub;
            _wb_LHASHITER_FOR_TRIE iter(m_phash);
            while (ppsub = iter.Next(key)) {
                (*ppsub)->Fill(d);
            }
            m_value = d;
        }
        DataT* GetData() { return &m_value; }
        _wb_LHASH_FOR_TRIE *GetHash() const { return m_phash; }
        bool IsDataLegal() { return !Map_noKeyP(m_value); }
        DataT* Find(const KeyT *p_pIndex, int nIndexLen, bool &bFound)
        {
            _wb_TRIE *psub = FindTrie(p_pIndex, nIndexLen, bFound);
            if (psub && psub->IsDataLegal()) {
                return psub->GetData();
            }
            bFound = false;
            return NULL;
        }
        DataT* Insert(const KeyT *p_pIndex, int nIndexLen, bool &bFound)
        {
            _wb_TRIE *psub = InsertTrie(p_pIndex, nIndexLen, bFound);
            return psub->GetData();
        }
        _wb_TRIE *FindTrie(const KeyT *p_pIndex, int nIndexLen, bool &bFound)
        {
            if (nIndexLen == 0) {
                bFound = true;
                return this;
            }

            if (!m_phash) {
                bFound = false;
                return NULL;
            }

            _wb_TRIE **ppsub = m_phash->Find(p_pIndex[0], bFound);
            if (!bFound || !ppsub) {
                bFound = false;
                return NULL;
            }

            return (*ppsub)->FindTrie(p_pIndex + 1, nIndexLen - 1, bFound);
        }
        _wb_TRIE *InsertTrie(const KeyT *p_pIndex, int nIndexLen, bool &bFound)
        {
            if (nIndexLen == 0) {
                return this;
            }

            if (!m_phash) {
                m_phash = new _wb_LHASH_FOR_TRIE;
            }
            _wb_TRIE **ppsub = m_phash->Insert(p_pIndex[0], bFound);
            if (!bFound) {
                *ppsub = new _wb_TRIE;
            }
            return (*ppsub)->InsertTrie(p_pIndex + 1, nIndexLen - 1, bFound);
        }
        DataT* Find(const KeyT *p_pIndex, int nIndexLen)
        {
            bool bFound;
            return Find(p_pIndex, nIndexLen, bFound);
        }
        DataT* Insert(const KeyT *p_pIndex, int nIndexLen)
        {
            bool bFound;
            return Insert(p_pIndex, nIndexLen, bFound);
        }
        _wb_TRIE *FindTrie(const KeyT *p_pIndex, int nIndexLen)
        {
            bool bFound;
            return FindTrie(p_pIndex, nIndexLen, bFound);
        }
        _wb_TRIE *InsertTrie(const KeyT *p_pIndex, int nIndexLen)
        {
            bool bFound;
            return InsertTrie(p_pIndex, nIndexLen, bFound);
        }
    };


    template <class KeyT, class DataT>
    class TrieIter
    {
        friend class TrieIter2<KeyT, DataT>;
    public:
        TrieIter(_wb_TRIE *ptrie, bool(*sort)(KeyT, KeyT) = 0)
            : m_Iter(ptrie->m_phash, sort) {};
        void Init() { m_Iter.Init(); };
        /* the returned trie must contain a legal value*/
        _wb_TRIE *Next(KeyT &key) {
            _wb_TRIE *p;
            while (p = Next_1(key)) {
                if (p->IsDataLegal())
                    break;
            }
            return p;
        }

    private:
        /* the returned trie may not contain a legal value */
        _wb_TRIE *Next_1(KeyT &key) {
            _wb_TRIE **pp = m_Iter.Next(key);
            return (pp) ? *pp : NULL;
        };
    private:
        _wb_LHASHITER_FOR_TRIE m_Iter; 
    };


    template <class KeyT, class DataT>
    class TrieIter2
    {
    public:
        TrieIter2(_wb_TRIE *ptrie, KeyT *pIndex, int level, bool(*sort)(KeyT, KeyT) = 0)
            : m_ptrie(ptrie), m_pIndex(pIndex), m_nLevel(level),
            m_iter(ptrie, sort), m_sort(sort){
            m_pSubIter2 = NULL;
        };
        void Init()
        {
            m_iter.Init();
            if (m_pSubIter2)
                delete m_pSubIter2;
            m_pSubIter2 = NULL;
        }
        /* The returned trie may not contain a legal value*/
        _wb_TRIE *Next()
        {
            if (m_nLevel == 0) {
                return m_ptrie;
            }
            else if (m_nLevel == 1) {
                return m_iter.Next(*m_pIndex);
            }

            while (1)
            {
                if (m_pSubIter2 == NULL) {
                    _wb_TRIE *pSub = m_iter.Next_1(*m_pIndex);
                    if (pSub == NULL)
                        return NULL;

                    m_pSubIter2 = new TrieIter2<KeyT, DataT>(pSub, m_pIndex + 1, m_nLevel - 1, m_sort);
                }

                _wb_TRIE *pNext = m_pSubIter2->Next();
                if (pNext == NULL) {
                    delete m_pSubIter2;
                    m_pSubIter2 = NULL;
                }
                else {
                    return pNext;
                }
            }

        }

    private:
        TrieIter<KeyT, DataT> m_iter;  
        TrieIter2<KeyT, DataT> *m_pSubIter2; 
        _wb_TRIE *m_ptrie; 
        KeyT *m_pIndex; 
        int m_nLevel; 
        bool(*m_sort)(KeyT, KeyT); 
    };
}

#endif