From Stdlib Require Import List NArith ZArith.
Import ListNotations.

Require Import skylabs.auto.cpp.proof.
Require Import skylabs.lang.cpp.cpp.
Require Import monad.proofs.exec_specs.
Require Import monad.proofs.execproofs.mip8.blake3model.
Require monad.proofs.libspecs.blake3.model.
Require Import monad.proofs.libspecs.blake3.blake3_impl_h_specs.

Import cQp_compat.

Local Open Scope Z_scope.

Definition blake3_hash_many_mode
    (flags flags_start flags_end : Z) : option hash_mode :=
  if (flags =? model.blake3_flag_value model.FlagDeriveKeyMaterial)
     && (flags_start =? 0)
     && (flags_end =? 0)
  then Some leaf_mode
  else
    if (flags =? 0)
       && (flags_start =? model.blake3_flag_value model.FlagChunkStart)
       && (flags_end =? model.blake3_flag_value model.FlagChunkEnd)
    then Some merge_mode
    else None.

Definition blake3_hash_many_outputs
    (mode : hash_mode) (inputs : list bytes64) : list N :=
  blake3_impl_h_specs.blake3_hash_many_outputs_from_params
    (params_for_hash_mode mode) inputs.

Lemma blake3_hash_many_call_params_leaf :
  blake3_impl_h_specs.blake3_hash_many_call_params
    mip8_leaf_iv_words 0 false
    (model.blake3_flag_value model.FlagDeriveKeyMaterial) 0 0 =
  leaf_hash_many_params.

Lemma blake3_hash_many_call_params_merge :
  blake3_impl_h_specs.blake3_hash_many_call_params
    model.blake3_iv_words 0 false 0
    (model.blake3_flag_value model.FlagChunkStart)
    (model.blake3_flag_value model.FlagChunkEnd) =
  merge_hash_many_params.

Lemma blake3_hash_many_outputs_range mode inputs :
  List.Forall (fun x ⇒ (x < 2 ^ 256)%N)
    (blake3_hash_many_outputs mode inputs).

Definition blake3_seal_digest
    (slot_bitmap : N) (root : option digest) : digest :=
  match root with
  | None ⇒ seal_empty (bytes16_of_N slot_bitmap)
  | Some root ⇒ seal_nonempty (bytes16_of_N slot_bitmap) root
  end.

Definition blake3_seal_model
    (slot_bitmap : N) (root : option digest) : N :=
  bytes32_to_N (blake3_seal_digest slot_bitmap root).

Opaque blake3_seal_digest blake3_seal_model.

