1- use std:: collections:: BTreeMap ;
1+ use std:: collections:: { BTreeMap , HashMap } ;
22
33use crate :: {
44 arithmetics:: { ceil_log2, next_pow2_instance_padding} ,
@@ -12,7 +12,7 @@ use crate::{
1212} ;
1313use ceno_zkvm:: {
1414 scheme:: { ZKVMChipProof , ZKVMProof } ,
15- structs:: { EccQuarkProof , TowerProofs } ,
15+ structs:: { EccQuarkProof , TowerProofs , ZKVMVerifyingKey } ,
1616} ;
1717use gkr_iop:: gkr:: { GKRProof , layer:: sumcheck_layer:: LayerProof } ;
1818use itertools:: Itertools ;
@@ -97,29 +97,62 @@ pub(crate) struct ZKVMProofInput {
9797 pub opening_proof : BasefoldProof ,
9898}
9999
100- impl From < ( usize , ZKVMProof < E , RecPcs > ) > for ZKVMProofInput {
101- fn from ( d : ( usize , ZKVMProof < E , RecPcs > ) ) -> Self {
100+ impl ZKVMProofInput {
101+ pub fn from_proof (
102+ shard_id : usize ,
103+ zkvm_proof : ZKVMProof < E , RecPcs > ,
104+ vk : & ZKVMVerifyingKey < E , RecPcs > ,
105+ ) -> Self {
106+ let mut chip_witin_num_vars: HashMap < usize , ( usize , usize ) > = HashMap :: new ( ) ; // (chip_id, (num_witin, num_fixed))
107+ let mut chip_indices = zkvm_proof
108+ . chip_proofs
109+ . keys ( )
110+ . copied ( )
111+ . collect :: < Vec < usize > > ( ) ;
112+ chip_indices. push ( 0 ) ;
113+
114+ let mut chip_idx: usize = 0 ;
115+ let mut chip_id: usize = chip_indices[ chip_idx] ;
116+ for ( i, ( _, chip_vk) ) in vk. circuit_vks . iter ( ) . enumerate ( ) {
117+ if chip_id == i {
118+ let composed_cs = chip_vk. get_cs ( ) ;
119+ chip_witin_num_vars. insert (
120+ chip_id,
121+ (
122+ composed_cs. zkvm_v1_css . num_witin as usize ,
123+ composed_cs. zkvm_v1_css . num_fixed ,
124+ ) ,
125+ ) ;
126+ chip_idx += 1 ;
127+ chip_id = chip_indices[ chip_idx] ;
128+ }
129+ }
130+
102131 ZKVMProofInput {
103- shard_id : d. 0 ,
104- raw_pi : d. 1 . raw_pi ,
105- pi_evals : d. 1 . pi_evals ,
106- chip_proofs : d
107- . 1
132+ shard_id,
133+ raw_pi : zkvm_proof. raw_pi ,
134+ pi_evals : zkvm_proof. pi_evals ,
135+ chip_proofs : zkvm_proof
108136 . chip_proofs
109137 . into_iter ( )
110138 . map ( |( chip_idx, proofs) | {
139+ let ( num_witin, num_fixed) = * chip_witin_num_vars
140+ . get ( & chip_idx)
141+ . expect ( "num_witin data should exist" ) ;
111142 (
112143 chip_idx,
113144 proofs
114145 . into_iter ( )
115- . map ( |proof| ZKVMChipProofInput :: from ( ( chip_idx, proof) ) )
146+ . map ( |proof| {
147+ ZKVMChipProofInput :: from ( ( chip_idx, proof, num_witin, num_fixed) )
148+ } )
116149 . collect :: < Vec < ZKVMChipProofInput > > ( )
117150 . into ( ) ,
118151 )
119152 } )
120153 . collect :: < BTreeMap < usize , ZKVMChipProofs > > ( ) ,
121- witin_commit : d . 1 . witin_commit . into ( ) ,
122- opening_proof : d . 1 . opening_proof . into ( ) ,
154+ witin_commit : zkvm_proof . witin_commit . into ( ) ,
155+ opening_proof : zkvm_proof . opening_proof . into ( ) ,
123156 }
124157 }
125158}
@@ -167,7 +200,8 @@ impl Hintable<InnerConfig> for ZKVMProofInput {
167200 . chip_proofs
168201 . iter ( )
169202 . flat_map ( |( _, proofs) | proofs. iter ( ) )
170- . map ( |proof| proof. sum_num_instances )
203+ . filter ( |proof| proof. num_witin > 0 )
204+ . map ( |proof| proof. num_vars )
171205 . collect :: < Vec < _ > > ( ) ;
172206 let witin_max_widths = self
173207 . chip_proofs
@@ -180,7 +214,7 @@ impl Hintable<InnerConfig> for ZKVMProofInput {
180214 . iter ( )
181215 . flat_map ( |( _, proofs) | proofs. iter ( ) )
182216 . filter ( |proof| !proof. fixed_in_evals . is_empty ( ) )
183- . map ( |proof| proof. sum_num_instances )
217+ . map ( |proof| proof. num_vars )
184218 . collect :: < Vec < _ > > ( ) ;
185219 let fixed_max_widths = self
186220 . chip_proofs
@@ -189,7 +223,12 @@ impl Hintable<InnerConfig> for ZKVMProofInput {
189223 . filter ( |proof| !proof. fixed_in_evals . is_empty ( ) )
190224 . map ( |proof| proof. fixed_in_evals . len ( ) )
191225 . collect :: < Vec < _ > > ( ) ;
192- let max_num_var = witin_num_vars. iter ( ) . copied ( ) . max ( ) . unwrap_or ( 0 ) ;
226+ let max_num_var = witin_num_vars
227+ . iter ( )
228+ . chain ( fixed_num_vars. iter ( ) )
229+ . copied ( )
230+ . max ( )
231+ . unwrap_or ( 0 ) ;
193232 let max_width = witin_max_widths
194233 . iter ( )
195234 . chain ( fixed_max_widths. iter ( ) )
@@ -319,7 +358,11 @@ impl Hintable<InnerConfig> for TowerProofInput {
319358
320359pub struct ZKVMChipProofInput {
321360 pub idx : usize ,
322- pub sum_num_instances : usize ,
361+ /// number of variables in the chip's multilinear polynomials
362+ pub num_vars : usize ,
363+
364+ pub num_witin : usize ,
365+ pub num_fixed : usize ,
323366
324367 // product constraints
325368 pub r_out_evals_len : usize ,
@@ -380,15 +423,30 @@ impl Hintable<InnerConfig> for ZKVMChipProofs {
380423 }
381424}
382425
383- impl From < ( usize , ZKVMChipProof < E > ) > for ZKVMChipProofInput {
384- fn from ( d : ( usize , ZKVMChipProof < E > ) ) -> Self {
426+ impl From < ( usize , ZKVMChipProof < E > , usize , usize ) > for ZKVMChipProofInput {
427+ fn from ( d : ( usize , ZKVMChipProof < E > , usize , usize ) ) -> Self {
385428 let idx = d. 0 ;
386429 let p = d. 1 ;
387- let sum_num_instances = p. num_instances . iter ( ) . sum ( ) ;
430+
431+ let num_vars = if p. gkr_iop_proof . is_some ( ) {
432+ let vars = p
433+ . gkr_iop_proof
434+ . as_ref ( )
435+ . unwrap ( )
436+ . 0
437+ . iter ( )
438+ . map ( |l| l. main . proof . proofs . len ( ) )
439+ . collect :: < Vec < usize > > ( ) ;
440+ vars[ 0 ]
441+ } else {
442+ 0
443+ } ;
388444
389445 Self {
390446 idx,
391- sum_num_instances,
447+ num_vars,
448+ num_witin : d. 2 ,
449+ num_fixed : d. 3 ,
392450 r_out_evals_len : p. r_out_evals . len ( ) ,
393451 w_out_evals_len : p. w_out_evals . len ( ) ,
394452 lk_out_evals_len : p. lk_out_evals . len ( ) ,
0 commit comments