diff --git a/crates/leanVm/src/air/mod.rs b/crates/leanVm/src/air/mod.rs index aa8f8a64..28174adf 100644 --- a/crates/leanVm/src/air/mod.rs +++ b/crates/leanVm/src/air/mod.rs @@ -1,3 +1,4 @@ pub mod constant; pub mod dot_product; +pub mod poseidon; pub mod vm; diff --git a/crates/leanVm/src/air/poseidon.rs b/crates/leanVm/src/air/poseidon.rs new file mode 100644 index 00000000..a99cdb52 --- /dev/null +++ b/crates/leanVm/src/air/poseidon.rs @@ -0,0 +1,126 @@ +use p3_field::PrimeCharacteristicRing; + +use crate::{ + constant::F, + poseidon2::{generate_trace_poseidon_16, generate_trace_poseidon_24}, + utils::pad_to_power_of_two, + witness::poseidon::{WitnessPoseidon16, WitnessPoseidon24}, +}; + +/// Build Poseidon Columns +/// +/// This function generates the execution traces for both Poseidon16 and Poseidon24 operations. +/// +/// ## Arguments +/// * `poseidons_16`: A slice of `WitnessPoseidon16` structs. +/// * `poseidons_24`: A slice of `WitnessPoseidon24` structs. +/// +/// ## Returns +/// A tuple containing two items: +/// 1. A `Vec>` for the Poseidon16 trace, where each inner vector is a column. +/// 2. A `Vec>` for the Poseidon24 trace, where each inner vector is a column. +pub fn build_poseidon_columns( + poseidons_16: &[WitnessPoseidon16], + poseidons_24: &[WitnessPoseidon24], +) -> (Vec>, Vec>) { + // Poseidon16 Trace Generation + // + // Extract the `input` state from each Poseidon16 witness. + let poseidon_16_data = poseidons_16.iter().map(|w| w.input).collect::>(); + // Generate the full execution trace for all Poseidon16 operations. + // This returns a `RowMajorMatrix`. + let witness_matrix_poseidon_16 = generate_trace_poseidon_16(poseidon_16_data); + + // Poseidon24 Trace Generation + // + // Extract the `input` state from each Poseidon24 witness. + let poseidon_24_data = poseidons_24.iter().map(|w| w.input).collect::>(); + // Generate the full execution trace for all Poseidon24 operations. + let witness_matrix_poseidon_24 = generate_trace_poseidon_24(poseidon_24_data); + + // Column Extraction + // + // To get the columns of a row-major matrix, we first transpose it. + let transposed_16 = witness_matrix_poseidon_16.transpose(); + // After transposing, the rows of the new matrix are the columns of the original. + // We can then iterate over the row slices of the transposed matrix and collect them. + let cols_16 = transposed_16.row_slices().map(<[F]>::to_vec).collect(); + + // Repeat the same process for the Poseidon24 trace. + let transposed_24 = witness_matrix_poseidon_24.transpose(); + let cols_24 = transposed_24.row_slices().map(<[F]>::to_vec).collect(); + + // Return the two sets of columns. + (cols_16, cols_24) +} + +/// All Poseidon16 Indexes +/// +/// This function collects all memory addresses used by the Poseidon16 precompile, +/// flattens them into a single vector, and pads it to the next power of two. +/// +/// ## Arguments +/// * `poseidons_16`: A slice of `WitnessPoseidon16` structs. +/// +/// ## Returns +/// A padded `Vec` containing the concatenated offsets of `addr_input_a`, `addr_input_b`, and `addr_output`. +#[must_use] +pub fn all_poseidon_16_indexes(poseidons_16: &[WitnessPoseidon16]) -> Vec { + // Collect all the indices into a single vector. + let all_indices: Vec = poseidons_16 + .iter() + .map(|p| F::from_usize(p.addr_input_a.offset)) + .chain( + poseidons_16 + .iter() + .map(|p| F::from_usize(p.addr_input_b.offset)), + ) + .chain( + poseidons_16 + .iter() + .map(|p| F::from_usize(p.addr_output.offset)), + ) + .collect(); + + // Pad the final vector with zeros to the next power of two. + pad_to_power_of_two(&all_indices) +} + +/// All Poseidon24 Indexes +/// +/// This function collects all memory addresses used by the Poseidon24 precompile, +/// flattens them into a single vector, and pads it to the next power of two. +/// +/// ## Arguments +/// * `poseidons_24`: A slice of `WitnessPoseidon24` structs. +/// +/// ## Returns +/// A padded `Vec` containing the concatenated and individually padded offsets of +/// `addr_input_a`, `addr_input_b`, and `addr_output`. +#[must_use] +pub fn all_poseidon_24_indexes(poseidons_24: &[WitnessPoseidon24]) -> Vec { + // Extract and pad the offsets for `addr_input_a`. + let indices_a = pad_to_power_of_two( + &poseidons_24 + .iter() + .map(|p| F::from_usize(p.addr_input_a.offset)) + .collect::>(), + ); + // Extract and pad the offsets for `addr_input_b`. + let indices_b = pad_to_power_of_two( + &poseidons_24 + .iter() + .map(|p| F::from_usize(p.addr_input_b.offset)) + .collect::>(), + ); + // Extract and pad the offsets for `addr_output`. + let indices_c = pad_to_power_of_two( + &poseidons_24 + .iter() + .map(|p| F::from_usize(p.addr_output.offset)) + .collect::>(), + ); + + // Concatenate the three padded vectors into a single flat vector. + [indices_a, indices_b, indices_c].concat() +} diff --git a/crates/leanVm/src/lang/simple_expr.rs b/crates/leanVm/src/lang/simple_expr.rs index 23b44a18..d9f20854 100644 --- a/crates/leanVm/src/lang/simple_expr.rs +++ b/crates/leanVm/src/lang/simple_expr.rs @@ -35,7 +35,8 @@ impl SimpleExpr { matches!(self, Self::Constant(_)) } - #[must_use] pub fn simplify_if_const(&self) -> Self { + #[must_use] + pub fn simplify_if_const(&self) -> Self { if let Self::Constant(constant) = self { return constant.try_naive_simplification().into(); } diff --git a/crates/leanVm/src/lib.rs b/crates/leanVm/src/lib.rs index 204e2157..84fb8e94 100644 --- a/crates/leanVm/src/lib.rs +++ b/crates/leanVm/src/lib.rs @@ -10,4 +10,5 @@ pub mod memory; pub mod parser; pub mod poseidon2; pub mod trace; +pub mod utils; pub mod witness; diff --git a/crates/leanVm/src/utils.rs b/crates/leanVm/src/utils.rs new file mode 100644 index 00000000..ac51e84d --- /dev/null +++ b/crates/leanVm/src/utils.rs @@ -0,0 +1,24 @@ +use p3_field::PrimeCharacteristicRing; + +use crate::constant::F; + +/// # Pad to Power of Two +/// +/// A utility function that pads a slice of field elements with zeros to the next +/// power-of-two length. +/// +/// ## Arguments +/// * `data`: A slice of field elements. +/// +/// ## Returns +/// A new `Vec` containing the original data, padded with zeros. +pub(crate) fn pad_to_power_of_two(data: &[F]) -> Vec { + // Calculate the next power-of-two length. + let next_power_of_two = data.len().next_power_of_two(); + // Create a new vector from the input slice. + let mut padded = data.to_vec(); + // Resize the vector to the target length, filling with zeros. + padded.resize(next_power_of_two, F::ZERO); + // Return the padded vector. + padded +}