path: root/src/crates/candelabra/src/profiler/info.rs

use std::collections::HashMap;
use std::str::FromStr;

use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};

use crate::{
    cost::{benchmark::OpName, Cost, CostModel, Estimator},
    types::ImplName,
};

/// The information we get from profiling.
/// Rather than keeping all results, we split them into 'similar enough' partitions,
/// with the idea that each partition will probably have the same best implementation.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct UsageProfile(pub Vec<ProfilerPartition>);

/// A vector of container lifetimes which have similar usage characteristics
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ProfilerPartition {
    pub occurences: f64,
    pub avg_n: f64,
    pub avg_op_counts: HashMap<OpName, f64>,
}

/// Lifetime of a single allocated collection.
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(
            iter.map(|contents| parse_output(&contents?))
                .fold(Ok(vec![]), partition_costs)?,
        ))
    }

    pub fn check_for_nsplit(
        &mut self,
        candidates: &HashMap<&String, CostModel>,
    ) -> Option<(ContainerSplitSpec, Cost)> {
        self.0.sort_by_key(|p| p.avg_n as usize);
        if self.0.is_empty() {
            return None;
        }

        let costs_by_partitions = candidates
            .iter()
            .map(|(name, model)| {
                (
                    name,
                    self.0
                        .iter()
                        .map(|p| p.estimate_cost(&model))
                        .collect::<Vec<_>>(),
                )
            })
            .collect::<Vec<(_, _)>>();

        let top_by_partition = (0..self.0.len())
            .map(|i| {
                costs_by_partitions.iter().fold(
                    ("".to_string(), f64::MAX),
                    |acc @ (_, val), (name, c)| {
                        if val < c[i] {
                            acc
                        } else {
                            (name.to_string(), c[i])
                        }
                    },
                )
            })
            .collect::<Vec<_>>();

        let split_idx = top_by_partition
            .iter()
            .enumerate()
            // TODO: fudge?
            .find(|(idx, (best, _))| *idx > 0 && *best != top_by_partition[idx - 1].0)
            .map(|(idx, _)| idx)?;

        let split_is_proper = top_by_partition.iter().enumerate().all(|(i, (best, _))| {
            if i >= split_idx {
                *best == top_by_partition[split_idx].0
            } else {
                *best == top_by_partition[0].0
            }
        });

        if !split_is_proper {
            return None;
        }

        // calculate cost of switching
        let before = &top_by_partition[0].0;
        let after = &top_by_partition[split_idx].0;
        let before_model = candidates.get(before).unwrap();
        let after_model = candidates.get(after).unwrap();

        let copy_n = self.0[split_idx].avg_n;
        let switching_cost = after_model.by_op.get("insert")?.estimatef(copy_n)
            + before_model.by_op.get("clear")?.estimatef(copy_n);

        // see if it's "worth it"
        let before_costs = &costs_by_partitions
            .iter()
            .find(|(name, _)| **name == before)
            .unwrap()
            .1;
        let after_costs = &costs_by_partitions
            .iter()
            .find(|(name, _)| **name == after)
            .unwrap()
            .1;
        let not_switching_cost = &before_costs[split_idx..].iter().sum::<f64>()
            - &after_costs[split_idx..].iter().sum::<f64>();

        if not_switching_cost < switching_cost {
            None
        } else {
            Some((
                ContainerSplitSpec {
                    before: before.to_string(),
                    threshold: copy_n as usize,
                    after: after.to_string(),
                },
                top_by_partition.iter().map(|(_, v)| v).sum(),
            ))
        }
    }

    /// Estimate the cost of using the implementation with the given cost model
    pub fn estimate_cost(&self, cost_model: &CostModel) -> f64 {
        self.0
            .iter()
            .map(|cl| cl.estimate_cost(cost_model))
            .sum::<f64>()
    }

    /// Get a breakdown of the cost by operation
    pub fn cost_breakdown<'a>(&self, cost_model: &'a CostModel) -> CostBreakdown<'a> {
        cost_model
            .by_op
            .iter()
            .map(|(op, estimator)| {
                (
                    op,
                    self.0
                        .iter()
                        .map(|cl| cl.op_cost(op, estimator))
                        .sum::<f64>(),
                )
            })
            .collect()
    }
}

