use core::ops::Range;
use super::{Felt, Word, ZERO};
use crate::field::BasedVectorSpace;
pub(crate) mod poseidon2;
pub(crate) mod rescue;
pub(crate) const STATE_WIDTH: usize = 12;
pub(crate) const RATE_RANGE: Range<usize> = 0..8;
pub(crate) const RATE_WIDTH: usize = RATE_RANGE.end - RATE_RANGE.start;
pub(crate) const RATE0_RANGE: Range<usize> = 0..4;
pub(crate) const RATE1_RANGE: Range<usize> = 4..8;
pub(crate) const CAPACITY_RANGE: Range<usize> = 8..12;
pub(crate) const DIGEST_RANGE: Range<usize> = 0..4;
const BINARY_CHUNK_SIZE: usize = 7;
pub(crate) trait AlgebraicSponge {
fn apply_permutation(state: &mut [Felt; STATE_WIDTH]);
fn hash_elements<E>(elements: &[E]) -> Word
where
E: BasedVectorSpace<Felt>,
{
let state = [ZERO; STATE_WIDTH];
hash_elements_internal::<E, Self>(elements, state)
}
fn hash(bytes: &[u8]) -> Word {
let mut state = [ZERO; STATE_WIDTH];
let num_field_elem = bytes.len().div_ceil(BINARY_CHUNK_SIZE);
state[CAPACITY_RANGE.start] =
Felt::from_u8((RATE_WIDTH + (num_field_elem % RATE_WIDTH)) as u8);
let mut buf = [0_u8; 8];
let mut current_chunk_idx = 0_usize;
let last_chunk_idx = if num_field_elem == 0 {
current_chunk_idx
} else {
num_field_elem - 1
};
let rate_pos = bytes.chunks(BINARY_CHUNK_SIZE).fold(0, |rate_pos, chunk| {
if current_chunk_idx != last_chunk_idx {
buf[..BINARY_CHUNK_SIZE].copy_from_slice(chunk);
} else {
buf.fill(0);
buf[..chunk.len()].copy_from_slice(chunk);
buf[chunk.len()] = 1;
}
current_chunk_idx += 1;
state[RATE_RANGE.start + rate_pos] = Felt::new_unchecked(u64::from_le_bytes(buf));
if rate_pos == RATE_WIDTH - 1 {
Self::apply_permutation(&mut state);
0
} else {
rate_pos + 1
}
});
if rate_pos != 0 {
state[RATE_RANGE.start + rate_pos..RATE_RANGE.end].fill(ZERO);
Self::apply_permutation(&mut state);
}
Word::new(state[DIGEST_RANGE].try_into().unwrap())
}
fn merge(values: &[Word; 2]) -> Word {
let mut state = [ZERO; STATE_WIDTH];
let it = Word::words_as_elements_iter(values.iter());
for (i, v) in it.enumerate() {
state[RATE_RANGE.start + i] = *v;
}
Self::apply_permutation(&mut state);
Word::new(state[DIGEST_RANGE].try_into().unwrap())
}
fn merge_many(values: &[Word]) -> Word {
let elements = Word::words_as_elements(values);
Self::hash_elements(elements)
}
fn merge_in_domain(values: &[Word; 2], domain: Felt) -> Word {
let mut state = [ZERO; STATE_WIDTH];
let it = Word::words_as_elements_iter(values.iter());
for (i, v) in it.enumerate() {
state[RATE_RANGE.start + i] = *v;
}
state[CAPACITY_RANGE.start + 1] = domain;
Self::apply_permutation(&mut state);
Word::new(state[DIGEST_RANGE].try_into().unwrap())
}
fn hash_elements_in_domain<E>(elements: &[E], domain: Felt) -> Word
where
E: BasedVectorSpace<Felt>,
{
let mut state = [ZERO; STATE_WIDTH];
state[CAPACITY_RANGE.start + 1] = domain;
hash_elements_internal::<E, Self>(elements, state)
}
}
fn hash_elements_internal<E, S>(elements: &[E], mut state: [Felt; STATE_WIDTH]) -> Word
where
E: BasedVectorSpace<Felt>,
S: AlgebraicSponge + ?Sized,
{
let total_len = elements
.iter()
.map(|elem| E::as_basis_coefficients_slice(elem).len())
.sum::<usize>();
state[CAPACITY_RANGE.start] = Felt::from_u8((total_len % RATE_WIDTH) as u8);
let mut i = 0;
for elem in elements.iter() {
for &felt in E::as_basis_coefficients_slice(elem) {
state[RATE_RANGE.start + i] = felt;
i += 1;
if i.is_multiple_of(RATE_WIDTH) {
S::apply_permutation(&mut state);
i = 0;
}
}
}
if i > 0 {
while i != RATE_WIDTH {
state[RATE_RANGE.start + i] = ZERO;
i += 1;
}
S::apply_permutation(&mut state);
}
Word::new(state[DIGEST_RANGE].try_into().unwrap())
}