fix adaptive containers for mappings

2 files changed, 57 insertions, 7 deletions

diff --git a/src/crates/library/src/adaptive.rs b/src/crates/library/src/adaptive.rs
index 3161970..6708a07 100644
--- a/src/crates/library/src/adaptive.rs
+++ b/src/crates/library/src/adaptive.rs
@@ -1,6 +1,6 @@
 use std::marker::PhantomData;
 
-use crate::traits::Container;
+use crate::traits::{Container, Mapping};
 use take_mut::take_or_recover;
 
 pub enum AdaptiveContainer<const THRESHOLD: usize, LO, HI, E> {
@@ -75,6 +75,61 @@ impl<const THRESHOLD: usize, LO: Default + Container<E>, HI: Default + Container
     }
 }
 
+impl<const THRESHOLD: usize, LO, HI, K, V> Mapping<K, V>
+    for AdaptiveContainer<THRESHOLD, LO, HI, (K, V)>
+where
+    LO: Mapping<K, V>,
+    HI: Mapping<K, V>,
+{
+    fn len(&self) -> usize {
+        get_inner!(self, c, c.len())
+    }
+
+    fn contains(&mut self, x: &K) -> bool {
+        get_inner!(self, c, c.contains(x))
+    }
+
+    fn insert(&mut self, key: K, val: V) -> Option<V> {
+        let res = get_inner!(self, c, c.insert(key, val));
+        if let Self::Low(l, _) = self {
+            if l.len() < THRESHOLD {
+                return res;
+            }
+            // adapt!
+            take_or_recover(self, Default::default, |old| {
+                let Self::Low(l, _) = old else {
+                    unreachable!();
+                };
+                Self::High(HI::from_iter(l))
+            });
+        }
+        res
+    }
+
+    fn get(&self, key: &K) -> Option<&V> {
+        get_inner!(self, c, c.get(key))
+    }
+
+    fn remove(&mut self, key: &K) -> Option<V> {
+        get_inner!(self, c, c.remove(key))
+    }
+
+    fn clear(&mut self) {
+        *self = Self::Low(LO::default(), PhantomData);
+    }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = (&'a K, &'a V)> + 'a
+    where
+        K: 'a,
+        V: 'a,
+    {
+        match self {
+            AdaptiveContainer::Low(lo, _) => AdaptiveContainerIter::Low(lo.iter(), PhantomData),
+            AdaptiveContainer::High(hi) => AdaptiveContainerIter::High(hi.iter()),
+        }
+    }
+}
+
 impl<const THRESHOLD: usize, LO, HI, E> FromIterator<E> for AdaptiveContainer<THRESHOLD, LO, HI, E>
 where
     LO: FromIterator<E>,
diff --git a/src/crates/primrose/src/codegen.rs b/src/crates/primrose/src/codegen.rs
index 82887bb..2794e16 100644
--- a/src/crates/primrose/src/codegen.rs
+++ b/src/crates/primrose/src/codegen.rs
@@ -137,16 +137,11 @@ fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
                 threshold,
                 after,
             } => {
-                let adaptive_container_type = if vars.contains(',') {
-                    "AdaptiveMappingContainer"
-                } else {
-                    "AdaptiveContainer"
-                };
                 format!(
                     r#"
 #[allow(non_snake_case)]
 fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
-    ::primrose_library::{adaptive_container_type}::<{threshold}, {before}<{vars}>, {after}<{vars}>, {vars}>::default()
+    ::primrose_library::AdaptiveContainer::<{threshold}, {before}<{vars}>, {after}<{vars}>, ({vars})>::default()
 }}
 "#
                 )