use hashmaps in more places for consistency

5 files changed, 122 insertions, 128 deletions

diff --git a/src/crates/candelabra/src/profiler.rs b/src/crates/candelabra/src/profiler.rs
index b6e010c..f12914b 100644
--- a/src/crates/candelabra/src/profiler.rs
+++ b/src/crates/candelabra/src/profiler.rs
@@ -17,6 +17,7 @@ use tempfile::tempdir;
 
 use crate::cache::{gen_tree_hash, FileCache};
 use crate::candidates::ConTypeName;
+use crate::cost::benchmark::OpName;
 use crate::cost::{Cost, CostModel};
 use crate::project::Project;
 use crate::{Paths, State};
@@ -33,49 +34,13 @@ pub(crate) struct CacheEntry {
 pub struct ProfilerInfo(pub Vec<CollectionLifetime>);
 
 /// Breakdown of a cost value
-#[derive(Clone, Debug)]
-pub struct CostBreakdown {
-    pub contains: Cost,
-    pub insert: Cost,
-    pub clear: Cost,
-    pub remove: Cost,
-    pub first: Cost,
-    pub last: Cost,
-    pub nth: Cost,
-    pub push: Cost,
-    pub pop: Cost,
-}
+pub type CostBreakdown<'a> = HashMap<&'a OpName, Cost>;
 
 /// Profiler info collected from the lifetime of a single collection instance
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct CollectionLifetime {
     pub n: usize,
-    pub contains: usize,
-    pub insert: usize,
-    pub clear: usize,
-    pub remove: usize,
-    pub first: usize,
-    pub last: usize,
-    pub nth: usize,
-    pub push: usize,
-    pub pop: usize,
-}
-
-macro_rules! agg_op_cost {
-    ($this:ident, $cost_model:ident, $op:ident) => {
-        $this
-            .0
-            .iter()
-            .map(|cl| {
-                $cost_model
-                    .by_op
-                    .get(stringify!($op))
-                    .unwrap()
-                    .estimatef(cl.avg_n())
-                    * cl.$op as f64
-            })
-            .sum()
-    };
+    pub op_counts: HashMap<OpName, usize>,
 }
 
 impl ProfilerInfo {
@@ -85,48 +50,34 @@ impl ProfilerInfo {
         sum / self.0.len() as f64
     }
 
-    pub fn cost_breakdown(&self, cost_model: &CostModel) -> CostBreakdown {
-        CostBreakdown {
-            contains: agg_op_cost!(self, cost_model, contains),
-            insert: agg_op_cost!(self, cost_model, insert),
-            clear: agg_op_cost!(self, cost_model, clear),
-            remove: agg_op_cost!(self, cost_model, remove),
-            first: agg_op_cost!(self, cost_model, first),
-            last: agg_op_cost!(self, cost_model, last),
-            nth: agg_op_cost!(self, cost_model, nth),
-            push: agg_op_cost!(self, cost_model, push),
-            pop: agg_op_cost!(self, cost_model, pop),
-        }
+    pub fn cost_breakdown<'a>(&self, cost_model: &'a CostModel) -> CostBreakdown<'a> {
+        cost_model
+            .by_op
+            .iter()
+            .map(|(op_name, estimator)| {
+                (
+                    op_name,
+                    self.0
+                        .iter()
+                        .map(|cl| estimator.estimatef(cl.avg_n()) * cl.op_count(op_name) as f64)
+                        .sum::<f64>()
+                        / self.0.len() as f64,
+                )
+            })
+            .collect()
     }
 }
 
 impl CollectionLifetime {
     pub fn avg_n(&self) -> f64 {
-        self.n as f64
-            / (self.contains
-                + self.insert
-                + self.clear
-                + self.remove
-                + self.first
-                + self.last
-                + self.nth
-                + self.push
-                + self.pop) as f64
+        self.n as f64 / (self.op_counts.values().sum::<usize>() as f64)
     }
 
     pub fn op_count(&self, op: &str) -> usize {
-        match op {
-            "contains" => self.contains,
-            "insert" => self.insert,
-            "clear" => self.clear,
-            "remove" => self.remove,
-            "first" => self.first,
-            "last" => self.last,
-            "nth" => self.nth,
-            "push" => self.push,
-            "pop" => self.pop,
-            _ => panic!("invalid op passed to op_count: {}", op),
-        }
+        *self
+            .op_counts
+            .get(op)
+            .expect("invalid op passed to op_count")
     }
 
     pub fn estimate_cost(&self, cost_model: &CostModel) -> f64 {
@@ -313,17 +264,44 @@ impl State {
 fn parse_output(contents: &str) -> Result<CollectionLifetime> {
     let mut lines = contents.lines().map(usize::from_str);
     let missing_line_err = || anyhow!("wrong number of lines in ");
-
-    Ok(CollectionLifetime {
-        n: lines.next().ok_or_else(missing_line_err)??,
-        contains: lines.next().ok_or_else(missing_line_err)??,
-        insert: lines.next().ok_or_else(missing_line_err)??,
-        clear: lines.next().ok_or_else(missing_line_err)??,
-        remove: lines.next().ok_or_else(missing_line_err)??,
-        first: lines.next().ok_or_else(missing_line_err)??,
-        last: lines.next().ok_or_else(missing_line_err)??,
-        nth: lines.next().ok_or_else(missing_line_err)??,
-        push: lines.next().ok_or_else(missing_line_err)??,
-        pop: lines.next().ok_or_else(missing_line_err)??,
-    })
+    let n = lines.next().ok_or_else(missing_line_err)??;
+    let mut op_counts = HashMap::new();
+    op_counts.insert(
+        "contains".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "insert".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "clear".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "remove".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "first".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "last".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "nth".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "push".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+    op_counts.insert(
+        "pop".to_string(),
+        lines.next().ok_or_else(missing_line_err)??,
+    );
+
+    Ok(CollectionLifetime { n, op_counts })
 }
diff --git a/src/crates/cli/src/estimate.rs b/src/crates/cli/src/estimate.rs
index 27948b8..ad7df17 100644
--- a/src/crates/cli/src/estimate.rs
+++ b/src/crates/cli/src/estimate.rs
@@ -3,9 +3,8 @@ use std::collections::HashMap;
 use anyhow::{anyhow, bail, Result};
 use argh::FromArgs;
 use log::info;
-use tabled::{builder::Builder, settings::Style};
 
-use crate::State;
+use crate::{util::print_table, State};
 
 /// Estimate the cost of a given set of assignments for a specific project, and detail how that was reached.
 #[derive(FromArgs)]
@@ -38,32 +37,23 @@ impl State {
         info!("Contribution to cost by operation:");
 
         let profile_info = self.inner.profiler_info(proj)?;
-        let mut builder = Builder::default();
-        builder.set_header([
-            "con type", "contains", "insert", "clear", "remove", "first", "last", "nth", "push",
-            "pop",
-        ]);
+        let mut acc = HashMap::new();
         for (con_type_name, impl_name) in assignments.iter() {
             let profiler = profile_info
                 .get(*con_type_name)
                 .ok_or_else(|| anyhow!("no profiling info for {} - wrong name?", con_type_name))?;
             let cost_model = self.inner.cost_model(&impl_name)?;
             let breakdown = profiler.cost_breakdown(&cost_model);
-            builder.push_record([
+            acc.insert(
                 *impl_name,
-                &breakdown.contains.to_string(),
-                &breakdown.insert.to_string(),
-                &breakdown.clear.to_string(),
-                &breakdown.remove.to_string(),
-                &breakdown.first.to_string(),
-                &breakdown.last.to_string(),
-                &breakdown.nth.to_string(),
-                &breakdown.push.to_string(),
-                &breakdown.pop.to_string(),
-            ]);
+                breakdown
+                    .into_iter()
+                    .map(|(n, i)| (n.to_string(), i))
+                    .collect::<HashMap<_, _>>(),
+            );
         }
 
-        println!("{}", builder.build().with(Style::sharp()));
+        print_table(acc);
 
         Ok(())
     }
diff --git a/src/crates/cli/src/main.rs b/src/crates/cli/src/main.rs
index a63f5d5..81d484d 100644
--- a/src/crates/cli/src/main.rs
+++ b/src/crates/cli/src/main.rs
@@ -9,6 +9,7 @@ mod library;
 mod model;
 mod profile;
 mod select;
+mod util;
 
 #[derive(FromArgs)]
 /// Find the best performing container type using primrose
diff --git a/src/crates/cli/src/profile.rs b/src/crates/cli/src/profile.rs
index 4568e5b..0971771 100644
--- a/src/crates/cli/src/profile.rs
+++ b/src/crates/cli/src/profile.rs
@@ -1,9 +1,8 @@
 use anyhow::Result;
 use argh::FromArgs;
 use log::info;
-use tabled::{builder::Builder, settings::Style};
 
-use crate::State;
+use crate::{util::print_table, State};
 
 /// Profile the selected projects and print the results
 #[derive(FromArgs)]
@@ -17,28 +16,12 @@ impl State {
 
             let all_info = self.inner.profiler_info(proj)?;
 
-            for (con_type, info) in all_info {
-                let mut builder = Builder::default();
-                builder.set_header([
-                    "N (cum.)", "contains", "insert", "clear", "remove", "first", "last", "nth",
-                    "push", "pop",
-                ]);
-                for info in info.0.iter() {
-                    builder.push_record([
-                        info.n.to_string(),
-                        info.contains.to_string(),
-                        info.insert.to_string(),
-                        info.clear.to_string(),
-                        info.remove.to_string(),
-                        info.first.to_string(),
-                        info.last.to_string(),
-                        info.nth.to_string(),
-                        info.push.to_string(),
-                        info.pop.to_string(),
-                    ]);
-                }
-                println!("{}:\n{}", con_type, builder.build().with(Style::sharp()));
-            }
+            print_table(all_info.iter().map(|(con_type_name, profile_info)| {
+                (
+                    con_type_name,
+                    profile_info.0.iter().flat_map(|cl| cl.op_counts.iter()),
+                )
+            }));
         }
         Ok(())
     }
diff --git a/src/crates/cli/src/util.rs b/src/crates/cli/src/util.rs
new file mode 100644
index 0000000..235771f
--- /dev/null
+++ b/src/crates/cli/src/util.rs
@@ -0,0 +1,42 @@
+use std::{fmt::Display, iter::once};
+
+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
+pub fn print_table<I1, I2, A, B, C>(map: I1)
+where
+    I1: IntoIterator<Item = (A, I2)>,
+    I2: IntoIterator<Item = (B, C)>,
+    A: Display,
+    B: Display,
+    C: Display,
+{
+    // Get headers (keys of first row)
+    let mut map = map.into_iter();
+    let Some((r1_name, r1_vals)) = map.next() else {
+        return; // nothing to print
+    };
+
+    let r1_vals = r1_vals.into_iter().collect::<Vec<_>>();
+    let mut header = vec!["".to_string()];
+    r1_vals
+        .iter()
+        .for_each(|(key, _)| header.push(key.to_string()));
+
+    let mut builder = Builder::new();
+    builder.set_header(header);
+
+    // Push 1st row
+    builder.push_record(
+        once(r1_name.to_string()).chain(r1_vals.iter().map(|(_, val)| val.to_string())),
+    );
+
+    // Rest of rows
+    for (key, vals) in map {
+        builder.push_record(
+            once(key.to_string()).chain(vals.into_iter().map(|(_, val)| val.to_string())),
+        );
+    }
+
+    println!("{}", builder.build().with(Style::sharp()));
+}