cpp-trie/trie.hpp

//
//  trie.hpp
//  trie
//
//  Created by Sam Jaffe on 6/16/17.
//

#pragma once

#include "pointers/const_propogating_ptr.hpp"
#include "pointers/not_null.hpp"
#include "pointers/value_ptr.hpp"
#include <map>

template <typename K, typename V, typename Compare = std::less<K>>
class trie;

namespace detail {
  template <typename Trie, typename Iter>
  class trie_iterator_base;
}

template <typename Trie, typename Iter>
class trie_iterator;
template <typename Trie, typename Iter>
class trie_reverse_iterator;

template <typename K, typename V, typename Compare>
class trie {
private:
  using self_t = trie<K, V, Compare>;
  using layer_t = value_ptr<self_t>;
  using backing_t = std::map<K, const_propogating_ptr<layer_t>, Compare>;
  template <typename KS>
  using is_collection_t = typename std::enable_if<std::is_same<K, std::decay_t<decltype(*std::begin(std::declval<KS>()))>>::value>::type;

public:
  using key_type = K;
  using mapped_type = V;
  using key_compare = Compare;
  
  using local_iterator = typename backing_t::iterator;
  using local_const_iterator = typename backing_t::const_iterator;
  using local_reverse_iterator = typename backing_t::reverse_iterator;
  using local_const_reverse_iterator = typename backing_t::const_reverse_iterator;
  
  using iterator = trie_iterator<self_t, local_iterator>;
  using const_iterator = trie_iterator<self_t const, local_const_iterator>;
  using post_iterator = trie_iterator<self_t, local_reverse_iterator>;
  using const_post_iterator = trie_iterator<self_t const, local_const_reverse_iterator>;
  using reverse_iterator = trie_reverse_iterator<self_t, local_reverse_iterator>;
  using const_reverse_iterator = trie_reverse_iterator<self_t const,  local_const_reverse_iterator>;
private:
  using impl_iterator = detail::trie_iterator_base<self_t, local_iterator>;
  using impl_const_iterator = detail::trie_iterator_base<self_t, local_const_iterator>;
public:
  trie() {}
  trie(trie const & other);
  trie(trie && other) { swap(other); }
  ~trie() { clear(); }
  self_t & operator=(mapped_type const & value);
  self_t & operator=(trie const & value);
  self_t & operator=(trie && value);

  operator mapped_type &() { return value(); }
  operator mapped_type const &() const { return value(); }
  mapped_type & value() { return value_; }
  mapped_type const & value() const { return value_; }
  
  template <typename KS, typename = is_collection_t<KS>>
  self_t & operator[](KS const & keys) {
    return *emplace(keys).first.stk.top();
  }
  self_t & operator[](key_type const & key) {
    return *emplace(key).first.stk.top();
  }
  self_t & operator[](std::initializer_list<key_type> keys) {
    return operator[]<std::initializer_list<key_type>>(keys);
  }
  
  template <typename KS>
  std::pair<iterator, bool> insert(KS const & keys, mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  std::pair<iterator, bool> insert(key_type const & key, mapped_type const & value) {
    return emplace_impl({key}, value);
  }
  std::pair<iterator, bool> insert(std::initializer_list<key_type> keys, mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  template <typename KS, typename... Args>
  std::pair<iterator, bool> emplace(KS && keys, Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type const & key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type && key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(std::initializer_list<key_type> keys, Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }
  
  iterator begin() { return {this}; }
  iterator end() { return {}; }
  const_iterator begin() const { return {this}; }
  const_iterator end() const { return {}; }
  const_iterator cbegin() const { return begin(); }
  const_iterator cend() const { return end(); }
  reverse_iterator rbegin() { return {this}; }
  reverse_iterator rend() { return {}; }
  const_reverse_iterator rbegin() const { return {this}; }
  const_reverse_iterator rend() const { return {}; }
  const_reverse_iterator crbegin() const { return rbegin(); }
  const_reverse_iterator crend() const { return rend(); }
  
  local_iterator local_begin() { return impl_.begin(); }
  local_iterator local_end() { return impl_.end(); }
  local_const_iterator local_begin() const { return impl_.begin(); }
  local_const_iterator local_end() const { return impl_.end(); }

  local_reverse_iterator local_rbegin() { return impl_.rbegin(); }
  local_reverse_iterator local_rend() { return impl_.rend(); }
  local_const_reverse_iterator local_rbegin() const { return impl_.rbegin(); }
  local_const_reverse_iterator local_rend() const { return impl_.rend(); }
  
  
  template <typename KS>
  iterator find(KS const & keys) { return find_impl<impl_iterator>(this, keys); }
  iterator find(key_type const & key) { return find_impl<impl_iterator>(this, {key}); }
  iterator find(std::initializer_list<key_type> keys) { return find_impl<impl_iterator>(this, keys); }
  template <typename KS>
  const_iterator find(KS const & keys) const { return find_impl<impl_const_iterator>(this, keys); }
  const_iterator find(key_type const & key) const { return find_impl<impl_const_iterator>(this, {key}); }
  const_iterator find(std::initializer_list<key_type> keys) const { return find_impl<impl_const_iterator>(this, keys); }
  
  template <typename KS>
  void erase(KS const & keys) { drop(find(keys)); }
  void erase(key_type const & key) { drop(find(key)); }
  void erase(std::initializer_list<key_type> keys) { drop(find(keys)); }
  
  void clear();
private:
  void drop(iterator it);
  template <typename Iter, typename Trie, typename KS>
  static Iter find_impl(Trie * tr, KS const & keys);
  template <typename Iter, typename Trie>
  static Iter find_impl(Trie * tr, std::initializer_list<key_type> keys) {
    return find_impl<Iter, Trie, std::initializer_list<key_type>>(tr, keys);
  }
  template <typename KS, typename... Args>
  std::pair<impl_iterator, bool> emplace_impl(KS && keys, Args &&... args);
  template <typename... Args>
  std::pair<impl_iterator, bool> emplace_impl(std::initializer_list<key_type> key, Args &&... args) {
    return emplace_impl<std::initializer_list<key_type>>(std::move(key), std::forward<Args>(args)...);
  }
  template <typename Key>
  void insert_impl(impl_iterator & out, bool & create, Key && key);
  
  friend bool operator==(trie const & lhs, trie const & rhs) {
    auto it1 = lhs.begin(), it2 = rhs.begin(), end1 = lhs.end(), end2 = lhs.end();
    for ( ; it1 != end1 && it2 != end2; ++it1, ++it2 ) {
      if (!it1.keys.empty() && !it2.keys.empty() && it1.keys.back() != it2.keys.back()) { return false; }
      else if (it1.stk.top()->value_ != it2.stk.top()->value_) { return false; }
    }
    return it1 == end1 && it2 == end2;
  }
  friend void swap(trie & lhs, trie & rhs) {
    using std::swap;
    swap(lhs.value_, rhs.value_);
    swap(lhs.impl_, rhs.impl_);
  }
  
  mapped_type value_{};
  backing_t impl_{};
};

#include "trie_impl.hpp"