include "../../../arch/x64/Vale.X64.InsBasic.vaf"
include "../../../arch/x64/Vale.X64.InsMem.vaf"
include "../../../arch/x64/Vale.X64.InsStack.vaf"
include "../../../arch/x64/Vale.X64.InsVector.vaf"
include "../../../arch/x64/Vale.X64.InsSha.vaf"
include "../../../lib/util/x64/Vale.X64.Stack.vaf"

include{:fstar}{:open} "FStar.Seq.Base"

include{:fstar}{:open} "Vale.Def.Words_s"
include{:fstar}{:open} "Vale.Def.Types_s"
include{:fstar}{:open} "Vale.Def.Words.Seq_s"
include{:fstar}{:open} "Vale.Arch.Types"
include{:fstar}{:open} "Spec.SHA2"
include{:fstar}{:open} "Vale.SHA.SHA_helpers"
include{:fstar}{:open} "Spec.Agile.Hash"
include{:fstar}{:open} "Spec.Hash.PadFinish"
include{:fstar}{:open} "Spec.Hash.Definitions"
include{:fstar}{:open} "Vale.X64.Machine_s"
include{:fstar}{:open} "Vale.X64.State"
include{:fstar}{:open} "Vale.X64.Decls"
include{:fstar}{:open} "Vale.X64.QuickCode"
include{:fstar}{:open} "Spec.Loops"
include{:fstar}{:open} "Vale.X64.CPU_Features_s"

module Vale.SHA.X64

#verbatim{:interface}{:implementation}
open Vale.Def.Opaque_s
open Vale.Def.Types_s
open Vale.Def.Words_s
open Vale.Def.Words.Seq_s
open FStar.Seq
open Vale.Arch.Types
open Vale.Arch.HeapImpl
open Vale.X64.Machine_s
open Vale.X64.Memory
open Vale.X64.Stack_i
open Vale.X64.State
open Vale.X64.Decls
open Vale.X64.InsBasic
open Vale.X64.InsMem
open Vale.X64.InsStack
open Vale.X64.InsVector
open Vale.X64.InsSha
open Vale.X64.QuickCode
open Vale.X64.QuickCodes
open Vale.SHA.SHA_helpers
open Spec.SHA2
open Spec.Agile.Hash
open Spec.Hash.PadFinish
open Spec.Hash.Definitions
open Spec.Loops
open Vale.X64.Stack
open Vale.X64.CPU_Features_s
#reset-options "--z3rlimit 40"
#endverbatim

procedure Preamble(ghost ctx_b:buffer128)
    {:quick}
    lets
        ctx @= rdi; tmp_reg @= rax;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        tmp_xmm @= xmm7; bswap @= xmm8;
    reads
        ctx; heap0; memLayout;
    modifies
        tmp_reg;
        Wi; abef; cdgh; tmp_xmm; bswap;
        efl;
    requires
        sse_enabled;
        validSrcAddrs128(heap0, ctx,  ctx_b,  2, memLayout, Secret);
    ensures
        // Why is this stored twice?
        tmp_xmm == bswap;
        bswap == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
        let abcd := buffer128_read(ctx_b, 0, heap0) in
        let efgh := buffer128_read(ctx_b, 1, heap0) in
        abef == Mkfour(efgh.lo1, efgh.lo0, abcd.lo1, abcd.lo0) /\ // LSB: FEBA
        cdgh == Mkfour(efgh.hi3, efgh.hi2, abcd.hi3, abcd.hi2) /\   // LSB: HGDC
        make_hash(abef, cdgh) == make_ordered_hash(abcd, efgh);
{
    // At the C level, ctx_b is an array of 8 32-bit values: ABCDEFGH
    // We load these out of memory using litte-endian loads,
    // but SHA spec is big endian, so we need to do some swaps
    let abcd := buffer128_read(ctx_b, 0, heap0);
    let efgh := buffer128_read(ctx_b, 1, heap0);
    Load128_buffer(heap0, abef, ctx,  0, Secret, ctx_b, 0);  // abef := LSB: ABCD
    Load128_buffer(heap0, cdgh, ctx, 16, Secret, ctx_b, 1);  // cdgh := LSB: EFGH
    InitPshufbStableMask(tmp_xmm, tmp_reg);

    Pshufd(Wi, abef, 0x1b);                           //   Wi := LSB: DCBA
    Pshufd(abef, abef, 0xb1);                         // abef := LSB: BADC
    Pshufd(cdgh, cdgh, 0x1b);                         // cdgh := LSB: HGFE

    Mov128(bswap, tmp_xmm);     // OpenSSL: offload   // bswap holds mask (why go via tmp?)
    Palignr8(abef, cdgh);                             // abef := LSB: FEBA == abef in big endian
    // OpenSSL uses punpcklqdq here:
    Shufpd(cdgh, Wi, 0);                              // cdgh := LSB: HGDC == cdgh in big endian
    assert equal(#(seq(Vale.SHA.SHA_helpers.word))(make_hash(abef, cdgh)), #(seq(Vale.SHA.SHA_helpers.word))(make_ordered_hash(abcd, efgh)));
}

procedure Loop_rounds_0_15(
        ghost in_b:buffer128,
        ghost k_b:buffer128,
        ghost offset:nat)
    {:quick}
    lets
        inp @= rsi; tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        tbl; heap0; memLayout;
    modifies
        inp;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm; abef_save; cdgh_save;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrsOffset128(heap0, inp,   in_b, offset, 4, memLayout, Secret);
        validSrcAddrs128(heap0, tbl,    k_b, 16, memLayout, Secret);
        inp + 0x40 < pow2_64;

        tmp_xmm == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);

        k_reqs(buffer128_as_seq(heap0, k_b));
    ensures
        inp == old(inp) + 0x40;
        let input_LE := slice(buffer128_as_seq(heap0, in_b), offset, offset+4) in
        let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in
        let block:block_w := quads_to_block(input_BE) in
        // Wi == Don't care
        make_hash(abef, cdgh) == repeat_range_vale(16, block, old(make_hash(abef, cdgh))) /\

        msg0 == ws_quad32(16, block) /\
        msg1 == add_wrap_quad32(ws_partial(20, block), ws_quad32(13, block)) /\
        msg2 == ws_partial(24, block) /\
        msg3 == ws_quad32(12, block);

        // tmp_xmm == Don't care
        abef_save == old(abef);
        cdgh_save == old(cdgh);

