openvm_stark_backend/transcript/
traits.rs

1use std::ops::Deref;
2
3use p3_field::{BasedVectorSpace, PrimeCharacteristicRing, PrimeField, PrimeField64};
4use p3_maybe_rayon::prelude::*;
5use tracing::instrument;
6
7use crate::StarkProtocolConfig;
8
9/// Unified trait describing the interface for the Fiat-Shamir transcript needed by the SWIRL
10/// protocol.
11pub trait FiatShamirTranscript<SC>: Clone + Send + Sync
12where
13    SC: StarkProtocolConfig,
14{
15    fn observe(&mut self, value: SC::F);
16    fn sample(&mut self) -> SC::F;
17
18    /// Implementations should pass through to [Self::observe], but no default implementation is
19    /// provided since an explicit conversion from `Digest` to array of `F` is required.
20    fn observe_commit(&mut self, digest: SC::Digest);
21
22    fn observe_ext(&mut self, value: SC::EF) {
23        // for i in 0..D
24        for &base_val in value.as_basis_coefficients_slice() {
25            self.observe(base_val);
26        }
27    }
28
29    fn sample_ext(&mut self) -> SC::EF {
30        SC::EF::from_basis_coefficients_fn(|_| self.sample())
31    }
32
33    /// Samples and returns an integer in the range [0, 2^bits).
34    ///
35    /// # Sampling bias
36    ///
37    /// The sampling is not uniform over `[0, 2^bits)`; it works by sampling a random field element,
38    /// using its canonical embedding, and then reducing it mod 2^bits. Writing the field order as
39    /// `p = c * 2^bits + r` with `0 <= r < 2^bits`, the `r` "heavy" residues `0..r` each occur with
40    /// probability `(c + 1) / p` while the remaining `2^bits - r` "light" residues occur with
41    /// probability `c / p`. Equivalently, residues `0..r` are favored by a factor of `(c + 1) / c`.
42    /// Note residue `0` is always heavy.
43    ///
44    /// For BabyBear, `p - 1 = 2^27 * 15`, so `p ≡ 1 (mod 2^bits)` and thus `r = 1` for every
45    /// `bits <= 27`—exactly one heavy residue (the all-zero one).
46    fn sample_bits(&mut self, bits: usize) -> u64 {
47        assert!(bits < (u32::BITS as usize));
48        assert!((1 << bits) < SC::F::ORDER_U64);
49        let rand_f: SC::F = self.sample();
50        let rand_u64 = rand_f.as_canonical_u64();
51        rand_u64 & ((1 << bits) - 1)
52    }
53
54    /// Checks a proof-of-work witness: observes `witness` into the transcript and accepts iff the
55    /// next [`Self::sample_bits`] is zero (probability `≈ 2^{-bits}` for a random witness, so
56    /// [`Self::grind`] tries `≈ 2^bits` witnesses to find one).
57    ///
58    /// The work is done with the transcript hash (the `observe`/`sample_bits` sponge permutation),
59    /// so `bits` counts hash evaluations; the attacker's wall-clock cost is `≈ 2^bits · cost(H)`
60    /// for that hash `H`. Soundness targets in bits are therefore comparable only across configs
61    /// that share a grinding hash.
62    #[must_use]
63    fn check_witness(&mut self, bits: usize, witness: SC::F) -> bool {
64        if bits == 0 {
65            return true;
66        }
67        self.observe(witness);
68        self.sample_bits(bits) == 0
69    }
70
71    /// Finds a proof-of-work witness `w` such that [`Self::check_witness`]`(bits, w)` holds, by
72    /// brute force over the base field.
73    #[instrument(name = "grind_pow", skip_all)]
74    fn grind(&mut self, bits: usize) -> SC::F {
75        assert!(bits < (u32::BITS as usize));
76        assert!((1 << bits) < SC::F::ORDER_U64);
77        // Trivial case: 0 bits mean no PoW is required and any witness is valid.
78        if bits == 0 {
79            return SC::F::ZERO;
80        }
81
82        let witness = (0..SC::F::ORDER_U64)
83            .into_par_iter()
84            .map(PrimeCharacteristicRing::from_u64)
85            .find_any(|witness| self.clone().check_witness(bits, *witness))
86            .expect("failed to find PoW witness");
87        assert!(self.check_witness(bits, witness));
88        witness
89    }
90}
91
92pub trait TranscriptHistory {
93    type F;
94    type State;
95
96    fn len(&self) -> usize;
97    fn into_log(self) -> TranscriptLog<Self::F, Self::State>;
98
99    fn is_empty(&self) -> bool {
100        self.len() == 0
101    }
102}
103
104/// Log of transcript history
105#[derive(Clone, Debug)]
106pub struct TranscriptLog<F, State> {
107    /// Every sampled or observed value F
108    values: Vec<F>,
109    /// True iff values[tidx] was a sampled value
110    is_sample: Vec<bool>,
111    /// Sponge state after every permutation; note that not all implementations of
112    /// TranscriptHistory will define this
113    perm_results: Vec<State>,
114}
115
116impl<F, State> Default for TranscriptLog<F, State> {
117    fn default() -> Self {
118        Self {
119            values: Vec::new(),
120            is_sample: Vec::new(),
121            perm_results: Vec::new(),
122        }
123    }
124}
125
126impl<F: Clone, State> TranscriptLog<F, State> {
127    pub fn new(values: Vec<F>, is_sample: Vec<bool>) -> Self {
128        debug_assert_eq!(values.len(), is_sample.len());
129        Self {
130            values,
131            is_sample,
132            perm_results: vec![],
133        }
134    }
135
136    pub fn values(&self) -> &[F] {
137        &self.values
138    }
139
140    pub fn values_mut(&mut self) -> &mut [F] {
141        &mut self.values
142    }
143
144    pub fn samples(&self) -> &[bool] {
145        &self.is_sample
146    }
147
148    pub fn samples_mut(&mut self) -> &mut [bool] {
149        &mut self.is_sample
150    }
151
152    pub fn push_observe(&mut self, value: F) {
153        self.values.push(value);
154        self.is_sample.push(false);
155    }
156
157    pub fn push_sample(&mut self, value: F) {
158        self.values.push(value);
159        self.is_sample.push(true);
160    }
161
162    pub fn push_perm_result(&mut self, state: State) {
163        self.perm_results.push(state);
164    }
165
166    pub fn extend_observe(&mut self, values: &[F]) {
167        self.values.extend_from_slice(values);
168        self.is_sample
169            .extend(core::iter::repeat_n(false, values.len()));
170    }
171
172    pub fn extend_sample(&mut self, values: &[F]) {
173        self.values.extend_from_slice(values);
174        self.is_sample
175            .extend(core::iter::repeat_n(true, values.len()));
176    }
177
178    pub fn extend_with_flags(&mut self, values: &[F], sample_flags: &[bool]) {
179        debug_assert_eq!(values.len(), sample_flags.len());
180        self.values.extend_from_slice(values);
181        self.is_sample.extend(sample_flags.iter().copied());
182    }
183
184    pub fn len(&self) -> usize {
185        self.values.len()
186    }
187
188    pub fn is_empty(&self) -> bool {
189        self.values.is_empty()
190    }
191
192    pub fn into_parts(self) -> (Vec<F>, Vec<bool>) {
193        (self.values, self.is_sample)
194    }
195
196    pub fn perm_results(&self) -> &Vec<State> {
197        &self.perm_results
198    }
199}
200
201impl<F, State> Deref for TranscriptLog<F, State> {
202    type Target = [F];
203
204    fn deref(&self) -> &Self::Target {
205        &self.values
206    }
207}
208
209/// Read-only transcript that replays a recorded log.
210#[derive(Clone, Debug)]
211pub struct ReadOnlyTranscript<'a, F, State> {
212    log: &'a TranscriptLog<F, State>,
213    position: usize,
214}
215
216impl<'a, F, State> ReadOnlyTranscript<'a, F, State> {
217    pub fn new(log: &'a TranscriptLog<F, State>, start_idx: usize) -> Self {
218        debug_assert!(start_idx <= log.len(), "start index out of bounds");
219        Self {
220            log,
221            position: start_idx,
222        }
223    }
224}
225
226impl<SC, F, const WIDTH: usize, const RATE: usize> FiatShamirTranscript<SC>
227    for ReadOnlyTranscript<'_, F, [F; WIDTH]>
228where
229    F: PrimeField,
230    SC: StarkProtocolConfig<F = F, Digest = [F; RATE]>,
231{
232    #[inline]
233    fn observe(&mut self, value: F) {
234        debug_assert!(
235            !self.log.samples()[self.position],
236            "expected observe at {}",
237            self.position
238        );
239        debug_assert_eq!(
240            self.log.values()[self.position],
241            value,
242            "value mismatch at {}",
243            self.position
244        );
245        self.position += 1;
246    }
247
248    #[inline]
249    fn sample(&mut self) -> F {
250        debug_assert!(
251            self.log.samples()[self.position],
252            "expected sample at {}",
253            self.position
254        );
255        let value = self.log.values()[self.position];
256        self.position += 1;
257        value
258    }
259
260    fn observe_commit(&mut self, digest: [F; RATE]) {
261        for x in digest {
262            FiatShamirTranscript::<SC>::observe(self, x);
263        }
264    }
265}
266
267impl<F, State> TranscriptHistory for ReadOnlyTranscript<'_, F, State>
268where
269    F: Clone,
270    State: Clone,
271{
272    type F = F;
273    type State = State;
274
275    fn len(&self) -> usize {
276        self.position
277    }
278
279    fn into_log(self) -> TranscriptLog<F, State> {
280        self.log.clone()
281    }
282}
openvm_stark_backend/transcript/traits.rs

openvm_stark_backend/transcript/
traits.rs
