Fix IMT skipping

.github/CODEOWNERS

@@ -1,10 +1,10 @@
-/src/app/zeko/* @L-as
-/src/app/zeko/circuits/ @L-as
-/src/app/zeko/circuits_as_prover/ @L-as
-/src/app/zeko/compile_simple/ @L-as
-/src/app/zeko/tests/ @L-as
-/src/app/zeko/v/ @L-as
+/src/app/zeko/* @MartinOndejka
+/src/app/zeko/circuits/ @MartinOndejka
+/src/app/zeko/circuits_as_prover/ @MartinOndejka
+/src/app/zeko/compile_simple/ @MartinOndejka
+/src/app/zeko/tests/ @MartinOndejka
+/src/app/zeko/v/ @MartinOndejka
 /src/app/zeko/da_layer/ @MartinOndejka
-/src/app/zeko/da_layer/core.ml @L-as
+/src/app/zeko/da_layer/core.ml @MartinOndejka
 /src/app/zeko/sequencer/ @MartinOndejka
 /src/app/lib/ @L-as @MartinOndejka

src/app/zeko/circuits/indexed_merkle_tree.ml

@@ -1,7 +1,6 @@
 open Core_kernel
 open Snark_params.Tick
 open Zeko_util
-open Checked.Let_syntax
 
 module Make (Inputs : sig
   module Key : SnarkType
@@ -28,28 +27,68 @@ struct
     type t = { hash_other : F.t; is_right : Boolean.t } [@@deriving snarky]
   end
 
-  module Path =
-    SnarkList
-      (PathStep)
-      (struct
-        let length = height
-      end)
+  module Path = struct
+    include
+      SnarkList
+        (PathStep)
+        (struct
+          let length = height
+        end)
+  end
 
   let hash_entry =
     var_to_hash ~init:Zeko_constants.indexed_merkle_tree_salt Entry.typ
 
-  let empty_path =
-    List.init height ~f:Fn.id
-    |> List.fold_map ~init:Field.zero ~f:(fun acc _ ->
-           let next =
-             Random_oracle.hash
-               ~init:
-                 (Hash_prefix_create.salt
-                    Zeko_constants.indexed_merkle_tree_merge_salt )
-               [| acc; acc |]
-           in
-           (next, constant Field.typ acc) )
-    |> snd
+  (** Two leaves are adjacent if and only if:
+      1. Their lowest common ancestor is as deep as possible
+      2. One leaf is the rightmost leaf of the left subtree
+      3. The other is the leftmost leaf of the right subtree
+    *)
+  let are_paths_neighbors ~(l : Path.var) ~(r : Path.var) =
+    let* _in_prefix, seen_div, valid =
+      foldl
+        (List.zip_exn l r |> List.rev)
+        ~init:(Boolean.true_, Boolean.false_, Boolean.true_)
+        ~f:(fun (in_prefix, seen_div, valid)
+                ( { PathStep.is_right = l_is_right; _ }
+                , { PathStep.is_right = r_is_right; _ } ) ->
+          let* same = Boolean.equal l_is_right r_is_right in
+
+          (* diverge happens exactly when we were in prefix and now differ *)
+          let* diverge_now = Boolean.(in_prefix && not same) in
+
+          (* At the divergence bit we require l=0 and r=1 *)
+          let* div_ok =
+            Boolean.Expr.(((not !l_is_right) && !r_is_right && !valid) |> eval)
+          in
+
+          (* After divergence we require l=1 and r=0 at every step *)
+          let* tail_ok =
+            Boolean.Expr.((!l_is_right && (not !r_is_right) && !valid) |> eval)
+          in
+
+          (* Update validity:
+             - before divergence: no constraint
+             - at divergence: enforce div_ok
+             - after divergence: enforce tail_ok
+          *)
+          let* if_diverged_then_tail_ok =
+            (* if already diverged earlier, enforce tail constraint *)
+            let* enforce_tail = Boolean.(seen_div && not diverge_now) in
+            if_ enforce_tail ~typ:Boolean.typ ~then_:tail_ok ~else_:valid
+          in
+          let* valid =
+            if_ diverge_now ~typ:Boolean.typ ~then_:div_ok
+              ~else_:if_diverged_then_tail_ok
+          in
+
+          (* Update flags *)
+          let* seen_div = Boolean.(seen_div || diverge_now) in
+          let*| in_prefix = Boolean.(in_prefix && same) in
+          (in_prefix, seen_div, valid) )
+    in
+    (* Must have diverged at least once, otherwise same leaf *)
+    Boolean.(valid && seen_div)
 
   (* TODO: consider different salt per level. *)
   (* NB: The first element in the list is the neighbor of init, and the next element
@@ -69,7 +108,7 @@ struct
     let* init = hash_entry entry in
     implied_root_raw init path
 
-  let add_key_var ~check ~x ~path_x ~y ~path_y ~z () =
+  let add_key_var ~check ~x ~path_x ~y_prev_hash ~path_y_prev ~y ~path_y ~z () =
     let* () =
       with_label __LOC__ (fun () -> assert_x_less_than_y_less_than_z ~x ~y ~z)
     in
@@ -86,18 +125,17 @@ struct
     in
     let* root_new = implied_root { key = y; next_key = z } path_y in
     let* root = if_ check ~typ:F.typ ~then_:root ~else_:root_new in
-    (* Check that no empty indices have been skipped. *)
-    let* is_y_most_left =
-      foldl (List.zip_exn path_y empty_path) ~init:Boolean.true_
-        ~f:(fun acc (PathStep.{ hash_other; is_right }, empty_hash) ->
-          let* is_valid_left =
-            Field.Checked.equal empty_hash hash_other >>= Boolean.( &&& ) acc
-          in
-          if_ is_right ~typ:Boolean.typ ~then_:acc ~else_:is_valid_left )
+    (* Check that the leaf before y is not empty *)
+    let* () =
+      assert_not_equal ~label:__LOC__ Field.typ y_prev_hash
+        Field.(constant typ zero)
     in
+    let* root_new' = implied_root_raw y_prev_hash path_y_prev in
+    let* () = assert_equal ~label:__LOC__ F.typ root_new root_new' in
+    (* Check that the path to leaf before y is really before y *)
+    let* are_neighbors = are_paths_neighbors ~l:path_y_prev ~r:path_y in
     let* () =
-      if_ check ~typ:Boolean.typ ~then_:is_y_most_left ~else_:Boolean.true_
-      >>= Boolean.Assert.is_true
+      assert_equal ~label:__LOC__ Boolean.typ are_neighbors Boolean.true_
     in
     Checked.return (`Before_adding_y root, `After_adding_y root_new)
 

src/app/zeko/circuits/indexed_merkle_tree.mli

@@ -51,6 +51,8 @@ end) : sig
        check:Boolean.var
     -> x:Key.var
     -> path_x:Path.var
+    -> y_prev_hash:F.var
+    -> path_y_prev:Path.var
     -> y:Key.var
     -> path_y:Path.var
     -> z:Key.var

src/app/zeko/circuits/rule_commit.ml

@@ -162,7 +162,9 @@ struct
       @@ fun () ->
       let input =
         let open Random_oracle.Input.Chunked in
-        Ledger_hash.var_to_field target_ledger |> field
+        append
+          (Ledger_hash.var_to_field target_ledger |> field)
+          (Account_set.to_input_var target_acc_set)
       in
       let* payload =
         make_checked (fun () ->

src/app/zeko/circuits/txn_state.ml

@@ -80,6 +80,8 @@ let constraint_constants : Genesis_constants.Constraint_constants.t =
 type update_acc_set_witness =
   { get_account_set_x : unit -> Token_id.t
   ; get_account_set_z : unit -> Token_id.t
+  ; get_account_set_y_prev_hash : unit -> Field.t
+  ; get_account_set_y_prev_path : unit -> Account_set.Path.t
   ; get_account_set_x_path : unit -> Account_set.Path.t
   ; get_account_set_y_path : unit -> Account_set.Path.t
   }
@@ -100,6 +102,14 @@ let update_acc_set accounts init ~witness =
         exists Account_set.Path.typ
           ~compute:(witness >>| fun x -> x.get_account_set_x_path ())
       in
+      let* y_prev_hash =
+        exists Field.typ
+          ~compute:(witness >>| fun x -> x.get_account_set_y_prev_hash ())
+      in
+      let* path_y_prev =
+        exists Account_set.Path.typ
+          ~compute:(witness >>| fun x -> x.get_account_set_y_prev_path ())
+      in
       let* path_y =
         exists Account_set.Path.typ
           ~compute:(witness >>| fun x -> x.get_account_set_y_path ())
@@ -110,8 +120,8 @@ let update_acc_set accounts init ~witness =
       in
       let* y = derive_token_id ~owner:account_id in
       let* `Before_adding_y set', `After_adding_y new_set =
-        Account_set.add_key_var ~x ~path_x ~y ~path_y ~z
-          ~check:is_empty_and_writeable ()
+        Account_set.add_key_var ~x ~path_x ~y_prev_hash ~path_y_prev ~y ~path_y
+          ~z ~check:is_empty_and_writeable ()
       in
       let*| () = assert_equal ~label:__LOC__ Account_set.typ set set' in
       new_set )

src/app/zeko/circuits/zeko_util.ml

@@ -405,6 +405,14 @@ let var_equal : ('var, 't) Typ.t -> 'var -> 'var -> Boolean.Expr.t Checked.t =
   in
   Boolean.Expr.all bools
 
+let assert_not_equal :
+    label:string -> ('var, 't) Typ.t -> 'var -> 'var -> unit Checked.t =
+ fun ~label typ x y ->
+  let* is_equal = var_equal typ x y in
+  let is_not_equal = Boolean.Expr.not is_equal in
+  let@ () = with_label label in
+  Boolean.Expr.assert_ is_not_equal
+
 module Checked32 = struct
   include Mina_numbers.Nat.Make32 ()
 

src/app/zeko/circuits/zeko_util.mli

@@ -184,6 +184,9 @@ val assert_equal_safer :
 
 val var_equal : ('var, 't) Typ.t -> 'var -> 'var -> Boolean.Expr.t Checked.t
 
+val assert_not_equal :
+  label:string -> ('var, 't) Typ.t -> 'var -> 'var -> unit Checked.t
+
 module Checked32 : sig
   include module type of Mina_numbers.Nat.Make32 ()
 

src/app/zeko/da_layer/cli.ml

@@ -16,12 +16,6 @@ let run_node =
          flag "--port"
            (optional_with_default 8080 int)
            ~doc:"int Port to listen on"
-       and node_to_sync =
-         flag "--da-node-to-sync" (optional string)
-           ~doc:"string Nodes to sync with"
-       and hash_to_sync =
-         flag "--hash-to-sync" (optional string)
-           ~doc:"string Hash to sync with in decimal string form"
        and testing_mode =
          flag "--random-sk" no_arg
            ~doc:"Run in testing mode, the signer key will be generated randomly"
@@ -45,20 +39,6 @@ let run_node =
          in
          let logger = Logger.create () in
          Stdout_log.setup log_json log_level ;
-         let sync_arg =
-           match (node_to_sync, hash_to_sync) with
-           | Some node_to_sync, Some hash_to_sync ->
-               Some
-                 ( Cli_lib.Flag.Types.
-                     { value = Core_kernel.Host_and_port.of_string node_to_sync
-                     ; name = "node-to-sync"
-                     }
-                 , Mina_base.Ledger_hash.of_decimal_string hash_to_sync )
-           | None, None ->
-               None
-           | _ ->
-               failwith "Both node-to-sync and hash-to-sync must be provided"
-         in
          let chain =
            match network_id with
            | "mainnet" ->
@@ -70,7 +50,7 @@ let run_node =
          in
          let%bind () =
            Deferred.ignore_m
-           @@ Da_layer.Node.create_server ~chain ~sync_arg ~logger ~port ~db_dir
+           @@ Da_layer.Node.create_server ~chain ~logger ~port ~db_dir
                 ~signer_sk:signer ~no_migrations ()
          in
          [%log info] "Server started on port %d" port ;