generate code for selected adaptive containers

12 files changed, 133 insertions, 214 deletions

diff --git a/src/crates/candelabra/src/candidates.rs b/src/crates/candelabra/src/candidates.rs
index 7bc1031..09b6ec6 100644
--- a/src/crates/candelabra/src/candidates.rs
+++ b/src/crates/candelabra/src/candidates.rs
@@ -5,14 +5,14 @@ use std::{collections::HashMap, fs::metadata, time::SystemTime};
 use anyhow::{bail, Context, Result};
 use camino::Utf8Path;
 use log::{debug, warn};
-use primrose::ContainerSelector;
+use primrose::{ConTypeName, ContainerSelector, ImplName};
 use serde::{Deserialize, Serialize};
 
 use crate::{
     cache::{gen_tree_hash, FileCache},
     paths::Paths,
     project::Project,
-    types::{ConTypeName, ConTypeTo, ImplName},
+    types::ConTypeTo,
     State,
 };
 
diff --git a/src/crates/candelabra/src/confirmation.rs b/src/crates/candelabra/src/confirmation.rs
index 7a593da..95267c6 100644
--- a/src/crates/candelabra/src/confirmation.rs
+++ b/src/crates/candelabra/src/confirmation.rs
@@ -2,7 +2,7 @@ use crate::types::ConTypeTo;
 use anyhow::{Context, Result};
 use camino::Utf8Path;
 use log::{debug, trace};
