stwo 2.2.0

Core library implementing the Circle STARK prover and verifier
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

//! Pre-allocated memory pools for the proving pipeline.
//!
//! The [`BaseColumnPool`] manages reusable [`Col<B, BaseField>`] buffers for polynomial evaluation,
//! avoiding repeated allocation/deallocation of large column buffers during proving.

use dashmap::DashMap;

use crate::core::fields::m31::BaseField;
use crate::prover::backend::{Col, Column, ColumnOps};

/// A pool of pre-allocated [`Col<B, BaseField>`] buffers, organized by log_size.
///
/// Used to avoid repeated allocation of evaluation buffers during polynomial commitment.
pub struct BaseColumnPool<B: ColumnOps<BaseField>> {
    /// Map from log_size -> stack of available buffers.
    pools: DashMap<u32, Vec<Col<B, BaseField>>>,
}

impl<B: ColumnOps<BaseField>> BaseColumnPool<B> {
    /// Creates a new empty base column pool.
    pub fn new() -> Self {
        Self {
            pools: DashMap::new(),
        }
    }

    /// Pre-allocates `count` zero-initialized buffers of size `1 << log_size`.
    pub fn reserve(&self, log_size: u32, count: usize) {
        let mut pool = self.pools.entry(log_size).or_default();
        for _ in 0..count {
            pool.push(Col::<B, BaseField>::zeros(1 << log_size));
        }
    }

    /// Takes a buffer from the pool for the given `log_size`.
    ///
    /// # Panics
    ///
    /// Panics if no buffer of the requested size is available.
    pub fn take(&self, log_size: u32) -> Col<B, BaseField> {
        self.pools
            .get_mut(&log_size)
            .and_then(|mut pool| pool.pop())
            .unwrap_or_else(|| {
                panic!("BaseColumnPool: no buffer available for log_size={log_size}")
            })
    }

    /// Takes a buffer from the pool, or allocates a new zero-initialized one if none is available.
    pub fn take_or_alloc(&self, log_size: u32) -> Col<B, BaseField> {
        self.pools
            .get_mut(&log_size)
            .and_then(|mut pool| pool.pop())
            .unwrap_or_else(|| unsafe { Col::<B, BaseField>::uninitialized(1 << log_size) })
    }

    /// Returns a buffer to the pool. The caller is responsible for ensuring the buffer's log_size
    /// matches.
    pub fn give_back(&self, log_size: u32, buf: Col<B, BaseField>) {
        debug_assert_eq!(buf.len(), 1 << log_size);
        self.pools.entry(log_size).or_default().push(buf);
    }
}

impl<B: ColumnOps<BaseField>> Default for BaseColumnPool<B> {
    fn default() -> Self {
        Self::new()
    }
}