openvm_stark_backend/
chip.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
use std::{
    cell::RefCell,
    rc::Rc,
    sync::{Arc, Mutex},
};

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

/// A chip is a stateful struct that stores the state necessary to
/// generate the trace of an AIR. This trait is for proving purposes
/// and has a generic [StarkGenericConfig] since it needs to know the STARK config.
pub trait Chip<SC: StarkGenericConfig>: ChipUsageGetter + Sized {
    fn air(&self) -> Arc<dyn AnyRap<SC>>;
    /// Generate all necessary input for proving a single AIR.
    fn generate_air_proof_input(self) -> AirProofInput<SC>;
    fn generate_air_proof_input_with_id(self, air_id: usize) -> (usize, AirProofInput<SC>) {
        (air_id, self.generate_air_proof_input())
    }
}

/// A trait to get chip usage information.
pub trait ChipUsageGetter {
    fn air_name(&self) -> String;
    /// If the chip has a state-independent trace height that is determined
    /// upon construction, return this height. This is used to distinguish
    /// "static" versus "dynamic" usage metrics.
    fn constant_trace_height(&self) -> Option<usize> {
        None
    }
    /// Height of used rows in the main trace.
    fn current_trace_height(&self) -> usize;
    /// Width of the main trace
    fn trace_width(&self) -> usize;
    /// For metrics collection
    fn current_trace_cells(&self) -> usize {
        self.trace_width() * self.current_trace_height()
    }
}

impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for RefCell<C> {
    fn air(&self) -> Arc<dyn AnyRap<SC>> {
        self.borrow().air()
    }
    fn generate_air_proof_input(self) -> AirProofInput<SC> {
        self.into_inner().generate_air_proof_input()
    }
}

impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Rc<C> {
    fn air(&self) -> Arc<dyn AnyRap<SC>> {
        self.as_ref().air()
    }
    fn generate_air_proof_input(self) -> AirProofInput<SC> {
        if let Some(c) = Rc::into_inner(self) {
            c.generate_air_proof_input()
        } else {
            panic!("Cannot generate AirProofInput while other chips still hold a reference");
        }
    }
}

impl<C: ChipUsageGetter> ChipUsageGetter for Rc<C> {
    fn air_name(&self) -> String {
        self.as_ref().air_name()
    }
    fn constant_trace_height(&self) -> Option<usize> {
        self.as_ref().constant_trace_height()
    }
    fn current_trace_height(&self) -> usize {
        self.as_ref().current_trace_height()
    }
    fn trace_width(&self) -> usize {
        self.as_ref().trace_width()
    }
}

impl<C: ChipUsageGetter> ChipUsageGetter for RefCell<C> {
    fn air_name(&self) -> String {
        self.borrow().air_name()
    }
    fn constant_trace_height(&self) -> Option<usize> {
        self.borrow().constant_trace_height()
    }
    fn current_trace_height(&self) -> usize {
        self.borrow().current_trace_height()
    }
    fn trace_width(&self) -> usize {
        self.borrow().trace_width()
    }
}

impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Arc<C> {
    fn air(&self) -> Arc<dyn AnyRap<SC>> {
        self.as_ref().air()
    }
    fn generate_air_proof_input(self) -> AirProofInput<SC> {
        if let Some(c) = Arc::into_inner(self) {
            c.generate_air_proof_input()
        } else {
            panic!("Cannot generate AirProofInput while other chips still hold a reference");
        }
    }
}

impl<C: ChipUsageGetter> ChipUsageGetter for Arc<C> {
    fn air_name(&self) -> String {
        self.as_ref().air_name()
    }
    fn constant_trace_height(&self) -> Option<usize> {
        self.as_ref().constant_trace_height()
    }
    fn current_trace_height(&self) -> usize {
        self.as_ref().current_trace_height()
    }
    fn trace_width(&self) -> usize {
        self.as_ref().trace_width()
    }
}

impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Mutex<C> {
    fn air(&self) -> Arc<dyn AnyRap<SC>> {
        self.lock().unwrap().air()
    }
    fn generate_air_proof_input(self) -> AirProofInput<SC> {
        self.into_inner().unwrap().generate_air_proof_input()
    }
}

impl<C: ChipUsageGetter> ChipUsageGetter for Mutex<C> {
    fn air_name(&self) -> String {
        self.lock().unwrap().air_name()
    }
    fn constant_trace_height(&self) -> Option<usize> {
        self.lock().unwrap().constant_trace_height()
    }
    fn current_trace_height(&self) -> usize {
        self.lock().unwrap().current_trace_height()
    }
    fn trace_width(&self) -> usize {
        self.lock().unwrap().trace_width()
    }
}