trie.hpp

#ifndef HALCHECK_LIB_TRIE_HPP
#define HALCHECK_LIB_TRIE_HPP
 
#include <halcheck/lib/functional.hpp>
#include <halcheck/lib/iterator.hpp>
#include <halcheck/lib/memory.hpp>
#include <halcheck/lib/type_traits.hpp>
 
#include <cstddef>
#include <functional>
#include <iterator>
#include <memory>
#include <type_traits>
#include <unordered_map>
#include <utility>
 
namespace halcheck { namespace lib {
 
template<typename K, typename V, typename Hash = std::hash<K>>
class trie {
private:
  static_assert(lib::is_invocable_r<std::size_t, Hash, const K &>(), "Hash must be a hasher for K");
  static_assert(std::is_default_constructible<V>(), "V must be default constructible");
 
  std::shared_ptr<const V> _value;
  std::shared_ptr<std::unordered_map<K, trie>> _children;
 
public:
  trie() = default;
 
  template<
      typename I,
      HALCHECK_REQUIRE(lib::is_input_iterator<I>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::iter_reference_t<I>, std::pair<const K, trie>>())>
  explicit trie(V value, I first, I last)
      : _value(new V(std::move(value))), _children(new std::unordered_map<K, trie>(std::move(first), std::move(last))) {
  }
 
  template<
      typename R,
      HALCHECK_REQUIRE(lib::is_input_range<R>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::range_reference_t<R>, std::pair<const K, trie>>())>
  explicit trie(V value, R &&range) : trie(std::move(value), lib::begin(range), lib::end(range)) {}
 
  const V &operator*() const {
    static const V fallback;
    if (_value)
      return *_value;
    else
      return fallback;
  }
 
  template<
      typename R,
      HALCHECK_REQUIRE(lib::is_range<R>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::range_reference_t<R>, K>())>
  const V &operator[](const R &range) const {
    return *drop(range);
  }
 
  template<
      typename R,
      HALCHECK_REQUIRE(lib::is_range<R>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::range_reference_t<R>, K>())>
  trie drop(const R &range) const {
    return drop(lib::begin(range), lib::end(range));
  }
 
  template<
      typename I,
      HALCHECK_REQUIRE(lib::is_input_iterator<I>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::iter_reference_t<I>, K>())>
  trie drop(I first, I last) const {
    auto output = *this;
    for (; first != last; ++first) {
      auto it = output.find(*first);
      if (it == output.end())
        return trie();
      else
        output = it->second;
    }
 
    return output;
  }
 
  trie drop(const K &key) const {
    auto it = find(key);
    if (it == end())
      return trie();
    else
      return it->second;
  }
 
  template<
      typename I,
      HALCHECK_REQUIRE(lib::is_input_iterator<I>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::iter_reference_t<I>, K>())>
  trie set(I first, I last, V value) const {
    auto output = *this;
    auto current = &output;
 
    for (; first != last; ++first) {
      current->_children = current->_children ? std::make_shared<std::unordered_map<K, trie, Hash>>(*current->_children)
                                              : std::make_shared<std::unordered_map<K, trie, Hash>>();
      current = &(*current->_children)[*first];
    }
 
    current->_value = std::make_shared<V>(std::move(value));
    return output;
  }
 
  template<
      typename R,
      HALCHECK_REQUIRE(lib::is_input_range<R>()),
      HALCHECK_REQUIRE(std::is_convertible<lib::range_reference_t<R>, K>())>
  trie set(const R &range, V value) const {
    return set(lib::begin(range), lib::end(range), std::move(value));
  }
 
  class iterator : public lib::iterator_interface<iterator> {
  public:
    using value_type = typename std::unordered_map<K, trie, Hash>::value_type;
    using difference_type = typename std::unordered_map<K, trie, Hash>::difference_type;
    using reference = typename std::unordered_map<K, trie, Hash>::const_reference;
    using pointer = typename std::unordered_map<K, trie, Hash>::const_pointer;
    using iterator_category = std::forward_iterator_tag;
 
    using lib::iterator_interface<iterator>::operator++;
 
    iterator() = default;
 
    explicit iterator(typename std::unordered_map<K, trie, Hash>::const_iterator base) : _base(std::move(base)) {}
 
    iterator &operator++() {
      ++_base;
      return *this;
    }
 
    reference operator*() const { return *_base; }
 
    pointer operator->() const { return _base.operator->(); }
 
  private:
    friend bool operator==(const iterator &lhs, const iterator &rhs) { return lhs._base == rhs._base; }
 
    typename std::unordered_map<K, trie, Hash>::const_iterator _base;
  };
 
  iterator begin() const { return _children ? iterator(_children->begin()) : iterator(); }
 
  iterator end() const { return _children ? iterator(_children->end()) : iterator(); }
 
  std::size_t size() const { return _children.size(); }
 
  bool empty() const { return _children.empty(); }
 
  iterator find(const K &key) const {
    if (_children)
      return iterator(_children->find(key));
    else
      return iterator();
  }
};
 
}} // namespace halcheck::lib
 
#endif