{
    // Create a ghost message block to pass to all of the sha256* instructions and lemmas
    let input_LE := slice(buffer128_as_seq(heap0, in_b), offset, offset+4);
    let input_BE := reverse_bytes_nat32_quad32_seq(input_LE);
    let block:block_w := quads_to_block(input_BE);
    lemma_quads_to_block(input_BE);

    // Prove that the original hash input starts the repeat_range_spec base case
    let hash_orig := make_hash(abef, cdgh);
    lemma_repeat_range_0_vale(block, hash_orig);

    let ks := buffer128_as_seq(heap0, k_b);

    // Load 4 128-bit message chunks == 16 32-bit chunks
    Load128_buffer(heap0, msg0, inp,  0, Secret, in_b, offset+0);
    Load128_buffer(heap0, msg1, inp, 16, Secret, in_b, offset+1);
    Load128_buffer(heap0, msg2, inp, 32, Secret, in_b, offset+2);
    PshufbStable(msg0, tmp_xmm);       // Convert msg0 to big endian
    assert msg0 == index(input_BE, 0);   // OBSERVE: TODO: Better trigger for lemma_quads_to_block?
    Load128_buffer(heap0, msg3, inp, 48, Secret, in_b, offset+3);

    Load128_buffer(heap0, Wi, tbl, 0, Secret, k_b, 0);  // Load K values from memory
    assert Wi == index(ks, 0);
    Paddd(Wi, msg0);            // Combine the round constant with the message block

    PshufbStable(msg1, tmp_xmm);      // Convert msg1 to big endian
    assert msg1 == index(input_BE, 1);   // OBSERVE: TODO: Better trigger for lemma_quads_to_block?
    Mov128(cdgh_save, cdgh);    // Save a copy of the original hash value to add at the end

    // Do two rounds of SHA, drawing on state in ABEF, CDGH, and WK in lower64 of XMM0
    SHA256_rnds2(cdgh, abef, 0, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);       // Move upper-64 into lower-64 of Wi, lining up the next two WK values
    Mov128(abef_save, abef);    // Save a copy of the original hash value to add at the end
    SHA256_rnds2(abef, cdgh, 2, block, hash_orig);

    // Repeat above, except we also do some message expansion
    Load128_buffer(heap0, Wi, tbl, 16, Secret, k_b, 1);  // Load K values from memory
    assert Wi == index(ks, 1);
    Paddd(Wi, msg1);
    PshufbStable(msg2, tmp_xmm);
    assert msg2 == index(input_BE, 2);   // OBSERVE: TODO: Better trigger for lemma_quads_to_block?
    SHA256_rnds2(cdgh, abef, 4, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Add64(inp, 0x40);
    SHA256_msg1(msg0, msg1, 16, block);     // Do a partial step of message expansion
    SHA256_rnds2(abef, cdgh, 6, block, hash_orig);

    // Repeat again
    Load128_buffer(heap0, Wi, tbl, 32, Secret, k_b, 2);  // Load K values from memory
    assert Wi == index(ks, 2);
    Paddd(Wi, msg2);
    PshufbStable(msg3, tmp_xmm);
    assert msg3 == index(input_BE, 3);   // OBSERVE: TODO: Better trigger for lemma_quads_to_block?
    SHA256_rnds2(cdgh, abef, 8, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Mov128(tmp_xmm, msg3);                  // We don't need the mask any more, but why save msg3?
    Palignr4(tmp_xmm, msg2);
    Paddd(msg0, tmp_xmm);
    SHA256_msg1(msg1, msg2, 20, block);     // Do another partial step of message expansion
    SHA256_rnds2(abef, cdgh, 10, block, hash_orig);

    // Repeat a final time
    Load128_buffer(heap0, Wi, tbl, 48, Secret, k_b, 3);  // Load K values from memory
    assert Wi == index(ks, 3);
    Paddd(Wi, msg3);
    SHA256_msg2(msg0, msg3, 16, block);      // Finalize the message expansion
    SHA256_rnds2(cdgh, abef, 12, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Mov128(tmp_xmm, msg0);                  // Why save msg0?
    Palignr4(tmp_xmm, msg3);
    Paddd(msg1, tmp_xmm);
    SHA256_msg1(msg2, msg3, 24, block);     // Do another partial step of message expansion
    SHA256_rnds2(abef, cdgh, 14, block, hash_orig);
}

procedure Loop_rounds_16_51_body(
        inline i:nat,
        ghost k_b:buffer128,
        ghost block:block_w,
        ghost hash_orig:hash256)
    {:quick}
    lets
        tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7;
    reads
        tbl; heap0; memLayout;
    modifies
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, tbl, k_b, 16, memLayout, Secret);

        4 <= i /\ i < 13;

        k_reqs(buffer128_as_seq(heap0, k_b));

        make_hash(abef, cdgh) == repeat_range_vale(4*i, block, hash_orig);

        msg0 == ws_quad32(4*i, block) /\
        msg1 == add_wrap_quad32(ws_partial(4*(i+1), block), ws_quad32(4*(#(nat)(i-1))+1, block)) /\
        msg2 == ws_partial(4*(i+2), block) /\
        msg3 == ws_quad32(4*(#(nat)(i-1)), block);

    ensures
        make_hash(abef, cdgh) == repeat_range_vale(4*(i+1), block, hash_orig);

        msg1 == ws_quad32(4*(i+1), block);
        msg2 == add_wrap_quad32(ws_partial(4*(i+2), block), ws_quad32(4*(i)+1, block));
        msg3 == ws_partial(4*(i+3), block);
        msg0 == ws_quad32(4*(i), block);
{
    let ks := buffer128_as_seq(heap0, k_b);
    Load128_buffer(heap0, Wi, tbl, 16*i, Secret, k_b, i);  // Load K values from memory
    assert Wi == index(ks, i);

    Paddd(Wi, msg0);
    SHA256_msg2(msg1, msg0, 4*(i+1), block);
    SHA256_rnds2(cdgh, abef, 4*i, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Mov128(tmp_xmm, msg1);
    Palignr4(tmp_xmm, msg0);
    Paddd(msg2, tmp_xmm);
    SHA256_msg1(msg3, msg0, 4*(i + 3), block);
    SHA256_rnds2(abef, cdgh, 4*i + 2, block, hash_orig);
}

// TODO: OpenSSL does this without any moves by renaming variables.
//       We can do the same once Vale quick-mode supports that,
//       or we can rewrite the procedure above in non-quick mode
procedure Msg_shift()
    {:quick}
    lets
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        Wi @= xmm0; tmp_xmm @= xmm7;
    modifies
        msg0; msg1; msg2; msg3; tmp_xmm; Wi;
    requires sse_enabled;
    ensures
        msg0 == old(msg1);
        msg1 == old(msg2);
        msg2 == old(msg3);
        msg3 == old(msg0);
{
    Mov128(tmp_xmm, msg0);
    Mov128(Wi, msg3);
    Mov128(msg0, msg1);
    Mov128(msg1, msg2);
    Mov128(msg3, tmp_xmm);
    Mov128(msg2, Wi);
}

// REVIEW: Vale wouldn't let me make this an exportOnly, so I duplicated it below :(
procedure Loop_rounds_16_51_recursive(
        inline i:nat,
        ghost k_b:buffer128,
        ghost block:block_w,
        ghost hash_orig:hash256)
    {:recursive}
    lets
        tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7;
    reads
        tbl; heap0; memLayout;
    modifies
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, tbl, k_b, 16, memLayout, Secret);

        4 <= i /\ i < 13;

        k_reqs(buffer128_as_seq(heap0, k_b));

        make_hash(abef, cdgh) == repeat_range_vale(4*4, block, hash_orig);

        msg0 == ws_quad32(4*4, block) /\
        msg1 == add_wrap_quad32(ws_partial(4*(4+1), block), ws_quad32(4*(4-1)+1, block)) /\
        msg2 == ws_partial(4*(4+2), block) /\
        msg3 == ws_quad32(4*(4-1), block);

    ensures
        make_hash(abef, cdgh) == repeat_range_vale(4*(i+1), block, hash_orig);

        msg0 == ws_quad32(4*(i+1), block);
        msg1 == add_wrap_quad32(ws_partial(4*(i+2), block), ws_quad32(4*(i)+1, block));
        msg2 == ws_partial(4*(i+3), block);
        msg3 == ws_quad32(4*(i), block);
{
    inline if (i > 4) {
        Loop_rounds_16_51_recursive((#(nat)(i-1)), k_b, block, hash_orig);
    }
    Loop_rounds_16_51_body(i, k_b, block, hash_orig);
    Msg_shift();
}

procedure Loop_rounds_16_51(
        ghost k_b:buffer128,
        ghost block:block_w,
        ghost hash_orig:hash256)
    {:quick exportOnly}
    lets
        tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7;
    reads
        tbl; heap0; memLayout;
    modifies
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, tbl, k_b, 16, memLayout, Secret);

        k_reqs(buffer128_as_seq(heap0, k_b));

        make_hash(abef, cdgh) == repeat_range_vale(4*4, block, hash_orig);

        msg0 == ws_quad32(4*4, block) /\
        msg1 == add_wrap_quad32(ws_partial(4*(4+1), block), ws_quad32(4*(4-1)+1, block)) /\
        msg2 == ws_partial(4*(4+2), block) /\
        msg3 == ws_quad32(4*(4-1), block);

    ensures
        let i := 12;
        make_hash(abef, cdgh) == repeat_range_vale(4*(i+1), block, hash_orig);

        msg0 == ws_quad32(4*(i+1), block);
        msg1 == add_wrap_quad32(ws_partial(4*(i+2), block), ws_quad32(4*(i)+1, block));
        msg2 == ws_partial(4*(i+3), block);
        msg3 == ws_quad32(4*(i), block);
{
    Loop_rounds_16_51_recursive(12, k_b, block, hash_orig);
}

procedure Loop_rounds_52_64(
        ghost k_b:buffer128,
        ghost block:block_w)
    {:quick}
    lets
        tbl @= rcx; num @= rdx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; bswap @= xmm8; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        tbl; bswap; abef_save; cdgh_save; heap0; memLayout;
    modifies
        num;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, tbl, k_b, 16, memLayout, Secret);

        num > 0;

        k_reqs(buffer128_as_seq(heap0, k_b));

        let hash_orig := make_hash(abef_save, cdgh_save) in
        make_hash(abef, cdgh) == repeat_range_vale(52, block, hash_orig);

        msg0 == ws_quad32(52, block);
        msg1 == add_wrap_quad32(ws_partial(56, block), ws_quad32(49, block));
        msg2 == ws_partial(60, block);
        msg3 == ws_quad32(48, block);
    ensures
        tmp_xmm == old(bswap);
        num == old(num) - 1;
        let hash_orig := make_hash(abef_save, cdgh_save) in
        make_hash(abef, cdgh) == update_block(hash_orig, block);
{
    let hash_orig := make_hash(abef_save, cdgh_save);

    let ks := buffer128_as_seq(heap0, k_b);
    Load128_buffer(heap0, Wi, tbl, 16*13, Secret, k_b, 13);  // Load K values from memory
    assert Wi == index(ks, 13);
    Paddd(Wi, msg0);
    SHA256_msg2(msg1, msg0, 4*(13+1), block);
    SHA256_rnds2(cdgh, abef, 4*13, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Mov128(tmp_xmm, msg1);
    Palignr4(tmp_xmm, msg0);
    SHA256_rnds2(abef, cdgh, 4*13+2, block, hash_orig);
    Paddd(msg2, tmp_xmm);

    Load128_buffer(heap0, Wi, tbl, 16*14, Secret, k_b, 14);  // Load K values from memory
    assert Wi == index(ks, 14);
    Paddd(Wi, msg1);
    SHA256_rnds2(cdgh, abef, 4*14, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    SHA256_msg2(msg2, msg1, 4*(14+1), block);
    Mov128(tmp_xmm, bswap);
    SHA256_rnds2(abef, cdgh, 4*14+2, block, hash_orig);

    Load128_buffer(heap0, Wi, tbl, 16*15, Secret, k_b, 15);  // Load K values from memory
    assert Wi == index(ks, 15);
    Paddd(Wi, msg2);
    SHA256_rnds2(cdgh, abef, 4*15, block, hash_orig);
    Pshufd(Wi, Wi, 0x0e);
    Sub64(num, 1);
    SHA256_rnds2(abef, cdgh, 4*15+2, block, hash_orig);
    //assert make_hash(abef, cdgh) == shuffle(SHA2_256, hash_orig, block);

    let abef_shuffle := abef;
    let cdgh_shuffle := cdgh;

    Paddd(cdgh, cdgh_save);
    Paddd(abef, abef_save);

    update_lemma(abef_shuffle, cdgh_shuffle, abef, cdgh, abef_save, cdgh_save, block);
}


procedure Loop_rounds(
        ghost in_b:buffer128,
        ghost k_b:buffer128,
        ghost offset:nat,
        ghost hash_orig:hash256)
    {:quick}
    lets
        num @= rdx; inp @= rsi; tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; bswap @= xmm8; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        tbl; bswap; heap0; memLayout;
    modifies
        inp; num;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm; abef_save; cdgh_save;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrsOffset128(heap0, inp,   in_b, offset, 4, memLayout, Secret);
        validSrcAddrs128(heap0, tbl,    k_b, 16, memLayout, Secret);
        inp + 0x40 < pow2_64;

        num > 0;

        tmp_xmm == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
        bswap == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);

        k_reqs(buffer128_as_seq(heap0, k_b));

        let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, offset) in
        let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in
        make_hash(abef, cdgh) == update_multi_quads(input_BE, hash_orig);
    ensures
        inp == old(inp) + 0x40;
        num == old(num) - 1;
        tmp_xmm == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
        bswap == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);

        let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, offset + 4) in
        let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in
        make_hash(abef, cdgh) == update_multi_quads(input_BE, hash_orig);

{
    // Create a ghost message block to pass to all of the sha256* instructions and lemmas
    let input_LE := slice(buffer128_as_seq(heap0, in_b), offset, offset+4);
    let input_BE := reverse_bytes_nat32_quad32_seq(input_LE);
    let block:block_w := quads_to_block(input_BE);

    let hash_init := make_hash(abef, cdgh);

    Loop_rounds_0_15(in_b, k_b, offset);
    Loop_rounds_16_51(k_b, block, hash_init);
    Loop_rounds_52_64(k_b, block);

    lemma_update_multi_quads(buffer128_as_seq(heap0, in_b), hash_orig, offset);
}

procedure Loop(
        ghost in_b:buffer128,
        ghost k_b:buffer128)
    {:quick}
    lets
        num @= rdx; inp @= rsi; tbl @= rcx;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; bswap @= xmm8; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        tbl; bswap; heap0; memLayout;
    modifies
        inp; num;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm; abef_save; cdgh_save;
        efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, inp,   in_b,  4*num, memLayout, Secret);
        validSrcAddrs128(heap0, tbl,    k_b, 16, memLayout, Secret);
        inp + 0x40*num < pow2_64;

        tmp_xmm == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
        bswap == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);

        k_reqs(buffer128_as_seq(heap0, k_b));
    ensures
        inp == old(inp) + 0x40*old(num);
        num == 0;
        let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, 4*old(num)) in
        let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in
        make_hash(abef, cdgh) == update_multi_quads(input_BE, old(make_hash(abef, cdgh)));
{
    let hash_orig:hash256 := make_hash(abef, cdgh);

    ghost var count:nat := 0;
    while (num > 0)
        invariant
            sha_enabled && sse_enabled;

            validSrcAddrs128(heap0, old(inp),   in_b,  4*old(num), memLayout, Secret);
            validSrcAddrs128(heap0,      tbl,    k_b,          16, memLayout, Secret);

            old(inp) + 0x40*old(num) < pow2_64;

            num == old(num) - count;
            inp == old(inp) + 0x40 * count;

            tmp_xmm == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
            bswap == Mkfour(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);

            k_reqs(buffer128_as_seq(heap0, k_b));

            let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, count*4) in
            let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in
            make_hash(abef, cdgh) == update_multi_quads(input_BE, hash_orig);
        decreases num;
    {
        Loop_rounds(in_b, k_b, count*4, hash_orig);
        count := count + 1;
    }
}

procedure Epilogue(
        ghost ctx_b:buffer128)
    {:quick}
    lets
        ctx @= rdi;
        abef @= xmm1; cdgh @= xmm2;
        tmp_xmm @= xmm7;
    reads
        ctx; memLayout;
    modifies
        abef; cdgh; tmp_xmm;
        heap0; efl;
    requires
        sse_enabled;
        validDstAddrs128(heap0, ctx,  ctx_b,  2, memLayout, Secret);
    ensures
        let abcd := old(Mkfour(abef.hi3, abef.hi2, cdgh.hi3, cdgh.hi2)) in
        let efgh := old(Mkfour(abef.lo1, abef.lo0, cdgh.lo1, cdgh.lo0)) in
        abcd == buffer128_read(ctx_b, 0, heap0) /\
        efgh == buffer128_read(ctx_b, 1, heap0) /\
        make_hash(old(abef), old(cdgh)) == make_ordered_hash(abcd, efgh);

        // Framing
        modifies_buffer128(ctx_b, old(heap0), heap0);
{
    // Initially:
    // abef == LSB: FEBA
    // cdgh == LSB: HGDC

    Pshufd(cdgh, cdgh, 0xb1);       // cdgh := LSB: GHCD
    Pshufd(tmp_xmm, abef, 0x1b);    // tmp_xmm := LSB: ABEF
    Pshufd(abef, abef, 0xb1);       // abef := LSB: EFAB
    // OpenSSL uses punpcklqdq here:
    Shufpd(abef, cdgh, 3);          // abef := LSB: ABCD
    Palignr8(cdgh, tmp_xmm);        // cdgh := LSB: EFGH

    Store128_buffer(heap0, ctx, abef, 0, Secret, ctx_b, 0);
    Store128_buffer(heap0, ctx, cdgh, 16, Secret, ctx_b, 1);
    let abcd := buffer128_read(ctx_b, 0, heap0);
    let efgh := buffer128_read(ctx_b, 1, heap0);
    assert equal(#(seq(Vale.SHA.SHA_helpers.word))(make_hash(old(abef), old(cdgh))), #(seq(Vale.SHA.SHA_helpers.word))(make_ordered_hash(abcd, efgh)));
}


procedure Sha_update(
        ghost ctx_b:buffer128,
        ghost in_b:buffer128,
        ghost k_b:buffer128)
    {:quick}
    lets
        tmp_reg @= rax; tbl @= rcx; num @= rdx; ctx @= rdi; inp @= rsi;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; bswap @= xmm8; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        ctx; tbl; memLayout;
    modifies
        inp; num; rax;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm; bswap; abef_save; cdgh_save;
        heap0; efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, inp,   in_b,  4*num, memLayout, Secret);
        validDstAddrs128(heap0, ctx,  ctx_b,  2, memLayout, Secret);
        validSrcAddrs128(heap0, tbl,    k_b, 16, memLayout, Secret);
        inp + 0x40*num < pow2_64;
        buffers_disjoint128(ctx_b, in_b);

        k_reqs(buffer128_as_seq(heap0, k_b));
    ensures
        inp == old(inp) + 0x40*old(num);
        let abcd  := old(buffer128_read(ctx_b, 0, heap0)) in
        let efgh  := old(buffer128_read(ctx_b, 1, heap0)) in
        let abcd' :=     buffer128_read(ctx_b, 0, heap0)  in
        let efgh' :=     buffer128_read(ctx_b, 1, heap0)  in

        let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, 4*old(num)) in
        let input_BE := reverse_bytes_nat32_quad32_seq(input_LE) in

        make_ordered_hash(abcd', efgh') == update_multi_quads(input_BE, make_ordered_hash(abcd, efgh));

        // Framing
        modifies_buffer128(ctx_b, old(heap0), heap0);
{
    Preamble(ctx_b);
    Loop(in_b, k_b);
    Epilogue(ctx_b);
}

