stwo 2.3.0

Core library implementing the Circle STARK prover and verifier
Documentation
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 {
    /// Computes the leaves of the Merkle tree. This is the core logic of the lifted Merkle
    /// commitment. The input columns are assumed to be in increasing order of length.
    ///
    /// The columns are interpreted as evaluations of polynomials in bit reversed order.
    /// For example, consider a polynomial that on the canonical circle domain of size 8 has
    /// evaluations (in natural order and bit reversed respectively):
    ///     a   a
    ///     b   e
    ///     c   c
    ///     d   g
    ///     e   b
    ///     f   f
    ///     g   d
    ///     h   h
    /// Then the evaluations of its lifted polynomial on the canonical circle domain of size 16 are
    /// (in natural and bit reversed order respectively):
    ///     a   a
    ///     b   e
    ///     c   a
    ///     d   e
    ///     a   c
    ///     b   g
    ///     c   c
    ///     d   g
    ///     e   b
    ///     f   f
    ///     g   b
    ///     h   f
    ///     e   d
    ///     f   h
    ///     g   d
    ///     h   h
    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)
                // We only clone when starting a column chunk of different size.
                .map(|idx| prev_layer[(idx >> (log_ratio + 1) << 1) + (idx & 1)].clone())
                .collect();

            // We chunk by 16 because it's the amount of M31 elements needed to trigger a
            // hash permutation, both in blake and in poseidon.
            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())
    }
}