refactor out nested hashmaps for mapping2d struct

7 files changed, 247 insertions, 89 deletions

diff --git a/src/crates/candelabra/src/candidates.rs b/src/crates/candelabra/src/candidates.rs
index e8693e0..7bc1031 100644
--- a/src/crates/candelabra/src/candidates.rs
+++ b/src/crates/candelabra/src/candidates.rs
@@ -3,7 +3,7 @@
 use std::{collections::HashMap, fs::metadata, time::SystemTime};
 
 use anyhow::{bail, Context, Result};
-use camino::{Utf8Path, Utf8PathBuf};
+use camino::Utf8Path;
 use log::{debug, warn};
 use primrose::ContainerSelector;
 use serde::{Deserialize, Serialize};
@@ -12,20 +12,15 @@ use crate::{
     cache::{gen_tree_hash, FileCache},
     paths::Paths,
     project::Project,
+    types::{ConTypeName, ConTypeTo, ImplName},
     State,
 };
 
-/// Names a container type we want to select.
-pub type ConTypeName = String;
+/// List of candidates for a whole project, by file and by type name
+pub type ProjectCandidateList = ConTypeTo<Vec<ImplName>>;
 
-/// Name of a container implementation we are considering
-pub type ImplName = String;
-
-/// A list of candidate container types
-pub type Candidates = HashMap<ConTypeName, Vec<ImplName>>;
-
-/// A list of candidates for each selection site, and each file in a given project
-pub type ProjectCandidateList = Vec<(Utf8PathBuf, Vec<(ConTypeName, Vec<ImplName>)>)>;
+/// Candidate container types for a single file
+type FileCandidates = HashMap<ConTypeName, Vec<ImplName>>;
 
 /// Info for getting & caching candidate types
 pub struct CandidatesStore {
@@ -38,7 +33,7 @@ pub struct CandidatesStore {
 pub struct CacheEntry {
     lib_hash: u64,
     mod_time: SystemTime,
-    value: Candidates,
+    value: FileCandidates,
 }
 
 impl CandidatesStore {
@@ -68,25 +63,23 @@ impl State {
     /// Run primrose on all files in the given project.
     /// Returns a list of all candidates for each container type in each file.
     pub fn project_candidate_list(&self, project: &Project) -> Result<ProjectCandidateList> {
-        let mut all_candidates = Vec::new();
+        let mut all_candidates = ProjectCandidateList::default();
         for file in project.find_primrose_files()? {
             let result = match self.candidates.store.find(&file)? {
                 Some(x) => x.value,
                 None => self.calc_candidates(&file)?,
             };
 
-            let mut typs = Vec::new();
             for (con_type_id, candidates) in result {
-                typs.push((con_type_id.clone(), candidates));
+                all_candidates.insert(&file, &con_type_id, candidates);
             }
-            all_candidates.push((file, typs));
         }
 
         Ok(all_candidates)
     }
 
     /// Find candidate types for every selection site in a given path
-    fn calc_candidates(&self, path: &Utf8Path) -> Result<Candidates> {
+    fn calc_candidates(&self, path: &Utf8Path) -> Result<FileCandidates> {
         let selector = ContainerSelector::from_path(
             path.as_std_path(),
             self.paths.library_src.as_std_path(),
@@ -94,7 +87,7 @@ impl State {
         )
         .with_context(|| format!("error getting container selector for {}", path))?;
 
-        let candidates: Candidates = selector
+        let candidates: FileCandidates = selector
             .find_all_candidates()?
             .into_iter()
             .map(|(k, v)| (k.to_string(), v))
diff --git a/src/crates/candelabra/src/confirmation.rs b/src/crates/candelabra/src/confirmation.rs
index 8949f56..7a593da 100644
--- a/src/crates/candelabra/src/confirmation.rs
+++ b/src/crates/candelabra/src/confirmation.rs
@@ -1,3 +1,4 @@
+use crate::types::ConTypeTo;
 use anyhow::{Context, Result};
 use camino::Utf8Path;
 use log::{debug, trace};
@@ -10,20 +11,24 @@ use std::{
 };
 
 use crate::{
-    candidates::{ConTypeName, ImplName},
     cost::{
         benchmark::{parse_criterion_output, tee_output},
         BenchmarkResult,
     },
+    select::Selection,
+    types::ConTypeName,
     Project, State,
 };
 
+/// A list of concrete selections for all container types in a project
+pub type Selections = ConTypeTo<Selection>;
+
 impl State {
     /// Run all benchmarks for the given project with the given container selections, returning the results
     pub fn run_benchmarks_with(
         &self,
         proj: &Project,
-        selections: &HashMap<&ConTypeName, &ImplName>,
+        selections: &Selections,
     ) -> Result<HashMap<String, BenchmarkResult>> {
         self.save_choices(proj, selections)
             .context("error setting up project")?;
@@ -51,18 +56,10 @@ impl State {
 
     /// Use the given selections for the container types in this project.
     /// Panics if a container type has not been selected.
-    pub fn save_choices(
-        &self,
-        proj: &Project,
-        selections: &HashMap<&ConTypeName, &ImplName>,
-    ) -> Result<()> {
+    pub fn save_choices(&self, proj: &Project, selections: &Selections) -> Result<()> {
         debug!("Saving choices for project {}", proj.name);
-        for (file, file_con_types) in self.project_candidate_list(proj)?.iter() {
-            // FIXME: this will cause type names to collide across files - need to rework this
-            let chosen = selections
-                .iter()
-                .filter(|(ctn, _)| file_con_types.iter().any(|(ctn2, _)| **ctn == ctn2))
-                .map(|(ctn, can)| (*ctn, can.as_str()));
+        for (file, _) in self.project_candidate_list(proj)?.iter_2d() {
+            let chosen = selections.iter_with_a(file);
 
             self.save_choices_file(file, chosen)
                 .with_context(|| format!("error setting up {}", file))?;
@@ -76,7 +73,7 @@ impl State {
     fn save_choices_file<'a>(
         &self,
         file: &Utf8Path,
-        choices: impl Iterator<Item = (&'a String, &'a str)>,
+        choices: impl Iterator<Item = &'a (ConTypeName, Selection)>,
     ) -> Result<()> {
         debug!("Saving choices for {}", file);
         let selector = ContainerSelector::from_path(
@@ -85,7 +82,15 @@ impl State {
             self.model_size,
         )?;
 
-        let new_code = selector.gen_replacement_file(choices);
+        let new_code = selector.gen_replacement_file(choices.map(|(ctn, sel)| {
+            (
+                ctn,
+                match sel {
+                    Selection::Singular(i) => i.as_str(),
+                    Selection::Split(_) => todo!(),
+                },
+            )
+        }));
         let new_path = file.to_string().replace(".pr", "");
 
         trace!("New code: {}", new_code);
diff --git a/src/crates/candelabra/src/lib.rs b/src/crates/candelabra/src/lib.rs
index c483e12..3a38265 100644
--- a/src/crates/candelabra/src/lib.rs
+++ b/src/crates/candelabra/src/lib.rs
@@ -1,3 +1,5 @@
+#![feature(impl_trait_in_assoc_type)]
+
 use anyhow::Result;
 use cache::FileCache;
 
@@ -11,6 +13,7 @@ mod confirmation;
 pub mod cost;
 pub mod profiler;
 pub mod select;
+pub mod types;
 
 mod paths;
 mod project;
@@ -40,7 +43,7 @@ impl State {
         Ok(Self {
             candidates: CandidatesStore::new(&paths)?,
             results: ResultsStore::new(&paths)?,
-            profiler_info_cache: Self::profiler_info_cache(&paths)?,
+            profiler_info_cache: Self::usage_profile_cache(&paths)?,
 
             model_size: 3, // TODO
             paths,
diff --git a/src/crates/candelabra/src/profiler/info.rs b/src/crates/candelabra/src/profiler/info.rs
index 5eb6734..488a8ba 100644
--- a/src/crates/candelabra/src/profiler/info.rs
+++ b/src/crates/candelabra/src/profiler/info.rs
@@ -5,15 +5,15 @@ use anyhow::{anyhow, Result};
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    candidates::ImplName,
     cost::{benchmark::OpName, Cost, CostModel, Estimator},
+    types::ImplName,
 };
 
 /// The information we get from profiling.
 /// Rather than keeping all results, we split them into 'similar enough' partitions,
 /// with the idea that each partition will probably have the same best implementation.
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
-pub struct ProfilerInfo(pub Vec<ProfilerPartition>);
+pub struct UsageProfile(pub Vec<ProfilerPartition>);
 
 /// A vector of container lifetimes which have similar usage characteristics
 #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
@@ -37,7 +37,7 @@ pub struct ContainerSplitSpec {
     pub after: ImplName,
 }
 
-impl ProfilerInfo {
+impl UsageProfile {
     pub fn from(iter: impl Iterator<Item = Result<String>>) -> Result<Self> {
         Ok(Self(
             iter.map(|contents| parse_output(&contents?))
@@ -299,7 +299,7 @@ mod tests {
         profiler::info::partition_costs,
     };
 
-    use super::{ProfilerInfo, ProfilerPartition};
+    use super::{ProfilerPartition, UsageProfile};
 
     const EPSILON: f64 = 1e-5;
     fn assert_feq(left: f64, right: f64, msg: &'static str) {
@@ -316,7 +316,7 @@ mod tests {
 
     #[test]
     fn test_cost_single_partition() {
-        let info = ProfilerInfo(vec![ProfilerPartition {
+        let info = UsageProfile(vec![ProfilerPartition {
             occurences: 1.0,
             avg_n: 100.0,
             avg_op_counts: {
@@ -340,7 +340,7 @@ mod tests {
 
     #[test]
     fn test_cost_multi_partitions_sums() {
-        let info = ProfilerInfo(vec![
+        let info = UsageProfile(vec![
             ProfilerPartition {
                 occurences: 1.0,
                 avg_n: 100.0,
@@ -375,7 +375,7 @@ mod tests {
 
     #[test]
     fn test_cost_multi_partitions_sums_weighted() {
-        let info = ProfilerInfo(vec![
+        let info = UsageProfile(vec![
             ProfilerPartition {
                 occurences: 2.0,
                 avg_n: 100.0,
diff --git a/src/crates/candelabra/src/profiler/mod.rs b/src/crates/candelabra/src/profiler/mod.rs
index 83c7954..014cf64 100644
--- a/src/crates/candelabra/src/profiler/mod.rs
+++ b/src/crates/candelabra/src/profiler/mod.rs
@@ -7,7 +7,6 @@ use camino::{Utf8Path, Utf8PathBuf};
 use log::{debug, log_enabled, trace, warn, Level};
 use primrose::ContainerSelector;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 use std::io::Write;
 use std::{
     fs::{read_dir, File},
@@ -17,22 +16,27 @@ use std::{
 use tempfile::tempdir;
 
 use crate::cache::{gen_tree_hash, FileCache};
-use crate::candidates::ConTypeName;
 use crate::cost::benchmark::tee_output;
 use crate::project::Project;
+use crate::types::{ConTypeName, ConTypeTo, ImplName, Mapping2D};
 use crate::{Paths, State};
 
-pub use self::info::{ContainerSplitSpec, ProfilerInfo, ProfilerPartition};
+pub use self::info::{ContainerSplitSpec, ProfilerPartition, UsageProfile};
+
+type ConTypeRef<'a> = (&'a Utf8PathBuf, &'a ConTypeName);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub(crate) struct CacheEntry {
     proj_hash: u64,
     proj_location: Utf8PathBuf,
-    info: HashMap<ConTypeName, ProfilerInfo>,
+    info: ProjectProfilingInfo,
 }
 
+/// Profiling
+type ProjectProfilingInfo = ConTypeTo<UsageProfile>;
+
 impl State {
-    pub(crate) fn profiler_info_cache(paths: &Paths) -> Result<FileCache<String, CacheEntry>> {
+    pub(crate) fn usage_profile_cache(paths: &Paths) -> Result<FileCache<String, CacheEntry>> {
         FileCache::new(
             paths.target_dir.join("candelabra").join("profiler_info"),
             |_, v: &CacheEntry| {
@@ -44,11 +48,11 @@ impl State {
 
     /// Get or calculate profiler info for the given project.
     /// Results are cached by the modification time of the project's source tree
-    pub fn profiler_info(&self, project: &Project) -> Result<HashMap<ConTypeName, ProfilerInfo>> {
+    pub fn usage_profile(&self, project: &Project) -> Result<ProjectProfilingInfo> {
         match self.profiler_info_cache.find(&project.name)? {
             Some(x) => Ok(x.info),
             None => {
-                let info = self.calc_profiler_info(project)?;
+                let info = self.calc_usage_profile(project)?;
 
                 let proj_hash = gen_tree_hash(&project.source_dir)
                     .context("Error generating project directory hash")?;
@@ -69,24 +73,24 @@ impl State {
     }
 
     /// Calculate profiler info for the given project.
-    fn calc_profiler_info(&self, project: &Project) -> Result<HashMap<ConTypeName, ProfilerInfo>> {
+    fn calc_usage_profile(&self, project: &Project) -> Result<ProjectProfilingInfo> {
         let candidate_list = self.project_candidate_list(project)?;
-        let con_types = candidate_list
+        let ct_refs = candidate_list
             .iter()
-            .flat_map(|(_, con_types)| con_types.iter())
-            .map(|(id, _)| id)
+            .map(|(f, ctn, _)| (f, ctn))
             .collect::<Vec<_>>();
 
-        self.project_profiling_prep(project, &con_types)?;
-        let mut acc = HashMap::new();
+        self.project_profiling_prep(project, &ct_refs)?;
+        let mut acc = ProjectProfilingInfo::default();
         for name in project.benchmarks.iter() {
-            for (con_type, new_results) in self
-                .profile_benchmark(project, name, &con_types)
+            for (f, ctn, new_results) in self
+                .profile_benchmark(project, name, &ct_refs)
                 .with_context(|| format!("Error profiling benchmark {}", name))?
+                .iter()
             {
-                acc.entry(con_type)
-                    .and_modify(|pi: &mut ProfilerInfo| pi.0.extend(new_results.0.iter().cloned()))
-                    .or_insert(new_results);
+                acc.entry_mut(f, ctn)
+                    .and_modify(|pi: &mut UsageProfile| pi.0.extend(new_results.0.iter().cloned()))
+                    .or_insert(new_results.clone());
             }
         }
 
@@ -94,9 +98,9 @@ impl State {
     }
 
     /// Prepare the given project to be profiled, by replacing all candidate types with the profiler wrapper.
-    fn project_profiling_prep(&self, project: &Project, con_types: &[&String]) -> Result<()> {
-        for (file, candidates) in self.project_candidate_list(project)? {
-            self.file_profiling_prep(&file, &candidates, con_types)
+    fn project_profiling_prep(&self, project: &Project, ct_refs: &[ConTypeRef<'_>]) -> Result<()> {
+        for (file, cts) in self.project_candidate_list(project)?.iter_2d() {
+            self.file_profiling_prep(file, cts, ct_refs)
                 .with_context(|| format!("error preparing {} for profiling", file))?;
         }
 
@@ -104,11 +108,11 @@ impl State {
     }
 
     /// Prepare the given file to be profiled, by replacing all candidate types with the profiler wrapper.
-    fn file_profiling_prep(
+    fn file_profiling_prep<'a>(
         &self,
         file: &Utf8Path,
-        candidates: &[(String, Vec<String>)],
-        con_types: &[&String],
+        candidates: impl Iterator<Item = (&'a ConTypeName, &'a Vec<ImplName>)>,
+        con_types: &[ConTypeRef<'_>],
     ) -> Result<()> {
         debug!("Setting up {} for profiling", file);
 
@@ -120,14 +124,16 @@ impl State {
         .context("error creating container selector")?;
 
         let chosen = candidates
-            .iter()
             .map(|(dest_name, impls)| (dest_name, &impls[0]))
             .collect::<Vec<_>>();
 
         let new_code = selector.gen_profiling_file(chosen.iter().map(|(d, c)| {
             (
                 *d,
-                con_types.iter().position(|id| id == d).unwrap(),
+                con_types
+                    .iter()
+                    .position(|(search_file, search_ctn)| file == *search_file && d == search_ctn)
+                    .unwrap(),
                 c.as_str(),
             )
         }));
@@ -149,8 +155,8 @@ impl State {
         &self,
         project: &Project,
         name: &str,
-        con_types: &[&String],
-    ) -> Result<HashMap<String, ProfilerInfo>> {
+        con_types: &[ConTypeRef<'_>],
+    ) -> Result<Mapping2D<Utf8PathBuf, ConTypeName, UsageProfile>> {
         let profiler_out_dir = tempdir()?;
         debug!(
             "Running benchmark {} with out dir {:?}",
@@ -172,21 +178,21 @@ impl State {
 
         tee_output(child)?;
 
-        let mut con_type_results = HashMap::new();
+        let mut con_type_results = Mapping2D::default();
         for dir in read_dir(&profiler_out_dir)? {
             // each directory has an index, corresponding to the container type name
             let dir = dir?;
-            let con_type: String = con_types[dir
+            let (f, ctn) = con_types[dir
                 .file_name()
                 .into_string()
                 .unwrap()
                 .parse::<usize>()
-                .unwrap()]
-            .to_string();
+                .unwrap()];
 
             con_type_results.insert(
-                con_type,
-                ProfilerInfo::from(read_dir(dir.path())?.map(|f| -> Result<String> {
+                f,
+                ctn,
+                UsageProfile::from(read_dir(dir.path())?.map(|f| -> Result<String> {
                     // read file contents
                     let mut contents = String::new();
                     File::open(f?.path())?.read_to_string(&mut contents)?;
diff --git a/src/crates/candelabra/src/select.rs b/src/crates/candelabra/src/select.rs
index a8653a4..12c359a 100644
--- a/src/crates/candelabra/src/select.rs
+++ b/src/crates/candelabra/src/select.rs
@@ -1,9 +1,9 @@
 use std::collections::HashMap;
 
 use crate::{
-    candidates::{ConTypeName, ImplName},
     cost::Cost,
     profiler::ContainerSplitSpec,
+    types::{ConTypeTo, ImplName},
     Project, State,
 };
 
@@ -15,15 +15,19 @@ pub enum Selection {
     Split(ContainerSplitSpec),
 }
 
+pub type ProjectSelections = ConTypeTo<Selection>;
+pub type ProjectCostEstimates = ConTypeTo<Vec<(Selection, Cost)>>;
+
 impl State {
     /// Select a container implementation for each container type in the given project
-    pub fn select(&self, project: &Project) -> Result<HashMap<ConTypeName, Selection>> {
+    pub fn select(&self, project: &Project) -> Result<ProjectSelections> {
         Ok(self
             .rank_candidates(project)?
             .into_iter()
-            .map(|(name, cs)| {
+            .map(|(f, ctn, cs)| {
                 (
-                    name,
+                    f,
+                    ctn,
                     cs.into_iter()
                         .reduce(
                             |acc @ (_, min), new @ (_, cost)| if min < cost { acc } else { new },
@@ -35,27 +39,24 @@ impl State {
             .collect())
     }
 
-    pub fn rank_candidates(
-        &self,
-        project: &Project,
-    ) -> Result<HashMap<ConTypeName, Vec<(Selection, Cost)>>> {
+    pub fn rank_candidates(&self, project: &Project) -> Result<ProjectCostEstimates> {
         // get all candidates
         let all_candidates = self.project_candidate_list(project)?;
 
         // get profiling information
-        let mut profiles = self.profiler_info(project)?;
+        let mut profiles = self.usage_profile(project)?;
 
-        let mut acc = HashMap::new();
-        let con_type_names = all_candidates.iter().flat_map(|(_, cs)| cs.iter());
-        for (con_type_name, candidates) in con_type_names {
+        let mut acc = ProjectCostEstimates::default();
+        for (f, ctn, candidates) in all_candidates.iter() {
             let mut costs = vec![];
-            let profile_info = profiles.get_mut(con_type_name).unwrap();
 
+            let profile_info = profiles.get_mut(f, ctn).unwrap();
             let cost_models = candidates
                 .iter()
                 .map(|name| Ok((name, self.cost_model(name)?)))
                 .collect::<Result<HashMap<_, _>>>()?;
 
+            // Get an estimate for each candidate
             for candidate in candidates {
                 let model = self.cost_model(candidate)?;
                 costs.push((
@@ -64,11 +65,12 @@ impl State {
                 ));
             }
 
+            // Also attempt to split at an n value
             if let Some((split, cost)) = profile_info.check_for_nsplit(&cost_models) {
                 costs.push((Selection::Split(split), cost));
             }
 
-            acc.insert(con_type_name.to_string(), costs);
+            acc.insert(f, ctn, costs);
         }
 
         Ok(acc)
diff --git a/src/crates/candelabra/src/types.rs b/src/crates/candelabra/src/types.rs
new file mode 100644
index 0000000..d16cc36
--- /dev/null
+++ b/src/crates/candelabra/src/types.rs
@@ -0,0 +1,149 @@
+use camino::Utf8PathBuf;
+use serde::{Deserialize, Serialize};
+
+/// Names a container type we want to select.
+pub type ConTypeName = String;
+
+/// Name of a container implementation we are considering
+pub type ImplName = String;
+
+/// Mapping from container type to some other type.
+/// This is super common so we have an alias for it
+pub type ConTypeTo<C> = Mapping2D<Utf8PathBuf, ConTypeName, C>;
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct Mapping2D<A, B, C>(Vec<(A, Vec<(B, C)>)>);
+
+impl<A: Clone + Eq, B: Clone + Eq, C> Mapping2D<A, B, C> {
+    fn with_a_or_create(&mut self, a: &A) -> &mut Vec<(B, C)> {
+        for (i, (ka, _)) in self.0.iter().enumerate() {
+            if ka == a {
+                return &mut self.0[i].1;
+            }
+        }
+
+        self.0.push((a.clone(), vec![]));
+        &mut self.0.last_mut().unwrap().1
+    }
+
+    pub fn get(&self, a: &A, b: &B) -> Option<&C> {
+        self.0
+            .iter()
+            .find(|(ka, _)| ka == a)
+            .and_then(|(_, bs)| bs.iter().find(|(kb, _)| kb == b).map(|(_, c)| c))
+    }
+
+    pub fn get_mut(&mut self, a: &A, b: &B) -> Option<&mut C> {
+        self.0
+            .iter_mut()
+            .find(|(ka, _)| ka == a)
+            .and_then(|(_, bs)| bs.iter_mut().find(|(kb, _)| kb == b).map(|(_, c)| c))
+    }
+
+    pub fn insert(&mut self, a: &A, b: &B, c: C) {
+        let bs = self.with_a_or_create(a);
+        if let Some((_, cv)) = bs.iter_mut().find(|(kb, _)| kb == b) {
+            *cv = c;
+        } else {
+            bs.push((b.clone(), c));
+        }
+    }
+
+    pub fn entry_mut<'a>(&'a mut self, a: &A, b: &'a B) -> EntryMut<'a, B, C> {
+        let as_ = self.with_a_or_create(a);
+        match as_.iter().position(|(kb, _)| kb == b) {
+            Some(idx) => EntryMut::Occupied(&mut as_[idx].1),
+            None => EntryMut::Vacant(b, as_),
+        }
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (&A, &B, &C)> {
+        self.0
+            .iter()
+            .flat_map(|(a, bs)| bs.iter().map(move |(b, c)| (a, b, c)))
+    }
+
+    pub fn keys(&self) -> impl Iterator<Item = (&A, &B)> {
+        self.0
+            .iter()
+            .flat_map(|(a, bs)| bs.iter().map(move |(b, _)| (a, b)))
+    }
+
+    pub fn iter_2d<'a>(
+        &'a self,
+    ) -> impl Iterator<Item = (&'a A, impl Iterator<Item = (&'a B, &'a C)>)> {
+        self.0
+            .iter()
+            .map(|(a, bs)| (a, bs.iter().map(|(b, c)| (b, c))))
+    }
+
+    pub fn into_iter_2d(self) -> impl Iterator<Item = (A, impl Iterator<Item = (B, C)>)> {
+        self.0
+            .into_iter()
+            .map(|(a, bs)| (a, bs.into_iter().map(|(b, c)| (b, c))))
+    }
+
+    pub(crate) fn iter_with_a(&self, target_a: &A) -> impl Iterator<Item = &(B, C)> {
+        self.0
+            .iter()
+            .find(|(a, _)| a == target_a)
+            .into_iter()
+            .flat_map(|(_, bs)| bs.iter())
+    }
+}
+
+impl<A: Clone, B, C> IntoIterator for Mapping2D<A, B, C> {
+    type Item = (A, B, C);
+
+    type IntoIter = impl Iterator<Item = (A, B, C)>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0
+            .into_iter()
+            .flat_map(|(a, bs)| bs.into_iter().map(move |(b, c)| (a.clone(), b, c)))
+    }
+}
+
+impl<A: Eq + Clone, B: Eq + Clone, C> FromIterator<(A, B, C)> for Mapping2D<A, B, C> {
+    fn from_iter<T: IntoIterator<Item = (A, B, C)>>(iter: T) -> Self {
+        let mut s = Self(vec![]);
+        for (a, b, c) in iter {
+            s.insert(&a, &b, c);
+        }
+        s
+    }
+}
+
+impl<A: Eq + Clone, B: Eq + Clone, C> FromIterator<((A, B), C)> for Mapping2D<A, B, C> {
+    fn from_iter<T: IntoIterator<Item = ((A, B), C)>>(iter: T) -> Self {
+        let mut s = Self(vec![]);
+        for ((a, b), c) in iter {
+            s.insert(&a, &b, c);
+        }
+        s
+    }
+}
+
+pub enum EntryMut<'a, B, C> {
+    Occupied(&'a mut C),
+    Vacant(&'a B, &'a mut Vec<(B, C)>),
+}
+
+impl<'a, B: Clone, C> EntryMut<'a, B, C> {
+    pub fn and_modify(&mut self, f: impl FnOnce(&mut C) -> ()) -> &mut Self {
+        match self {
+            EntryMut::Occupied(c) => f(*c),
+            EntryMut::Vacant(_, _) => (),
+        };
+
+        self
+    }
+    pub fn or_insert(&mut self, new: C) -> &mut Self {
+        match self {
+            EntryMut::Occupied(c) => **c = new,
+            EntryMut::Vacant(b, bs) => bs.push(((**b).clone(), new)),
+        };
+
+        self
+    }
+}