1use alloc::vec::Vec;
  2
  3use winter_math::{fields::f64::BaseElement, FieldElement, ToElements};
  4use winter_prover::{
  5    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
  6    TransitionConstraintDegree,
  7};
  8
  9use crate::{
 10    hash::{ARK1, ARK2, MDS, STATE_WIDTH},
 11    Word, ZERO,
 12};
 13
 14// CONSTANTS
 15// ================================================================================================
 16
 17pub const HASH_CYCLE_LEN: usize = 8;
 18
 19// AIR
 20// ================================================================================================
 21
 22pub struct RescueAir {
 23    context: AirContext<BaseElement>,
 24    pub_key: Word,
 25}
 26
 27impl Air for RescueAir {
 28    type BaseField = BaseElement;
 29    type PublicInputs = PublicInputs;
 30
 31    type GkrProof = ();
 32    type GkrVerifier = ();
 33
 34    // CONSTRUCTOR
 35    // --------------------------------------------------------------------------------------------
 36    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
 37        let degrees = vec![
 38            // Apply RPO rounds.
 39            TransitionConstraintDegree::new(7),
 40            TransitionConstraintDegree::new(7),
 41            TransitionConstraintDegree::new(7),
 42            TransitionConstraintDegree::new(7),
 43            TransitionConstraintDegree::new(7),
 44            TransitionConstraintDegree::new(7),
 45            TransitionConstraintDegree::new(7),
 46            TransitionConstraintDegree::new(7),
 47            TransitionConstraintDegree::new(7),
 48            TransitionConstraintDegree::new(7),
 49            TransitionConstraintDegree::new(7),
 50            TransitionConstraintDegree::new(7),
 51        ];
 52        assert_eq!(STATE_WIDTH, trace_info.width());
 53        let context = AirContext::new(trace_info, degrees, 12, options);
 54        let context = context.set_num_transition_exemptions(1);
 55        RescueAir { context, pub_key: pub_inputs.pub_key }
 56    }
 57
 58    fn context(&self) -> &AirContext<Self::BaseField> {
 59        &self.context
 60    }
 61
 62    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
 63        &self,
 64        frame: &EvaluationFrame<E>,
 65        periodic_values: &[E],
 66        result: &mut [E],
 67    ) {
 68        let current = frame.current();
 69        let next = frame.next();
 70        // expected state width is 12 field elements
 71        debug_assert_eq!(STATE_WIDTH, current.len());
 72        debug_assert_eq!(STATE_WIDTH, next.len());
 73
 74        enforce_rpo_round(frame, result, periodic_values);
 75    }
 76
 77    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
 78        let initial_step = 0;
 79        let last_step = self.trace_length() - 1;
 80        vec![
 81            // Assert that the capacity as well as the second half of the rate portion of the state
 82            // are initialized to `ZERO`.The first half of the rate is unconstrained as it will
 83            // contain the secret key
 84            Assertion::single(0, initial_step, Self::BaseField::ZERO),
 85            Assertion::single(1, initial_step, Self::BaseField::ZERO),
 86            Assertion::single(2, initial_step, Self::BaseField::ZERO),
 87            Assertion::single(3, initial_step, Self::BaseField::ZERO),
 88            Assertion::single(8, initial_step, Self::BaseField::ZERO),
 89            Assertion::single(9, initial_step, Self::BaseField::ZERO),
 90            Assertion::single(10, initial_step, Self::BaseField::ZERO),
 91            Assertion::single(11, initial_step, Self::BaseField::ZERO),
 92            // Assert that the public key is the correct one
 93            Assertion::single(4, last_step, self.pub_key[0]),
 94            Assertion::single(5, last_step, self.pub_key[1]),
 95            Assertion::single(6, last_step, self.pub_key[2]),
 96            Assertion::single(7, last_step, self.pub_key[3]),
 97        ]
 98    }
 99
100    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
101        get_round_constants()
102    }
103}
104
105pub struct PublicInputs {
106    pub(crate) pub_key: Word,
107    pub(crate) msg: Word,
108}
109
110impl PublicInputs {
111    pub fn new(pub_key: Word, msg: Word) -> Self {
112        Self { pub_key, msg }
113    }
114}
115
116impl ToElements<BaseElement> for PublicInputs {
117    fn to_elements(&self) -> Vec<BaseElement> {
118        let mut res = self.pub_key.to_vec();
119        res.extend_from_slice(self.msg.as_ref());
120        res
121    }
122}
123
124// HELPER EVALUATORS
125// ------------------------------------------------------------------------------------------------
126
127/// Enforces constraints for a single round of the Rescue Prime Optimized hash functions.
128pub fn enforce_rpo_round<E: FieldElement + From<BaseElement>>(
129    frame: &EvaluationFrame<E>,
130    result: &mut [E],
131    ark: &[E],
132) {
133    // compute the state that should result from applying the first 5 operations of the RPO round to
134    // the current hash state.
135    let mut step1 = [E::ZERO; STATE_WIDTH];
136    step1.copy_from_slice(frame.current());
137
138    apply_mds(&mut step1);
139    // add constants
140    for i in 0..STATE_WIDTH {
141        step1[i] += ark[i];
142    }
143    apply_sbox(&mut step1);
144    apply_mds(&mut step1);
145    // add constants
146    for i in 0..STATE_WIDTH {
147        step1[i] += ark[STATE_WIDTH + i];
148    }
149
150    // compute the state that should result from applying the inverse of the last operation of the
151    // RPO round to the next step of the computation.
152    let mut step2 = [E::ZERO; STATE_WIDTH];
153    step2.copy_from_slice(frame.next());
154    apply_sbox(&mut step2);
155
156    // make sure that the results are equal.
157    for i in 0..STATE_WIDTH {
158        result[i] = step2[i] - step1[i]
159    }
160}
161
162#[inline(always)]
163fn apply_sbox<E: FieldElement + From<BaseElement>>(state: &mut [E; STATE_WIDTH]) {
164    state.iter_mut().for_each(|v| {
165        let t2 = v.square();
166        let t4 = t2.square();
167        *v *= t2 * t4;
168    });
169}
170
171#[inline(always)]
172fn apply_mds<E: FieldElement + From<BaseElement>>(state: &mut [E; STATE_WIDTH]) {
173    let mut result = [E::ZERO; STATE_WIDTH];
174    result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
175        state.iter().zip(mds_row).for_each(|(&s, m)| {
176            *r += E::from(m) * s;
177        });
178    });
179    *state = result
180}
181
182/// Returns RPO round constants arranged in column-major form.
183pub fn get_round_constants() -> Vec<Vec<BaseElement>> {
184    let mut constants = Vec::new();
185    for _ in 0..(STATE_WIDTH * 2) {
186        constants.push(vec![ZERO; HASH_CYCLE_LEN]);
187    }
188
189    #[allow(clippy::needless_range_loop)]
190    for i in 0..HASH_CYCLE_LEN - 1 {
191        for j in 0..STATE_WIDTH {
192            constants[j][i] = ARK1[i][j];
193            constants[j + STATE_WIDTH][i] = ARK2[i][j];
194        }
195    }
196
197    constants
198}