impl ProfilerPartition {
    pub fn avg_op_count(&self, op: &str) -> f64 {
        *self
            .avg_op_counts
            .get(op)
            .expect("invalid op passed to op_count")
    }

    pub fn estimate_cost(&self, cost_model: &CostModel) -> f64 {
        cost_model
            .by_op
            .iter()
            .map(|(op, estimator)| self.op_cost(op, estimator))
            .sum::<f64>()
    }

    pub fn op_cost(&self, op: &str, estimator: &Estimator) -> f64 {
        estimator.estimatef(self.avg_n) * self.avg_op_count(op) * self.occurences
    }

    /// Get a breakdown of the cost by operation
    pub fn cost_breakdown<'a>(&'a self, cost_model: &'a CostModel) -> CostBreakdown<'a> {
        self.avg_op_counts
            .keys()
            .flat_map(|op| Some((op, self.op_cost(op, cost_model.by_op.get(op)?))))
            .collect()
    }
    fn add_lifetime(&mut self, (n, ops): (f64, HashMap<String, usize>)) {
        self.avg_n = self.avg_n + (n - self.avg_n) / (self.occurences + 1.0);
        for (op, count) in ops {
            let count = count as f64;
            self.avg_op_counts
                .entry(op)
                .and_modify(|avg| *avg = *avg + (count - *avg) / (self.occurences + 1.0))
                .or_insert(count);
        }
        self.occurences += 1.0;
    }
}

/// Attempt to compress an iterator of collection lifetimes into as few partitions as possible
fn partition_costs(
    acc: Result<Vec<ProfilerPartition>>,
    cl: Result<CollectionLifetime>,
) -> Result<Vec<ProfilerPartition>> {
    // error short circuiting
    let (mut acc, (n, ops)) = (acc?, cl?);

    // attempt to find a partition with a close enough n value
    let (closest_idx, closest_delta) =
        acc.iter()
            .enumerate()
            .fold((0, f64::MAX), |acc @ (_, val), (idx, partition)| {
                let delta = (partition.avg_n - n).abs();
                if delta < val {
                    (idx, delta)
                } else {
                    acc
                }
            });

    if closest_delta < 100.0 {
        acc[closest_idx].add_lifetime((n, ops));
    } else {
        // add a new partition
        acc.push(ProfilerPartition {
            occurences: 1.0,
            avg_n: n,
            avg_op_counts: ops.into_iter().map(|(k, v)| (k, v as f64)).collect(),
        })
    }

    Ok(acc)
}