procedure Sha_update_bytes(
        ghost ctx_b:buffer128,
        ghost in_b:buffer128,
        ghost k_b:buffer128)
    {:quick}
    {:restartProver}
    {:options z3rlimit(80)}
    lets
        tmp_reg @= rax; tbl @= rcx; num @= rdx; ctx @= rdi; inp @= rsi;
        Wi @= xmm0; abef @= xmm1; cdgh @= xmm2;
        msg0 @= xmm3; msg1 @= xmm4; msg2 @= xmm5; msg3 @= xmm6;
        tmp_xmm @= xmm7; bswap @= xmm8; abef_save @= xmm9; cdgh_save @= xmm10;
    reads
        ctx; tbl; memLayout;
    modifies
        inp; num; rax;
        Wi; abef; cdgh; msg0; msg1; msg2; msg3; tmp_xmm; bswap; abef_save; cdgh_save;
        heap0; efl;
    requires
        sha_enabled && sse_enabled;
        validSrcAddrs128(heap0, inp,   in_b,  4*num, memLayout, Secret);
        validDstAddrs128(heap0, ctx,  ctx_b,  2, memLayout, Secret);
        validSrcAddrs128(heap0, tbl,    k_b, 16, memLayout, Secret);
        inp + 0x40*num < pow2_64;
        buffers_disjoint128(ctx_b, in_b);
        buffer_length(ctx_b) == 2;
        buffer_length(in_b) == 4*num;

        k_reqs(buffer128_as_seq(heap0, k_b));
    ensures
        inp == old(inp) + 0x40*old(num);

        let hash_in  := old(le_bytes_to_hash(le_seq_quad32_to_bytes(buffer128_as_seq(heap0, ctx_b))));
        let hash_out :=     le_bytes_to_hash(le_seq_quad32_to_bytes(buffer128_as_seq(heap0, ctx_b)));

        let input_LE := seq_nat8_to_seq_uint8(le_seq_quad32_to_bytes(buffer128_as_seq(heap0, in_b))) in
        length(input_LE) % 64 == 0 /\
        hash_out == update_multi_opaque_vale(hash_in, input_LE);

        // Framing
        modifies_buffer128(ctx_b, old(heap0), heap0);
{
    Sha_update(ctx_b, in_b, k_b);
    let old_ctx:seq(four(nat32)) := old(buffer128_as_seq(heap0, ctx_b));
    let new_ctx:seq(four(nat32)) :=     buffer128_as_seq(heap0, ctx_b);
    lemma_hash_to_bytes(old_ctx);
    lemma_hash_to_bytes(new_ctx);
    let hash_in  := le_bytes_to_hash(le_seq_quad32_to_bytes(old_ctx));
    let hash_out := le_bytes_to_hash(le_seq_quad32_to_bytes(new_ctx));

    // Lots of OBSERVE below
//    assert old(buffer128_read(ctx_b, 0, heap0)) == index(old_ctx, 0);
//    assert old(buffer128_read(ctx_b, 1, heap0)) == index(old_ctx, 1);
//    assert     buffer128_read(ctx_b, 0, heap0)  == index(new_ctx, 0);
//    assert     buffer128_read(ctx_b, 1, heap0)  == index(new_ctx, 1);

    let input_LE := slice(buffer128_as_seq(heap0, in_b), 0, 4*old(num));
    let input_BE := reverse_bytes_nat32_quad32_seq(input_LE);
    assert hash_out == update_multi_quads(input_BE, hash_in);
    lemma_update_multi_equiv_vale(hash_in, hash_out, input_LE, input_BE, le_seq_quad32_to_bytes(input_LE),
                                  seq_nat8_to_seq_uint8(le_seq_quad32_to_bytes(input_LE)));
}

