let key_expansion_st (a: algorithm { a = AES_128 \/ a = AES_256 }) =
  (input_key_b:uint8_p) ->
  (output_key_expansion_b:uint8_p) ->
  Stack unit
    (requires fun h0 ->
      B.disjoint input_key_b output_key_expansion_b /\

      B.live h0 input_key_b /\ B.live h0 output_key_expansion_b /\

      B.length input_key_b = key_length a /\
      B.length output_key_expansion_b = key_offset a /\

      aesni_enabled && avx_enabled && sse_enabled)
    (ensures fun h0 _ h1 ->
      B.modifies (B.loc_buffer output_key_expansion_b) h0 h1 /\

      (let key = seq_nat8_to_seq_nat32_LE (seq_uint8_to_seq_nat8 (B.as_seq h0 input_key_b)) in
      Seq.equal (B.as_seq h1 output_key_expansion_b)
         (seq_nat8_to_seq_uint8 (le_seq_quad32_to_bytes (key_to_round_keys_LE a key)))))
let key_expansion_st (a: algorithm{a = AES_128 \/ a = AES_256}) =
input_key_b: uint8_p -> output_key_expansion_b: uint8_p
  -> Stack unit
      (requires
        fun h0 ->
          B.disjoint input_key_b output_key_expansion_b /\ B.live h0 input_key_b /\
          B.live h0 output_key_expansion_b /\ B.length input_key_b = key_length a /\
          B.length output_key_expansion_b = key_offset a /\
          aesni_enabled && avx_enabled && sse_enabled)
      (ensures
        fun h0 _ h1 ->
          B.modifies (B.loc_buffer output_key_expansion_b) h0 h1 /\
          (let key = seq_nat8_to_seq_nat32_LE (seq_uint8_to_seq_nat8 (B.as_seq h0 input_key_b)) in
            Seq.equal (B.as_seq h1 output_key_expansion_b)
              (seq_nat8_to_seq_uint8 (le_seq_quad32_to_bytes (key_to_round_keys_LE a key)))))
Vale.Wrapper.X64.AES.fsti.checked
Vale.X64.CPU_Features_s.fst.checked Vale.Interop.Base.fst.checked Vale.Def.Words_s.fsti.checked Vale.Def.Words.Seq_s.fsti.checked Vale.Def.Types_s.fst.checked Vale.AES.AES_s.fst.checked prims.fst.checked LowStar.BufferView.Up.fsti.checked LowStar.BufferView.Down.fsti.checked LowStar.Buffer.fst.checked FStar.UInt8.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.fsti.checked FStar.Mul.fst.checked FStar.HyperStack.ST.fsti.checked FStar.HyperStack.fst.checked
Vale.Wrapper.X64.AES.fsti
Vale.AES.AES_common_s.algorithm Prims.l_or Prims.b2t Prims.op_Equality Vale.AES.AES_common_s.AES_128 Vale.AES.AES_common_s.AES_256 Vale.Wrapper.X64.AES.uint8_p Prims.unit FStar.Monotonic.HyperStack.mem Prims.l_and LowStar.Monotonic.Buffer.disjoint FStar.UInt8.t LowStar.Buffer.trivial_preorder LowStar.Monotonic.Buffer.live Prims.int LowStar.Monotonic.Buffer.length Vale.Wrapper.X64.AES.key_length Vale.Wrapper.X64.AES.key_offset Prims.op_AmpAmp Vale.X64.CPU_Features_s.aesni_enabled Vale.X64.CPU_Features_s.avx_enabled Vale.X64.CPU_Features_s.sse_enabled LowStar.Monotonic.Buffer.modifies LowStar.Monotonic.Buffer.loc_buffer FStar.Seq.Base.equal LowStar.Monotonic.Buffer.as_seq Vale.Def.Words.Seq_s.seq_nat8_to_seq_uint8 Vale.Def.Types_s.le_seq_quad32_to_bytes Vale.AES.AES_s.key_to_round_keys_LE FStar.Seq.Base.seq Vale.Def.Words_s.nat32 Vale.Def.Words.Seq_s.seq_nat8_to_seq_nat32_LE Vale.Def.Words.Seq_s.seq_uint8_to_seq_nat8
module Vale.Wrapper.X64.AES

open Vale.X64.CPU_Features_s
open FStar.HyperStack.ST
module B = LowStar.Buffer
module HS = FStar.HyperStack
module DV = LowStar.BufferView.Down
module UV = LowStar.BufferView.Up
open FStar.Mul
open Vale.Def.Words_s
open Vale.Def.Words.Seq_s
open Vale.AES.AES_s
open Vale.Interop.Base
open Vale.Def.Types_s

unfold
let uint8_p = B.buffer UInt8.t

let length_aux (b:uint8_p) : Lemma
  (requires B.length b = 176)
  (ensures DV.length (get_downview b) % 16 = 0) =
    let db = get_downview b in
    DV.length_eq db

let length_aux2 (b:uint8_p) : Lemma
  (requires B.length b = 240)
  (ensures DV.length (get_downview b) % 16 = 0) =
    let db = get_downview b in
    DV.length_eq db

inline_for_extraction unfold
let key_offset (a: algorithm { a = AES_128 \/ a = AES_256 }) =
  match a with
  | AES_128 -> 176
  | AES_256 -> 240

inline_for_extraction unfold
let key_length (a: algorithm { a = AES_128 \/ a = AES_256 }) =
  match a with
  | AES_128 -> 16
  | AES_256 -> 32

inline_for_extraction
val key_expansion_st : a:
      Vale.AES.AES_common_s.algorithm
        {a = Vale.AES.AES_common_s.AES_128 \/ a = Vale.AES.AES_common_s.AES_256}
  -> Type0

    a:
      Vale.AES.AES_common_s.algorithm
        {a = Vale.AES.AES_common_s.AES_128 \/ a = Vale.AES.AES_common_s.AES_256}
  -> Type0