11use ceno_emul:: KeccakSpec ;
22use ff_ext:: ExtensionField ;
3+ use gkr_iop:: gkr:: GKRCircuitWitness ;
34use std:: {
45 collections:: { BTreeMap , BTreeSet , HashMap } ,
56 sync:: Arc ,
@@ -99,6 +100,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
99100 let mut commitments = BTreeMap :: new ( ) ;
100101 let mut wits = BTreeMap :: new ( ) ;
101102 let mut structural_wits = BTreeMap :: new ( ) ;
103+ let mut keccak_gkr_wit = Some ( witnesses. keccak_gkr_wit . clone ( ) ) ;
102104
103105 let commit_to_traces_span = entered_span ! ( "commit_to_traces" , profiling_1 = true ) ;
104106 // commit to opcode circuits first and then commit to table circuits, sorted by name
@@ -180,10 +182,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
180182 && cs. r_table_expressions . is_empty ( )
181183 && cs. w_table_expressions . is_empty ( ) ;
182184
183- if * circuit_name == <LargeEcallDummy < E , KeccakSpec > as Instruction < E > >:: name ( ) {
184- // unimplemented!("keccak impl wip");
185- ( )
186- } else if is_opcode_circuit {
185+ if is_opcode_circuit {
187186 tracing:: debug!(
188187 "opcode circuit {} has {} witnesses, {} reads, {} writes, {} lookups" ,
189188 circuit_name,
@@ -197,7 +196,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
197196 let gkr_iop_pk = if * circuit_name
198197 == <LargeEcallDummy < E , KeccakSpec > as Instruction < E > >:: name ( )
199198 {
200- Some ( self . pk . keccak_pk . clone ( ) )
199+ Some ( ( & self . pk . keccak_pk , keccak_gkr_wit . take ( ) . unwrap ( ) ) )
201200 } else {
202201 None
203202 } ;
@@ -206,7 +205,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
206205 circuit_name,
207206 & self . pk . pp ,
208207 pk,
209- & gkr_iop_pk,
208+ gkr_iop_pk,
210209 witness,
211210 wits_commit,
212211 & pi,
@@ -266,7 +265,10 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
266265 name : & str ,
267266 pp : & PCS :: ProverParam ,
268267 circuit_pk : & ProvingKey < E , PCS > ,
269- gkr_iop_pk : & Option < GKRIOPProvingKey < E , PCS , KeccakGKRIOP < E > > > ,
268+ gkr_iop_pk : Option < (
269+ & GKRIOPProvingKey < E , PCS , KeccakGKRIOP < E > > ,
270+ GKRCircuitWitness < E > ,
271+ ) > ,
270272 witnesses : Vec < ArcMultilinearExtension < ' _ , E > > ,
271273 wits_commit : PCS :: CommitmentWithWitness ,
272274 pi : & [ ArcMultilinearExtension < ' _ , E > ] ,
@@ -279,11 +281,6 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
279281 let log2_num_instances = ceil_log2 ( next_pow2_instances) ;
280282 let ( chip_record_alpha, _) = ( challenges[ 0 ] , challenges[ 1 ] ) ;
281283
282- // if let Some(gkr_iop_pk) = gkr_iop_pk {
283- // let mut gkr_iop_pk = gkr_iop_pk.clone();
284- // unimplemented!("cannot fully handle GKRIOP component yet")
285- // }
286-
287284 // sanity check
288285 assert_eq ! ( witnesses. len( ) , cs. num_witin as usize ) ;
289286 assert ! (
@@ -331,6 +328,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
331328
332329 // infer all tower witness after last layer
333330 let span = entered_span ! ( "tower_witness_r_layers" ) ;
331+
334332 let r_wit_layers = infer_tower_product_witness (
335333 log2_num_instances + log2_r_count,
336334 r_records_last_layer,
@@ -455,6 +453,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
455453 rt_tower[ ..log2_num_instances] . to_vec ( ) ,
456454 ) ;
457455
456+ dbg ! ( rt_r. len( ) , rt_w. len( ) , rt_lk. len( ) ) ;
457+
458458 let num_threads = optimal_sumcheck_threads ( log2_num_instances) ;
459459 let alpha_pow = get_challenge_pows (
460460 MAINCONSTRAIN_SUMCHECK_BATCH_SIZE + cs. assert_zero_sumcheck_expressions . len ( ) ,
@@ -481,6 +481,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
481481 ) ;
482482 }
483483
484+ dbg ! ( sel_r. len( ) ) ;
485+
484486 let mut sel_w = build_eq_x_r_vec ( & rt_w[ log2_w_count..] ) ;
485487 if num_instances < sel_w. len ( ) {
486488 sel_w. splice (
@@ -630,6 +632,11 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
630632 distrinct_zerocheck_terms_set. len( ) + 1 // +1 from sel_non_lc_zero_sumcheck
631633 }
632634 ) ;
635+
636+ dbg ! ( r_counts_per_instance) ;
637+ dbg ! ( w_counts_per_instance) ;
638+ dbg ! ( lk_counts_per_instance) ;
639+
633640 let mut main_sel_evals_iter = main_sel_evals. into_iter ( ) ;
634641 main_sel_evals_iter. next ( ) ; // skip sel_r
635642 let r_records_in_evals = ( 0 ..r_counts_per_instance)
@@ -653,6 +660,12 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
653660 distrinct_zerocheck_terms_set. len( ) + 1
654661 }
655662 ) ;
663+
664+ // if let Some(gkr_iop_pk) = gkr_iop_pk {
665+ // let mut gkr_iop_pk = gkr_iop_pk.clone();
666+ // unimplemented!("cannot fully handle GKRIOP component yet")
667+ // }
668+
656669 let input_open_point = main_sel_sumcheck_proofs. point . clone ( ) ;
657670 assert ! ( input_open_point. len( ) == log2_num_instances) ;
658671 exit_span ! ( main_sel_span) ;
@@ -689,6 +702,19 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
689702 exit_span ! ( pcs_open_span) ;
690703 let wits_commit = PCS :: get_pure_commitment ( & wits_commit) ;
691704
705+ if let Some ( ( gkr_iop_pk, gkr_wit) ) = gkr_iop_pk {
706+ let mut gkr_iop_pk = gkr_iop_pk. clone ( ) ;
707+ let gkr_circuit = gkr_iop_pk. vk . get_state ( ) . chip . gkr_circuit ( ) ;
708+
709+ // let mut prover_transcript = transcript::BasicTranscript::<E>::new(b"protocol");
710+ let out_evals = vec ! [ ] ;
711+
712+ let gkr_iop:: gkr:: GKRProverOutput { gkr_proof, .. } = gkr_circuit
713+ . prove ( gkr_wit, & out_evals, & vec ! [ ] , transcript)
714+ . expect ( "Failed to prove phase" ) ;
715+ // unimplemented!("cannot fully handle GKRIOP component yet")
716+ }
717+
692718 Ok ( ZKVMOpcodeProof {
693719 num_instances,
694720 record_r_out_evals,
0 commit comments