procedure Sha_update_bytes_stdcall(
        inline win:bool,
        ghost ctx_b:buffer128,
        ghost in_b:buffer128,
        ghost num_val:nat64,
        ghost k_b:buffer128)
    {:public}
    {:quick}
    {:exportSpecs}
    {:restartProver}
    {:options z3rlimit(80)}
    lets
        ctx_ptr     := if win then rcx else rdi;
        in_ptr      := if win then rdx else rsi;
        num         := if win then r8  else rdx;
        k_ptr       := if win then r9  else rcx;
    modifies
        rax; rbx; rcx; rdx; rsi; rdi; rbp; rsp; r8; r9; r10; r11; r12; r13; r14; r15;
        xmm0; xmm1; xmm2; xmm3; xmm4; xmm5; xmm6; xmm7; xmm8; xmm9; xmm10; xmm11; xmm12; xmm13; xmm14; xmm15;
        efl; heap0; memLayout; stack; stackTaint;
    requires
        rsp == init_rsp(stack);
        is_initial_heap(memLayout, mem);

        sha_enabled && sse_enabled;
        locs_disjoint(list(loc_buffer(ctx_b), loc_buffer(in_b))) \/ ctx_b == in_b;
        locs_disjoint(list(loc_buffer(ctx_b), loc_buffer(k_b))) \/ ctx_b == k_b;
        locs_disjoint(list(loc_buffer(in_b), loc_buffer(k_b))) \/ in_b == k_b;
        validDstAddrs128(mem, ctx_ptr, ctx_b, 2, memLayout, Secret);
        validSrcAddrs128(mem, in_ptr, in_b, 4 * num, memLayout, Secret);
        validSrcAddrs128(mem, k_ptr, k_b, 16, memLayout, Secret);
        num_val == num;

        in_ptr + 0x40 * num < pow2_64;
        buffers_disjoint128(ctx_b, in_b);
        buffer_length(ctx_b) == 2;
        buffer_length(in_b) == 4 * num;

        k_reqs(buffer128_as_seq(mem, k_b));

    ensures
        let hash_in  := old(le_bytes_to_hash(le_seq_quad32_to_bytes(buffer128_as_seq(mem, ctx_b))));
        let hash_out :=     le_bytes_to_hash(le_seq_quad32_to_bytes(buffer128_as_seq(mem, ctx_b)));

        let input_LE := seq_nat8_to_seq_uint8(le_seq_quad32_to_bytes(buffer128_as_seq(mem, in_b))) in
        length(input_LE) % 64 == 0 /\
        hash_out == update_multi_transparent(hash_in, #(Vale.SHA.SHA_helpers.bytes_blocks)(input_LE));

        // Framing
        modifies_mem(loc_buffer(ctx_b), old(mem), mem);

        rsp == old(rsp);

        win ==>  rbx == old(rbx);
        win ==>  rbp == old(rbp);
        win ==>  rdi == old(rdi);
        win ==>  rsi == old(rsi);
        win ==>  r12 == old(r12);
        win ==>  r13 == old(r13);
        win ==>  r14 == old(r14);
        win ==>  r15 == old(r15);
        !win ==>  rbx == old(rbx);
        !win ==>  rbp == old(rbp);
        !win ==>  r12 == old(r12);
        !win ==>  r13 == old(r13);
        !win ==>  r14 == old(r14);
        !win ==>  r15 == old(r15);
        win ==>  xmm6 == old(xmm6);
        win ==>  xmm7 == old(xmm7);
        win ==>  xmm8 == old(xmm8);
        win ==>  xmm9 == old(xmm9);
        win ==>  xmm10 == old(xmm10);
        win ==>  xmm11 == old(xmm11);
        win ==>  xmm12 == old(xmm12);
        win ==>  xmm13 == old(xmm13);
        win ==>  xmm14 == old(xmm14);
        win ==>  xmm15 == old(xmm15);
{
    CreateHeaplets(list(
        declare_buffer128(in_b, 0, Secret, Immutable),
        declare_buffer128(k_b, 0, Secret, Immutable),
        declare_buffer128(ctx_b, 0, Secret, Mutable)));

    // Save registers
    inline if (win)
    {
        PushXmm_Secret(xmm15, rax);
        PushXmm_Secret(xmm14, rax);
        PushXmm_Secret(xmm13, rax);
        PushXmm_Secret(xmm12, rax);
        PushXmm_Secret(xmm11, rax);
        PushXmm_Secret(xmm10, rax);
        PushXmm_Secret(xmm9,  rax);
        PushXmm_Secret(xmm8,  rax);
        PushXmm_Secret(xmm7,  rax);
        PushXmm_Secret(xmm6,  rax);
        Push_Secret(r15);
        Push_Secret(r14);
        Push_Secret(r13);
        Push_Secret(r12);
        Push_Secret(rsi);
        Push_Secret(rdi);
        Push_Secret(rbp);
        Push_Secret(rbx);
    }
    else
    {
        Push_Secret(r15);
        Push_Secret(r14);
        Push_Secret(r13);
        Push_Secret(r12);
        Push_Secret(rsi);
        Push_Secret(rdi);
        Push_Secret(rbp);
        Push_Secret(rbx);
    }

    inline if (win)
    {
        Mov64(rdi, rcx);
        Mov64(rsi, rdx);
        Mov64(rdx, r8);
        Mov64(rcx, r9);
    }

    Sha_update_bytes(ctx_b, in_b, k_b);

    // Restore registers
    inline if (win) {
        Pop_Secret(rbx);
        Pop_Secret(rbp);
        Pop_Secret(rdi);
        Pop_Secret(rsi);
        Pop_Secret(r12);
        Pop_Secret(r13);
        Pop_Secret(r14);
        Pop_Secret(r15);

        PopXmm_Secret(xmm6,  rax, old(xmm6));
        PopXmm_Secret(xmm7,  rax, old(xmm7));
        PopXmm_Secret(xmm8,  rax, old(xmm8));
        PopXmm_Secret(xmm9,  rax, old(xmm9));
        PopXmm_Secret(xmm10, rax, old(xmm10));
        PopXmm_Secret(xmm11, rax, old(xmm11));
        PopXmm_Secret(xmm12, rax, old(xmm12));
        PopXmm_Secret(xmm13, rax, old(xmm13));
        PopXmm_Secret(xmm14, rax, old(xmm14));
        PopXmm_Secret(xmm15, rax, old(xmm15));
    }
    else
    {
        Pop_Secret(rbx);
        Pop_Secret(rbp);
        Pop_Secret(rdi);
        Pop_Secret(rsi);
        Pop_Secret(r12);
        Pop_Secret(r13);
        Pop_Secret(r14);
        Pop_Secret(r15);
    }

    let hash_in  := old(le_bytes_to_hash(le_seq_quad32_to_bytes(buffer128_as_seq(heap0, ctx_b))));
    let input_LE := seq_nat8_to_seq_uint8(le_seq_quad32_to_bytes(buffer128_as_seq(heap0, in_b)));
    lemma_update_multi_opaque_vale_is_update_multi(hash_in, input_LE);

    DestroyHeaplets();
}