openvm_stark_backend/transcript/
duplex_sponge.rs

1use p3_challenger::CanObserve;
2use p3_field::PrimeField;
3use p3_symmetric::CryptographicPermutation;
4
5use crate::{FiatShamirTranscript, StarkProtocolConfig, TranscriptHistory, TranscriptLog};
6
7/// Permutation-based duplex sponge in overwrite mode.
8///
9/// "Duplex" refers to being able to alternately absorb (observe) and squeeze
10/// (sample), rather than a single absorb phase followed by a single squeeze
11/// phase.
12///
13/// This variant operates in *overwrite mode*, meaning new inputs overwrite
14/// state elements directly (instead of, e.g., being added in).
15#[derive(Clone, Debug)]
16pub struct DuplexSponge<F, P, const WIDTH: usize, const RATE: usize> {
17    perm: P,
18    /// Poseidon2 state
19    state: [F; WIDTH],
20    /// Invariant to be preserved: 0 <= absorb_idx < RATE
21    absorb_idx: usize,
22    /// Invariant to be preserved: 0 <= sample_idx <= RATE
23    sample_idx: usize,
24    /// True iff last sample/observe triggered a permutation
25    last_op_perm: bool,
26}
27
28impl<F: Default + Copy, P, const WIDTH: usize, const RATE: usize> From<P>
29    for DuplexSponge<F, P, WIDTH, RATE>
30{
31    fn from(perm: P) -> Self {
32        Self {
33            perm,
34            state: [F::default(); WIDTH],
35            absorb_idx: 0,
36            sample_idx: 0,
37            last_op_perm: false,
38        }
39    }
40}
41
42impl<F, P, const WIDTH: usize, const RATE: usize> DuplexSponge<F, P, WIDTH, RATE> {
43    pub fn state(&self) -> &[F; WIDTH] {
44        &self.state
45    }
46
47    pub fn absorb_idx(&self) -> usize {
48        self.absorb_idx
49    }
50
51    pub fn sample_idx(&self) -> usize {
52        self.sample_idx
53    }
54}
55
56impl<F, P, const WIDTH: usize, const RATE: usize> DuplexSponge<F, P, WIDTH, RATE>
57where
58    F: Copy,
59    P: CryptographicPermutation<[F; WIDTH]>,
60{
61    /// Absorb a single value (overwrite mode).
62    pub fn absorb(&mut self, value: F) {
63        self.state[self.absorb_idx] = value;
64        self.absorb_idx += 1;
65        self.last_op_perm = self.absorb_idx == RATE;
66        if self.last_op_perm {
67            self.perm.permute_mut(&mut self.state);
68            self.absorb_idx = 0;
69            self.sample_idx = RATE;
70        }
71    }
72
73    /// Squeeze a single value. Permutes if there are pending absorbs.
74    pub fn squeeze(&mut self) -> F {
75        self.last_op_perm = self.absorb_idx != 0 || self.sample_idx == 0;
76        if self.last_op_perm {
77            self.perm.permute_mut(&mut self.state);
78            self.absorb_idx = 0;
79            self.sample_idx = RATE;
80        }
81        self.sample_idx -= 1;
82        self.state[self.sample_idx]
83    }
84}
85
86// This implementation **must** be equivalent to Plonky3's `DuplexChallenger`.
87impl<SC, F, P, const WIDTH: usize, const RATE: usize> FiatShamirTranscript<SC>
88    for DuplexSponge<F, P, WIDTH, RATE>
89where
90    F: PrimeField,
91    SC: StarkProtocolConfig<Digest = [F; RATE], F = F>,
92    P: CryptographicPermutation<[F; WIDTH]> + Send + Sync,
93{
94    fn observe(&mut self, value: F) {
95        self.absorb(value);
96    }
97
98    fn sample(&mut self) -> F {
99        self.squeeze()
100    }
101
102    fn observe_commit(&mut self, digest: [F; RATE]) {
103        CanObserve::observe(self, digest);
104    }
105}
106
107impl<F, P, const WIDTH: usize, const RATE: usize> CanObserve<F> for DuplexSponge<F, P, WIDTH, RATE>
108where
109    F: Copy,
110    P: CryptographicPermutation<[F; WIDTH]>,
111{
112    fn observe(&mut self, value: F) {
113        self.absorb(value);
114    }
115}
116
117impl<F, P, const WIDTH: usize, const RATE: usize, const N: usize> CanObserve<[F; N]>
118    for DuplexSponge<F, P, WIDTH, RATE>
119where
120    F: Copy,
121    P: CryptographicPermutation<[F; WIDTH]>,
122{
123    fn observe(&mut self, values: [F; N]) {
124        for x in values {
125            self.absorb(x);
126        }
127    }
128}
129
130#[derive(Clone)]
131pub struct DuplexSpongeRecorder<F, P, const WIDTH: usize, const RATE: usize> {
132    pub inner: DuplexSponge<F, P, WIDTH, RATE>,
133    pub log: TranscriptLog<F, [F; WIDTH]>,
134}
135
136impl<F: Default + Copy, P, const WIDTH: usize, const RATE: usize> From<P>
137    for DuplexSpongeRecorder<F, P, WIDTH, RATE>
138{
139    fn from(perm: P) -> Self {
140        let inner = DuplexSponge::from(perm);
141        let mut log = TranscriptLog::default();
142        log.push_perm_result([F::default(); WIDTH]);
143        Self { inner, log }
144    }
145}
146
147impl<SC, F, P, const WIDTH: usize, const RATE: usize> FiatShamirTranscript<SC>
148    for DuplexSpongeRecorder<F, P, WIDTH, RATE>
149where
150    F: PrimeField,
151    SC: StarkProtocolConfig<Digest = [F; RATE], F = F>,
152    P: CryptographicPermutation<[F; WIDTH]> + Send + Sync,
153{
154    fn observe(&mut self, x: F) {
155        CanObserve::observe(&mut self.inner, x);
156        self.log.push_observe(x);
157        if self.inner.last_op_perm {
158            self.log.push_perm_result(self.inner.state);
159        }
160    }
161
162    fn sample(&mut self) -> F {
163        let x = FiatShamirTranscript::<SC>::sample(&mut self.inner);
164        self.log.push_sample(x);
165        if self.inner.last_op_perm {
166            self.log.push_perm_result(self.inner.state);
167        }
168        x
169    }
170
171    fn observe_commit(&mut self, digest: [F; RATE]) {
172        for x in digest {
173            FiatShamirTranscript::<SC>::observe(self, x);
174        }
175    }
176}
177
178impl<F, P, const WIDTH: usize, const RATE: usize> TranscriptHistory
179    for DuplexSpongeRecorder<F, P, WIDTH, RATE>
180{
181    type F = F;
182    type State = [F; WIDTH];
183
184    fn len(&self) -> usize {
185        self.log.len()
186    }
187
188    fn into_log(self) -> TranscriptLog<F, [F; WIDTH]> {
189        self.log
190    }
191}
192
193/// [DuplexSpongeRecorder] that checks the live transcript logs against a provided transcript log.
194/// For testing usage.
195#[derive(Clone)]
196pub struct DuplexSpongeValidator<F, P, const WIDTH: usize, const RATE: usize> {
197    pub inner: DuplexSpongeRecorder<F, P, WIDTH, RATE>,
198    pub idx: usize,
199    log: TranscriptLog<F, [F; WIDTH]>,
200}
201
202impl<F: Default + Copy, P, const WIDTH: usize, const RATE: usize>
203    DuplexSpongeValidator<F, P, WIDTH, RATE>
204{
205    pub fn new(perm: P, log: TranscriptLog<F, [F; WIDTH]>) -> Self {
206        debug_assert_eq!(log.len(), log.samples().len());
207        Self {
208            inner: perm.into(),
209            idx: 0,
210            log,
211        }
212    }
213}
214
215impl<SC, F, P, const WIDTH: usize, const RATE: usize> FiatShamirTranscript<SC>
216    for DuplexSpongeValidator<F, P, WIDTH, RATE>
217where
218    F: PrimeField,
219    SC: StarkProtocolConfig<Digest = [F; RATE], F = F>,
220    P: CryptographicPermutation<[F; WIDTH]> + Send + Sync,
221{
222    fn observe(&mut self, x: F) {
223        debug_assert!(self.idx < self.log.len(), "transcript replay overflow");
224        assert!(!self.log.samples()[self.idx]);
225        let exp_x = self.log[self.idx];
226        assert_eq!(x, exp_x);
227        self.idx += 1;
228        FiatShamirTranscript::<SC>::observe(&mut self.inner, x);
229    }
230
231    fn sample(&mut self) -> F {
232        debug_assert!(self.idx < self.log.len(), "transcript replay overflow");
233        assert!(self.log.samples()[self.idx]);
234        let x = FiatShamirTranscript::<SC>::sample(&mut self.inner);
235        let exp_x = self.log[self.idx];
236        self.idx += 1;
237        assert_eq!(x, exp_x);
238        x
239    }
240
241    fn observe_commit(&mut self, digest: [F; RATE]) {
242        for x in digest {
243            FiatShamirTranscript::<SC>::observe(self, x);
244        }
245    }
246}
247
248impl<F, P, const WIDTH: usize, const RATE: usize> TranscriptHistory
249    for DuplexSpongeValidator<F, P, WIDTH, RATE>
250{
251    type F = F;
252    type State = [F; WIDTH];
253
254    fn len(&self) -> usize {
255        self.inner.len()
256    }
257
258    fn into_log(self) -> TranscriptLog<F, [F; WIDTH]> {
259        debug_assert_eq!(self.inner.len(), self.log.len());
260        debug_assert_eq!(
261            self.inner.len(),
262            self.idx,
263            "transcript replay ended with {} of {} entries consumed",
264            self.idx,
265            self.inner.len()
266        );
267        debug_assert_eq!(
268            self.log.len(),
269            self.idx,
270            "transcript replay ended with {} of {} entries consumed",
271            self.idx,
272            self.log.len()
273        );
274        self.inner.into_log()
275    }
276}
openvm_stark_backend/transcript/duplex_sponge.rs

openvm_stark_backend/transcript/
duplex_sponge.rs
