2 files changed, 62 insertions, 38 deletions

diff --git a/src/crates/cli/src/profiler/mod.rs b/src/crates/cli/src/profiler/mod.rs
index 63c0ff4..d7e7c6b 100644
--- a/src/crates/cli/src/profiler/mod.rs
+++ b/src/crates/cli/src/profiler/mod.rs
@@ -55,18 +55,10 @@ impl State {
 
         let chosen = candidates
             .iter()
-            .map(|(dest_name, impls)| {
-                (
-                    dest_name,
-                    format!(
-                        "::primrose_traits::profiler::ProfilerWrapper::<{}>",
-                        &impls[0]
-                    ),
-                )
-            })
+            .map(|(dest_name, impls)| (dest_name, &impls[0]))
             .collect::<Vec<_>>();
 
-        let new_code = selector.gen_replacement_file(chosen.iter().map(|(d, c)| (*d, c.as_str())));
+        let new_code = selector.gen_profiling_file(chosen.iter().map(|(d, c)| (*d, c.as_str())));
 
         let new_path = file.to_string().replace(".pr", "");
 
@@ -82,6 +74,11 @@ impl State {
 
     fn profile_benchmark(&self, project: &Project, name: &str) -> Result<ProfilerInfo> {
         let profiler_out_dir = tempdir()?;
+        debug!(
+            "Running benchmark {} with out dir {:?}",
+            name, profiler_out_dir
+        );
+
         let output = Command::new("cargo")
             .current_dir(&project.source_dir)
             .args(["bench", "--bench", name])
@@ -95,7 +92,7 @@ impl State {
         }
 
         let mut info = ProfilerInfo::default();
-        for file in read_dir(profiler_out_dir)? {
+        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)?;
diff --git a/src/crates/primrose/src/codegen.rs b/src/crates/primrose/src/codegen.rs
index 206f034..e60cd46 100644
--- a/src/crates/primrose/src/codegen.rs
+++ b/src/crates/primrose/src/codegen.rs
@@ -19,56 +19,83 @@ impl ContainerSelector {
         selections: T,
     ) -> String {
         let mut result = String::new();
-
-        result += CODEGEN;
-        result += &self.bounds_decl;
-        result += CODEGENEND;
-
+        result += &codegen_block(&self.bounds_decl);
         for (tag_id, selection) in selections {
             result += &self.gen_replacement_code(tag_id, selection);
         }
 
         // rest of the rust code, minus the spec stuff
-        for block in self.blocks.iter().filter(|block| block.is_code_block()) {
-            match block {
-                Block::CodeBlock(code, _) => result += code,
-                _ => unreachable!(),
-            };
+        add_all_code(&mut result, self.blocks.iter());
+        result
+    }
+
+    /// 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)>>(
+        &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);
         }
 
+        // rest of the rust code, minus the spec stuff
+        add_all_code(&mut result, self.blocks.iter());
         result
     }
-
     /// Generate replacement code for the given tag, choosing the given library spec.
     /// 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_code(&self, tag_id: &String, selection: &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");
         };
 
-        // generated code at top
-        let mut result = String::new();
-        result += CODEGEN;
-        result += &gen_output_code(tag_id, elem_ty, selection);
-        result += CODEGENEND;
-
-        result
+        codegen_block(&format!(
+            r#"
+#[allow(non_snake_case)]
+fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
+    {selection}::<{elem_ty}>::default()
+}}
+"#
+        ))
     }
-}
 
-// TODO: Constructing a box like this is inefficient, and might affect performance of some programs
-fn gen_output_code(s: &str, elem_type: &str, chosen: &str) -> String {
-    format!(
-        r#"
+    /// 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 {
+        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");
+        };
+
+        codegen_block(&format!(
+            r#"
 #[allow(non_snake_case)]
-fn _{s}<{elem_type}: PartialEq>() -> {s}<{elem_type}> {{
-    {chosen}::<{elem_type}>::default()
+fn _{tag_id}<{elem_ty}: PartialEq>() -> {tag_id}<{elem_ty}> {{
+    ::primrose_library::ProfilerWrapper::<{inner}<{elem_ty}>>::default()
 }}
 "#
-    )
+        ))
+    }
+}
+
+fn codegen_block(x: &str) -> String {
+    format!("{}{}{}", CODEGEN, x, CODEGENEND)
+}
+
+fn add_all_code<'a>(result: &mut String, blocks: impl Iterator<Item = &'a Block>) {
+    for block in blocks.filter(|block| block.is_code_block()) {
+        match block {
+            Block::CodeBlock(code, _) => *result += code,
+            _ => unreachable!(),
+        };
+    }
 }
 
 pub fn process_bound_decl(ctx: &InforMap) -> String {