Proof

open Protocol
open Protocol.Tx_rollup_l2_storage_sig
open Lwt_syntax

exception Error of Environment.Error_monad.error

module Prover_storage :
  STORAGE
    with type t = Tezos_context_memory.Context_binary.tree
     and type 'a m = 'a Lwt.t = struct
  type t = Tezos_context_memory.Context_binary.tree

  type 'a m = 'a Lwt.t

  module Syntax = struct
    include Lwt.Syntax

    let return = Lwt.return

    let fail e = Lwt.fail (Error e)

    let catch (m : 'a m) k h =
      Lwt.catch
        (fun () -> m >>= k)
        (function Error e -> h e | e -> Lwt.fail e)

    let list_fold_left_m = Lwt_list.fold_left_s
  end

  let path k = [Bytes.to_string k]

  let get store key =
    Tezos_context_memory.Context_binary.Tree.find store (path key)

  let set store key value =
    Tezos_context_memory.Context_binary.Tree.add store (path key) value

  let remove store key =
    Tezos_context_memory.Context_binary.Tree.remove store (path key)
end

module Storage :
  STORAGE
    with type t = Tezos_context_memory.Context_binary.t
     and type 'a m = 'a Lwt.t = struct
  type t = Tezos_context_memory.Context_binary.t

  type 'a m = 'a Lwt.t

  module Syntax = struct
    include Lwt.Syntax

    let return = Lwt.return

    let fail e = Lwt.fail (Error e)

    let catch (m : 'a m) k h =
      Lwt.catch
        (fun () -> m >>= k)
        (function Error e -> h e | e -> Lwt.fail e)

    let list_fold_left_m = Lwt_list.fold_left_s
  end

  let path k = [Bytes.to_string k]

  let get store key = Tezos_context_memory.Context_binary.find store (path key)

  let set store key value =
    Tezos_context_memory.Context_binary.add store (path key) value

  let remove store key =
    Tezos_context_memory.Context_binary.remove store (path key)
end

module Prover_context = Tx_rollup_l2_context.Make (Prover_storage)
module Context = Tx_rollup_l2_context.Make (Storage)
module Prover_apply = Tx_rollup_l2_apply.Make (Prover_context)
module Apply = Tx_rollup_l2_apply.Make (Context)

type proof =
  Tezos_context_memory.Context_binary.Proof.stream
  Tezos_context_memory.Context_binary.Proof.t

let produce_proof repo
    (node_key : Tezos_context_memory.Context_binary.kinded_key) f =
  Tezos_context_memory.Context_binary.produce_stream_proof
    repo
    node_key
    (fun t ->
      let* length = Tezos_context_memory.Context_binary.Tree.length t [] in
      let+ p = f t in
      (p, length))

let produce_ctxt () =
  let open Context.Syntax in
  let empty_ctxt = Tezos_context_memory.Context_binary.create () in
  (* Add stuff to the ctxt here *)
  return empty_ctxt

let run_transaction ctxt msg =
  let open Prover_context.Syntax in
  let* (ctxt, _) = Prover_apply.apply_message ctxt msg in
  return ctxt

(** Tx_rollup_l2_message.t *)
let msg = assert false

let _ =
  Lwt_main.run
  @@
  let open Context.Syntax in
  let* ctxt = produce_ctxt () in

  let index = Tezos_context_memory.Context_binary.index ctxt in
  let time = Time.Protocol.of_seconds 0L in

  let* h = Tezos_context_memory.Context_binary.commit ~time ctxt in
  let* store = Tezos_context_memory.Context_binary.checkout_exn index h in

  let* hash =
    let+ tree_opt = Tezos_context_memory.Context_binary.find_tree store [] in
    match tree_opt with
    | Some t -> Tezos_context_memory.Context_binary.Tree.hash t
    | None -> assert false
  in

  let+ (proof, _) =
    Tezos_context_memory.Context_binary.produce_stream_proof
      index
      (`Node hash)
      (fun ctxt ->
        let* length = Tezos_context_memory.Context_binary.Tree.length ctxt [] in
        let+ ctxt = run_transaction ctxt msg in
        (ctxt, length))
  in

  let bytes =
    Data_encoding.Binary.to_bytes_exn
      Tezos_context_helpers.Context.Proof_encoding.V2.Tree2
      .stream_proof_encoding
      proof
  in

  Printf.printf "bytes: %d\n" (Bytes.length bytes)