openvm_stark_backend/
chip.rs

Help
1use std::{
2    cell::RefCell,
3    rc::Rc,
4    sync::{Arc, Mutex},
5};
6
7use crate::{config::StarkGenericConfig, prover::types::AirProofInput, rap::AnyRap};
8
9/// A chip is a stateful struct that stores the state necessary to
10/// generate the trace of an AIR. This trait is for proving purposes
11/// and has a generic [StarkGenericConfig] since it needs to know the STARK config.
12pub trait Chip<SC: StarkGenericConfig>: ChipUsageGetter + Sized {
13    fn air(&self) -> Arc<dyn AnyRap<SC>>;
14    /// Generate all necessary input for proving a single AIR.
15    fn generate_air_proof_input(self) -> AirProofInput<SC>;
16    fn generate_air_proof_input_with_id(self, air_id: usize) -> (usize, AirProofInput<SC>) {
17        (air_id, self.generate_air_proof_input())
18    }
19}
20
21/// A trait to get chip usage information.
22pub trait ChipUsageGetter {
23    fn air_name(&self) -> String;
24    /// If the chip has a state-independent trace height that is determined
25    /// upon construction, return this height. This is used to distinguish
26    /// "static" versus "dynamic" usage metrics.
27    fn constant_trace_height(&self) -> Option<usize> {
28        None
29    }
30    /// Height of used rows in the main trace.
31    fn current_trace_height(&self) -> usize;
32    /// Width of the main trace
33    fn trace_width(&self) -> usize;
34    /// For metrics collection
35    fn current_trace_cells(&self) -> usize {
36        self.trace_width() * self.current_trace_height()
37    }
38}
39
40impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for RefCell<C> {
41    fn air(&self) -> Arc<dyn AnyRap<SC>> {
42        self.borrow().air()
43    }
44    fn generate_air_proof_input(self) -> AirProofInput<SC> {
45        self.into_inner().generate_air_proof_input()
46    }
47}
48
49impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Rc<C> {
50    fn air(&self) -> Arc<dyn AnyRap<SC>> {
51        self.as_ref().air()
52    }
53    fn generate_air_proof_input(self) -> AirProofInput<SC> {
54        if let Some(c) = Rc::into_inner(self) {
55            c.generate_air_proof_input()
56        } else {
57            panic!("Cannot generate AirProofInput while other chips still hold a reference");
58        }
59    }
60}
61
62impl<C: ChipUsageGetter> ChipUsageGetter for Rc<C> {
63    fn air_name(&self) -> String {
64        self.as_ref().air_name()
65    }
66    fn constant_trace_height(&self) -> Option<usize> {
67        self.as_ref().constant_trace_height()
68    }
69    fn current_trace_height(&self) -> usize {
70        self.as_ref().current_trace_height()
71    }
72    fn trace_width(&self) -> usize {
73        self.as_ref().trace_width()
74    }
75}
76
77impl<C: ChipUsageGetter> ChipUsageGetter for RefCell<C> {
78    fn air_name(&self) -> String {
79        self.borrow().air_name()
80    }
81    fn constant_trace_height(&self) -> Option<usize> {
82        self.borrow().constant_trace_height()
83    }
84    fn current_trace_height(&self) -> usize {
85        self.borrow().current_trace_height()
86    }
87    fn trace_width(&self) -> usize {
88        self.borrow().trace_width()
89    }
90}
91
92impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Arc<C> {
93    fn air(&self) -> Arc<dyn AnyRap<SC>> {
94        self.as_ref().air()
95    }
96    fn generate_air_proof_input(self) -> AirProofInput<SC> {
97        if let Some(c) = Arc::into_inner(self) {
98            c.generate_air_proof_input()
99        } else {
100            panic!("Cannot generate AirProofInput while other chips still hold a reference");
101        }
102    }
103}
104
105impl<C: ChipUsageGetter> ChipUsageGetter for Arc<C> {
106    fn air_name(&self) -> String {
107        self.as_ref().air_name()
108    }
109    fn constant_trace_height(&self) -> Option<usize> {
110        self.as_ref().constant_trace_height()
111    }
112    fn current_trace_height(&self) -> usize {
113        self.as_ref().current_trace_height()
114    }
115    fn trace_width(&self) -> usize {
116        self.as_ref().trace_width()
117    }
118}
119
120impl<SC: StarkGenericConfig, C: Chip<SC>> Chip<SC> for Mutex<C> {
121    fn air(&self) -> Arc<dyn AnyRap<SC>> {
122        self.lock().unwrap().air()
123    }
124    fn generate_air_proof_input(self) -> AirProofInput<SC> {
125        self.into_inner().unwrap().generate_air_proof_input()
126    }
127}
128
129impl<C: ChipUsageGetter> ChipUsageGetter for Mutex<C> {
130    fn air_name(&self) -> String {
131        self.lock().unwrap().air_name()
132    }
133    fn constant_trace_height(&self) -> Option<usize> {
134        self.lock().unwrap().constant_trace_height()
135    }
136    fn current_trace_height(&self) -> usize {
137        self.lock().unwrap().current_trace_height()
138    }
139    fn trace_width(&self) -> usize {
140        self.lock().unwrap().trace_width()
141    }
142}
openvm_stark_backend/chip.rs

openvm_stark_backend/
chip.rs
