fix cli tables getting messed up by key differences

2 files changed, 26 insertions, 24 deletions

diff --git a/src/crates/cli/src/estimate.rs b/src/crates/cli/src/estimate.rs
index 0365388..9618fb6 100644
--- a/src/crates/cli/src/estimate.rs
+++ b/src/crates/cli/src/estimate.rs
@@ -34,9 +34,11 @@ impl State {
             .assignments
             .iter()
             .step_by(2)
-            .zip(args.assignments.iter().skip(1))
+            .zip(args.assignments.iter().skip(1).step_by(2))
             .collect::<HashMap<_, _>>();
 
+        info!("Using assignments: {:?}", &assignments);
+
         // get breakdown by operation for each assignment
         info!("Contribution to cost by operation:");
 
@@ -49,7 +51,7 @@ impl State {
             let cost_model = self.inner.cost_model(&impl_name)?;
             let breakdown = profiler.cost_breakdown(&cost_model);
             acc.insert(
-                *impl_name,
+                *con_type_name,
                 breakdown
                     .into_iter()
                     .map(|(n, i)| (n.to_string(), i))
diff --git a/src/crates/cli/src/util.rs b/src/crates/cli/src/util.rs
index 235771f..499f4ba 100644
--- a/src/crates/cli/src/util.rs
+++ b/src/crates/cli/src/util.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, iter::once};
+use std::{collections::HashSet, fmt::Display, hash::Hash, iter::once};
 
 use tabled::{builder::Builder, settings::Style};
 
@@ -8,34 +8,34 @@ where
     I1: IntoIterator<Item = (A, I2)>,
     I2: IntoIterator<Item = (B, C)>,
     A: Display,
-    B: Display,
+    B: Display + Eq + Hash,
     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
-    };
+    // Collect everything up
+    let map = map
+        .into_iter()
+        .map(|(k, v)| (k, v.into_iter().collect::<Vec<(_, _)>>()))
+        .collect::<Vec<(_, _)>>();
 
-    let r1_vals = r1_vals.into_iter().collect::<Vec<_>>();
-    let mut header = vec!["".to_string()];
-    r1_vals
+    // Get keys
+    let keys = map
         .iter()
-        .for_each(|(key, _)| header.push(key.to_string()));
+        .flat_map(|(_, vs)| vs.iter().map(|(k, _)| k))
+        .collect::<HashSet<_>>() // dedup
+        .into_iter()
+        .collect::<Vec<_>>(); // consistent order
 
     let mut builder = Builder::new();
-    builder.set_header(header);
+    builder.set_header(once("".to_string()).chain(keys.iter().map(|s| s.to_string())));
 
-    // 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())),
-        );
+    // Add rows
+    for (key, vals) in map.iter() {
+        builder.push_record(once(key.to_string()).chain(keys.iter().map(|k| {
+            vals.iter()
+                .find(|(search, _)| search == *k)
+                .map(|(_, v)| v.to_string())
+                .unwrap_or("".to_string())
+        })));
     }
 
     println!("{}", builder.build().with(Style::sharp()));