refactor out nested hashmaps for mapping2d struct

13 files changed, 323 insertions, 158 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
+    }
+}
diff --git a/src/crates/cli/src/candidates.rs b/src/crates/cli/src/candidates.rs
index 764f45f..6d76da3 100644
--- a/src/crates/cli/src/candidates.rs
+++ b/src/crates/cli/src/candidates.rs
@@ -2,7 +2,7 @@ use anyhow::Result;
 use argh::FromArgs;
 use log::info;
 
-use crate::State;
+use crate::{display::print_table, State};
 
 /// List selection sites and their candidates
 #[derive(FromArgs)]
@@ -15,14 +15,15 @@ impl State {
             info!("Processing project {}", &proj.name);
 
             let candidates = self.inner.project_candidate_list(proj)?;
-            for (file, sites) in candidates {
-                info!("{}:", file);
-                for (name, candidates) in sites {
-                    info!("  {}:", name);
-                    for candidate in candidates {
-                        info!("    {}", candidate);
-                    }
-                }
+            for (file, con_types) in candidates.iter_2d() {
+                print_table(con_types.flat_map(|(ctn, candidates)| {
+                    candidates.iter().map(move |v| {
+                        (
+                            ctn,
+                            [("implementation", v.as_str()), ("file", file.as_str())],
+                        )
+                    })
+                }))
             }
         }
         Ok(())
diff --git a/src/crates/cli/src/display.rs b/src/crates/cli/src/display.rs
index 2debede..b48e370 100644
--- a/src/crates/cli/src/display.rs
+++ b/src/crates/cli/src/display.rs
@@ -1,6 +1,6 @@
 use std::{collections::HashSet, fmt::Display, hash::Hash, iter::once};
 
-use candelabra::profiler::ProfilerInfo;
+use candelabra::profiler::UsageProfile;
 use tabled::{builder::Builder, settings::Style};
 
 // Print the given 2D map as a table, where the first key is the left-most column, and the second key the column index
@@ -42,8 +42,8 @@ where
     println!("{}", builder.build().with(Style::sharp()));
 }
 