/// Parse the output of the profiler
fn parse_output(contents: &str) -> Result<(f64, HashMap<OpName, usize>)> {
    let mut lines = contents.lines().map(usize::from_str);
    let missing_line_err = || anyhow!("wrong number of lines in ");
    let n = lines.next().ok_or_else(missing_line_err)??;
    let mut op_counts = HashMap::new();
    op_counts.insert(
        "contains".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "insert".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "clear".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "remove".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "first".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "last".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "nth".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "push".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "pop".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );
    op_counts.insert(
        "get".to_string(),
        lines.next().ok_or_else(missing_line_err)??,
    );

    Ok((n as f64, op_counts))
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use crate::{
        cost::{CostModel, Estimator},
        profiler::info::partition_costs,
    };

    use super::{ProfilerPartition, UsageProfile};

    const EPSILON: f64 = 1e-5;
    fn assert_feq(left: f64, right: f64, msg: &'static str) {
        assert!((left - right).abs() < EPSILON, "{}", msg);
    }

    fn linear_estimator() -> Estimator {
        Estimator {
            coeffs: [0.0, 1.0, 0.0, 0.0],
            transform_x: (0.0, 1.0),
            transform_y: (0.0, 1.0),
        }
    }

    #[test]
    fn test_cost_single_partition() {
        let info = UsageProfile(vec![ProfilerPartition {
            occurences: 1.0,
            avg_n: 100.0,
            avg_op_counts: {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), 100.0);
                map
            },
        }]);

        let model = CostModel {
            by_op: {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), linear_estimator());
                map
            },
        };

        let cost = dbg!(info.estimate_cost(&model));
        assert_feq(cost, 10_000.0, "per op = 100 * 100 ops");
    }

    #[test]
    fn test_cost_multi_partitions_sums() {
        let info = UsageProfile(vec![
            ProfilerPartition {
                occurences: 1.0,
                avg_n: 100.0,
                avg_op_counts: {
                    let mut map = HashMap::new();
                    map.insert("insert".to_string(), 100.0);
                    map
                },
            },
            ProfilerPartition {
                occurences: 1.0,
                avg_n: 10.0,
                avg_op_counts: {
                    let mut map = HashMap::new();
                    map.insert("insert".to_string(), 10.0);
                    map
                },
            },
        ]);

        let model = CostModel {
            by_op: {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), linear_estimator());
                map
            },
        };

        let cost = dbg!(info.estimate_cost(&model));
        assert_feq(cost, 10_100.0, "100ns/op * 100 ops + 10ns/op * 10 ops");
    }

    #[test]
    fn test_cost_multi_partitions_sums_weighted() {
        let info = UsageProfile(vec![
            ProfilerPartition {
                occurences: 2.0,
                avg_n: 100.0,
                avg_op_counts: {
                    let mut map = HashMap::new();
                    map.insert("insert".to_string(), 100.0);
                    map
                },
            },
            ProfilerPartition {
                occurences: 1.0,
                avg_n: 10.0,
                avg_op_counts: {
                    let mut map = HashMap::new();
                    map.insert("insert".to_string(), 10.0);
                    map
                },
            },
        ]);

        let model = CostModel {
            by_op: {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), linear_estimator());
                map
            },
        };

        let cost = dbg!(info.estimate_cost(&model));
        assert_feq(cost, 20_100.0, "100ns/op * 100 ops * 2 + 10ns/op * 10 ops");
    }

    #[test]
    fn test_partition_costs_merges_duplicates() {
        let cl = (100.0, {
            let mut map = HashMap::new();
            map.insert("insert".to_string(), 10);
            map
        });
        let outp = vec![Ok(cl.clone()), Ok(cl)]
            .into_iter()
            .fold(Ok(vec![]), partition_costs)
            .unwrap();

        assert_eq!(outp.len(), 1, "merged duplicates");
        assert_eq!(outp[0].occurences, 2.0, "weight updated");
        assert_feq(outp[0].avg_n, 100.0, "average n correct");
        assert_feq(
            *outp[0].avg_op_counts.get("insert").unwrap(),
            10.0,
            "average n correct",
        );
    }

    #[test]
    fn test_partition_costs_merges_close() {
        let outp = vec![
            Ok((100.0, {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), 50);
                map
            })),
            Ok((110.0, {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), 100);
                map
            })),
        ]
        .into_iter()
        .fold(Ok(vec![]), partition_costs)
        .unwrap();

        assert_eq!(outp.len(), 1, "merged duplicates");
        assert_eq!(outp[0].occurences, 2.0, "weight updated");
        assert_feq(outp[0].avg_n, 105.0, "average n correct");
        assert_feq(
            *outp[0].avg_op_counts.get("insert").unwrap(),
            75.0,
            "average n correct",
        );
    }
    #[test]
    fn test_partition_costs_keeps_separate() {
        let outp = vec![
            Ok((100.0, {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), 10);
                map
            })),
            Ok((999999.0, {
                let mut map = HashMap::new();
                map.insert("insert".to_string(), 10);
                map
            })),
        ]
        .into_iter()
        .fold(Ok(vec![]), partition_costs)
        .unwrap();

        assert_eq!(
            outp.len(),
            2,
            "large difference in n values causes partition"
        );
    }
}