1  // Copyright (c) 2019-2025 Alpha-Delta Network Inc.
  2  // This file is part of the alphavm library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  mod hash;
 17  mod hash_many;
 18  mod hash_to_group;
 19  mod hash_to_scalar;
 20  mod prf;
 21  
 22  #[cfg(test)]
 23  use alphavm_circuit_types::environment::assert_scope;
 24  #[cfg(test)]
 25  use alphavm_utilities::{TestRng, Uniform};
 26  
 27  use crate::{Elligator2, Hash, HashMany, HashToGroup, HashToScalar, PRF};
 28  use alphavm_circuit_types::{Field, Group, Scalar, environment::prelude::*};
 29  
 30  /// Poseidon2 is a cryptographic hash function of input rate 2.
 31  pub type Poseidon2<E> = Poseidon<E, 2>;
 32  /// Poseidon4 is a cryptographic hash function of input rate 4.
 33  pub type Poseidon4<E> = Poseidon<E, 4>;
 34  /// Poseidon8 is a cryptographic hash function of input rate 8.
 35  pub type Poseidon8<E> = Poseidon<E, 8>;
 36  
 37  const CAPACITY: usize = 1;
 38  
 39  /// The mode structure for duplex sponges.
 40  #[derive(PartialEq, Eq, Clone, Debug)]
 41  pub enum DuplexSpongeMode {
 42      /// The sponge is currently absorbing data.
 43      Absorbing {
 44          /// The next position of the state to be XOR-ed when absorbing.
 45          next_absorb_index: usize,
 46      },
 47      /// The sponge is currently squeezing data out.
 48      Squeezing {
 49          /// The next position of the state to be outputted when squeezing.
 50          next_squeeze_index: usize,
 51      },
 52  }
 53  
 54  #[derive(Clone)]
 55  pub struct Poseidon<E: Environment, const RATE: usize> {
 56      /// The domain separator for the Poseidon hash function.
 57      domain: Field<E>,
 58      /// The number of rounds in a full-round operation.
 59      full_rounds: usize,
 60      /// The number of rounds in a partial-round operation.
 61      partial_rounds: usize,
 62      /// The exponent used in S-boxes.
 63      alpha: Field<E>,
 64      /// The additive round keys. These are added before each MDS matrix application to make it an affine shift.
 65      /// They are indexed by `ark[round_number][state_element_index]`
 66      ark: Vec<Vec<Field<E>>>,
 67      /// The Maximally Distance Separating (MDS) matrix.
 68      mds: Vec<Vec<Field<E>>>,
 69  }
 70  
 71  impl<E: Environment, const RATE: usize> Inject for Poseidon<E, RATE> {
 72      type Primitive = console::Poseidon<E::Network, RATE>;
 73  
 74      fn new(_mode: Mode, poseidon: Self::Primitive) -> Self {
 75          // Initialize the domain separator.
 76          let domain = Field::constant(poseidon.domain());
 77  
 78          // Initialize the Poseidon parameters.
 79          let parameters = poseidon.parameters();
 80          let full_rounds = parameters.full_rounds;
 81          let partial_rounds = parameters.partial_rounds;
 82          let alpha = Field::constant(console::Field::from_u128(parameters.alpha as u128));
 83          // Cache the bits for the field element.
 84          alpha.to_bits_le();
 85          let ark = parameters
 86              .ark
 87              .iter()
 88              .take(full_rounds + partial_rounds)
 89              .map(|round| {
 90                  round.iter().take(RATE + 1).copied().map(|field| Field::constant(console::Field::new(field))).collect()
 91              })
 92              .collect();
 93          let mds = parameters
 94              .mds
 95              .iter()
 96              .take(RATE + 1)
 97              .map(|round| {
 98                  round.iter().take(RATE + 1).copied().map(|field| Field::constant(console::Field::new(field))).collect()
 99              })
100              .collect();
101  
102          Self { domain, full_rounds, partial_rounds, alpha, ark, mds }
103      }
104  }