Skip to content

Commit 5aef784

Browse files
roynalnarutoroynalnaruto
authored
Refactor proving task build process (#1757)
1 parent 9aae8ec commit 5aef784

File tree

4 files changed

+68
-77
lines changed

4 files changed

+68
-77
lines changed

‎crates/libzkp/src/tasks.rs‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -63,10 +63,9 @@ pub fn gen_universal_chunk_task(
6363
task: ChunkProvingTask,
6464
) -> eyre::Result<(B256, ChunkProofMetadata, ProvingTask)> {
6565
let chunk_total_gas = task.stats().total_gas_used;
66-
let (chunk_info, pi_hash) = task.precheck_and_build_metadata()?;
67-
let proving_task = task.try_into()?;
66+
let (proving_task, chunk_info, chunk_pi_hash) = task.into_proving_task_with_precheck()?;
6867
Ok((
69-
pi_hash,
68+
chunk_pi_hash,
7069
ChunkProofMetadata {
7170
chunk_info,
7271
chunk_total_gas,
@@ -79,8 +78,7 @@ pub fn gen_universal_chunk_task(
7978
pub fn gen_universal_batch_task(
8079
task: BatchProvingTask,
8180
) -> eyre::Result<(B256, BatchProofMetadata, ProvingTask)> {
82-
let (batch_info, batch_pi_hash) = task.precheck_and_build_metadata()?;
83-
let proving_task = task.try_into()?;
81+
let (proving_task, batch_info, batch_pi_hash) = task.into_proving_task_with_precheck()?;
8482
Ok((
8583
batch_pi_hash,
8684
BatchProofMetadata { batch_info },
@@ -92,8 +90,7 @@ pub fn gen_universal_batch_task(
9290
pub fn gen_universal_bundle_task(
9391
task: BundleProvingTask,
9492
) -> eyre::Result<(B256, BundleProofMetadata, ProvingTask)> {
95-
let (bundle_info, bundle_pi_hash) = task.precheck_and_build_metadata()?;
96-
let proving_task = task.try_into()?;
93+
let (proving_task, bundle_info, bundle_pi_hash) = task.into_proving_task_with_precheck()?;
9794
Ok((
9895
bundle_pi_hash,
9996
BundleProofMetadata {

‎crates/libzkp/src/tasks/batch.rs‎

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -112,33 +112,31 @@ pub struct BatchProvingTask {
112112
pub fork_name: String,
113113
}
114114

115-
impl TryFrom<BatchProvingTask> for ProvingTask {
116-
type Error = eyre::Error;
117-
118-
fn try_from(value: BatchProvingTask) -> Result<Self> {
119-
let witness = value.build_guest_input(value.version.into());
120-
let serialized_witness = if crate::witness_use_legacy_mode(&value.fork_name)? {
115+
impl BatchProvingTask {
116+
pub fn into_proving_task_with_precheck(self) -> Result<(ProvingTask, BatchInfo, B256)> {
117+
let (witness, metadata, batch_pi_hash) = self.precheck()?;
118+
let serialized_witness = if crate::witness_use_legacy_mode(&self.fork_name)? {
121119
let legacy_witness = LegacyBatchWitness::from(witness);
122120
to_rkyv_bytes::<RancorError>(&legacy_witness)?.into_vec()
123121
} else {
124122
super::encode_task_to_witness(&witness)?
125123
};
126124

127-
Ok(ProvingTask {
128-
identifier: value.batch_header.batch_hash().to_string(),
129-
fork_name: value.fork_name,
130-
aggregated_proofs: value
125+
let proving_task = ProvingTask {
126+
identifier: self.batch_header.batch_hash().to_string(),
127+
fork_name: self.fork_name,
128+
aggregated_proofs: self
131129
.chunk_proofs
132130
.into_iter()
133131
.map(|w_proof| w_proof.proof.into_stark_proof().expect("expect root proof"))
134132
.collect(),
135133
serialized_witness: vec![serialized_witness],
136134
vk: Vec::new(),
137-
})
135+
};
136+
137+
Ok((proving_task, metadata, batch_pi_hash))
138138
}
139-
}
140139

141-
impl BatchProvingTask {
142140
fn build_guest_input(&self, version: Version) -> BatchWitness {
143141
tracing::info!(
144142
"Handling batch task for input, version byte {}, Version data: {:?}",
@@ -281,7 +279,7 @@ impl BatchProvingTask {
281279
}
282280
}
283281

284-
pub fn precheck_and_build_metadata(&self) -> Result<(BatchInfo, B256)> {
282+
pub fn precheck(&self) -> Result<(BatchWitness, BatchInfo, B256)> {
285283
// for every aggregation task, there are two steps needed to build the metadata:
286284
// 1. generate data for metadata from the witness
287285
// 2. validate every adjacent proof pair
@@ -294,7 +292,7 @@ impl BatchProvingTask {
294292
)?;
295293
let pi_hash = metadata.pi_hash_by_version(version);
296294

297-
Ok((metadata, pi_hash))
295+
Ok((witness, metadata, pi_hash))
298296
}
299297
}
300298

‎crates/libzkp/src/tasks/bundle.rs‎

Lines changed: 26 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,30 @@ pub struct BundleProvingTask {
2525
}
2626

2727
impl BundleProvingTask {
28+
pub fn into_proving_task_with_precheck(self) -> Result<(ProvingTask, BundleInfo, B256)> {
29+
let (witness, bundle_info, bundle_pi_hash) = self.precheck()?;
30+
let serialized_witness = if crate::witness_use_legacy_mode(&self.fork_name)? {
31+
let legacy = LegacyBundleWitness::from(witness);
32+
to_rkyv_bytes::<RancorError>(&legacy)?.into_vec()
33+
} else {
34+
super::encode_task_to_witness(&witness)?
35+
};
36+
37+
let proving_task = ProvingTask {
38+
identifier: self.identifier(),
39+
fork_name: self.fork_name,
40+
aggregated_proofs: self
41+
.batch_proofs
42+
.into_iter()
43+
.map(|w_proof| w_proof.proof.into_stark_proof().expect("expect root proof"))
44+
.collect(),
45+
serialized_witness: vec![serialized_witness],
46+
vk: Vec::new(),
47+
};
48+
49+
Ok((proving_task, bundle_info, bundle_pi_hash))
50+
}
51+
2852
fn identifier(&self) -> String {
2953
assert!(!self.batch_proofs.is_empty(), "{BUNDLE_SANITY_MSG}",);
3054

@@ -59,7 +83,7 @@ impl BundleProvingTask {
5983
}
6084
}
6185

62-
pub fn precheck_and_build_metadata(&self) -> Result<(BundleInfo, B256)> {
86+
fn precheck(&self) -> Result<(BundleWitness, BundleInfo, B256)> {
6387
// for every aggregation task, there are two steps needed to build the metadata:
6488
// 1. generate data for metadata from the witness
6589
// 2. validate every adjacent proof pair
@@ -72,32 +96,6 @@ impl BundleProvingTask {
7296
)?;
7397
let pi_hash = metadata.pi_hash_by_version(version);
7498

75-
Ok((metadata, pi_hash))
76-
}
77-
}
78-
79-
impl TryFrom<BundleProvingTask> for ProvingTask {
80-
type Error = eyre::Error;
81-
82-
fn try_from(value: BundleProvingTask) -> Result<Self> {
83-
let witness = value.build_guest_input(value.version.into());
84-
let serialized_witness = if crate::witness_use_legacy_mode(&value.fork_name)? {
85-
let legacy = LegacyBundleWitness::from(witness);
86-
to_rkyv_bytes::<RancorError>(&legacy)?.into_vec()
87-
} else {
88-
super::encode_task_to_witness(&witness)?
89-
};
90-
91-
Ok(ProvingTask {
92-
identifier: value.identifier(),
93-
fork_name: value.fork_name,
94-
aggregated_proofs: value
95-
.batch_proofs
96-
.into_iter()
97-
.map(|w_proof| w_proof.proof.into_stark_proof().expect("expect root proof"))
98-
.collect(),
99-
serialized_witness: vec![serialized_witness],
100-
vk: Vec::new(),
101-
})
99+
Ok((witness, metadata, pi_hash))
102100
}
103101
}

‎crates/libzkp/src/tasks/chunk.rs‎

Lines changed: 25 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -94,28 +94,6 @@ pub struct ChunkDetails {
9494
pub total_gas_used: u64,
9595
}
9696

97-
impl TryFrom<ChunkProvingTask> for ProvingTask {
98-
type Error = eyre::Error;
99-
100-
fn try_from(value: ChunkProvingTask) -> Result<Self> {
101-
let witness = value.build_guest_input(value.version.into());
102-
let serialized_witness = if crate::witness_use_legacy_mode(&value.fork_name)? {
103-
let legacy_witness = LegacyChunkWitness::from(witness);
104-
to_rkyv_bytes::<RancorError>(&legacy_witness)?.into_vec()
105-
} else {
106-
super::encode_task_to_witness(&witness)?
107-
};
108-
109-
Ok(ProvingTask {
110-
identifier: value.identifier(),
111-
fork_name: value.fork_name,
112-
aggregated_proofs: Vec::new(),
113-
serialized_witness: vec![serialized_witness],
114-
vk: Vec::new(),
115-
})
116-
}
117-
}
118-
11997
impl ChunkProvingTask {
12098
pub fn stats(&self) -> ChunkDetails {
12199
let num_blocks = self.block_witnesses.len();
@@ -137,6 +115,26 @@ impl ChunkProvingTask {
137115
}
138116
}
139117

118+
pub fn into_proving_task_with_precheck(self) -> Result<(ProvingTask, ChunkInfo, B256)> {
119+
let (witness, chunk_info, chunk_pi_hash) = self.precheck()?;
120+
let serialized_witness = if crate::witness_use_legacy_mode(&self.fork_name)? {
121+
let legacy_witness = LegacyChunkWitness::from(witness);
122+
to_rkyv_bytes::<RancorError>(&legacy_witness)?.into_vec()
123+
} else {
124+
super::encode_task_to_witness(&witness)?
125+
};
126+
127+
let proving_task = ProvingTask {
128+
identifier: self.identifier(),
129+
fork_name: self.fork_name,
130+
aggregated_proofs: Vec::new(),
131+
serialized_witness: vec![serialized_witness],
132+
vk: Vec::new(),
133+
};
134+
135+
Ok((proving_task, chunk_info, chunk_pi_hash))
136+
}
137+
140138
fn identifier(&self) -> String {
141139
assert!(!self.block_witnesses.is_empty(), "{CHUNK_SANITY_MSG}",);
142140

@@ -180,13 +178,13 @@ impl ChunkProvingTask {
180178
self.block_witnesses[0].states.push(node);
181179
}
182180

183-
pub fn precheck_and_build_metadata(&self) -> Result<(ChunkInfo, B256)> {
181+
fn precheck(&self) -> Result<(ChunkWitness, ChunkInfo, B256)> {
184182
let version = Version::from(self.version);
185183
let witness = self.build_guest_input(version);
186-
let ret = ChunkInfo::try_from(witness).map_err(|e| eyre::eyre!("{e}"))?;
187-
assert_eq!(ret.post_msg_queue_hash, self.post_msg_queue_hash);
188-
let pi_hash = ret.pi_hash_by_version(version);
189-
Ok((ret, pi_hash))
184+
let chunk_info = ChunkInfo::try_from(witness.clone()).map_err(|e| eyre::eyre!("{e}"))?;
185+
assert_eq!(chunk_info.post_msg_queue_hash, self.post_msg_queue_hash);
186+
let chunk_pi_hash = chunk_info.pi_hash_by_version(version);
187+
Ok((witness, chunk_info, chunk_pi_hash))
190188
}
191189

192190
/// this method check the validate of current task (there may be missing storage node)

0 commit comments

Comments
 (0)