Section with_Sigma.
  Context `{Sigma : cpp_logic} {CU : genv} {hh : HasOwn mpredI fracR}.

  Definition bytes32_ty : type := blake3_impl_h_specs.bytes32_ty.

  Definition blake3_input_pointer_kind : Type :=
    blake3_impl_h_specs.blake3_input_pointer_kind.

  Definition Blake3ConstInputPointers : blake3_input_pointer_kind :=
    blake3_impl_h_specs.Blake3ConstInputPointers.

  Definition Blake3MutableInputPointers : blake3_input_pointer_kind :=
    blake3_impl_h_specs.Blake3MutableInputPointers.

  Definition blake3_input_ptr_ty_for :=
    blake3_impl_h_specs.blake3_input_ptr_ty_for.

  Definition blake3_input_ptr_ty :=
    blake3_impl_h_specs.blake3_input_ptr_ty.

  Definition blake3_input_ptr_store_ty_for :=
    blake3_impl_h_specs.blake3_input_ptr_store_ty_for.

  Definition blake3_input_ptr_store_ty :=
    blake3_impl_h_specs.blake3_input_ptr_store_ty.

  Definition byte_value :=
    blake3_impl_h_specs.byte_value.

  Definition bytes32_byte_values :=
    blake3_impl_h_specs.bytes32_byte_values.

  Definition bytes64_byte_values :=
    blake3_impl_h_specs.bytes64_byte_values.

  Definition RawDigestBytesR :=
    blake3_impl_h_specs.RawDigestBytesR.

  Definition Blake3BlockR :=
    blake3_impl_h_specs.Blake3BlockR.

  Definition Blake3BlocksAtR :=
    blake3_impl_h_specs.Blake3BlocksAtR.

  Definition Blake3InputBlocksRFor :=
    blake3_impl_h_specs.Blake3InputBlocksRFor.

  Definition Blake3InputBlocksR :=
    blake3_impl_h_specs.Blake3InputBlocksR.

  Definition Blake3InputBlocksExactRFor :=
    blake3_impl_h_specs.Blake3InputBlocksExactRFor.

  Definition Blake3OutputWordsR :=
    blake3_impl_h_specs.Blake3OutputWordsR.

  Definition Blake3OutputBytesR :=
    blake3_impl_h_specs.Blake3OutputBytesR.

  Definition blake3_hash_many_outputs_from_params_range :=
    blake3_impl_h_specs.blake3_hash_many_outputs_from_params_range.

  Definition Blake3KeyWordsR :=
    blake3_impl_h_specs.Blake3KeyWordsR.

  Definition Blake3ConstKeyWordsR :=
    blake3_impl_h_specs.Blake3ConstKeyWordsR.

  Definition storage_get_leaf_iv_lambda_name : name :=
    Nscoped
      (Nscoped (Nscoped (Nglobal (Nid "monad")) Nanonymous)
         (Nfunction function_qualifiers.N "get_leaf_iv" nil))
      (Nanon 0).

  Definition storage_get_leaf_iv_lambda_call_name : name :=
    Nscoped storage_get_leaf_iv_lambda_name
      (Nop function_qualifiers.Nc OOCall nil).

  Definition storage_get_leaf_iv_static_iv_name : name :=
    Nscoped storage_get_leaf_iv_lambda_call_name (Nid "iv").

  Definition storage_blake3_hash_many_spec : mpred :=
    blake3_impl_h_specs.blake3_hash_many_spec.

  Definition SealRootArgR (rootp : ptr) (root : option digest) : mpred :=
    match root with
    | None ⇒ [| rootp = nullptr |]
    | Some root ⇒
        rootp |-> type_ptrR (Tarray Tuchar 32)
        ** rootp |-> RawDigestBytesR 1 root
    end.

  Definition blake3_seal_block_bytesR
      (base : ptr) (slot_bitmap : N) (bitmap_bytes : list N)
      (root : option digest) : mpred :=
    match root with
    | Some root ⇒
        base |-> arrayLR Tuchar 0 16
          (fun byte : val ⇒ primR Tuchar 1$m byte)
          (map Vint (Z.of_N <$> bitmap_bytes))
        ** base .[ Tuchar ! 16 ] |-> arrayLR Tuchar 0 32
          (fun byte : val ⇒ primR Tuchar 1$m byte)
          (bytes32_byte_values root)
        ** base |-> arrayLR Tuchar 48 64
          (fun byte : val ⇒ primR Tuchar 1$m byte)
          (replicateN 16 (Vint 0))
    | None ⇒
        base |-> arrayLR Tuchar 0 16
          (fun byte : val ⇒ primR Tuchar 1$m byte)
          (map Vint (Z.of_N <$> bitmap_bytes))
        ** base |-> arrayLR Tuchar 16 64
          (fun byte : val ⇒ primR Tuchar 1$m byte)
          (replicateN 48 (Vint 0))
    end.

  Axiom blake3_seal_blockR :
    ∀ (base : ptr) slot_bitmap bitmap_bytes root,
      decodes_uint bitmap_bytes (Z.of_N slot_bitmap) →
      length bitmap_bytes = N.to_nat 16 →
      blake3_seal_block_bytesR base slot_bitmap bitmap_bytes root
      |--
      base |-> Blake3BlockR 1
        (model.bytes64_of_bytes
           (seal_input_bytes
              (bytes16_of_N slot_bitmap) root)).

  Lemma blake3_seal_nonempty_blockR :
    ∀ (base : ptr) slot_bitmap bitmap_bytes root,
      decodes_uint bitmap_bytes (Z.of_N slot_bitmap) →
      length bitmap_bytes = N.to_nat 16 →
      base |-> arrayLR Tuchar 0 16
        (fun byte : val ⇒ primR Tuchar 1$m byte)
        (map Vint (Z.of_N <$> bitmap_bytes))
      ** base .[ Tuchar ! 16 ] |-> arrayLR Tuchar 0 32
        (fun byte : val ⇒ primR Tuchar 1$m byte)
        (bytes32_byte_values root)
      ** base |-> arrayLR Tuchar 48 64
        (fun byte : val ⇒ primR Tuchar 1$m byte)
        (replicateN 16 (Vint 0))
      |--
      base |-> Blake3BlockR 1
        (model.bytes64_of_bytes
           (seal_input_bytes
              (bytes16_of_N slot_bitmap) (Some root))).

  Lemma blake3_seal_empty_blockR :
    ∀ (base : ptr) slot_bitmap bitmap_bytes,
      decodes_uint bitmap_bytes (Z.of_N slot_bitmap) →
      length bitmap_bytes = N.to_nat 16 →
      base |-> arrayLR Tuchar 0 16
        (fun byte : val ⇒ primR Tuchar 1$m byte)
        (map Vint (Z.of_N <$> bitmap_bytes))
      ** base |-> arrayLR Tuchar 16 64
        (fun byte : val ⇒ primR Tuchar 1$m byte)
        (replicateN 48 (Vint 0))
      |--
      base |-> Blake3BlockR 1
        (model.bytes64_of_bytes
           (seal_input_bytes
              (bytes16_of_N slot_bitmap) None)).
End with_Sigma.