sp1_recursion_core/air/
block.rs

1use p3_air::AirBuilder;
2use p3_field::{AbstractField, ExtensionField, Field};
3use serde::{Deserialize, Serialize};
4use sp1_derive::AlignedBorrow;
5use sp1_stark::air::{BinomialExtension, ExtensionAirBuilder, SP1AirBuilder};
6
7use std::ops::{Index, IndexMut};
8
9use crate::runtime::D;
10
11/// The smallest unit of memory that can be read and written to.
12#[derive(
13    AlignedBorrow, Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize,
14)]
15#[repr(C)]
16pub struct Block<T>(pub [T; D]);
17
18pub trait BlockBuilder: AirBuilder {
19    fn assert_block_eq<Lhs: Into<Self::Expr>, Rhs: Into<Self::Expr>>(
20        &mut self,
21        lhs: Block<Lhs>,
22        rhs: Block<Rhs>,
23    ) {
24        for (l, r) in lhs.0.into_iter().zip(rhs.0) {
25            self.assert_eq(l, r);
26        }
27    }
28}
29
30impl<AB: AirBuilder> BlockBuilder for AB {}
31
32impl<T> Block<T> {
33    pub fn map<F, U>(self, f: F) -> Block<U>
34    where
35        F: FnMut(T) -> U,
36    {
37        Block(self.0.map(f))
38    }
39
40    pub fn ext<E>(&self) -> E
41    where
42        T: Field,
43        E: ExtensionField<T>,
44    {
45        E::from_base_slice(&self.0)
46    }
47}
48
49impl<T: Clone> Block<T> {
50    pub fn as_extension<AB: ExtensionAirBuilder<Var = T>>(&self) -> BinomialExtension<AB::Expr> {
51        let arr: [AB::Expr; 4] = self.0.clone().map(|x| AB::Expr::zero() + x);
52        BinomialExtension(arr)
53    }
54
55    pub fn as_extension_from_base<AB: SP1AirBuilder<Var = T>>(
56        &self,
57        base: AB::Expr,
58    ) -> BinomialExtension<AB::Expr> {
59        let mut arr: [AB::Expr; 4] = self.0.clone().map(|_| AB::Expr::zero());
60        arr[0] = base;
61
62        BinomialExtension(arr)
63    }
64}
65
66impl<T> From<[T; D]> for Block<T> {
67    fn from(arr: [T; D]) -> Self {
68        Self(arr)
69    }
70}
71
72impl<T: AbstractField> From<T> for Block<T> {
73    fn from(value: T) -> Self {
74        Self([value, T::zero(), T::zero(), T::zero()])
75    }
76}
77
78impl<T: Copy> From<&[T]> for Block<T> {
79    fn from(slice: &[T]) -> Self {
80        let arr: [T; D] = slice.try_into().unwrap();
81        Self(arr)
82    }
83}
84
85impl<T, I> Index<I> for Block<T>
86where
87    [T]: Index<I>,
88{
89    type Output = <[T] as Index<I>>::Output;
90
91    #[inline]
92    fn index(&self, index: I) -> &Self::Output {
93        Index::index(&self.0, index)
94    }
95}
96
97impl<T, I> IndexMut<I> for Block<T>
98where
99    [T]: IndexMut<I>,
100{
101    #[inline]
102    fn index_mut(&mut self, index: I) -> &mut Self::Output {
103        IndexMut::index_mut(&mut self.0, index)
104    }
105}
106
107impl<T> IntoIterator for Block<T> {
108    type Item = T;
109    type IntoIter = std::array::IntoIter<T, D>;
110
111    fn into_iter(self) -> Self::IntoIter {
112        self.0.into_iter()
113    }
114}
sp1_recursion_core/air/block.rs

sp1_recursion_core/air/
block.rs
