use itertools::Itertools;
#[cfg(feature = "parallel")]
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SECURE_EXTENSION_DEGREE;
use crate::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
use crate::core::vcs_lifted::verifier::PACKED_LEAF_SIZE;
use crate::parallel_iter;
use crate::prover::backend::{Col, Column, CpuBackend};
use crate::prover::vcs_lifted::ops::{MerkleOpsLifted, PackLeavesOps};
impl<H: MerkleHasherLifted> MerkleOpsLifted<H> for CpuBackend {
fn build_leaves(columns: &[&Vec<BaseField>], lifting_log_size: u32) -> Vec<H::Hash> {
let hasher = H::default();
if columns.is_empty() {
return vec![hasher.finalize()];
}
assert!(columns[0].len() >= 2, "A column must be of length >= 2.");
let mut prev_layer: Vec<H> = vec![hasher; 2];
let mut prev_layer_log_size: u32 = 1;
for (log_size, group) in columns.iter().group_by(|c| c.len().ilog2()).into_iter() {
let log_ratio = log_size - prev_layer_log_size;
prev_layer = (0..1 << log_size)
.map(|idx| prev_layer[(idx >> (log_ratio + 1) << 1) + (idx & 1)].clone())
.collect();
for chunk in &group.into_iter().chunks(16) {
let vec = chunk.into_iter().collect_vec();
prev_layer.iter_mut().enumerate().for_each(|(i, hasher)| {
hasher.update_leaf(&vec.iter().map(|v| v[i]).collect_vec());
})
}
prev_layer_log_size = log_size;
}
let log_ratio = lifting_log_size - prev_layer_log_size;
if log_ratio > 0 {
prev_layer = (0..1 << lifting_log_size)
.map(|idx| prev_layer[(idx >> (log_ratio + 1) << 1) + (idx & 1)].clone())
.collect();
}
prev_layer.into_iter().map(|x| x.finalize()).collect()
}
fn build_next_layer(prev_layer: &Vec<H::Hash>) -> Vec<H::Hash> {
let log_size: u32 = prev_layer.len().ilog2() - 1;
parallel_iter!(0..(1 << log_size))
.map(|i| H::hash_children((prev_layer[2 * i], prev_layer[2 * i + 1])))
.collect()
}
}
impl PackLeavesOps for CpuBackend {
fn pack_leaves_input(
values: &[&Col<Self, BaseField>; SECURE_EXTENSION_DEGREE],
) -> [Col<Self, BaseField>; SECURE_EXTENSION_DEGREE * PACKED_LEAF_SIZE] {
let len_m31 = values[0].len();
assert!(values.iter().all(|c| c.len() == len_m31));
assert!(len_m31.is_multiple_of(PACKED_LEAF_SIZE));
let packed_len = len_m31 / PACKED_LEAF_SIZE;
let cpu_columns: [Vec<BaseField>; SECURE_EXTENSION_DEGREE] =
core::array::from_fn(|coord| values[coord].to_cpu());
let mut packed_cpu: [Vec<BaseField>; SECURE_EXTENSION_DEGREE * PACKED_LEAF_SIZE] =
core::array::from_fn(|_| Vec::with_capacity(packed_len));
for packed_row in 0..packed_len {
let row_start = packed_row * PACKED_LEAF_SIZE;
for offset in 0..PACKED_LEAF_SIZE {
for coord in 0..SECURE_EXTENSION_DEGREE {
packed_cpu[coord + offset * SECURE_EXTENSION_DEGREE]
.push(cpu_columns[coord][row_start + offset]);
}
}
}
packed_cpu.map(|column| column.into_iter().collect())
}
}