Skip to content

Commit 9f02efc

Browse files
ythamytham
committed
Update to handle input byte lengths
1 parent 6822893 commit 9f02efc

File tree

5 files changed

+213
-142
lines changed

5 files changed

+213
-142
lines changed

‎bin/afs/src/commands/cache/mod.rs

Lines changed: 56 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,14 @@ use afs_test_utils::{
66
config::{self, baby_bear_poseidon2::BabyBearPoseidon2Config},
77
page_config::PageConfig,
88
};
9-
use alloy_primitives::U256;
9+
use alloy_primitives::{U256, U512};
1010
use clap::Parser;
1111
use color_eyre::eyre::Result;
12-
use logical_interface::{afs_interface::AfsInterface, mock_db::MockDb};
12+
use logical_interface::{
13+
afs_interface::AfsInterface,
14+
mock_db::MockDb,
15+
types::{Data, Index},
16+
};
1317
use p3_baby_bear::BabyBear;
1418
use p3_util::log2_strict_usize;
1519

@@ -46,37 +50,23 @@ impl CacheCommand {
4650
/// Execute the `cache` command
4751
pub fn execute(&self, config: &PageConfig) -> Result<()> {
4852
println!("Caching table {} from {}", self.table_id, self.db_file_path);
49-
50-
let start = Instant::now();
51-
let mut db = MockDb::from_file(&self.db_file_path);
52-
let height = config.page.height;
53-
let mut interface = AfsInterface::<U256, U256>::new(&mut db);
54-
let page = interface
55-
.get_table(self.table_id.clone())
56-
.unwrap()
57-
.to_page(height);
58-
59-
assert!(height > 0);
60-
61-
let checker_trace_degree = config.page.max_rw_ops * 4;
62-
63-
let max_log_degree = log2_strict_usize(checker_trace_degree)
64-
.max(log2_strict_usize(height))
65-
.max(8);
66-
67-
let trace = page.gen_trace::<BabyBear>();
68-
let engine = config::baby_bear_poseidon2::default_engine(max_log_degree);
69-
let prover = MultiTraceStarkProver::new(&engine.config);
70-
let trace_builder = TraceCommitmentBuilder::<BabyBearPoseidon2Config>::new(prover.pcs());
71-
let trace_prover_data = trace_builder.committer.commit(vec![trace]);
72-
let encoded_data = bincode::serialize(&trace_prover_data).unwrap();
73-
let path = self.output_folder.clone() + "/" + &self.table_id + ".cache.bin";
74-
write_bytes(&encoded_data, path).unwrap();
75-
76-
let duration = start.elapsed();
77-
println!("Cached table {} in {:?}", self.table_id, duration);
78-
79-
Ok(())
53+
println!(
54+
"Index bytes: {}, Data bytes: {}",
55+
config.page.index_bytes, config.page.data_bytes
56+
);
57+
match (config.page.index_bytes, config.page.data_bytes) {
58+
(2, 2) => cache::<u16, u16>(self, config),
59+
(4, 4) => cache::<u32, u32>(self, config),
60+
(8, 8) => cache::<u64, u64>(self, config),
61+
(16, 16) => cache::<u128, u128>(self, config),
62+
(32, 32) => cache::<U256, U256>(self, config),
63+
(32, 64) => cache::<U256, U512>(self, config),
64+
(32, 128) => cache::<U256, [U256; 4]>(self, config),
65+
(32, 256) => cache::<U256, [U256; 8]>(self, config),
66+
(32, 512) => cache::<U256, [U256; 16]>(self, config),
67+
(32, 1024) => cache::<U256, [U256; 32]>(self, config),
68+
_ => panic!("Unsupported index/data size"),
69+
}
8070
}
8171

8272
pub fn read_page_file(&self) -> Result<Page> {
@@ -92,3 +82,36 @@ impl CacheCommand {
9282
Ok(())
9383
}
9484
}
85+
86+
fn cache<I: Index, D: Data>(cmd: &CacheCommand, config: &PageConfig) -> Result<()> {
87+
let start = Instant::now();
88+
let mut db = MockDb::from_file(&cmd.db_file_path);
89+
let height = config.page.height;
90+
let mut interface = AfsInterface::<I, D>::new(&mut db);
91+
let page = interface
92+
.get_table(cmd.table_id.clone())
93+
.unwrap()
94+
.to_page(height);
95+
96+
assert!(height > 0);
97+
98+
let checker_trace_degree = config.page.max_rw_ops * 4;
99+
100+
let max_log_degree = log2_strict_usize(checker_trace_degree)
101+
.max(log2_strict_usize(height))
102+
.max(8);
103+
104+
let trace = page.gen_trace::<BabyBear>();
105+
let engine = config::baby_bear_poseidon2::default_engine(max_log_degree);
106+
let prover = MultiTraceStarkProver::new(&engine.config);
107+
let trace_builder = TraceCommitmentBuilder::<BabyBearPoseidon2Config>::new(prover.pcs());
108+
let trace_prover_data = trace_builder.committer.commit(vec![trace]);
109+
let encoded_data = bincode::serialize(&trace_prover_data).unwrap();
110+
let path = cmd.output_folder.clone() + "/" + &cmd.table_id + ".cache.bin";
111+
write_bytes(&encoded_data, path).unwrap();
112+
113+
let duration = start.elapsed();
114+
println!("Cached table {} in {:?}", cmd.table_id, duration);
115+
116+
Ok(())
117+
}

‎bin/afs/src/commands/prove/mod.rs

Lines changed: 125 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ use afs_test_utils::{
1616
engine::StarkEngine,
1717
page_config::{PageConfig, PageMode},
1818
};
19-
use alloy_primitives::U256;
19+
use alloy_primitives::{U256, U512};
2020
use clap::Parser;
2121
use color_eyre::eyre::Result;
2222
use logical_interface::{
@@ -85,9 +85,8 @@ impl ProveCommand {
8585
/// Execute the `prove` command
8686
pub fn execute(&self, config: &PageConfig) -> Result<()> {
8787
let start = Instant::now();
88-
let prefix = create_prefix(config);
8988
match config.page.mode {
90-
PageMode::ReadWrite => self.execute_rw(config, prefix)?,
89+
PageMode::ReadWrite => self.execute_rw(config)?,
9190
PageMode::ReadOnly => panic!(),
9291
}
9392

@@ -97,127 +96,148 @@ impl ProveCommand {
9796
Ok(())
9897
}
9998

100-
pub fn execute_rw(&self, config: &PageConfig, prefix: String) -> Result<()> {
99+
pub fn execute_rw(&self, config: &PageConfig) -> Result<()> {
101100
println!("Proving ops file: {}", self.afi_file_path);
102-
let instructions = AfsInputInstructions::from_file(&self.afi_file_path)?;
103-
let mut db = MockDb::from_file(&self.db_file_path);
104-
let idx_len = (config.page.index_bytes + 1) / 2;
105-
let data_len = (config.page.data_bytes + 1) / 2;
106-
let height = config.page.height;
107-
let codec = FixedBytesCodec::new(config.page.index_bytes, config.page.data_bytes);
108-
let mut interface = AfsInterface::<U256, U256>::new(&mut db);
109-
let table_id = instructions.header.table_id;
110-
let page_init = interface
111-
.get_table(table_id.clone())
112-
.unwrap()
113-
.to_page(height);
114-
let zk_ops = instructions
115-
.operations
116-
.iter()
117-
.enumerate()
118-
.map(|(i, op)| afi_op_conv(op, table_id.clone(), &mut interface, i + 1, &codec))
119-
.collect::<Vec<_>>();
101+
println!(
102+
"Index bytes: {}, Data bytes: {}",
103+
config.page.index_bytes, config.page.data_bytes
104+
);
105+
match (config.page.index_bytes, config.page.data_bytes) {
106+
(2, 2) => prove::<u16, u16>(self, config),
107+
(4, 4) => prove::<u32, u32>(self, config),
108+
(8, 8) => prove::<u64, u64>(self, config),
109+
(16, 16) => prove::<u128, u128>(self, config),
110+
(32, 32) => prove::<U256, U256>(self, config),
111+
(32, 64) => prove::<U256, U512>(self, config),
112+
(32, 128) => prove::<U256, [U256; 4]>(self, config),
113+
(32, 256) => prove::<U256, [U256; 8]>(self, config),
114+
(32, 512) => prove::<U256, [U256; 16]>(self, config),
115+
(32, 1024) => prove::<U256, [U256; 32]>(self, config),
116+
_ => panic!("Unsupported index/data size"),
117+
}
118+
}
119+
}
120120

121-
assert!(height > 0);
122-
let page_bus_index = 0;
123-
let range_bus_index = 2;
124-
let ops_bus_index = 3;
121+
fn prove<I: Index, D: Data>(cmd: &ProveCommand, config: &PageConfig) -> Result<()> {
122+
let instructions = AfsInputInstructions::from_file(&cmd.afi_file_path)?;
123+
let mut db = MockDb::from_file(&cmd.db_file_path);
124+
let mut interface = AfsInterface::<I, D>::new(&mut db);
125+
let idx_len = (config.page.index_bytes + 1) / 2;
126+
let data_len = (config.page.data_bytes + 1) / 2;
127+
let height = config.page.height;
128+
let codec = FixedBytesCodec::new(config.page.index_bytes, config.page.data_bytes);
129+
let table_id = instructions.header.table_id;
130+
let page_init = interface
131+
.get_table(table_id.clone())
132+
.unwrap()
133+
.to_page(height);
134+
let zk_ops = instructions
135+
.operations
136+
.iter()
137+
.enumerate()
138+
.map(|(i, op)| afi_op_conv(op, table_id.clone(), &mut interface, i + 1, &codec))
139+
.collect::<Vec<_>>();
125140

126-
let checker_trace_degree = config.page.max_rw_ops * 4;
141+
assert!(height > 0);
142+
let page_bus_index = 0;
143+
let range_bus_index = 2;
144+
let ops_bus_index = 3;
127145

128-
let idx_limb_bits = config.page.bits_per_fe;
146+
let checker_trace_degree = config.page.max_rw_ops * 4;
129147

130-
let max_log_degree = log2_strict_usize(checker_trace_degree)
131-
.max(log2_strict_usize(height))
132-
.max(8);
148+
let idx_limb_bits = config.page.bits_per_fe;
133149

134-
let idx_decomp = 8;
150+
let max_log_degree = log2_strict_usize(checker_trace_degree)
151+
.max(log2_strict_usize(height))
152+
.max(8);
135153

136-
let mut page_controller: PageController<BabyBearPoseidon2Config> = PageController::new(
137-
page_bus_index,
138-
range_bus_index,
139-
ops_bus_index,
140-
idx_len,
141-
data_len,
142-
idx_limb_bits,
143-
idx_decomp,
144-
);
145-
let ops_sender = ExecutionAir::new(ops_bus_index, idx_len, data_len);
146-
let engine = config::baby_bear_poseidon2::default_engine(max_log_degree);
147-
let prover = MultiTraceStarkProver::new(&engine.config);
148-
let mut trace_builder = TraceCommitmentBuilder::new(prover.pcs());
154+
let idx_decomp = 8;
149155

150-
let init_prover_data_encoded =
151-
read_from_path(self.cache_folder.clone() + "/" + &table_id + ".cache.bin").unwrap();
152-
let init_prover_data: ProverTraceData<BabyBearPoseidon2Config> =
153-
bincode::deserialize(&init_prover_data_encoded).unwrap();
156+
let mut page_controller: PageController<BabyBearPoseidon2Config> = PageController::new(
157+
page_bus_index,
158+
range_bus_index,
159+
ops_bus_index,
160+
idx_len,
161+
data_len,
162+
idx_limb_bits,
163+
idx_decomp,
164+
);
165+
let ops_sender = ExecutionAir::new(ops_bus_index, idx_len, data_len);
166+
let engine = config::baby_bear_poseidon2::default_engine(max_log_degree);
167+
let prover = MultiTraceStarkProver::new(&engine.config);
168+
let mut trace_builder = TraceCommitmentBuilder::new(prover.pcs());
154169

155-
let span = tracing::info_span!("Trace generation");
156-
let _enter = span.enter();
157-
let (page_traces, mut prover_data) = page_controller.load_page_and_ops(
158-
&page_init,
159-
zk_ops.clone(),
160-
checker_trace_degree,
161-
&mut trace_builder.committer,
162-
);
163-
let offline_checker_trace = page_controller.offline_checker_trace();
164-
let final_page_aux_trace = page_controller.final_page_aux_trace();
165-
let range_checker_trace = page_controller.range_checker_trace();
170+
let init_prover_data_encoded =
171+
read_from_path(cmd.cache_folder.clone() + "/" + &table_id + ".cache.bin").unwrap();
172+
let init_prover_data: ProverTraceData<BabyBearPoseidon2Config> =
173+
bincode::deserialize(&init_prover_data_encoded).unwrap();
166174

167-
// Generating trace for ops_sender and making sure it has height num_ops
168-
let ops_sender_trace =
169-
ops_sender.generate_trace_testing(&zk_ops, config.page.max_rw_ops, 1);
170-
drop(_enter);
175+
let span = tracing::info_span!("Trace generation");
176+
let _enter = span.enter();
177+
let (page_traces, mut prover_data) = page_controller.load_page_and_ops(
178+
&page_init,
179+
zk_ops.clone(),
180+
checker_trace_degree,
181+
&mut trace_builder.committer,
182+
);
183+
let offline_checker_trace = page_controller.offline_checker_trace();
184+
let final_page_aux_trace = page_controller.final_page_aux_trace();
185+
let range_checker_trace = page_controller.range_checker_trace();
171186

172-
// Clearing the range_checker counts
173-
page_controller.update_range_checker(idx_decomp);
187+
// Generating trace for ops_sender and making sure it has height num_ops
188+
let ops_sender_trace = ops_sender.generate_trace_testing(&zk_ops, config.page.max_rw_ops, 1);
189+
drop(_enter);
174190

175-
trace_builder.clear();
191+
// Clearing the range_checker counts
192+
page_controller.update_range_checker(idx_decomp);
176193

177-
trace_builder.load_cached_trace(page_traces[0].clone(), init_prover_data);
178-
trace_builder.load_cached_trace(page_traces[1].clone(), prover_data.remove(0));
179-
trace_builder.load_trace(final_page_aux_trace);
180-
trace_builder.load_trace(offline_checker_trace.clone());
181-
trace_builder.load_trace(range_checker_trace);
182-
trace_builder.load_trace(ops_sender_trace);
194+
trace_builder.clear();
183195

184-
tracing::info_span!("Trace commitment").in_scope(|| trace_builder.commit_current());
185-
let encoded_pk =
186-
read_from_path(self.keys_folder.clone() + "/" + &prefix + ".partial.pk").unwrap();
187-
let partial_pk: MultiStarkPartialProvingKey<BabyBearPoseidon2Config> =
188-
bincode::deserialize(&encoded_pk).unwrap();
189-
let partial_vk = partial_pk.partial_vk();
190-
let main_trace_data = trace_builder.view(
191-
&partial_vk,
192-
vec![
193-
&page_controller.init_chip,
194-
&page_controller.final_chip,
195-
&page_controller.offline_checker,
196-
&page_controller.range_checker.air,
197-
&ops_sender,
198-
],
199-
);
196+
trace_builder.load_cached_trace(page_traces[0].clone(), init_prover_data);
197+
trace_builder.load_cached_trace(page_traces[1].clone(), prover_data.remove(0));
198+
trace_builder.load_trace(final_page_aux_trace);
199+
trace_builder.load_trace(offline_checker_trace.clone());
200+
trace_builder.load_trace(range_checker_trace);
201+
trace_builder.load_trace(ops_sender_trace);
200202

201-
let pis = vec![vec![]; partial_vk.per_air.len()];
203+
let prefix = create_prefix(config);
202204

203-
let prover = engine.prover();
205+
tracing::info_span!("Trace commitment").in_scope(|| trace_builder.commit_current());
206+
let encoded_pk =
207+
read_from_path(cmd.keys_folder.clone() + "/" + &prefix + ".partial.pk").unwrap();
208+
let partial_pk: MultiStarkPartialProvingKey<BabyBearPoseidon2Config> =
209+
bincode::deserialize(&encoded_pk).unwrap();
210+
let partial_vk = partial_pk.partial_vk();
211+
let main_trace_data = trace_builder.view(
212+
&partial_vk,
213+
vec![
214+
&page_controller.init_chip,
215+
&page_controller.final_chip,
216+
&page_controller.offline_checker,
217+
&page_controller.range_checker.air,
218+
&ops_sender,
219+
],
220+
);
204221

205-
let mut challenger = engine.new_challenger();
206-
let proof = prover.prove(&mut challenger, &partial_pk, main_trace_data, &pis);
207-
let encoded_proof: Vec<u8> = bincode::serialize(&proof).unwrap();
208-
let table = interface.get_table(table_id.clone()).unwrap();
209-
if !self.silent {
210-
println!("Table ID: {}", table_id);
211-
println!("{:?}", table.metadata);
212-
for (index, data) in table.body.iter() {
213-
println!("{:?}: {:?}", index, data);
214-
}
222+
let pis = vec![vec![]; partial_vk.per_air.len()];
223+
224+
let prover = engine.prover();
225+
226+
let mut challenger = engine.new_challenger();
227+
let proof = prover.prove(&mut challenger, &partial_pk, main_trace_data, &pis);
228+
let encoded_proof: Vec<u8> = bincode::serialize(&proof).unwrap();
229+
let table = interface.get_table(table_id.clone()).unwrap();
230+
if !cmd.silent {
231+
println!("Table ID: {}", table_id);
232+
println!("{:?}", table.metadata);
233+
for (index, data) in table.body.iter() {
234+
println!("{:?}: {:?}", index, data);
215235
}
216-
let proof_path = self.db_file_path.clone() + ".prove.bin";
217-
write_bytes(&encoded_proof, proof_path).unwrap();
218-
db.save_to_file(&(self.db_file_path.clone() + ".0"))?;
219-
Ok(())
220236
}
237+
let proof_path = cmd.db_file_path.clone() + ".prove.bin";
238+
write_bytes(&encoded_proof, proof_path).unwrap();
239+
db.save_to_file(&(cmd.db_file_path.clone() + ".0"))?;
240+
Ok(())
221241
}
222242

223243
fn afi_op_conv<I, D>(

‎bin/afs/src/commands/verify/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,8 @@ impl VerifyCommand {
8585
.max(8);
8686

8787
let idx_decomp = 8;
88+
8889
println!("Verifying proof file: {}", self.proof_file);
89-
// verify::verify_ops(&self.proof_file).await?;
9090
let encoded_vk =
9191
read_from_path(self.keys_folder.clone() + "/" + &prefix + ".partial.vk").unwrap();
9292
let partial_vk: MultiStarkPartialVerifyingKey<BabyBearPoseidon2Config> =

0 commit comments

Comments
 (0)