Refactor zkapp and add test with more account updates

Author: sebastiencs

ledger/src/proofs/merge.rs

@@ -53,7 +53,7 @@ pub fn read_witnesses() -> std::io::Result<Vec<Fp>> {
 }
 
 fn merge_main(
-    statement: Statement<SokDigest>,
+    statement: &Statement<SokDigest>,
     proofs: &[v2::LedgerProofProdStableV2; 2],
     w: &mut Witness<Fp>,
 ) -> (Statement<SokDigest>, Statement<SokDigest>) {
@@ -179,11 +179,11 @@ pub fn generate_merge_proof(params: MergeParams, w: &mut Witness<Fp>) -> WrapPro
     } = params;
 
     let sok_digest = message.digest();
-    let statement_with_sok = statement.with_digest(sok_digest);
+    let statement_with_sok = Rc::new(statement.with_digest(sok_digest));
 
-    w.exists(&statement_with_sok);
+    w.exists(&*statement_with_sok);
 
-    let (s1, s2) = merge_main(statement_with_sok.clone(), proofs, w);
+    let (s1, s2) = merge_main(&statement_with_sok, proofs, w);
 
     let proofs: [&v2::PicklesProofProofsVerified2ReprStableV2; 2] = {
         let [p1, p2] = proofs;
@@ -228,7 +228,7 @@ pub fn generate_merge_proof(params: MergeParams, w: &mut Witness<Fp>) -> WrapPro
         proof,
     } = step::<StepMergeProof, MERGE_N_PREVIOUS_PROOFS>(
         StepParams {
-            app_state: Rc::new(statement_with_sok.clone()),
+            app_state: Rc::clone(&statement_with_sok) as _,
             rule,
             for_step_datas,
             indexes,
@@ -253,7 +253,7 @@ pub fn generate_merge_proof(params: MergeParams, w: &mut Witness<Fp>) -> WrapPro
 
     wrap::<WrapMergeProof>(
         WrapParams {
-            app_state: Rc::new(statement_with_sok),
+            app_state: statement_with_sok,
             proof: &proof,
             step_statement: statement,
             prev_evals: &prev_evals,

ledger/src/proofs/transaction.rs

@@ -4669,22 +4669,18 @@ mod tests {
             zkapp_step_proof_prover,
         } = make_provers();
 
-        let mut witnesses: Witness<Fp> = Witness::new::<StepZkappOptSignedOptSignedProof>();
-        let LedgerProof { proof, .. } = generate_zkapp_proof(
-            ZkappParams {
-                statement: &statement,
-                tx_witness: &tx_witness,
-                message: &message,
-                step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
-                step_opt_signed_prover: &zkapp_step_opt_signed_prover,
-                step_proof_prover: &zkapp_step_proof_prover,
-                merge_step_prover: &merge_step_prover,
-                tx_wrap_prover: &tx_wrap_prover,
-                opt_signed_path: Some("zkapp_opt_signed"),
-                proved_path: None,
-            },
-            &mut witnesses,
-        );
+        let LedgerProof { proof, .. } = generate_zkapp_proof(ZkappParams {
+            statement: &statement,
+            tx_witness: &tx_witness,
+            message: &message,
+            step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
+            step_opt_signed_prover: &zkapp_step_opt_signed_prover,
+            step_proof_prover: &zkapp_step_proof_prover,
+            merge_step_prover: &merge_step_prover,
+            tx_wrap_prover: &tx_wrap_prover,
+            opt_signed_path: Some("zkapp_opt_signed"),
+            proved_path: None,
+        });
 
         let proof_json = serde_json::to_vec(&proof.proof).unwrap();
         let sum = dbg!(sha256_sum(&proof_json));
@@ -4719,22 +4715,18 @@ mod tests {
             zkapp_step_proof_prover,
         } = make_provers();
 
-        let mut witnesses: Witness<Fp> = Witness::new::<StepZkappProofProof>();
-        let LedgerProof { proof, .. } = generate_zkapp_proof(
-            ZkappParams {
-                statement: &statement,
-                tx_witness: &tx_witness,
-                message: &message,
-                step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
-                step_opt_signed_prover: &zkapp_step_opt_signed_prover,
-                step_proof_prover: &zkapp_step_proof_prover,
-                merge_step_prover: &merge_step_prover,
-                tx_wrap_prover: &tx_wrap_prover,
-                opt_signed_path: Some("zkapp_proof"),
-                proved_path: Some("zkapp_proof2"),
-            },
-            &mut witnesses,
-        );
+        let LedgerProof { proof, .. } = generate_zkapp_proof(ZkappParams {
+            statement: &statement,
+            tx_witness: &tx_witness,
+            message: &message,
+            step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
+            step_opt_signed_prover: &zkapp_step_opt_signed_prover,
+            step_proof_prover: &zkapp_step_proof_prover,
+            merge_step_prover: &merge_step_prover,
+            tx_wrap_prover: &tx_wrap_prover,
+            opt_signed_path: Some("zkapp_proof"),
+            proved_path: Some("zkapp_proof2"),
+        });
 
         let proof_json = serde_json::to_vec(&proof.proof).unwrap();
         let sum = dbg!(sha256_sum(&proof_json));
@@ -4824,66 +4816,37 @@ mod tests {
             zkapp_step_proof_prover,
         } = make_provers();
 
-        // zkapp proof with signature authorization
-        {
-            let data = std::fs::read(base_dir.join("zkapp_0_rampup4.bin")).unwrap();
-            let (statement, tx_witness, message) = extract_request(&data);
-
-            let mut witnesses: Witness<Fp> = Witness::new::<StepZkappOptSignedOptSignedProof>();
-            let LedgerProof { proof, .. } = generate_zkapp_proof(
-                ZkappParams {
-                    statement: &statement,
-                    tx_witness: &tx_witness,
-                    message: &message,
-                    step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
-                    step_opt_signed_prover: &zkapp_step_opt_signed_prover,
-                    step_proof_prover: &zkapp_step_proof_prover,
-                    merge_step_prover: &merge_step_prover,
-                    tx_wrap_prover: &tx_wrap_prover,
-                    opt_signed_path: Some("zkapp_opt_signed"),
-                    proved_path: None,
-                },
-                &mut witnesses,
-            );
-
-            let proof_json = serde_json::to_vec(&proof.proof).unwrap();
-            let sum = dbg!(sha256_sum(&proof_json));
-
-            assert_eq!(
-                sum,
-                "6e9bb6ed613cf0aa737188e0e8ddde7438211ca54c02e89aff32816c181caca9"
-            );
-        }
+        #[rustfmt::skip]
+        let zkapp_cases = [
+            // zkapp proof with signature authorization
+            ("zkapp_0_rampup4.bin", Some("zkapp_opt_signed"), None, "6e9bb6ed613cf0aa737188e0e8ddde7438211ca54c02e89aff32816c181caca9"),
+            // zkapp proof with proof authorization
+            ("zkapp_10_rampup4.bin", Some("zkapp_proof"), Some("zkapp_proof2"), "e2ca355ce4ed5aaf379e992c0c8c5b1c4ac1687546ceac5a5c6c9c4994002249"),
+            // zkapp with multiple account updates
+            ("zkapp_2_0_rampup4.bin", None, None, "03153d1c5b934e00c7102d3683f27572b6e8bfe0335817cb822d701c83415930"),
+        ];
 
-        // zkapp proof with proof authorization
-        {
-            let data = std::fs::read(base_dir.join("zkapp_10_rampup4.bin")).unwrap();
+        for (file, opt_signed_path, proved_path, expected_sum) in zkapp_cases {
+            let data = std::fs::read(base_dir.join(file)).unwrap();
             let (statement, tx_witness, message) = extract_request(&data);
 
-            let mut witnesses: Witness<Fp> = Witness::new::<StepZkappProofProof>();
-            let LedgerProof { proof, .. } = generate_zkapp_proof(
-                ZkappParams {
-                    statement: &statement,
-                    tx_witness: &tx_witness,
-                    message: &message,
-                    step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
-                    step_opt_signed_prover: &zkapp_step_opt_signed_prover,
-                    step_proof_prover: &zkapp_step_proof_prover,
-                    merge_step_prover: &merge_step_prover,
-                    tx_wrap_prover: &tx_wrap_prover,
-                    opt_signed_path: Some("zkapp_proof"),
-                    proved_path: Some("zkapp_proof2"),
-                },
-                &mut witnesses,
-            );
+            let LedgerProof { proof, .. } = generate_zkapp_proof(ZkappParams {
+                statement: &statement,
+                tx_witness: &tx_witness,
+                message: &message,
+                step_opt_signed_opt_signed_prover: &zkapp_step_opt_signed_opt_signed_prover,
+                step_opt_signed_prover: &zkapp_step_opt_signed_prover,
+                step_proof_prover: &zkapp_step_proof_prover,
+                merge_step_prover: &merge_step_prover,
+                tx_wrap_prover: &tx_wrap_prover,
+                opt_signed_path,
+                proved_path,
+            });
 
             let proof_json = serde_json::to_vec(&proof.proof).unwrap();
             let sum = dbg!(sha256_sum(&proof_json));
 
-            assert_eq!(
-                sum,
-                "e2ca355ce4ed5aaf379e992c0c8c5b1c4ac1687546ceac5a5c6c9c4994002249"
-            );
+            assert_eq!(sum, expected_sum);
         }
 
         // Block proof

ledger/src/proofs/zkapp.rs

@@ -1197,7 +1197,18 @@ fn zkapp_main(
                 };
                 finish(v, (&mut global, &mut local))
             }
-            IsStart::Yes => todo!(),
+            IsStart::Yes => {
+                assert!(local.stack_frame.data.calls.data.is_empty());
+
+                let v = match start_zkapp_command {
+                    [] => unreachable!(),
+                    [p, ps @ ..] => {
+                        start_zkapp_command = ps;
+                        StartOrSkip::Start(p)
+                    }
+                };
+                finish(v, (&mut global, &mut local))
+            }
         };
 
         new_acc.unwrap() // TODO: Remove unwrap
@@ -1282,7 +1293,7 @@ fn snapp_proof_data<'a>(
 
 fn of_zkapp_command_segment_exn<StepConstants, WrapConstants>(
     statement: Statement<SokDigest>,
-    witness: &ZkappCommandSegmentWitness,
+    zkapp_witness: &ZkappCommandSegmentWitness,
     spec: &SegmentBasic,
     step_prover: &Prover<Fp>,
     tx_wrap_prover: &Prover<Fq>,
@@ -1302,7 +1313,7 @@ where
         w.ocaml_aux = read_witnesses(path);
     };
 
-    let (zkapp_input, must_verify) = zkapp_main(statement.clone(), witness, &s, &mut w);
+    let (zkapp_input, must_verify) = zkapp_main(statement.clone(), zkapp_witness, &s, &mut w);
 
     let StepProof {
         statement: step_statement,
@@ -1323,7 +1334,7 @@ where
             &mut w,
         ),
         Proved => {
-            let (proof, vk) = snapp_proof_data(witness).unwrap();
+            let (proof, vk) = snapp_proof_data(zkapp_witness).unwrap();
             let proof = proof.into();
 
             let dlog_plonk_index_cvar = vk.wrap_index.to_cvar(CircuitVar::Var);
@@ -1614,7 +1625,7 @@ impl From<&WrapProof> for v2::PicklesProofProofsVerified2ReprStableV2 {
 
 fn of_zkapp_command_segment(
     statement: Statement<SokDigest>,
-    witness: &ZkappCommandSegmentWitness,
+    zkapp_witness: &ZkappCommandSegmentWitness,
     spec: &SegmentBasic,
     step_opt_signed_opt_signed_prover: &Prover<Fp>,
     step_opt_signed_prover: &Prover<Fp>,
@@ -1651,7 +1662,7 @@ fn of_zkapp_command_segment(
 
     of_zkapp_command_segment_exn(
         statement,
-        witness,
+        zkapp_witness,
         spec,
         step_prover,
         tx_wrap_prover,
@@ -1660,7 +1671,7 @@ fn of_zkapp_command_segment(
     )
 }
 
-pub fn generate_zkapp_proof(params: ZkappParams, w: &mut Witness<Fp>) -> LedgerProof {
+pub fn generate_zkapp_proof(params: ZkappParams) -> LedgerProof {
     let ZkappParams {
         statement,
         tx_witness,
@@ -1703,11 +1714,11 @@ pub fn generate_zkapp_proof(params: ZkappParams, w: &mut Witness<Fp>) -> LedgerP
     });
 
     let of_zkapp_command_segment = |statement: Statement<SokDigest>,
-                                    witness: &ZkappCommandSegmentWitness<'_>,
+                                    zkapp_witness: &ZkappCommandSegmentWitness<'_>,
                                     spec: &SegmentBasic| {
         of_zkapp_command_segment(
             statement,
-            witness,
+            zkapp_witness,
             spec,
             step_opt_signed_opt_signed_prover,
             step_opt_signed_prover,
@@ -1720,23 +1731,32 @@ pub fn generate_zkapp_proof(params: ZkappParams, w: &mut Witness<Fp>) -> LedgerP
 
     let sok_digest = message.digest();
     let mut witnesses_specs_stmts = witnesses_specs_stmts.into_iter().rev();
-    let (witness, spec, statement) = witnesses_specs_stmts.next().unwrap(); // last one
+    let (zkapp_witness, spec, statement) = witnesses_specs_stmts.next().unwrap(); // last one
 
-    let mut first_proof =
-        of_zkapp_command_segment(statement.with_digest(sok_digest.clone()), &witness, &spec);
+    let mut first_proof = of_zkapp_command_segment(
+        statement.with_digest(sok_digest.clone()),
+        &zkapp_witness,
+        &spec,
+    );
 
-    witnesses_specs_stmts.fold(first_proof, |prev_proof, (witness, spec, statement)| {
-        let curr_proof =
-            of_zkapp_command_segment(statement.with_digest(sok_digest.clone()), &witness, &spec);
+    witnesses_specs_stmts.fold(
+        first_proof,
+        |prev_proof, (zkapp_witness, spec, statement)| {
+            let curr_proof = of_zkapp_command_segment(
+                statement.with_digest(sok_digest.clone()),
+                &zkapp_witness,
+                &spec,
+            );
 
-        merge_zkapp_proofs(
-            prev_proof,
-            curr_proof,
-            message,
-            merge_step_prover,
-            tx_wrap_prover,
-        )
-    })
+            merge_zkapp_proofs(
+                prev_proof,
+                curr_proof,
+                message,
+                merge_step_prover,
+                tx_wrap_prover,
+            )
+        },
+    )
 }
 
 fn merge_zkapp_proofs(