text: #pragma once
  #include "utils.h"
  #include "map_ops.h"
  
  // *******************************************
  //   AUGMENTED MAP OPERATIONS
  // *******************************************
  
  template<class Map>
  struct augmented_ops : Map {
    using Entry = typename Map::Entry;
    using node = typename Map::node;
    using ET = typename Map::ET;
    using GC = typename Map::GC;
    using K = typename Map::K;
    using aug_t = typename Entry::aug_t;
  
    struct aug_sum_t {
      aug_t result;
      aug_sum_t() : result(Entry::get_empty()) {}
      void add_entry(ET e) {
        result = Entry::combine(result, Entry::from_entry(e));
      }
      void add_aug_val(aug_t av) {
        result = Entry::combine(result, av);
      }
    };
  
    template<class aug>
    // the sum right of or at key
    static void aug_sum_right(node* b, const K& key, aug& a) {
      while (b) {
        if (!Map::comp(Map::get_key(b), key)) {
  	a.add_entry(Map::get_entry(b));
  	if (b->rc) a.add_aug_val(Map::aug_val(b->rc));
  	b = b->lc;
        } else b = b->rc;
      }
    }
  
    template<class aug>
    // the sum left of or at key
    static void aug_sum_left(node* b, const K& key, aug& a) {
      while (b) {
        if (!Map::comp(key, Map::get_key(b))) {
  	a.add_entry(Map::get_entry(b));
  	if (b->lc) a.add_aug_val(Map::aug_val(b->lc));
  	b = b->rc;
        } else b = b->lc;
      }
    }
  
    template<class aug>
    static void aug_sum_range(node* b, const K& key_left, const K& key_right, aug& a) {
      node* r = Map::range_root(b, key_left, key_right);
      if (r) {
        // add in left side (right of or at key_left)
        aug_sum_right(r->lc, key_left, a);  
        a.add_entry(Map::get_entry(r));   // add in middle
        // add in right side (left of or at key_right)
        aug_sum_left(r->rc, key_right, a);  
      }
    }
  
    template<typename Func>
    static node* aug_select(node* b, const Func& f) {
      if (b == NULL) return NULL;
      if (f(Map::aug_val(b->lc))) {
        if (f(Entry::from_entry(Map::get_entry(b))))
  	return aug_select(b->rc, f);
        return b;
      } return aug_select(b->lc, f);
    }
  
    template<class Func>
    static node* aug_filter(node* b, const Func& f, bool extra_ptr = false) {
      if (!b) return NULL;
      if (!Map::aug_val(b)) return NULL;
      bool copy = extra_ptr || (b->ref_cnt > 1);
      
      auto P = utils::fork<node*>(Map::size(b) >= utils::node_limit,
  	        [&]() {return aug_filter(b->lc, f, copy);},
  		[&]() {return aug_filter(b->rc, f, copy);});
  
      if (f(Entry::from_entry(Map::get_entry(b)))) {
  	return Map::node_join(P.first, P.second, GC::copy_if(b, copy, extra_ptr));
      } else {
        GC::dec_if(b, copy, extra_ptr);
        return Map::join2(P.first, P.second);
      }
    }
  
    template <class Func>
    static node* insert_lazy(node* b, const ET& e, const Func& f) {
      aug_t av = Entry::from_entry(e);
      auto g = [&] (const aug_t& a) { return Entry::combine(av,a);};
  
      auto lazy_join = [&] (node* l, node* r, node* m) {
        m->rc = r; m->lc = l;
        if (Map::is_balanced(m)) {
  	Map::lazy_update(m,g);
  	return m;
        } else return Map::node_join(l,r,m);
      };
  
      return Map::insert_j(b, e, f, lazy_join, false);
    }
  
  };
  