1use std::sync::Arc;
2
3use crate::{MemInput, MemModule, MemPreviousSegment};
4use mem_common::{MEM_BYTES_BITS, SEGMENT_ADDR_MAX_RANGE};
5
6use fields::PrimeField64;
7use pil_std_lib::Std;
8use proofman_common::{AirInstance, FromTrace, ProofmanResult};
9use zisk_common::SegmentId;
10use zisk_core::{INPUT_ADDR, MAX_INPUT_SIZE};
11use zisk_pil::InputDataAirValues;
12#[cfg(not(feature = "packed"))]
13use zisk_pil::InputDataTrace;
14#[cfg(feature = "packed")]
15use zisk_pil::InputDataTracePacked;
16
17#[cfg(feature = "packed")]
18type InputDataTraceType<F> = InputDataTracePacked<F>;
19
20#[cfg(not(feature = "packed"))]
21type InputDataTraceType<F> = InputDataTrace<F>;
22
23pub const INPUT_DATA_W_ADDR_INIT: u32 = INPUT_ADDR as u32 >> MEM_BYTES_BITS;
24pub const INPUT_DATA_W_ADDR_END: u32 = (INPUT_ADDR + MAX_INPUT_SIZE - 1) as u32 >> MEM_BYTES_BITS;
25
26#[allow(clippy::assertions_on_constants)]
27const _: () = {
28 assert!(
29 INPUT_ADDR + MAX_INPUT_SIZE - 1 <= 0xFFFF_FFFF,
30 "INPUT_DATA memory exceeds the 32-bit addressable range"
31 );
32 assert!(
33 (MAX_INPUT_SIZE - 1) <= (128 << 20),
34 "INPUT_DATA is too large. Input size must be <= 128MB"
35 );
36};
47}
48
69}
70
71impl<F: PrimeField64> MemModule<F> for InputDataSM<F> {
72 fn get_addr_range(&self) -> (u32, u32) {
73 (INPUT_DATA_W_ADDR_INIT, INPUT_DATA_W_ADDR_END)
74 }
75 fn is_dual(&self) -> bool {
76 false
77 }
78
79 // TODO PRE: proxy calculate if exists jmp on step out-of-range, adding internal inputs
80 // memory only need to process these special inputs, but inputs no change. At end of
81 // inputs proxy add an extra internal input to jump to last address
82
83 /// Finalizes the witness accumulation process and triggers the proof generation.
84 ///
85 /// This method is invoked by the executor when no further witness data remains to be added.
86 ///
87 /// # Parameters
88 ///
89 /// - `mem_inputs`: A slice of all `ZiskRequiredMemory` inputs
90 fn compute_witness(
91 &self,
92 mem_ops: &[MemInput],
93 segment_id: SegmentId,
94 is_last_segment: bool,
95 previous_segment: &MemPreviousSegment,
96 trace_buffer: Vec<F>,
97 ) -> ProofmanResult<AirInstance<F>> {
98 let mut trace = InputDataTraceType::<F>::new_from_vec(trace_buffer)?;
99
100 let num_rows = InputDataTraceType::<F>::NUM_ROWS;
101 debug_assert!(
102 !mem_ops.is_empty() && mem_ops.len() <= num_rows,
103 "InputDataSM: mem_ops.len()={} out of range {}",
104 mem_ops.len(),
105 num_rows
106 );
107
108 let mut range_check_data: Vec<u32> = vec![0; 1 << 16];
109
110 // range of instance
111 self.std.range_check(
112 self.range_id,
113 (previous_segment.addr - INPUT_DATA_W_ADDR_INIT) as i64,
114 1,
115 );
116
117 let mut max_range_distance_count = 0;
118
119 let mut last_addr: u32 = previous_segment.addr;
120 let mut last_step: u64 = previous_segment.step;
121 let mut last_value: u64 = previous_segment.value;
122 let mut i = 0;
123
124 for mem_op in mem_ops.iter() {
125 let distance = mem_op.addr - last_addr;
126
127 if i >= num_rows {
128 break;
129 }
130
131 if distance > SEGMENT_ADDR_MAX_RANGE as u32 {
132 let mut internal_reads = (distance - 1) / SEGMENT_ADDR_MAX_RANGE as u32;
133
134 let incomplete = (i + internal_reads as usize) >= num_rows;
135 if incomplete {
136 internal_reads = (num_rows - i) as u32;
137 }
138
139 trace[i].set_addr_changes(true);
140 last_addr += SEGMENT_ADDR_MAX_RANGE as u32;
141 max_range_distance_count += 1;
142 trace[i].set_addr(last_addr);
143
144 // the step, value of internal reads isn't relevant
145 last_step = 0;
146 trace[i].set_step(0);
147 trace[i].set_sel(false);
148
149 // setting value to zero, is not relevant for internal reads
150 last_value = 0;
151 for j in 0..4 {
152 trace[i].set_value_word(j, 0);
153 }
154 i += 1;
155
156 for _j in 1..internal_reads {
157 trace[i] = trace[i - 1];
158 last_addr += SEGMENT_ADDR_MAX_RANGE as u32;
159 max_range_distance_count += 1;
160 trace[i].set_addr(last_addr);
161
162 i += 1;
163 }
164 range_check_data[0] += 4 * internal_reads;
165 if incomplete {
166 break;
167 }
168 }
169
170 trace[i].set_addr(mem_op.addr);
171 trace[i].set_step(mem_op.step);
172 trace[i].set_sel(true);
173 trace[i].set_is_free_read(mem_op.addr == INPUT_DATA_W_ADDR_INIT);
174
175 let value = mem_op.value;
176 let value_words = self.get_u16_values(value);
177 for j in 0..4 {
178 range_check_data[value_words[j] as usize] += 1;
179 trace[i].set_value_word(j, value_words[j]);
180 }
181
182 let addr_changes = last_addr != mem_op.addr;
183 if addr_changes {
184 trace[i].set_addr_changes(true);
185 self.std.range_check(self.range_id, (mem_op.addr - last_addr - 1) as i64, 1);
186 } else {
187 trace[i].set_addr_changes(false);
188 }
189
190 last_addr = mem_op.addr;
191 last_step = mem_op.step;
192 last_value = mem_op.value;
193 i += 1;
194 }
195 let count = i;
196
197 // STEP3. Add dummy rows to the output vector to fill the remaining rows
198 //PADDING: At end of memory fill with same addr, incrementing step, same value, sel = 0
199 let last_row_idx = count - 1;
200 let addr = trace[last_row_idx].get_addr();
201 let is_free_read = last_addr == INPUT_DATA_W_ADDR_INIT;
202
203 let padding_size = num_rows - count;
204 for i in count..num_rows {
205 last_step += 1;
206
207 trace[i].set_addr(addr);
208 trace[i].set_step(last_step);
209 trace[i].set_sel(false);
210 for j in 0..4 {
211 let value = trace[last_row_idx].get_value_word(j);
212 trace[i].set_value_word(j, value);
213 }
214 trace[i].set_is_free_read(is_free_read);
215
216 trace[i].set_addr_changes(false);
217
218 // address doesn't change in padding rows, no range check is required
219 }
220
221 self.std.range_check(
222 self.range_id,
223 SEGMENT_ADDR_MAX_RANGE as i64,
224 max_range_distance_count,
225 );
226 self.std.range_check(self.range_id, (INPUT_DATA_W_ADDR_END - last_addr) as i64, 1);
227
228 // range of chunks
229 for j in 0..4 {
230 let value = trace[last_row_idx].get_value_word(j);
231 range_check_data[value as usize] += padding_size as u32;
232 }
233 self.std.range_checks(self.range_chunks_id, range_check_data);
234
235 let mut air_values = InputDataAirValues::<F>::new();
236 air_values.segment_id = F::from_usize(segment_id.into());
237 air_values.is_first_segment = F::from_bool(segment_id == 0);
238 air_values.is_last_segment = F::from_bool(is_last_segment);
239 air_values.previous_segment_step = F::from_u64(previous_segment.step);
240 air_values.previous_segment_addr = F::from_u32(previous_segment.addr);
241 air_values.segment_last_addr = F::from_u32(last_addr);
242 air_values.segment_last_step = F::from_u64(last_step);
243
244 air_values.previous_segment_value[0] = F::from_u32(previous_segment.value as u32);
245 air_values.previous_segment_value[1] = F::from_u32((previous_segment.value >> 32) as u32);
246
247 air_values.segment_last_value[0] = F::from_u32(last_value as u32);
248 air_values.segment_last_value[1] = F::from_u32((last_value >> 32) as u32);
249
250 Ok(AirInstance::new_from_trace(FromTrace::new(&mut trace).with_air_values(&mut air_values)))
251 }
252}