Skip to main content

sp1_hypercube/shape/
ordered.rs

1use std::{cmp::Reverse, collections::BTreeSet, fmt};
2
3use itertools::Itertools;
4use serde::{Deserialize, Serialize};
5use slop_matrix::{dense::RowMajorMatrix, Matrix};
6
7/// A way to keep track of the log2 heights of some set of chips and in canonical order.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd, Ord, Eq, Hash)]
9pub struct OrderedShape {
10    /// The inner data.
11    pub inner: Vec<(String, usize)>,
12}
13
14impl OrderedShape {
15    /// Create an [`OrderedShape`] from a set of traces.
16    #[must_use]
17    pub fn from_traces<V: Clone + Send + Sync>(traces: &[(String, RowMajorMatrix<V>)]) -> Self {
18        traces
19            .iter()
20            .map(|(name, trace)| (name.clone(), trace.height().ilog2() as usize))
21            .sorted_by_key(|(_, height)| *height)
22            .collect()
23    }
24
25    /// Create an [`OrderedShape`] from a set of log2 heights.
26    #[must_use]
27    pub fn from_log2_heights(traces: &[(String, usize)]) -> Self {
28        traces
29            .iter()
30            .map(|(name, height)| (name.clone(), *height))
31            .sorted_by_key(|(_, height)| *height)
32            .collect()
33    }
34}
35
36impl FromIterator<(String, usize)> for OrderedShape {
37    fn from_iter<T: IntoIterator<Item = (String, usize)>>(iter: T) -> Self {
38        let set = iter
39            .into_iter()
40            .map(|(name, log_degree)| (Reverse(log_degree), name))
41            .collect::<BTreeSet<_>>();
42        Self {
43            inner: set.into_iter().map(|(Reverse(log_degree), name)| (name, log_degree)).collect(),
44        }
45    }
46}
47
48impl IntoIterator for OrderedShape {
49    type Item = (String, usize);
50
51    type IntoIter = <Vec<(String, usize)> as IntoIterator>::IntoIter;
52
53    fn into_iter(self) -> Self::IntoIter {
54        self.inner.into_iter()
55    }
56}
57
58impl fmt::Display for OrderedShape {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        writeln!(f, "OrderedShape:")?;
61        for (name, log_degree) in &self.inner {
62            writeln!(f, "{name}: {log_degree}")?;
63        }
64        Ok(())
65    }
66}