lean2/src/util/trie.h

/*
Copyright (c) 2014 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.

Author: Leonardo de Moura
*/
#pragma once
#include <algorithm>
#include <string>
#include "util/rb_map.h"
#include "util/optional.h"

namespace lean {
template<typename Key, typename Val, typename KeyCMP>
class trie : public KeyCMP {
    struct cell;
    struct node {
        cell * m_ptr;
        node():m_ptr(new cell()) { m_ptr->inc_ref(); }
        node(cell * ptr):m_ptr(ptr) { if (m_ptr) ptr->inc_ref(); }
        node(node const & s):m_ptr(s.m_ptr) { if (m_ptr) m_ptr->inc_ref(); }
        node(node && s):m_ptr(s.m_ptr) { s.m_ptr = nullptr; }
        ~node() { if (m_ptr) m_ptr->dec_ref(); }
        node & operator=(node const & n) { LEAN_COPY_REF(n); }
        node & operator=(node&& n) { LEAN_MOVE_REF(n); }
        cell * operator->() const { lean_assert(m_ptr); return m_ptr; }
        bool is_shared() const { return m_ptr && m_ptr->get_rc() > 1; }
        friend void swap(node & n1, node & n2) { std::swap(n1.m_ptr, n2.m_ptr); }
        node steal() { node r; swap(r, *this); return r; }
    };

    struct cell {
        rb_map<Key, node, KeyCMP> m_children;
        optional<Val>             m_value;
        MK_LEAN_RC();
        void dealloc() { delete this; }
        cell():m_rc(0) {}
        cell(Val const & v):m_value(v), m_rc(0) {}
        cell(cell const & s):m_children(s.m_children), m_value(s.m_value), m_rc(0) {}
    };

    static node ensure_unshared(node && n) {
        if (n.is_shared())
            return node(new cell(*n.m_ptr));
        else
            return n;
    }

    template<typename It>
    static node insert(node && n, It const & begin, It const & end, Val const & v) {
        node h = ensure_unshared(n.steal());
        if (begin == end) {
            h->m_value = v;
            return h;
        } else {
            Key k = *begin;
            node const * c = h->m_children.find(k);
            It it(begin); it++;
            if (c == nullptr) {
                h->m_children.insert(k, insert(node(), it, end, v));
            } else {
                node n(*c);
                h->m_children.erase(k);
                h->m_children.insert(k, insert(n.steal(), it, end, v));
            }
            return h;
        }
    }

    node m_root;
    trie(node const & n):m_root(n) {}
public:
    trie() {}
    trie(trie const & s):m_root(s.m_root) {}
    trie(trie && s):m_root(s.m_root) {}

    trie & operator=(trie const & s) { m_root = s.m_root; return *this; }
    trie & operator=(trie && s) { m_root = s.m_root; return *this; }

    template<typename It>
    optional<Val> find(It const & begin, It const & end) const {
        node const * n = &m_root;
        for (It it = begin; it != end; ++it) {
            n = (*n)->m_children.find(*it);
            if (!n)
                return optional<Val>();
        }
        return (*n)->m_value;
    }

    template<typename It>
    void insert(It const & begin, It const & end, Val const & v) {
        m_root = insert(m_root.steal(), begin, end, v);
    }

    optional<trie> find(Key const & k) const {
        node const * c = m_root->m_children.find(k);
        if (c)
            return optional<trie>(trie(*c));
        else
            return optional<trie>();
    }
};

struct char_cmp { int operator()(char c1, char c2) const { return c1 < c2 ? -1 : (c1 == c2 ? 0 : 1); } };

template<typename Val>
using ctrie = trie<char, Val, char_cmp>;

template<typename Val>
inline ctrie<Val> insert(ctrie<Val> const & t, std::string const & str, Val const & v) {
    ctrie<Val> r(t);
    r.insert(str.begin(), str.end(), v);
    return r;
}

template<typename Val>
inline ctrie<Val> insert(ctrie<Val> const & t, char const * str, Val const & v) {
    ctrie<Val> r(t);
    r.insert(str, str+strlen(str), v);
    return r;
}

template<typename Val>
optional<Val> find(ctrie<Val> const & t, std::string const & str) { return t.find(str.begin(), str.end()); }

template<typename Val>
optional<Val> find(ctrie<Val> const & t, char const * str) { return t.find(str, str + strlen(str)); }
}