-pub fn display_profiler_info(profile_info: ProfilerInfo) {
-    print_table(profile_info.0.into_iter().enumerate().map(|(i, p)| {
+pub fn display_usage_info(profile_info: &UsageProfile) {
+    print_table(profile_info.0.iter().enumerate().map(|(i, p)| {
         (
             i,
             [
@@ -51,7 +51,7 @@ pub fn display_profiler_info(profile_info: ProfilerInfo) {
                 ("occurences".to_string(), p.occurences),
             ]
             .into_iter()
-            .chain(p.avg_op_counts),
+            .chain(p.avg_op_counts.clone()),
         )
     }))
 }
diff --git a/src/crates/cli/src/estimate.rs b/src/crates/cli/src/estimate.rs
index eafe101..be5dc5c 100644
--- a/src/crates/cli/src/estimate.rs
+++ b/src/crates/cli/src/estimate.rs
@@ -1,7 +1,9 @@
-use std::{collections::HashMap, iter::once};
+use std::iter::once;
 
 use anyhow::{anyhow, bail, Result};
 use argh::FromArgs;
+use candelabra::{select::Selection, types::Mapping2D};
+use cargo_metadata::camino::Utf8PathBuf;
 use log::info;
 
 use crate::{display::print_table, State};
@@ -10,7 +12,7 @@ use crate::{display::print_table, State};
 #[derive(FromArgs)]
 #[argh(subcommand, name = "estimate")]
 pub struct Args {
-    /// the assignments, written <type name> <implementation> ...
+    /// the assignments, written <file path> <type name> <implementation> ...
     #[argh(positional, greedy)]
     assignments: Vec<String>,
 
@@ -27,69 +29,70 @@ impl State {
         let proj = &self.projects[0];
 
         // parse assignment list
-        if args.assignments.len() % 2 != 0 {
-            bail!("assignments list must have even length");
+        if args.assignments.len() % 3 != 0 {
+            bail!("assignments list must have length divisible by 3");
         }
         let assignments = args
             .assignments
             .iter()
-            .step_by(2)
-            .zip(args.assignments.iter().skip(1).step_by(2))
-            .collect::<HashMap<_, _>>();
+            .step_by(3)
+            .map(|p| Utf8PathBuf::from(p))
+            .zip(args.assignments.iter().skip(1).step_by(3))
+            .zip(args.assignments.iter().skip(2).step_by(3))
+            .collect::<Mapping2D<_, _, _>>();
 
         info!("Using assignments: {:?}", &assignments);
 
         // get breakdown by operation for each assignment
         info!("Summary breakdown:");
 
-        let profile_info = self.inner.profiler_info(proj)?;
+        let profile_info = self.inner.usage_profile(proj)?;
         let cost_models = profile_info
             .keys()
-            .map(|ctn| {
+            .map(|(file, ctn)| {
                 (
+                    file,
                     ctn,
                     assignments
-                        .get(ctn)
+                        .get(file, &ctn)
                         .ok_or_else(|| anyhow!("missing assignment for {}", ctn)),
                 )
             })
-            .map(|(ctn, impl_name)| Ok((ctn, self.inner.cost_model(impl_name?)?)))
-            .collect::<Result<HashMap<_, _>>>()?;
+            .map(|(f, ctn, impl_name)| Ok((f, ctn, self.inner.cost_model(impl_name?)?)))
+            .collect::<Result<Mapping2D<_, _, _>>>()?;
 
-        let mut acc = HashMap::new();
-        for con_type_name in assignments.keys() {
-            let profiler = profile_info
-                .get(*con_type_name)
-                .ok_or_else(|| anyhow!("no profiling info for {} - wrong name?", con_type_name))?;
+        let mut acc = Mapping2D::default();
+        for (f, ctn) in assignments.keys() {
+            let ctn_alias = format!("{} / {}", f, ctn);
 
-            let cost_model = cost_models.get(con_type_name).unwrap();
+            let cost_model = cost_models.get(&f, ctn).unwrap();
+            let profiler = profile_info
+                .get(f, *ctn)
+                .ok_or_else(|| anyhow!("no profiling info for {} / {} - wrong name?", f, ctn))?;
 
             let breakdown = profiler.cost_breakdown(&cost_model);
-            acc.insert(
-                *con_type_name,
-                breakdown
-                    .into_iter()
-                    .map(|(n, i)| (n.to_string(), i))
-                    .collect::<HashMap<_, _>>(),
-            );
+            for (op, i) in breakdown {
+                acc.insert(&ctn_alias, op, i);
+            }
         }
 
-        print_table(acc);
+        print_table(acc.into_iter_2d());
 
         info!("Breakdown by split");
         profile_info
             .iter()
-            .map(|(ctn, p)| (ctn, cost_models.get(ctn).unwrap(), p))
-            .map(|(ctn, cost_model, p)| {
+            .map(|(f, ctn, p)| (f, ctn, cost_models.get(&f, &ctn).unwrap(), p))
+            .map(|(f, ctn, cost_model, p)| {
                 (
+                    f,
                     ctn,
                     p.0.iter().enumerate().map(|(i, partition)| {
                         (i, partition.cost_breakdown(cost_model).into_iter())
                     }),
                 )
             })
-            .for_each(|(ctn, table)| {
-                info!("{}:", ctn);
+            .for_each(|(f, ctn, table)| {
+                info!("{} / {}:", f, ctn);
                 print_table(table);
             });
 
@@ -101,7 +104,15 @@ impl State {
 
         print_table(
             self.inner
-                .run_benchmarks_with(proj, &assignments)?
+                .run_benchmarks_with(
+                    proj,
+                    &assignments
+                        .into_iter()
+                        .map(|(f, ctn, sel)| {
+                            (f, ctn.to_string(), Selection::Singular(sel.to_string()))
+                        })
+                        .collect(),
+                )?
                 .into_iter()
                 .map(|(name, results)| (name, once(("avg", results.avg)))),
         );
diff --git a/src/crates/cli/src/model.rs b/src/crates/cli/src/model.rs
index b1dcb67..d6824b6 100644
--- a/src/crates/cli/src/model.rs
+++ b/src/crates/cli/src/model.rs
@@ -36,21 +36,6 @@ impl State {
 
         println!("{}", builder.build().with(Style::sharp()));
 
-        // Table of example cost estimates
-        let mut builder = Builder::default();
-        builder.set_header(["op", "n = 8", "n = 64", "n = 512", "n = 4096"]);
-        for (k, v) in model.by_op.iter() {
-            builder.push_record(&[
-                k.to_string(),
-                format!("{0:.3}", v.estimate(8)),
-                format!("{0:.3}", v.estimate(64)),
-                format!("{0:.3}", v.estimate(512)),
-                format!("{0:.3}", v.estimate(4096)),
-            ]);
-        }
-
-        println!("{}", builder.build().with(Style::sharp()));
-
         Ok(())
     }
 }
diff --git a/src/crates/cli/src/profile.rs b/src/crates/cli/src/profile.rs
index efb1571..af437ba 100644
--- a/src/crates/cli/src/profile.rs
+++ b/src/crates/cli/src/profile.rs
@@ -2,7 +2,7 @@ use anyhow::Result;
 use argh::FromArgs;
 use log::info;
 
-use crate::{display::display_profiler_info, State};
+use crate::{display::display_usage_info, State};
 
 /// Profile the selected projects and print the results
 #[derive(FromArgs)]
@@ -14,11 +14,11 @@ impl State {
         for proj in self.projects.iter() {
             info!("Processing project {}", &proj.name);
 
-            let all_info = self.inner.profiler_info(proj)?;
+            let all_info = self.inner.usage_profile(proj)?;
 
-            for (con_type_name, profile_info) in all_info {
-                info!("{}:", con_type_name);
-                display_profiler_info(profile_info);
+            for (f, ctn, usage_profile) in all_info.iter() {
+                info!("{} / {}:", f, ctn);
+                display_usage_info(usage_profile);
             }
         }
         Ok(())
diff --git a/src/crates/cli/src/select.rs b/src/crates/cli/src/select.rs
index bb613f5..740b988 100644
--- a/src/crates/cli/src/select.rs
+++ b/src/crates/cli/src/select.rs
@@ -26,8 +26,8 @@ impl State {
             let costs = self.inner.rank_candidates(proj).unwrap();
 
             let mut builder = Builder::default();
-            builder.set_header(["name", "implementation", "estimated cost"]);
-            for (con_type_name, candidates) in costs.iter() {
+            builder.set_header(["name", "implementation", "estimated cost", "file"]);
+            for (f, ctn, candidates) in costs.iter() {
                 for (candidate, cost) in candidates.iter() {
                     let name = match candidate {
                         Selection::Singular(x) => x.to_string(),
@@ -39,9 +39,10 @@ impl State {
                         }
                     };
                     builder.push_record([
-                        con_type_name.as_str(),
+                        ctn.as_str(),
                         name.as_str(),
                         cost.to_string().as_str(),
+                        f.as_str(),
                     ]);
                 }
             }
@@ -57,15 +58,24 @@ impl State {
                 .inner
                 .project_candidate_list(proj)?
                 .into_iter()
-                .flat_map(|(_, v)| v.into_iter())
+                .map(|(f, ctn, v)| ((f, ctn), v))
                 .collect();
             let possible_assignments = nary_cartesian_product(&candidates);
+
             let mut assignments_results = HashMap::new();
             for assignment in possible_assignments.iter() {
                 info!("Running benchmarks with {:?}", &assignment);
                 assignments_results.insert(
                     format!("{:?}", &assignment),
-                    self.inner.run_benchmarks_with(proj, assignment)?,
+                    self.inner.run_benchmarks_with(
+                        proj,
+                        &assignment
+                            .into_iter()
+                            .map(|((f, ctn), i)| {
+                                (f.clone(), ctn.clone(), Selection::Singular(i.to_string()))
+                            })
+                            .collect(),
+                    )?,
                 );
             }