1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

//! Mina Poseidon hasher
//!
//! An implementation of Mina's hasher based on the poseidon arithmetic sponge
//!
use std::marker::PhantomData;

use crate::DomainParameter;
use mina_curves::pasta::Fp;
use mina_poseidon::{
    constants::{PlonkSpongeConstantsKimchi, PlonkSpongeConstantsLegacy, SpongeConstants},
    pasta,
    poseidon::{ArithmeticSponge, ArithmeticSpongeParams, Sponge, SpongeState},
};

use super::{domain_prefix_to_field, Hashable, Hasher};

/// Poseidon hasher context
//
//  The arithmetic sponge parameters are large and costly to initialize,
//  so we only want to do this once and then re-use the Poseidon context
//  for many hashes. Also, following approach of the mina code we store
//  a backup of the initialized sponge state for efficient reuse.
pub struct Poseidon<SC: SpongeConstants, H: Hashable> {
    sponge: ArithmeticSponge<Fp, SC>,
    sponge_state: SpongeState,
    state: Vec<Fp>,
    phantom: PhantomData<H>,
}

impl<SC: SpongeConstants, H: Hashable> Poseidon<SC, H> {
    fn new(domain_param: H::D, sponge_params: &'static ArithmeticSpongeParams<Fp>) -> Self {
        let mut poseidon = Poseidon::<SC, H> {
            sponge: ArithmeticSponge::<Fp, SC>::new(sponge_params),
            sponge_state: SpongeState::Absorbed(0),
            state: vec![],
            phantom: PhantomData,
        };

        poseidon.init(domain_param);

        poseidon
    }
}

/// Poseidon hasher type with legacy plonk sponge constants
pub type PoseidonHasherLegacy<H> = Poseidon<PlonkSpongeConstantsLegacy, H>;

/// Create a legacy hasher context
pub(crate) fn new_legacy<H: Hashable>(domain_param: H::D) -> PoseidonHasherLegacy<H> {
    Poseidon::<PlonkSpongeConstantsLegacy, H>::new(domain_param, pasta::fp_legacy::static_params())
}

/// Poseidon hasher type with experimental kimchi plonk sponge constants
pub type PoseidonHasherKimchi<H> = Poseidon<PlonkSpongeConstantsKimchi, H>;

/// Create an experimental kimchi hasher context
pub(crate) fn new_kimchi<H: Hashable>(domain_param: H::D) -> PoseidonHasherKimchi<H> {
    Poseidon::<PlonkSpongeConstantsKimchi, H>::new(domain_param, pasta::fp_kimchi::static_params())
}

impl<SC: SpongeConstants, H: Hashable> Hasher<H> for Poseidon<SC, H>
where
    H::D: DomainParameter,
{
    fn reset(&mut self) -> &mut dyn Hasher<H> {
        // Efficient reset
        self.sponge.sponge_state = self.sponge_state.clone();
        self.sponge.state = self.state.clone();

        self
    }

    fn init(&mut self, domain_param: H::D) -> &mut dyn Hasher<H> {
        // Set sponge initial state and save it so the hasher context can be reused efficiently
        // N.B. Mina sets the sponge's initial state by hashing the input type's domain bytes
        self.sponge.reset();

        if let Some(domain_string) = H::domain_string(domain_param) {
            self.sponge
                .absorb(&[domain_prefix_to_field::<Fp>(domain_string)]);
            self.sponge.squeeze();
        }

        // Save initial state for efficient reset
        self.sponge_state = self.sponge.sponge_state.clone();
        self.state = self.sponge.state.clone();

        self
    }

    fn update(&mut self, input: &H) -> &mut dyn Hasher<H> {
        self.sponge.absorb(&input.to_roinput().to_fields());

        self
    }

    fn digest(&mut self) -> Fp {
        let output = self.sponge.squeeze();
        self.sponge.reset();
        output
    }
}