feat(profiler): be aware of multiple container types per project

4 files changed, 110 insertions, 57 deletions

diff --git a/src/crates/candelabra/src/profiler.rs b/src/crates/candelabra/src/profiler.rs
index b52a96b..9f56f54 100644
--- a/src/crates/candelabra/src/profiler.rs
+++ b/src/crates/candelabra/src/profiler.rs
@@ -5,6 +5,7 @@ use camino::{Utf8Path, Utf8PathBuf};
 use log::{debug, trace, warn};
 use primrose::ContainerSelector;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use std::io::Write;
 use std::str::FromStr;
 use std::{
@@ -15,6 +16,7 @@ use std::{
 use tempfile::tempdir;
 
 use crate::cache::{gen_tree_hash, FileCache};
+use crate::candidates::ConTypeName;
 use crate::project::Project;
 use crate::{Paths, State};
 
@@ -22,7 +24,7 @@ use crate::{Paths, State};
 pub(crate) struct CacheEntry {
     proj_hash: u64,
     proj_location: Utf8PathBuf,
-    info: ProfilerInfo,
+    info: HashMap<ConTypeName, ProfilerInfo>,
 }
 
 /// The information we get from profiling.
@@ -57,7 +59,7 @@ 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<ProfilerInfo> {
+    pub fn profiler_info(&self, project: &Project) -> Result<HashMap<ConTypeName, ProfilerInfo>> {
         match self.profiler_info_cache.find(&project.name)? {
             Some(x) => Ok(x.info),
             None => {
@@ -82,24 +84,34 @@ impl State {
     }
 
     /// Calculate profiler info for the given project.
-    fn calc_profiler_info(&self, project: &Project) -> Result<ProfilerInfo> {
-        self.project_profiling_prep(project)?;
-        let mut acc = Vec::new();
+    fn calc_profiler_info(&self, project: &Project) -> Result<HashMap<ConTypeName, ProfilerInfo>> {
+        let candidate_list = self.project_candidate_list(project)?;
+        let con_types = candidate_list
+            .iter()
+            .flat_map(|(_, con_types)| con_types.iter())
+            .map(|(id, _)| id)
+            .collect::<Vec<_>>();
+
+        self.project_profiling_prep(project, &con_types)?;
+        let mut acc = HashMap::new();
         for name in project.benchmarks.iter() {
-            acc.extend(
-                self.profile_benchmark(project, name)
-                    .with_context(|| format!("Error profiling benchmark {}", name))?
-                    .0,
-            );
+            for (con_type, new_results) in self
+                .profile_benchmark(project, name, &con_types)
+                .with_context(|| format!("Error profiling benchmark {}", name))?
+            {
+                acc.entry(con_type)
+                    .and_modify(|pi: &mut ProfilerInfo| pi.0.extend(new_results.0.iter().cloned()))
+                    .or_insert(new_results);
+            }
         }
 
-        Ok(ProfilerInfo(acc))
+        Ok(acc)
     }
 
     /// Prepare the given project to be profiled, by replacing all candidate types with the profiler wrapper.
-    fn project_profiling_prep(&self, project: &Project) -> Result<()> {
+    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)
+            self.file_profiling_prep(&file, &candidates, con_types)
                 .with_context(|| format!("error preparing {} for profiling", file))?;
         }
 
@@ -111,6 +123,7 @@ impl State {
         &self,
         file: &Utf8Path,
         candidates: &[(String, Vec<String>)],
+        con_types: &[&String],
     ) -> Result<()> {
         debug!("Setting up {} for profiling", file);
 
@@ -126,7 +139,13 @@ impl State {
             .map(|(dest_name, impls)| (dest_name, &impls[0]))
             .collect::<Vec<_>>();
 
-        let new_code = selector.gen_profiling_file(chosen.iter().map(|(d, c)| (*d, c.as_str())));
+        let new_code = selector.gen_profiling_file(chosen.iter().map(|(d, c)| {
+            (
+                *d,
+                con_types.iter().position(|id| id == d).unwrap(),
+                c.as_str(),
+            )
+        }));
 
         let new_path = file.to_string().replace(".pr", "");
 
@@ -141,7 +160,12 @@ impl State {
     }
 
     /// Run the given benchmark on the project, and parse the resulting profiling information.
-    fn profile_benchmark(&self, project: &Project, name: &str) -> Result<ProfilerInfo> {
+    fn profile_benchmark(
+        &self,
+        project: &Project,
+        name: &str,
+        con_types: &[&String],
+    ) -> Result<HashMap<String, ProfilerInfo>> {
         let profiler_out_dir = tempdir()?;
         debug!(
             "Running benchmark {} with out dir {:?}",
@@ -160,16 +184,29 @@ impl State {
             bail!("Error running benchmark");
         }
 
-        let mut acc = Vec::default();
-        for file in read_dir(&profiler_out_dir)? {
-            let file = file?;
-            let mut contents = String::new();
-            File::open(file.path())?.read_to_string(&mut contents)?;
+        let mut con_type_results = HashMap::new();
+        for dir in read_dir(&profiler_out_dir)? {
+            let dir = dir?;
+            let con_type: String = con_types[dir
+                .file_name()
+                .into_string()
+                .unwrap()
+                .parse::<usize>()
+                .unwrap()]
+            .to_string();
+            let mut acc = Vec::default();
+            for file in read_dir(dir.path())? {
+                let file = file?;
+                let mut contents = String::new();
+                File::open(file.path())?.read_to_string(&mut contents)?;
+
+                acc.push(parse_output(&contents)?);
+            }
 
-            acc.push(parse_output(&contents)?);
+            con_type_results.insert(con_type, ProfilerInfo(acc));
         }
 
-        Ok(ProfilerInfo(acc))
+        Ok(con_type_results)
     }
 }
 
diff --git a/src/crates/cli/src/profile.rs b/src/crates/cli/src/profile.rs
index cc74db1..4568e5b 100644
--- a/src/crates/cli/src/profile.rs
+++ b/src/crates/cli/src/profile.rs
@@ -15,28 +15,30 @@ impl State {
         for proj in self.projects.iter() {
             info!("Processing project {}", &proj.name);
 
-            let info = self.inner.profiler_info(proj)?;
+            let all_info = self.inner.profiler_info(proj)?;
 
-            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(),
+            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()));
             }
-            println!("{}", builder.build().with(Style::sharp()));
         }
         Ok(())
     }
diff --git a/src/crates/library/src/profiler.rs b/src/crates/library/src/profiler.rs
index 2b6fc43..952d0a0 100644
--- a/src/crates/library/src/profiler.rs
+++ b/src/crates/library/src/profiler.rs
@@ -1,8 +1,13 @@
-use std::{env::var, fs::File, io::Write, time::SystemTime};
+use std::{
+    env::var,
+    fs::{create_dir, metadata, File},
+    io::Write,
+    time::SystemTime,
+};
 
 use crate::traits::{Container, Indexable, Stack};
 
-pub struct ProfilerWrapper<T> {
+pub struct ProfilerWrapper<const ID: usize, T> {
     inner: T,
     sum_ns: usize,
     n_contains: usize,
@@ -16,7 +21,7 @@ pub struct ProfilerWrapper<T> {
     n_pop: usize,
 }
 
-impl<T: Default> Default for ProfilerWrapper<T> {
+impl<const ID: usize, T: Default> Default for ProfilerWrapper<ID, T> {
     fn default() -> Self {
         Self {
             inner: T::default(),
@@ -34,7 +39,7 @@ impl<T: Default> Default for ProfilerWrapper<T> {
     }
 }
 
-impl<T: Container<E>, E> Container<E> for ProfilerWrapper<T> {
+impl<const ID: usize, T: Container<E>, E> Container<E> for ProfilerWrapper<ID, T> {
     fn len(&mut self) -> usize {
         self.inner.len()
     }
@@ -68,7 +73,7 @@ impl<T: Container<E>, E> Container<E> for ProfilerWrapper<T> {
     }
 }
 
-impl<T: Indexable<E> + Container<E>, E> Indexable<E> for ProfilerWrapper<T> {
+impl<const ID: usize, T: Indexable<E> + Container<E>, E> Indexable<E> for ProfilerWrapper<ID, T> {
     fn first(&mut self) -> Option<&E> {
         self.sum_ns += self.inner.len();
         self.n_first += 1;
@@ -88,7 +93,7 @@ impl<T: Indexable<E> + Container<E>, E> Indexable<E> for ProfilerWrapper<T> {
     }
 }
 
-impl<T: Stack<E> + Container<E>, E> Stack<E> for ProfilerWrapper<T> {
+impl<const ID: usize, T: Stack<E> + Container<E>, E> Stack<E> for ProfilerWrapper<ID, T> {
     fn push(&mut self, elt: E) {
         self.sum_ns += self.inner.len();
         self.n_push += 1;
@@ -102,14 +107,18 @@ impl<T: Stack<E> + Container<E>, E> Stack<E> for ProfilerWrapper<T> {
     }
 }
 
-impl<T> Drop for ProfilerWrapper<T> {
+impl<const ID: usize, T> Drop for ProfilerWrapper<ID, T> {
     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 mut f = File::create(format!("{}/{}", base_dir, unix_time.as_nanos())).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.sum_ns).unwrap();
         writeln!(f, "{}", self.n_contains).unwrap();
diff --git a/src/crates/primrose/src/codegen.rs b/src/crates/primrose/src/codegen.rs
index e60cd46..d72b0db 100644
--- a/src/crates/primrose/src/codegen.rs
+++ b/src/crates/primrose/src/codegen.rs
@@ -32,14 +32,14 @@ impl ContainerSelector {
     /// Generate replacement code for profiling the whole file, with the given `(tag_id, inner)` pairs.
     /// This will generate invalid code if any selection is invalid, or panic if any `tag_id` is invalid.
     /// Returns the original file with the generated code in place of the original specification.
-    pub fn gen_profiling_file<'a, T: Iterator<Item = (&'a String, &'a str)>>(
+    pub fn gen_profiling_file<'a, T: Iterator<Item = (&'a String, usize, &'a str)>>(
         &self,
         inners: T,
     ) -> String {
         let mut result = String::new();
         result += &codegen_block(&self.bounds_decl);
-        for (tag_id, selection) in inners {
-            result += &self.gen_replacement_profiling_code(tag_id, selection);
+        for (tag_id, id, selection) in inners {
+            result += &self.gen_replacement_profiling_code(tag_id, id, selection);
         }
 
         // rest of the rust code, minus the spec stuff
@@ -58,7 +58,7 @@ impl ContainerSelector {
         codegen_block(&format!(
             r#"
 #[allow(non_snake_case)]
-fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
+fn _{tag_id}<{elem_ty}: PartialEq + Ord>() -> {tag_id}<{elem_ty}> {{
     {selection}::<{elem_ty}>::default()
 }}
 "#
@@ -68,7 +68,12 @@ fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
     /// Generate replacement code for profiling the given tag, with the given inner implementation.
     /// This will generate invalid code if selection is invalid, or panic if `tag_id` is invalid.
     /// Returns only the required code, and not the rest of the code originally in the file.
-    pub fn gen_replacement_profiling_code(&self, tag_id: &String, inner: &str) -> String {
+    pub fn gen_replacement_profiling_code(
+        &self,
+        tag_id: &String,
+        id: usize,
+        inner: &str,
+    ) -> String {
         let tag = self.analyser.ctx().get(tag_id).expect("invalid tag_id");
         let Tag::Con(elem_ty, _, _) = tag else {
             panic!("tag_id was not Tag::Con");
@@ -77,8 +82,8 @@ fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
         codegen_block(&format!(
             r#"
 #[allow(non_snake_case)]
-fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
-    ::primrose_library::ProfilerWrapper::<{inner}<{elem_ty}>>::default()
+fn _{tag_id}<{elem_ty}: PartialEq + Ord>() -> {tag_id}<{elem_ty}> {{
+    ::primrose_library::ProfilerWrapper::<{id}, {inner}<{elem_ty}>>::default()
 }}
 "#
         ))
@@ -120,7 +125,7 @@ pub fn process_bound_decl(ctx: &InforMap) -> String {
 fn gen_trait_code(typ_name: &str, elem_type: &str, traits: &str) -> String {
     format!(
         r#"
-pub type {typ_name}<{elem_type}: PartialEq> = impl {traits} + Default;
+pub type {typ_name}<{elem_type}: PartialEq + Ord> = impl {traits} + Default;
 "#
     )
 }