some fixes for mappings

7 files changed, 156 insertions, 40 deletions

diff --git a/src/crates/candelabra/src/cost/benchmark.rs b/src/crates/candelabra/src/cost/benchmark.rs
index 9c5c382..439d42a 100644
--- a/src/crates/candelabra/src/cost/benchmark.rs
+++ b/src/crates/candelabra/src/cost/benchmark.rs
@@ -204,12 +204,17 @@ primrose-library = {{ path = \"{}\" }}
     // Add benchmarks for implemented traits
     let implemented_traits = lib_spec.interface_provide_map.keys();
     for tr in implemented_traits {
-        benchmark_statements += &format!(
-            "candelabra_benchmarker::benchmark_{}::<{}<{}>, _>(&mut rng, &NS);",
+        if tr == "Mapping" {
+            benchmark_statements += &format!(
+            "candelabra_benchmarker::benchmark_{0}::<{name}<{ELEM_TYPE}, {ELEM_TYPE}>, _, _>(&mut rng, &NS);",
             tr.to_lowercase(),
-            name,
-            ELEM_TYPE,
         );
+        } else {
+            benchmark_statements += &format!(
+                "candelabra_benchmarker::benchmark_{0}::<{name}<{ELEM_TYPE}>, _>(&mut rng, &NS);",
+                tr.to_lowercase(),
+            );
+        }
     }
 
     // Write the benchmarking source, using our generated benchmarker code.
diff --git a/src/crates/library/src/btreemap.rs b/src/crates/library/src/btreemap.rs
index d3f5144..4cd490e 100644
--- a/src/crates/library/src/btreemap.rs
+++ b/src/crates/library/src/btreemap.rs
@@ -57,7 +57,7 @@ impl<K: Ord + Hash, V> Mapping<K, V> for BTreeMap<K, V> {
     (define (get xs k) (cdr (assoc k xs)))
     (define (post-get xs k r) (equal? r (get xs k)))
     *ENDLIBSPEC*/
-    fn get(&mut self, key: &K) -> Option<&V> {
+    fn get(&self, key: &K) -> Option<&V> {
         BTreeMap::get(self, key)
     }
 
diff --git a/src/crates/library/src/hashmap.rs b/src/crates/library/src/hashmap.rs
index 58bf564..e14074c 100644
--- a/src/crates/library/src/hashmap.rs
+++ b/src/crates/library/src/hashmap.rs
@@ -57,7 +57,7 @@ impl<K: Ord + Hash, V> Mapping<K, V> for HashMap<K, V> {
     (define (get xs k) (cdr (assoc k xs)))
     (define (post-get xs k r) (equal? r (get xs k)))
     *ENDLIBSPEC*/
-    fn get(&mut self, key: &K) -> Option<&V> {
+    fn get(&self, key: &K) -> Option<&V> {
         HashMap::get(self, key)
     }
 
diff --git a/src/crates/library/src/lib.rs b/src/crates/library/src/lib.rs
index 0a62c75..825ffad 100644
--- a/src/crates/library/src/lib.rs
+++ b/src/crates/library/src/lib.rs
@@ -5,7 +5,7 @@
 pub mod traits;
 
 mod profiler;
-pub use profiler::ProfilerWrapper;
+pub use profiler::{MappingProfilerWrapper, ProfilerWrapper};
 
 mod adaptive;
 pub use adaptive::AdaptiveContainer;
diff --git a/src/crates/library/src/profiler.rs b/src/crates/library/src/profiler.rs
index 2f1af37..5391cb7 100644
--- a/src/crates/library/src/profiler.rs
+++ b/src/crates/library/src/profiler.rs
@@ -21,7 +21,7 @@ pub struct ProfilerWrapper<const ID: usize, T, E> {
     n_nth: Cell<usize>,
     n_push: usize,
     n_pop: usize,
-    n_get: usize,
+    n_get: Cell<usize>,
     _d: PhantomData<E>,
 }
 
@@ -39,7 +39,7 @@ impl<const ID: usize, T: Default, E> Default for ProfilerWrapper<ID, T, E> {
             n_nth: Cell::new(0),
             n_push: 0,
             n_pop: 0,
-            n_get: 0,
+            n_get: Cell::new(0),
             _d: Default::default(),
         }
     }
@@ -52,13 +52,6 @@ impl<const ID: usize, T: Container<E>, E> ProfilerWrapper<ID, T, E> {
     }
 }
 
-impl<const ID: usize, T: Mapping<K, V>, K, V> ProfilerWrapper<ID, T, (K, V)> {
-    fn add_n_map(&self) {
-        self.max_n
-            .set(self.max_n.get().max(self.inner.as_ref().unwrap().len()));
-    }
-}
-
 impl<const ID: usize, T: Container<E>, E> Container<E> for ProfilerWrapper<ID, T, E> {
     fn len(&self) -> usize {
         self.inner.as_ref().unwrap().len()
@@ -129,46 +122,151 @@ impl<const ID: usize, T: Stack<E> + Container<E>, E> Stack<E> for ProfilerWrappe
     }
 }
 
-impl<const ID: usize, T: Mapping<K, V>, K, V> Mapping<K, V> for ProfilerWrapper<ID, T, (K, V)> {
+impl<const ID: usize, T, E> Drop for ProfilerWrapper<ID, T, E> {
+    fn drop(&mut self) {
+        let base_dir =
+            var("PROFILER_OUT_DIR").expect("ProfilerWrapper used without environment variable");
+        let unix_time = SystemTime::now()
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .unwrap();
+        let dir = format!("{}/{}", base_dir, ID);
+        if metadata(&dir).is_err() {
+            create_dir(&dir).unwrap();
+        }
+        let mut f = File::create(format!("{}/{}", dir, unix_time.as_nanos())).unwrap();
+
+        writeln!(f, "{}", self.max_n.get()).unwrap();
+        writeln!(f, "{}", self.n_contains).unwrap();
+        writeln!(f, "{}", self.n_insert).unwrap();
+        writeln!(f, "{}", self.n_clear).unwrap();
+        writeln!(f, "{}", self.n_remove).unwrap();
+        writeln!(f, "{}", self.n_first.get()).unwrap();
+        writeln!(f, "{}", self.n_last.get()).unwrap();
+        writeln!(f, "{}", self.n_nth.get()).unwrap();
+        writeln!(f, "{}", self.n_push).unwrap();
+        writeln!(f, "{}", self.n_pop).unwrap();
+        writeln!(f, "{}", self.n_get.get()).unwrap();
+    }
+}
+
+impl<const ID: usize, T: IntoIterator, E> IntoIterator for ProfilerWrapper<ID, T, E> {
+    type Item = T::Item;
+
+    type IntoIter = T::IntoIter;
+
+    fn into_iter(mut self) -> Self::IntoIter {
+        self.inner.take().unwrap().into_iter()
+    }
+}
+
+impl<const ID: usize, C: FromIterator<E> + Container<E>, E> FromIterator<E>
+    for ProfilerWrapper<ID, C, E>
+{
+    fn from_iter<T: IntoIterator<Item = E>>(iter: T) -> Self {
+        let c = C::from_iter(iter);
+        Self {
+            max_n: Cell::new(c.len()),
+            inner: Some(c),
+            n_contains: 0,
+            n_insert: 0,
+            n_clear: 0,
+            n_remove: 0,
+            n_first: Cell::new(0),
+            n_last: Cell::new(0),
+            n_nth: Cell::new(0),
+            n_push: 0,
+            n_pop: 0,
+            n_get: Cell::new(0),
+            _d: PhantomData,
+        }
+    }
+}
+
+pub struct MappingProfilerWrapper<const ID: usize, T, E> {
+    inner: Option<T>,
+    max_n: Cell<usize>,
+    n_contains: usize,
+    n_insert: usize,
+    n_clear: usize,
+    n_remove: usize,
+    n_first: Cell<usize>,
+    n_last: Cell<usize>,
+    n_nth: Cell<usize>,
+    n_push: usize,
+    n_pop: usize,
+    n_get: Cell<usize>,
+    _d: PhantomData<E>,
+}
+
+impl<const ID: usize, T: Default, E> Default for MappingProfilerWrapper<ID, T, E> {
+    fn default() -> Self {
+        Self {
+            inner: Some(T::default()),
+            max_n: Cell::new(0),
+            n_contains: 0,
+            n_insert: 0,
+            n_clear: 0,
+            n_remove: 0,
+            n_first: Cell::new(0),
+            n_last: Cell::new(0),
+            n_nth: Cell::new(0),
+            n_push: 0,
+            n_pop: 0,
+            n_get: Cell::new(0),
+            _d: Default::default(),
+        }
+    }
+}
+
+impl<const ID: usize, T: Mapping<K, V>, K, V> MappingProfilerWrapper<ID, T, (K, V)> {
+    fn add_n(&self) {
+        self.max_n
+            .set(self.max_n.get().max(self.inner.as_ref().unwrap().len()));
+    }
+}
+
+impl<const ID: usize, T: Mapping<K, V>, K, V> Mapping<K, V>
+    for MappingProfilerWrapper<ID, T, (K, V)>
+{
     fn len(&self) -> usize {
         self.inner.as_ref().unwrap().len()
     }
 
     fn contains(&mut self, x: &K) -> bool {
-        self.add_n_map();
+        self.add_n();
         self.n_contains += 1;
         self.inner.as_mut().unwrap().contains(x)
     }
 
     fn insert(&mut self, key: K, val: V) -> Option<V> {
-        self.add_n_map();
+        self.add_n();
         self.n_insert += 1;
         self.inner.as_mut().unwrap().insert(key, val)
     }
 
-    fn get(&mut self, key: &K) -> Option<&V> {
-        self.add_n_map();
-        self.n_get += 1;
-        self.inner.as_mut().unwrap().get(key)
+    fn get(&self, key: &K) -> Option<&V> {
+        self.add_n();
+        self.n_get.set(self.n_get.get() + 1);
+        self.inner.as_ref().unwrap().get(key)
     }
 
     fn remove(&mut self, key: &K) -> Option<V> {
-        self.add_n_map();
+        self.add_n();
         self.n_remove += 1;
         self.inner.as_mut().unwrap().remove(key)
     }
 
     fn clear(&mut self) {
-        self.add_n_map();
+        self.add_n();
         self.n_clear += 1;
         self.inner.as_mut().unwrap().clear()
     }
 }
 
-impl<const ID: usize, T, E> Drop for ProfilerWrapper<ID, T, E> {
+impl<const ID: usize, T, E> Drop for MappingProfilerWrapper<ID, T, E> {
     fn drop(&mut self) {
-        let base_dir =
-            var("PROFILER_OUT_DIR").expect("ProfilerWrapper used without environment variable");
+        let base_dir = var("PROFILER_OUT_DIR")
+            .expect("MappingProfilerWrapper used without environment variable");
         let unix_time = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
             .unwrap();
@@ -188,11 +286,11 @@ impl<const ID: usize, T, E> Drop for ProfilerWrapper<ID, T, E> {
         writeln!(f, "{}", self.n_nth.get()).unwrap();
         writeln!(f, "{}", self.n_push).unwrap();
         writeln!(f, "{}", self.n_pop).unwrap();
-        writeln!(f, "{}", self.n_get).unwrap();
+        writeln!(f, "{}", self.n_get.get()).unwrap();
     }
 }
 
-impl<const ID: usize, T: IntoIterator, E> IntoIterator for ProfilerWrapper<ID, T, E> {
+impl<const ID: usize, T: IntoIterator, E> IntoIterator for MappingProfilerWrapper<ID, T, E> {
     type Item = T::Item;
 
     type IntoIter = T::IntoIter;
@@ -202,10 +300,10 @@ impl<const ID: usize, T: IntoIterator, E> IntoIterator for ProfilerWrapper<ID, T
     }
 }
 
-impl<const ID: usize, C: FromIterator<E> + Container<E>, E> FromIterator<E>
-    for ProfilerWrapper<ID, C, E>
+impl<const ID: usize, C: Mapping<K, V>, K, V> FromIterator<(K, V)>
+    for MappingProfilerWrapper<ID, C, (K, V)>
 {
-    fn from_iter<T: IntoIterator<Item = E>>(iter: T) -> Self {
+    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
         let c = C::from_iter(iter);
         Self {
             max_n: Cell::new(c.len()),
@@ -219,7 +317,7 @@ impl<const ID: usize, C: FromIterator<E> + Container<E>, E> FromIterator<E>
             n_nth: Cell::new(0),
             n_push: 0,
             n_pop: 0,
-            n_get: 0,
+            n_get: Cell::new(0),
             _d: PhantomData,
         }
     }
diff --git a/src/crates/library/src/traits.rs b/src/crates/library/src/traits.rs
index 7d75f13..39a1874 100644
--- a/src/crates/library/src/traits.rs
+++ b/src/crates/library/src/traits.rs
@@ -31,7 +31,7 @@ pub trait Indexable<T> {
     fn nth(&self, n: usize) -> Option<&T>;
 }
 
-pub trait Mapping<K, V>: Default + IntoIterator<Item = (K, V)> {
+pub trait Mapping<K, V>: Default + IntoIterator<Item = (K, V)> + FromIterator<(K, V)> {
     /// Get the current number of elements
     fn len(&self) -> usize;
 
@@ -47,7 +47,7 @@ pub trait Mapping<K, V>: Default + IntoIterator<Item = (K, V)> {
     fn insert(&mut self, key: K, val: V) -> Option<V>;
 
     /// Get the value for the given key
-    fn get(&mut self, key: &K) -> Option<&V>;
+    fn get(&self, key: &K) -> Option<&V>;
 
     /// Remove the first occurence of the given item
     fn remove(&mut self, key: &K) -> Option<V>;
diff --git a/src/crates/primrose/src/codegen.rs b/src/crates/primrose/src/codegen.rs
index 2eccf8d..e7766bf 100644
--- a/src/crates/primrose/src/codegen.rs
+++ b/src/crates/primrose/src/codegen.rs
@@ -66,14 +66,27 @@ fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
 }}
 "#
             ),
-            ContainerSelection::Profile(inner, idx) => format!(
-                r#"
+            ContainerSelection::Profile(inner, idx) => {
+                let wrapper_name = if self
+                    .lib_specs
+                    .get(inner)
+                    .expect("invalid inner implementation of profiled container")
+                    .interface_provide_map
+                    .contains_key("Mapping")
+                {
+                    "MappingProfilerWrapper"
+                } else {
+                    "ProfilerWrapper"
+                };
+                format!(
+                    r#"
 #[allow(non_snake_case)]
 fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
-    ::primrose_library::ProfilerWrapper::<{idx}, {inner}<{vars}>, _>::default()
+    ::primrose_library::{wrapper_name}::<{idx}, {inner}<{vars}>, _>::default()
 }}
 "#
-            ),
+                )
+            }
             ContainerSelection::Split {
                 before,
                 threshold,