sp1_stark/
word.rs

1use std::ops::{Index, IndexMut};
2
3use crate::air::SP1AirBuilder;
4use arrayref::array_ref;
5use itertools::Itertools;
6use p3_air::AirBuilder;
7use p3_field::{AbstractField, Field};
8use serde::{Deserialize, Serialize};
9use sp1_derive::AlignedBorrow;
10use sp1_primitives::consts::WORD_SIZE;
11use std::array::IntoIter;
12
13/// An array of four bytes to represent a 32-bit value.
14///
15/// We use the generic type `T` to represent the different representations of a byte, ranging from
16/// a `u8` to a `AB::Var` or `AB::Expr`.
17#[derive(
18    AlignedBorrow, Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize,
19)]
20#[repr(C)]
21pub struct Word<T>(pub [T; WORD_SIZE]);
22
23impl<T> Word<T> {
24    /// Applies `f` to each element of the word.
25    pub fn map<F, S>(self, f: F) -> Word<S>
26    where
27        F: FnMut(T) -> S,
28    {
29        Word(self.0.map(f))
30    }
31
32    /// Extends a variable to a word.
33    pub fn extend_var<AB: SP1AirBuilder<Var = T>>(var: T) -> Word<AB::Expr> {
34        Word([AB::Expr::zero() + var, AB::Expr::zero(), AB::Expr::zero(), AB::Expr::zero()])
35    }
36}
37
38impl<T: AbstractField> Word<T> {
39    /// Extends a variable to a word.
40    pub fn extend_expr<AB: SP1AirBuilder<Expr = T>>(expr: T) -> Word<AB::Expr> {
41        Word([AB::Expr::zero() + expr, AB::Expr::zero(), AB::Expr::zero(), AB::Expr::zero()])
42    }
43
44    /// Returns a word with all zero expressions.
45    #[must_use]
46    pub fn zero<AB: SP1AirBuilder<Expr = T>>() -> Word<T> {
47        Word([AB::Expr::zero(), AB::Expr::zero(), AB::Expr::zero(), AB::Expr::zero()])
48    }
49}
50
51impl<F: Field> Word<F> {
52    /// Converts a word to a u32.
53    pub fn to_u32(&self) -> u32 {
54        u32::from_le_bytes(self.0.map(|x| x.to_string().parse::<u8>().unwrap()))
55    }
56}
57
58impl<V: Copy> Word<V> {
59    /// Reduces a word to a single variable.
60    pub fn reduce<AB: AirBuilder<Var = V>>(&self) -> AB::Expr {
61        let base = [1, 1 << 8, 1 << 16, 1 << 24].map(AB::Expr::from_canonical_u32);
62        self.0.iter().enumerate().map(|(i, x)| base[i].clone() * *x).sum()
63    }
64}
65
66impl<T> Index<usize> for Word<T> {
67    type Output = T;
68
69    fn index(&self, index: usize) -> &Self::Output {
70        &self.0[index]
71    }
72}
73
74impl<T> IndexMut<usize> for Word<T> {
75    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
76        &mut self.0[index]
77    }
78}
79
80impl<F: AbstractField> From<u32> for Word<F> {
81    fn from(value: u32) -> Self {
82        Word(value.to_le_bytes().map(F::from_canonical_u8))
83    }
84}
85
86impl<T> IntoIterator for Word<T> {
87    type Item = T;
88    type IntoIter = IntoIter<T, WORD_SIZE>;
89
90    fn into_iter(self) -> Self::IntoIter {
91        self.0.into_iter()
92    }
93}
94
95impl<T: Clone> FromIterator<T> for Word<T> {
96    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
97        let elements = iter.into_iter().take(WORD_SIZE).collect_vec();
98
99        Word(array_ref![elements, 0, WORD_SIZE].clone())
100    }
101}
sp1_stark/word.rs

sp1_stark/
word.rs
