feat(cli): WIP curve fitting with fitme

6 files changed, 72 insertions, 10 deletions

diff --git a/src/crates/cli/Cargo.toml b/src/crates/cli/Cargo.toml
index c6dc852..d91b34c 100644
--- a/src/crates/cli/Cargo.toml
+++ b/src/crates/cli/Cargo.toml
@@ -18,4 +18,5 @@ camino = "1.1.6"
 cargo_metadata = "0.18.1"
 argh = "0.1.12"
 glob = "0.3.1"
-tempfile = "3"
\ No newline at end of file
+tempfile = "3"
+fitme = "1.1.0"
\ No newline at end of file
diff --git a/src/crates/cli/src/cost/benchmark.rs b/src/crates/cli/src/cost/benchmark.rs
index a4824d9..49e1919 100644
--- a/src/crates/cli/src/cost/benchmark.rs
+++ b/src/crates/cli/src/cost/benchmark.rs
@@ -3,7 +3,7 @@
 use std::{
     fs::{copy, create_dir, File},
     io::Write,
-    process::Command,
+    process::{Command, Stdio},
 };
 
 use anyhow::{bail, Context, Result};
@@ -18,7 +18,7 @@ use crate::paths::Paths;
 pub const ELEM_TYPE: &str = "usize";
 
 /// String representation of the array of N values we use for benchmarking
-pub const NS: &str = "[1, 10, 100]";
+pub const NS: &str = "[65536]";
 
 /// Run benchmarks for the given container type, returning the results.
 /// Panics if the given name is not in the library specs.
diff --git a/src/crates/cli/src/cost/fit.rs b/src/crates/cli/src/cost/fit.rs
new file mode 100644
index 0000000..7e17379
--- /dev/null
+++ b/src/crates/cli/src/cost/fit.rs
@@ -0,0 +1,48 @@
+//! Fitting a curve to benchmark results
+
+use std::time::Duration;
+
+use candelabra_benchmarker::Observation;
+use fitme::{expr::v1::Eq, Data, Equation, Fit, Headers, Output};
+
+/// Fit a curve to the given set of observations.
+pub fn fit(results: &Vec<Observation>) -> impl Estimator {
+    let headers = Headers::from_iter(&["N", "T"]);
+    let eq = Eq::parse("m * T + c", &headers).unwrap();
+    let fit = fitme::fit(
+        eq.clone(),
+        Data::new(
+            headers,
+            results
+                .into_iter()
+                .map(|(n, results)| dbg!([*n as f64, as_millis_f64(&results.avg)])),
+        )
+        .unwrap(),
+        "N",
+    )
+    .unwrap();
+
+    dbg!(&fit.parameter_names);
+    dbg!(&fit.parameter_values);
+
+    (eq, fit)
+}
+
+/// Can estimate a duration for a given `n`.
+pub trait Estimator {
+    /// Estimate the duration for `n`.
+    fn estimate(&self, n: usize) -> Duration;
+}
+
+impl Estimator for (Eq, Fit) {
+    fn estimate(&self, n: usize) -> Duration {
+        todo!()
+    }
+}
+
+fn as_millis_f64(d: &Duration) -> f64 {
+    let millis = d.as_millis() as f64;
+    let exp = 10.0_f64.powf(6.0);
+    let remainder_nanos = d.as_nanos() as f64 - (millis * exp);
+    millis + (remainder_nanos / exp)
+}
diff --git a/src/crates/cli/src/cost/mod.rs b/src/crates/cli/src/cost/mod.rs
index eda5957..c407d9e 100644
--- a/src/crates/cli/src/cost/mod.rs
+++ b/src/crates/cli/src/cost/mod.rs
@@ -1,5 +1,6 @@
 //! Generating, caching, and using cost models
 pub mod benchmark;
+pub mod fit;
 
 use anyhow::{anyhow, Context, Result};
 
diff --git a/src/crates/cli/src/main.rs b/src/crates/cli/src/main.rs
index 00c85cc..bf9a084 100644
--- a/src/crates/cli/src/main.rs
+++ b/src/crates/cli/src/main.rs
@@ -2,10 +2,14 @@ use std::collections::HashSet;
 
 use anyhow::{anyhow, Context, Result};
 use argh::FromArgs;
-use log::info;
+use log::{debug, info};
 use project::Project;
 
-use crate::{candidates::CandidatesStore, cost::ResultsStore, paths::Paths};
+use crate::{
+    candidates::CandidatesStore,
+    cost::{fit::fit, ResultsStore},
+    paths::Paths,
+};
 
 mod cache;
 mod candidates;
@@ -52,8 +56,16 @@ fn main() -> Result<()> {
     }
 
     info!("Found all candidate types. Running benchmarks");
-    for typ in seen_types.into_iter() {
-        dbg!(benchmarks.get(&typ).context("Error running benchmark")?);
+    for typ in seen_types
+        .into_iter()
+        .filter(|x| x == "primrose_library::EagerSortedVec")
+    {
+        let results = benchmarks.get(&typ).context("Error running benchmark")?;
+
+        for (op, results) in results.by_op.iter().filter(|(k, _)| **k == "insert") {
+            debug!("Fitting curve for op {}", op);
+            fit(results);
+        }
     }
 
     Ok(())
diff --git a/src/crates/cli/src/paths.rs b/src/crates/cli/src/paths.rs
index 639e1c5..2b44400 100644
--- a/src/crates/cli/src/paths.rs
+++ b/src/crates/cli/src/paths.rs
@@ -14,9 +14,9 @@ pub struct Paths {
 impl Paths {
     fn from_base(base: Utf8PathBuf) -> Self {
         Paths {
-            library_crate: base.join("crates").join("primrose-library"),
-            library_src: base.join("crates").join("primrose-library").join("src"),
-            benchmarker_crate: base.join("crates").join("candelabra-benchmarker"),
+            library_crate: base.join("crates").join("library"),
+            library_src: base.join("crates").join("library").join("src"),
+            benchmarker_crate: base.join("crates").join("benchmarker"),
             target_dir: base.join("target"),
             base,
         }