sp1_recursion_executor/
shape.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    marker::PhantomData,
4};
5
6use serde::{Deserialize, Serialize};
7use slop_algebra::{Field, PrimeField32};
8use sp1_hypercube::{air::MachineAir, ChipDimensions};
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
11pub struct RecursionShape<F> {
12    heights: BTreeMap<String, usize>,
13    _marker: PhantomData<F>,
14}
15
16impl<F: PrimeField32> RecursionShape<F> {
17    pub const fn new(heights: BTreeMap<String, usize>) -> Self {
18        Self { heights, _marker: PhantomData }
19    }
20
21    pub fn height<A>(&self, air: &A) -> Option<usize>
22    where
23        F: Field,
24        A: MachineAir<F>,
25    {
26        self.heights.get(air.name()).copied()
27    }
28
29    pub fn height_of_name(&self, name: &str) -> Option<usize> {
30        self.heights.get(name).copied()
31    }
32
33    pub fn insert<A>(&mut self, air: &A, height: usize)
34    where
35        F: Field,
36        A: MachineAir<F>,
37    {
38        self.heights.insert(air.name().to_string(), height);
39    }
40
41    pub fn insert_with_name(&mut self, name: &str, height: usize) {
42        self.heights.insert(name.to_string(), height);
43    }
44
45    pub const fn empty() -> Self {
46        Self { heights: BTreeMap::new(), _marker: PhantomData }
47    }
48
49    pub fn preprocessed_chip_information<A>(
50        &self,
51        chips: &BTreeSet<A>,
52    ) -> BTreeMap<String, ChipDimensions<F>>
53    where
54        F: Field,
55        A: MachineAir<F>,
56    {
57        chips
58            .iter()
59            .filter_map(|chip| {
60                self.height(chip).map(|height| {
61                    (
62                        chip.name().to_string(),
63                        ChipDimensions {
64                            height: F::from_canonical_u32(height as u32),
65                            num_polynomials: F::from_canonical_u32(chip.preprocessed_width() as u32),
66                        },
67                    )
68                })
69            })
70            .collect()
71    }
72}
73
74impl<F: Field, A: MachineAir<F>> FromIterator<(A, usize)> for RecursionShape<F> {
75    fn from_iter<T: IntoIterator<Item = (A, usize)>>(iter: T) -> Self {
76        RecursionShape {
77            heights: iter
78                .into_iter()
79                .map(|(air, height)| (air.name().to_string(), height))
80                .collect(),
81            _marker: PhantomData,
82        }
83    }
84}
85
86impl<F: Field> IntoIterator for RecursionShape<F> {
87    type Item = (String, usize);
88    type IntoIter = <BTreeMap<String, usize> as IntoIterator>::IntoIter;
89
90    fn into_iter(self) -> Self::IntoIter {
91        self.heights.into_iter()
92    }
93}
94
95impl<'a, F: Field> IntoIterator for &'a RecursionShape<F> {
96    type Item = (&'a String, &'a usize);
97    type IntoIter = <&'a BTreeMap<String, usize> as IntoIterator>::IntoIter;
98
99    fn into_iter(self) -> Self::IntoIter {
100        self.heights.iter()
101    }
102}
sp1_recursion_executor/shape.rs

sp1_recursion_executor/
shape.rs
