1use std::{cmp::Ordering, fmt::Debug};
  2
  3use crate::{Bias, Dimension, Item, KeyedItem, SeekTarget, SumTree, Summary};
  4
  5pub struct TreeMap<K, V>(SumTree<MapEntry<K, V>>)
  6where
  7    K: Clone + Debug + Default,
  8    V: Clone + Debug;
  9
 10#[derive(Clone)]
 11pub struct MapEntry<K, V> {
 12    key: K,
 13    value: V,
 14}
 15
 16#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
 17pub struct MapKey<K>(K);
 18
 19#[derive(Clone, Debug, Default)]
 20pub struct MapKeyRef<'a, K>(Option<&'a K>);
 21
 22impl<K: Clone + Debug + Default + Ord, V: Clone + Debug + Default> TreeMap<K, V> {
 23    pub fn get<'a>(&self, key: &'a K) -> Option<&V> {
 24        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
 25        let key = MapKeyRef(Some(key));
 26        cursor.seek(&key, Bias::Left, &());
 27        if key.cmp(cursor.start(), &()) == Ordering::Equal {
 28            Some(&cursor.item().unwrap().value)
 29        } else {
 30            None
 31        }
 32    }
 33
 34    pub fn insert(&mut self, key: K, value: V) {
 35        self.0.insert_or_replace(MapEntry { key, value }, &());
 36    }
 37
 38    pub fn remove<'a>(&mut self, key: &'a K) -> Option<V> {
 39        let mut removed = None;
 40        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
 41        let key = MapKeyRef(Some(key));
 42        let mut new_tree = cursor.slice(&key, Bias::Left, &());
 43        if key.cmp(cursor.start(), &()) == Ordering::Equal {
 44            removed = Some(cursor.item().unwrap().value.clone());
 45            cursor.next(&());
 46        }
 47        new_tree.push_tree(cursor.suffix(&()), &());
 48        drop(cursor);
 49        self.0 = new_tree;
 50        removed
 51    }
 52}
 53
 54impl<K, V> Item for MapEntry<K, V>
 55where
 56    K: Clone + Debug + Default + Clone,
 57    V: Clone,
 58{
 59    type Summary = MapKey<K>;
 60
 61    fn summary(&self) -> Self::Summary {
 62        todo!()
 63    }
 64}
 65
 66impl<K, V> KeyedItem for MapEntry<K, V>
 67where
 68    K: Clone + Debug + Default + Ord,
 69    V: Clone,
 70{
 71    type Key = MapKey<K>;
 72
 73    fn key(&self) -> Self::Key {
 74        MapKey(self.key.clone())
 75    }
 76}
 77
 78impl<K> Summary for MapKey<K>
 79where
 80    K: Clone + Debug + Default,
 81{
 82    type Context = ();
 83
 84    fn add_summary(&mut self, summary: &Self, _: &()) {
 85        *self = summary.clone()
 86    }
 87}
 88
 89impl<'a, K> Dimension<'a, MapKey<K>> for MapKeyRef<'a, K>
 90where
 91    K: Clone + Debug + Default + Ord,
 92{
 93    fn add_summary(&mut self, summary: &'a MapKey<K>, _: &()) {
 94        self.0 = Some(&summary.0)
 95    }
 96}
 97
 98impl<'a, K> SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>> for MapKeyRef<'_, K>
 99where
100    K: Clone + Debug + Default + Ord,
101{
102    fn cmp(&self, cursor_location: &MapKeyRef<K>, _: &()) -> Ordering {
103        self.0.cmp(&cursor_location.0)
104    }
105}