From 5c146c4de76ffa5192e69b15378bef0c507f72a6 Mon Sep 17 00:00:00 2001
From: georgeee <george.agapov@gmail.com>
Date: Tue, 21 Nov 2023 15:07:36 +0100
Subject: [PATCH] Parametrize usage of masks in tool

---
 src/app/cli/src/init/client.ml                | 22 +++++++++--
 .../cli/src/init/test_ledger_application.ml   | 37 +++++++++++--------
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/src/app/cli/src/init/client.ml b/src/app/cli/src/init/client.ml
index 3a08cebc532..9657c876aec 100644
--- a/src/app/cli/src/init/client.ml
+++ b/src/app/cli/src/init/client.ml
@@ -2242,25 +2242,39 @@ let test_ledger_application =
        flag "--ledger-path" ~doc:"FILE directory with ledger DB"
          (required string)
      and num_txs =
-       flag "--num-txs" ~doc:"NN Number of transactions to create"
+       flag "--num-txs"
+         ~doc:"NN Number of transactions to create after preparatory rounds"
          (required int)
+     and num_txs_per_round =
+       flag "--num-txs-per-round"
+         ~doc:
+           "NN Number of transactions to create per preparatory round \
+            (default: 3)"
+         (optional int)
+     and rounds =
+       flag "--rounds" ~doc:"NN Number of preparatory rounds (default: 290)"
+         (optional int)
      and first_partition_slots =
        flag "--first-partition-slots"
-         ~doc:"NN Number of slots in first partition of scan state"
+         ~doc:
+           "NN Number of slots in first partition of scan state (default: 128)"
          (optional int)
      and no_new_stack =
        flag "--old-stack" ~doc:"Use is_new_stack: false (scan state)" no_arg
      and has_second_partition =
        flag "--has-second-partition"
          ~doc:"Assume there is a second partition (scan state)" no_arg
-     in
+     and no_masks = flag "--no-masks" ~doc:"Do not create masks" no_arg in
      Cli_lib.Exceptions.handle_nicely
      @@ fun () ->
      let first_partition_slots =
        Option.value ~default:128 first_partition_slots
      in
+     let num_txs_per_round = Option.value ~default:3 num_txs_per_round in
+     let rounds = Option.value ~default:290 rounds in
      Test_ledger_application.test ~privkey_path ~ledger_path ~prev_block_path
-       ~first_partition_slots ~no_new_stack ~has_second_partition num_txs )
+       ~first_partition_slots ~no_new_stack ~has_second_partition
+       ~num_txs_per_round ~rounds ~no_masks num_txs )
 
 let itn_create_accounts =
   Command.async ~summary:"Fund new accounts for incentivized testnet"
diff --git a/src/app/cli/src/init/test_ledger_application.ml b/src/app/cli/src/init/test_ledger_application.ml
index b2e16b6ae7d..9cf8d3360a8 100644
--- a/src/app/cli/src/init/test_ledger_application.ml
+++ b/src/app/cli/src/init/test_ledger_application.ml
@@ -58,7 +58,8 @@ let mk_tx acc_creation_fee keypair nonce =
   Transaction_snark.For_tests.multiple_transfers multispec
 
 let apply_txs ~first_partition_slots ~no_new_stack ~has_second_partition
-    ~num_txs ~prev_protocol_state ~(keypair : Signature_lib.Keypair.t) ~ledger =
+    ~num_txs ~prev_protocol_state ~(keypair : Signature_lib.Keypair.t) ~i ledger
+    =
   let init_nonce =
     let account_id = Account_id.of_public_key keypair.public_key in
     let loc =
@@ -119,8 +120,8 @@ let apply_txs ~first_partition_slots ~no_new_stack ~has_second_partition
   with
   | Ok (b, _, _, _, _) ->
       printf
-        !"Result of application: %B (took %s)\n%!"
-        b
+        !"Result of application %d: %B (took %s)\n%!"
+        i b
         Time.(Span.to_string @@ diff (now ()) start)
   | Error e ->
       eprintf
@@ -129,9 +130,10 @@ let apply_txs ~first_partition_slots ~no_new_stack ~has_second_partition
       exit 1
 
 let test ~privkey_path ~ledger_path ~prev_block_path ~first_partition_slots
-    ~no_new_stack ~has_second_partition num_txs =
+    ~no_new_stack ~has_second_partition ~num_txs_per_round ~rounds ~no_masks
+    num_txs_final =
   let%bind keypair = read_privkey privkey_path in
-  let ledger =
+  let init_ledger =
     Ledger.create ~directory_name:ledger_path
       ~depth:constraint_constants.ledger_depth ()
   in
@@ -142,14 +144,17 @@ let test ~privkey_path ~ledger_path ~prev_block_path ~first_partition_slots
   let prev_protocol_state =
     Mina_block.header prev_block |> Mina_block.Header.protocol_state
   in
-  Deferred.ignore_m
-    (Deferred.List.fold
-       (List.init 290 ~f:(Fn.const ()))
-       ~init:ledger
-       ~f:(fun ledger () ->
-         let%map () =
-           apply_txs ~first_partition_slots ~no_new_stack ~has_second_partition
-             ~num_txs ~prev_protocol_state ~keypair ~ledger
-         in
-         Ledger.register_mask ledger
-           (Ledger.Mask.create ~depth:constraint_constants.ledger_depth ()) ) )
+  let apply =
+    apply_txs ~first_partition_slots ~no_new_stack ~has_second_partition
+      ~prev_protocol_state ~keypair
+  in
+  let mask_handler ledger =
+    if no_masks then Fn.const ledger
+    else
+      Fn.compose (Ledger.register_mask ledger)
+      @@ Ledger.Mask.create ~depth:constraint_constants.ledger_depth
+  in
+  Deferred.List.fold (List.init rounds ~f:ident) ~init:init_ledger
+    ~f:(fun ledger i ->
+      apply ~num_txs:num_txs_per_round ~i ledger >>| mask_handler ledger )
+  >>= apply ~num_txs:num_txs_final ~i:rounds