openvm_stark_backend/prover/
helper.rs

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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
use std::sync::Arc;

use itertools::izip;
use p3_matrix::{dense::RowMajorMatrix, Matrix};

use crate::{
    config::{StarkGenericConfig, Val},
    prover::types::{AirProofInput, AirProofRawInput},
    rap::AnyRap,
};

/// Test helper trait for AirProofInput
/// Don't use this trait in production code
pub trait AirProofInputTestHelper<SC: StarkGenericConfig> {
    fn cached_traces_no_pis(
        air: Arc<dyn AnyRap<SC>>,
        cached_traces: Vec<RowMajorMatrix<Val<SC>>>,
        common_trace: RowMajorMatrix<Val<SC>>,
    ) -> Self;
}

impl<SC: StarkGenericConfig> AirProofInputTestHelper<SC> for AirProofInput<SC> {
    fn cached_traces_no_pis(
        air: Arc<dyn AnyRap<SC>>,
        cached_traces: Vec<RowMajorMatrix<Val<SC>>>,
        common_trace: RowMajorMatrix<Val<SC>>,
    ) -> Self {
        Self {
            air,
            cached_mains_pdata: vec![],
            raw: AirProofRawInput {
                cached_mains: cached_traces.into_iter().map(Arc::new).collect(),
                common_main: Some(common_trace),
                public_values: vec![],
            },
        }
    }
}
impl<SC: StarkGenericConfig> AirProofInput<SC> {
    pub fn simple(
        air: Arc<dyn AnyRap<SC>>,
        trace: RowMajorMatrix<Val<SC>>,
        public_values: Vec<Val<SC>>,
    ) -> Self {
        Self {
            air,
            cached_mains_pdata: vec![],
            raw: AirProofRawInput {
                cached_mains: vec![],
                common_main: Some(trace),
                public_values,
            },
        }
    }
    pub fn simple_no_pis(air: Arc<dyn AnyRap<SC>>, trace: RowMajorMatrix<Val<SC>>) -> Self {
        Self::simple(air, trace, vec![])
    }

    pub fn multiple_simple(
        airs: Vec<Arc<dyn AnyRap<SC>>>,
        traces: Vec<RowMajorMatrix<Val<SC>>>,
        public_values: Vec<Vec<Val<SC>>>,
    ) -> Vec<Self> {
        izip!(airs, traces, public_values)
            .map(|(air, trace, pis)| AirProofInput::simple(air, trace, pis))
            .collect()
    }

    pub fn multiple_simple_no_pis(
        airs: Vec<Arc<dyn AnyRap<SC>>>,
        traces: Vec<RowMajorMatrix<Val<SC>>>,
    ) -> Vec<Self> {
        izip!(airs, traces)
            .map(|(air, trace)| AirProofInput::simple_no_pis(air, trace))
            .collect()
    }
    /// Return the height of the main trace.
    pub fn main_trace_height(&self) -> usize {
        if self.raw.cached_mains.is_empty() {
            // An AIR must have a main trace.
            self.raw.common_main.as_ref().unwrap().height()
        } else {
            self.raw.cached_mains[0].height()
        }
    }
}