-use primrose::ContainerSelector;
+use primrose::{ConTypeName, ContainerSelection, ContainerSelector};
 use std::{
     collections::HashMap,
     fs::File,
@@ -15,13 +15,11 @@ use crate::{
         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>;
+pub type Selections = ConTypeTo<ContainerSelection>;
 
 impl State {
     /// Run all benchmarks for the given project with the given container selections, returning the results
@@ -61,7 +59,7 @@ impl State {
         for (file, _) in self.project_candidate_list(proj)?.iter_2d() {
             let chosen = selections.iter_with_a(file);
 
-            self.save_choices_file(file, chosen)
+            self.save_choices_file(file, chosen.map(|(ctn, sel)| (ctn, sel)))
                 .with_context(|| format!("error setting up {}", file))?;
         }
 
@@ -73,7 +71,7 @@ impl State {
     fn save_choices_file<'a>(
         &self,
         file: &Utf8Path,
-        choices: impl Iterator<Item = &'a (ConTypeName, Selection)>,
+        choices: impl Iterator<Item = (&'a ConTypeName, &'a ContainerSelection)>,
     ) -> Result<()> {
         debug!("Saving choices for {}", file);
         let selector = ContainerSelector::from_path(
@@ -82,15 +80,7 @@ impl State {
             self.model_size,
         )?;
 
-        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_code = selector.gen_replacement_file(choices);
         let new_path = file.to_string().replace(".pr", "");
 
         trace!("New code: {}", new_code);
diff --git a/src/crates/candelabra/src/profiler/info.rs b/src/crates/candelabra/src/profiler/info.rs
index 488a8ba..f061e78 100644
--- a/src/crates/candelabra/src/profiler/info.rs
+++ b/src/crates/candelabra/src/profiler/info.rs
@@ -2,12 +2,11 @@ use std::collections::HashMap;
 use std::str::FromStr;
 
 use anyhow::{anyhow, Result};
+use log::debug;
+use primrose::ContainerSelection;
 use serde::{Deserialize, Serialize};
 
-use crate::{
-    cost::{benchmark::OpName, Cost, CostModel, Estimator},
-    types::ImplName,
-};
+use crate::cost::{benchmark::OpName, Cost, CostModel, Estimator};
 
 /// The information we get from profiling.
 /// Rather than keeping all results, we split them into 'similar enough' partitions,
@@ -29,14 +28,6 @@ type CollectionLifetime = (f64, HashMap<OpName, usize>);
 /// Breakdown of a cost value by operation
 pub type CostBreakdown<'a> = HashMap<&'a OpName, Cost>;
 
-/// A single result of container selection
-#[derive(Clone, Debug)]
-pub struct ContainerSplitSpec {
-    pub before: ImplName,
-    pub threshold: usize,
-    pub after: ImplName,
-}
-
 impl UsageProfile {
     pub fn from(iter: impl Iterator<Item = Result<String>>) -> Result<Self> {
         Ok(Self(
@@ -48,12 +39,16 @@ impl UsageProfile {
     pub fn check_for_nsplit(
         &mut self,
         candidates: &HashMap<&String, CostModel>,
-    ) -> Option<(ContainerSplitSpec, Cost)> {
+    ) -> Option<(ContainerSelection, Cost)> {
+        debug!("Checking for nsplit");
+
         self.0.sort_by_key(|p| p.avg_n as usize);
         if self.0.is_empty() {
             return None;
         }
 
+        debug!("Partitions: {:?}", self.0);
+
         let costs_by_partitions = candidates
             .iter()
             .map(|(name, model)| {
@@ -67,6 +62,8 @@ impl UsageProfile {
             })
             .collect::<Vec<(_, _)>>();
 
+        debug!("Costs by partitions: {:?}", costs_by_partitions);
+
         let top_by_partition = (0..self.0.len())
             .map(|i| {
                 costs_by_partitions.iter().fold(
@@ -82,6 +79,8 @@ impl UsageProfile {
             })
             .collect::<Vec<_>>();
 
+        debug!("Top by partition: {:?}", top_by_partition);
+
         let split_idx = top_by_partition
             .iter()
             .enumerate()
@@ -97,6 +96,11 @@ impl UsageProfile {
             }
         });
 
+        debug!(
+            "With split index {}, split proper is {}",
+            split_idx, split_is_proper
+        );
+
         if !split_is_proper {
             return None;
         }
@@ -125,11 +129,14 @@ impl UsageProfile {
         let not_switching_cost = &before_costs[split_idx..].iter().sum::<f64>()
             - &after_costs[split_idx..].iter().sum::<f64>();
 
+        debug!("Estimated switching cost: {}", switching_cost);
+        debug!("Estimated not switching cost: {}", not_switching_cost);
+
         if not_switching_cost < switching_cost {
             None
         } else {
             Some((
-                ContainerSplitSpec {
+                ContainerSelection::Split {
                     before: before.to_string(),
                     threshold: copy_n as usize,
                     after: after.to_string(),
diff --git a/src/crates/candelabra/src/profiler/mod.rs b/src/crates/candelabra/src/profiler/mod.rs
index 014cf64..8cd163e 100644
--- a/src/crates/candelabra/src/profiler/mod.rs
+++ b/src/crates/candelabra/src/profiler/mod.rs
@@ -5,7 +5,7 @@ mod info;
 use anyhow::{Context, Result};
 use camino::{Utf8Path, Utf8PathBuf};
 use log::{debug, log_enabled, trace, warn, Level};
-use primrose::ContainerSelector;
+use primrose::{ConTypeName, ContainerSelection, ContainerSelector, ImplName};
 use serde::{Deserialize, Serialize};
 use std::io::Write;
 use std::{
@@ -18,10 +18,10 @@ use tempfile::tempdir;
 use crate::cache::{gen_tree_hash, FileCache};
 use crate::cost::benchmark::tee_output;
 use crate::project::Project;
-use crate::types::{ConTypeName, ConTypeTo, ImplName, Mapping2D};
+use crate::types::{ConTypeTo, Mapping2D};
 use crate::{Paths, State};
 
-pub use self::info::{ContainerSplitSpec, ProfilerPartition, UsageProfile};
+pub use self::info::{ProfilerPartition, UsageProfile};
 
 type ConTypeRef<'a> = (&'a Utf8PathBuf, &'a ConTypeName);
 
@@ -124,19 +124,23 @@ impl State {
         .context("error creating container selector")?;
 
         let chosen = candidates
-            .map(|(dest_name, impls)| (dest_name, &impls[0]))
+            .map(|(dest_name, impls)| {
+                (
+                    dest_name,
+                    ContainerSelection::Profile(
+                        impls[0].clone(),
+                        con_types
+                            .iter()
+                            .position(|(search_file, search_ctn)| {
+                                file == *search_file && dest_name == *search_ctn
+                            })
+                            .unwrap(),
+                    ),
+                )
+            })
             .collect::<Vec<_>>();
 
-        let new_code = selector.gen_profiling_file(chosen.iter().map(|(d, c)| {
-            (
-                *d,
-                con_types
-                    .iter()
-                    .position(|(search_file, search_ctn)| file == *search_file && d == search_ctn)
-                    .unwrap(),
-                c.as_str(),
-            )
-        }));
+        let new_code = selector.gen_replacement_file(chosen.iter().map(|(ctn, sel)| (*ctn, sel)));
 
         let new_path = file.to_string().replace(".pr", "");
 
diff --git a/src/crates/candelabra/src/select.rs b/src/crates/candelabra/src/select.rs
index 12c359a..70b0f4e 100644
--- a/src/crates/candelabra/src/select.rs
+++ b/src/crates/candelabra/src/select.rs
@@ -1,22 +1,12 @@
 use std::collections::HashMap;
 
-use crate::{
-    cost::Cost,
-    profiler::ContainerSplitSpec,
-    types::{ConTypeTo, ImplName},
-    Project, State,
-};
+use crate::{cost::Cost, types::ConTypeTo, Project, State};
 
 use anyhow::Result;
+use primrose::ContainerSelection;
 
-#[derive(Clone, Debug)]
-pub enum Selection {
-    Singular(ImplName),
-    Split(ContainerSplitSpec),
-}
-
-pub type ProjectSelections = ConTypeTo<Selection>;
-pub type ProjectCostEstimates = ConTypeTo<Vec<(Selection, Cost)>>;
+pub type ProjectSelections = ConTypeTo<ContainerSelection>;
+pub type ProjectCostEstimates = ConTypeTo<Vec<(ContainerSelection, Cost)>>;
 
 impl State {
     /// Select a container implementation for each container type in the given project
@@ -60,14 +50,14 @@ impl State {
             for candidate in candidates {
                 let model = self.cost_model(candidate)?;
                 costs.push((
-                    Selection::Singular(candidate.clone()),
+                    ContainerSelection::Singular(candidate.to_string()),
                     profile_info.estimate_cost(&model),
                 ));
             }
 
             // 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));
+                costs.push((split, cost));
             }
 
             acc.insert(f, ctn, costs);
diff --git a/src/crates/candelabra/src/types.rs b/src/crates/candelabra/src/types.rs
index d16cc36..3738cd6 100644
--- a/src/crates/candelabra/src/types.rs
+++ b/src/crates/candelabra/src/types.rs
@@ -1,12 +1,7 @@
 use camino::Utf8PathBuf;
+use primrose::ConTypeName;
 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>;
diff --git a/src/crates/cli/src/estimate.rs b/src/crates/cli/src/estimate.rs
index 0ba0615..86e8d15 100644
--- a/src/crates/cli/src/estimate.rs
+++ b/src/crates/cli/src/estimate.rs
@@ -2,9 +2,10 @@ use std::iter::once;
 
 use anyhow::{anyhow, bail, Result};
 use argh::FromArgs;
-use candelabra::{select::Selection, types::Mapping2D};
+use candelabra::types::Mapping2D;
 use cargo_metadata::camino::Utf8PathBuf;
 use log::info;
+use primrose::ContainerSelection;
 
 use crate::State;
 
@@ -109,7 +110,11 @@ impl State {
                     &assignments
                         .into_iter()
                         .map(|(f, ctn, sel)| {
-                            (f, ctn.to_string(), Selection::Singular(sel.to_string()))
+                            (
+                                f,
+                                ctn.to_string(),
+                                ContainerSelection::Singular(sel.to_string()),
+                            )
                         })
                         .collect(),
                 )?
diff --git a/src/crates/cli/src/select.rs b/src/crates/cli/src/select.rs
index 740b988..1578c1a 100644
--- a/src/crates/cli/src/select.rs
+++ b/src/crates/cli/src/select.rs
@@ -2,9 +2,8 @@ use std::collections::HashMap;
 
 use anyhow::Result;
 use argh::FromArgs;
-use candelabra::select::Selection;
 use log::info;
-use primrose::tools::nary_cartesian_product;
+use primrose::{tools::nary_cartesian_product, ContainerSelection};
 use tabled::{builder::Builder, settings::Style};
 
 use crate::State;
@@ -30,13 +29,15 @@ impl State {
             for (f, ctn, candidates) in costs.iter() {
                 for (candidate, cost) in candidates.iter() {
                     let name = match candidate {
-                        Selection::Singular(x) => x.to_string(),
-                        Selection::Split(split) => {
-                            format!(
-                                "{} until n={}, then {}",
-                                split.before, split.threshold, split.after
-                            )
+                        ContainerSelection::Singular(x) => x.to_string(),
+                        ContainerSelection::Split {
+                            before,
+                            threshold,
+                            after,
+                        } => {
+                            format!("{} until n={}, then {}", before, threshold, after)
                         }
+                        _ => unreachable!(),
                     };
                     builder.push_record([
                         ctn.as_str(),
@@ -72,7 +73,11 @@ impl State {
                         &assignment
                             .into_iter()
                             .map(|((f, ctn), i)| {
-                                (f.clone(), ctn.clone(), Selection::Singular(i.to_string()))
+                                (
+                                    f.clone(),
+                                    ctn.clone(),
+                                    ContainerSelection::Singular(i.to_string()),
+                                )
                             })
                             .collect(),
                     )?,
diff --git a/src/crates/primrose/src/codegen.rs b/src/crates/primrose/src/codegen.rs
index fac55a6..2eccf8d 100644
--- a/src/crates/primrose/src/codegen.rs
+++ b/src/crates/primrose/src/codegen.rs
@@ -4,17 +4,30 @@ use crate::{
     description::{InforMap, Tag},
     parser::Block,
     selector::ContainerSelector,
+    ImplName,
 };
 
 const CODEGEN: &str = "/*CODEGEN*/\n";
 const CODEGENEND: &str = "/*ENDCODEGEN*/\n";
 const TRAITCRATE: &str = "primrose_library::traits::";
 
+/// A selected concrete type that we can use
+#[derive(Clone, Debug)]
+pub enum ContainerSelection {
+    Singular(ImplName),
+    Profile(ImplName, usize),
+    Split {
+        before: ImplName,
+        threshold: usize,
+        after: ImplName,
+    },
+}
+
 impl ContainerSelector {
     /// Generate replacement code for the whole file, with the given `(tag_id, selection)` 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_replacement_file<'a, T: Iterator<Item = (&'a String, &'a str)>>(
+    pub fn gen_replacement_file<'a, T: Iterator<Item = (&'a String, &'a ContainerSelection)>>(
         &self,
         selections: T,
     ) -> String {
@@ -29,27 +42,14 @@ impl ContainerSelector {
         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, usize, &'a str)>>(
-        &self,
-        inners: T,
-    ) -> String {
-        let mut result = String::new();
-        result += &codegen_block(&self.bounds_decl);
-        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
-        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 {
+    pub(crate) fn gen_replacement_code(
+        &self,
+        tag_id: &String,
+        selection: &ContainerSelection,
+    ) -> 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");
@@ -57,40 +57,38 @@ impl ContainerSelector {
         let bounds = to_trait_bounds(elem_ty);
         let vars = elem_ty.join(", ");
 
-        codegen_block(&format!(
-            r#"
+        let inner = match selection {
+            ContainerSelection::Singular(inner) => format!(
+                r#"
 #[allow(non_snake_case)]
 fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
-    {selection}::<{vars}>::default()
+    {inner}::<{vars}>::default()
 }}
 "#
-        ))
-    }
-
-    /// 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,
-        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");
-        };
-        let bounds = to_trait_bounds(elem_ty);
-        let vars = elem_ty.join(", ");
-
-        codegen_block(&format!(
-            r#"
+            ),
+            ContainerSelection::Profile(inner, idx) => format!(
+                r#"
 #[allow(non_snake_case)]
 fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
-    ::primrose_library::ProfilerWrapper::<{id}, {inner}<{vars}>, _>::default()
+    ::primrose_library::ProfilerWrapper::<{idx}, {inner}<{vars}>, _>::default()
 }}
 "#
-        ))
+            ),
+            ContainerSelection::Split {
+                before,
+                threshold,
+                after,
+            } => format!(
+                r#"
+#[allow(non_snake_case)]
+fn _{tag_id}<{bounds}>() -> {tag_id}<{vars}> {{
+    ::primrose_library::AdaptiveContainer::<{threshold}, {before}<{vars}>, {after}<{vars}>::default()
+}}
+"#
+            ),
+        };
+
+        codegen_block(&inner)
     }
 }
 
diff --git a/src/crates/primrose/src/lib.rs b/src/crates/primrose/src/lib.rs
index a283f28..a891084 100644
--- a/src/crates/primrose/src/lib.rs
+++ b/src/crates/primrose/src/lib.rs
@@ -13,6 +13,7 @@ pub mod tools;
 
 mod codegen;
 mod selector;
+pub use codegen::ContainerSelection;
 pub use selector::ContainerSelector;
 
 mod library_specs;
@@ -21,3 +22,9 @@ pub use spec_map::LibSpecs;
 
 mod error;
 pub use error::Error;
+
+/// Names a container type we want to select.
+pub type ConTypeName = String;
+
+/// Name of a container implementation we are considering
+pub type ImplName = String;
diff --git a/src/crates/primrose/src/main.rs b/src/crates/primrose/src/main.rs
deleted file mode 100644
index d307942..0000000
--- a/src/crates/primrose/src/main.rs
+++ /dev/null
@@ -1,91 +0,0 @@
-use log::info;
-use primrose::tools::nary_cartesian_product;
-use primrose::{ContainerSelector, Error};
-use std::collections::HashMap;
-use std::error::Error as StdError;
-use std::path::Path;
-use std::{env, fs, io::Write};
-
-const LIB: &str = "./crates/primrose-library/src/";
-
-fn main() -> Result<(), Box<dyn StdError>> {
-    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
-
-    let (file_name, output_name, model_size) =
-        parse_args().map_err(Into::<Box<dyn StdError>>::into)?;
-
-    info!(
-        "Running on {}, outputting to {}, with model size {}",
-        file_name, output_name, model_size
-    );
-    run(file_name, output_name, model_size)?;
-
-    info!("Yippeeeeee!!");
-    Ok(())
-}
-
-fn parse_args() -> Result<(String, String, usize), &'static str> {
-    let args: Vec<String> = env::args().collect();
-    if args.len() == 1 {
-        Ok((
-            "./spec_code/example_unique.rs".to_string(),
-            "default".to_string(),
-            5,
-        ))
-    } else if args.len() == 4 {
-        let model_size_input = match args.get(3).unwrap().parse::<u64>() {
-            Ok(val) => val,
-            Err(_) => {
-                return Err("Invalid model size");
-            }
-        };
-        let model_size = model_size_input as usize;
-
-        Ok((
-            "./spec_code/".to_string() + args.get(1).unwrap(),
-            args.get(2).unwrap().to_string(),
-            model_size,
-        ))
-    } else {
-        return Err("Usage: <file_name> [output_path] [model_size]");
-    }
-}
-
-/// Read input from the filename, calculate all valid implementations, and output them to seperate files in output_path
-fn run(input: String, output_name: String, model_size: usize) -> Result<(), Error> {
-    info!("Generating candidate code outputs");
-    let gen_code = gen_outputs(input, model_size)?;
-
-    let output_path = format!("./gen_code/{}/", output_name);
-    fs::create_dir_all(&output_path).expect("error creating output directory");
-
-    info!("Writing {} different outputs", gen_code.len());
-    for (i, code) in gen_code.iter().enumerate() {
-        let mut output = fs::File::create(format!("{}/{}{}.rs", output_path, output_name, i))
-            .expect("error creating output file");
-        write!(output, "{}", &code).expect("error writing output file");
-    }
-
-    Ok(())
-}
-
-/// Process the given file, returning code for all possible types
-fn gen_outputs(filename: String, model_size: usize) -> Result<Vec<String>, Error> {
-    let selector = ContainerSelector::from_path(Path::new(&filename), Path::new(LIB), model_size)?;
-
-    let mut candidates = HashMap::new();
-    for tag in selector.container_tags() {
-        let found = selector.find_candidates(tag)?;
-        candidates.insert(tag, found);
-    }
-
-    Ok(nary_cartesian_product(&candidates)
-        .into_iter()
-        .map(|selections| {
-            selections
-                .into_iter()
-                .map(|(tag_name, typ_name)| selector.gen_replacement_code(tag_name, typ_name))
-                .collect::<String>()
-        })
-        .collect())
-}
diff --git a/src/tests/Cargo.lock b/src/tests/Cargo.lock
index 3b7b6dc..183a743 100644
--- a/src/tests/Cargo.lock
+++ b/src/tests/Cargo.lock
@@ -308,6 +308,9 @@ dependencies = [
 [[package]]
 name = "primrose-library"
 version = "0.1.0"
+dependencies = [
+ "take_mut",
+]
 
 [[package]]
 name = "proc-macro2"
@@ -474,6 +477,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "take_mut"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f764005d11ee5f36500a149ace24e00e3da98b0158b3e2d53a7495660d3f4d60"
+
+[[package]]
 name = "textwrap"
 version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"