1use core::marker::PhantomData;
  2
  3use rand_chacha::ChaCha20Rng;
  4use winter_air::{
  5    AuxRandElements, ConstraintCompositionCoefficients, PartitionOptions, ZkParameters,
  6};
  7use winter_crypto::{ElementHasher, SaltedMerkleTree};
  8use winter_math::{fields::f64::BaseElement, FieldElement};
  9use winter_prover::{
 10    matrix::ColMatrix, CompositionPoly, CompositionPolyTrace, DefaultConstraintCommitment,
 11    DefaultConstraintEvaluator, DefaultTraceLde, ProofOptions, Prover, StarkDomain, Trace,
 12    TraceInfo, TracePolyTable, TraceTable,
 13};
 14
 15use super::air::{PublicInputs, RescueAir, HASH_CYCLE_LEN};
 16use crate::{
 17    hash::{rpo::Rpo256, STATE_WIDTH},
 18    rand::RpoRandomCoin,
 19    Word, ZERO,
 20};
 21
 22// PROVER
 23// ================================================================================================
 24
 25/// A prover for the RPO STARK-based signature scheme.
 26///
 27/// The signature is based on the the one-wayness of the RPO hash function but it is generic over
 28/// the hash function used for instantiating the random oracle for the BCS transform.
 29pub(crate) struct RpoSignatureProver<H: ElementHasher + Sync> {
 30    message: Word,
 31    options: ProofOptions,
 32    _hasher: PhantomData<H>,
 33}
 34
 35impl<H: ElementHasher + Sync> RpoSignatureProver<H> {
 36    pub(crate) fn new(message: Word, options: ProofOptions) -> Self {
 37        Self { message, options, _hasher: PhantomData }
 38    }
 39
 40    pub(crate) fn build_trace(&self, sk: Word) -> TraceTable<BaseElement> {
 41        let mut trace = TraceTable::new(STATE_WIDTH, HASH_CYCLE_LEN);
 42
 43        trace.fill(
 44            |state| {
 45                // initialize first half of the rate portion of the state with the secret key
 46                state[0] = ZERO;
 47                state[1] = ZERO;
 48                state[2] = ZERO;
 49                state[3] = ZERO;
 50                state[4] = sk[0];
 51                state[5] = sk[1];
 52                state[6] = sk[2];
 53                state[7] = sk[3];
 54                state[8] = ZERO;
 55                state[9] = ZERO;
 56                state[10] = ZERO;
 57                state[11] = ZERO;
 58            },
 59            |step, state| {
 60                Rpo256::apply_round(
 61                    state.try_into().expect("should not fail given the size of the array"),
 62                    step,
 63                );
 64            },
 65        );
 66        trace
 67    }
 68}
 69
 70impl<H: ElementHasher> Prover for RpoSignatureProver<H>
 71where
 72    H: ElementHasher<BaseField = BaseElement> + Sync,
 73{
 74    type BaseField = BaseElement;
 75    type Air = RescueAir;
 76    type Trace = TraceTable<BaseElement>;
 77    type HashFn = Rpo256;
 78    type VC = SaltedMerkleTree<Self::HashFn, Self::ZkPrng>;
 79    type RandomCoin = RpoRandomCoin;
 80    type TraceLde<E: FieldElement<BaseField = Self::BaseField>> =
 81        DefaultTraceLde<E, Self::HashFn, Self::VC>;
 82    type ConstraintCommitment<E: FieldElement<BaseField = Self::BaseField>> =
 83        DefaultConstraintCommitment<E, Self::HashFn, Self::ZkPrng, Self::VC>;
 84    type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> =
 85        DefaultConstraintEvaluator<'a, Self::Air, E>;
 86    type ZkPrng = ChaCha20Rng;
 87
 88    fn get_pub_inputs(&self, trace: &Self::Trace) -> PublicInputs {
 89        let last_step = trace.length() - 1;
 90        // Note that the message is not part of the execution trace but is part of the public
 91        // inputs. This is explained in the reference description of the DSA and intuitively
 92        // it is done in order to make sure that the message is part of the Fiat-Shamir
 93        // transcript and hence binds the proof/signature to the message
 94        PublicInputs {
 95            pub_key: [
 96                trace.get(4, last_step),
 97                trace.get(5, last_step),
 98                trace.get(6, last_step),
 99                trace.get(7, last_step),
100            ],
101            msg: self.message,
102        }
103    }
104
105    fn options(&self) -> &ProofOptions {
106        &self.options
107    }
108
109    fn new_trace_lde<E: FieldElement<BaseField = Self::BaseField>>(
110        &self,
111        trace_info: &TraceInfo,
112        main_trace: &ColMatrix<Self::BaseField>,
113        domain: &StarkDomain<Self::BaseField>,
114        partition_option: PartitionOptions,
115        zk_parameters: Option<ZkParameters>,
116        prng: &mut Option<Self::ZkPrng>,
117    ) -> (Self::TraceLde<E>, TracePolyTable<E>) {
118        DefaultTraceLde::new(trace_info, main_trace, domain, partition_option, zk_parameters, prng)
119    }
120
121    fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(
122        &self,
123        air: &'a Self::Air,
124        aux_rand_elements: Option<AuxRandElements<E>>,
125        composition_coefficients: ConstraintCompositionCoefficients<E>,
126    ) -> Self::ConstraintEvaluator<'a, E> {
127        DefaultConstraintEvaluator::new(air, aux_rand_elements, composition_coefficients)
128    }
129
130    fn build_constraint_commitment<E: FieldElement<BaseField = Self::BaseField>>(
131        &self,
132        composition_poly_trace: CompositionPolyTrace<E>,
133        num_constraint_composition_columns: usize,
134        domain: &StarkDomain<Self::BaseField>,
135        partition_options: PartitionOptions,
136        zk_parameters: Option<ZkParameters>,
137        prng: &mut Option<Self::ZkPrng>,
138    ) -> (Self::ConstraintCommitment<E>, CompositionPoly<E>) {
139        DefaultConstraintCommitment::new(
140            composition_poly_trace,
141            num_constraint_composition_columns,
142            domain,
143            partition_options,
144            zk_parameters,
145            prng,
146        )
147    }
148}