Require Export Fiat.Common.Coq__8_4__8_5__Compat.
Require Import
        Coq.Strings.String
        Coq.Vectors.Vector
        Coq.ZArith.ZArith.

Require Import
        Fiat.Common.SumType
        Fiat.Common.EnumType
        Fiat.Common.BoundedLookup
        Fiat.Common.ilist
        Fiat.Computation
        Fiat.QueryStructure.Specification.Representation.Notations
        Fiat.QueryStructure.Specification.Representation.Heading
        Fiat.QueryStructure.Specification.Representation.Tuple
        Fiat.Narcissus.BinLib.AlignedByteString
        Fiat.Narcissus.BinLib.AlignWord
        Fiat.Narcissus.BinLib.AlignedList
        Fiat.Narcissus.BinLib.AlignedDecoders
        Fiat.Narcissus.BinLib.AlignedDecodeMonad
        Fiat.Narcissus.BinLib.AlignedEncodeMonad
        Fiat.Narcissus.Common.Specs
        Fiat.Narcissus.Common.WordFacts
        Fiat.Narcissus.Common.ComposeCheckSum
        Fiat.Narcissus.Common.ComposeIf
        Fiat.Narcissus.Common.ComposeOpt
        Fiat.Narcissus.Formats
        Fiat.Narcissus.BaseFormats
        Fiat.Narcissus.Stores.EmptyStore.

Require Import Bedrock.Word.

Definition decode_IPChecksum
  : ByteString -> CacheDecode -> option (() * ByteString * CacheDecode) :=
  decode_unused_word (sz := 16).

Definition encode_word {sz} (w : word sz) : ByteString :=
  encode_word' sz w ByteString_id.

Fixpoint Vector_checksum_bound n {sz} (bytes :ByteBuffer.t sz) acc : InternetChecksum.W16 :=
  match n, bytes with
  | 0, _ => acc
  | _, Vector.nil => acc
  | S 0, Vector.cons x _ _ => InternetChecksum.add_bytes_into_checksum x (wzero _) acc
  | _, Vector.cons x _ Vector.nil => InternetChecksum.add_bytes_into_checksum x (wzero _) acc
  | S (S n'), Vector.cons x _ (Vector.cons y _ t) =>
    (Vector_checksum_bound n' t (InternetChecksum.add_bytes_into_checksum x y acc))
  end.

Definition ByteBuffer_checksum_bound' n {sz} (bytes : ByteBuffer.t sz) : InternetChecksum.W16 :=
  InternetChecksum.ByteBuffer_fold_left_pair InternetChecksum.add_bytes_into_checksum n bytes (wzero _) (wzero _).

Lemma ByteBuffer_checksum_bound'_ok' :
  forall n {sz} (bytes :ByteBuffer.t sz) acc,
    Vector_checksum_bound n bytes acc =
    InternetChecksum.ByteBuffer_fold_left_pair InternetChecksum.add_bytes_into_checksum n bytes acc (wzero _).
Proof.
  fix IH 3.
  destruct bytes as [ | hd sz [ | hd' sz' tl ] ]; intros; simpl.
  - destruct n as [ | [ | ] ]; reflexivity.
  - destruct n as [ | [ | ] ]; reflexivity.
  - destruct n as [ | [ | ] ]; simpl; try reflexivity.
    rewrite IH; reflexivity.
Qed.

Lemma ByteBuffer_checksum_bound'_ok :
  forall n {sz} (bytes :ByteBuffer.t sz),
    Vector_checksum_bound n bytes (wzero _) = ByteBuffer_checksum_bound' n bytes.
Proof.
  intros; apply ByteBuffer_checksum_bound'_ok'.
Qed.

Definition IPChecksum_Valid_dec (n : nat) (b : ByteString)
  : {IPChecksum_Valid n b} + {~IPChecksum_Valid n b} := weq _ _.

Definition calculate_IPChecksum {S} {sz}
  : AlignedEncodeM (S := S) sz :=
  (fun v =>
     (let checksum := InternetChecksum.ByteBuffer_checksum_bound 20 v in
      (fun v idx s => SetByteAt (n := sz) 10 v 0 (wnot (split2 8 8 checksum)) ) >>
                                                                                (fun v idx s => SetByteAt (n := sz) 11 v 0 (wnot (split1 8 8 checksum)))) v)%AlignedEncodeM.

Definition splitLength (len: word 16) : Vector.t (word 8) 2 :=
  Vector.cons _ (split2 8 8 len) _ (Vector.cons _ (split1 8 8 len) _ (Vector.nil _)).

Definition Pseudo_Checksum_Valid
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (udpLength : word 16)
           (protoCode : word 8)
           (n : nat) (* Number of /bits/ in checksum; needed by
                        ByteString2ListOfChar *)
           (b : ByteString)
  := onesComplement (wzero 8 :: protoCode ::
                           to_list srcAddr ++ to_list destAddr ++ to_list (splitLength udpLength)
                           ++ (ByteString2ListOfChar n b)
                    )%list
     = wones 16.

Import VectorNotations.

Definition pseudoHeader_checksum
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (measure : nat)
           (udpLength : word 16)
           (protoCode : word 8)
           {sz} (packet: ByteBuffer.t sz) :=
  InternetChecksum.ByteBuffer_checksum_bound (12 + measure)
                                             (srcAddr ++ destAddr ++ [wzero 8; protoCode] ++ (splitLength udpLength) ++ packet).

Infix "^1+" := (InternetChecksum.OneC_plus) (at level 50, left associativity).

Import InternetChecksum.

Definition pseudoHeader_checksum'
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (measure : nat)
           (udpLength : word 16)
           (protoCode : word 8)
           {sz} (packet: ByteBuffer.t sz) :=
  ByteBuffer_checksum srcAddr ^1+
                               ByteBuffer_checksum destAddr ^1+
                                                             zext protoCode 8 ^1+
                                                                               udpLength ^1+
                                                                                InternetChecksum.ByteBuffer_checksum_bound measure packet.

Lemma OneC_plus_wzero_l :
  forall w, OneC_plus (wzero 16) w = w.
Proof. reflexivity. Qed.

Lemma OneC_plus_wzero_r :
  forall w, OneC_plus w (wzero 16) = w.
Proof.
  intros; rewrite OneC_plus_comm; reflexivity.
Qed.

Lemma Buffer_fold_left16_acc_oneC_plus :
  forall {sz} (packet: ByteBuffer.t sz) acc n,
    ByteBuffer_fold_left16 add_w16_into_checksum n packet acc =
    OneC_plus
      (ByteBuffer_fold_left16 add_w16_into_checksum n packet (wzero 16))
      acc.
Proof.
  fix IH 2.
  unfold ByteBuffer_fold_left16 in *.
  destruct packet as [ | hd sz [ | hd' sz' tl ] ]; intros; simpl.
  - destruct n as [ | [ | ] ]; reflexivity.
  - destruct n as [ | [ | ] ]; simpl; unfold add_bytes_into_checksum, add_w16_into_checksum;
      try rewrite OneC_plus_wzero_l, OneC_plus_comm; reflexivity.
  - destruct n as [ | [ | ] ]; simpl; unfold add_bytes_into_checksum, add_w16_into_checksum;
      try rewrite OneC_plus_wzero_l, OneC_plus_comm; try reflexivity.
    rewrite (IH _ tl (hd' +^+ hd ^1+ acc)).
    rewrite (IH _ tl (hd' +^+ hd)).
    rewrite OneC_plus_assoc.
    reflexivity.
Qed.

Lemma Vector_destruct_S :
  forall {A sz} (v: Vector.t A (S sz)),
  exists hd tl, v = hd :: tl.
Proof.
  repeat eexists.
  apply VectorSpec.eta.
Defined.

Lemma Vector_destruct_O :
  forall {A} (v: Vector.t A 0),
    v = [].
Proof.
  intro; apply Vector.case0; reflexivity.
Qed.

Ltac explode_vector :=
  unfold ByteBuffer.t in *;
  lazymatch goal with
  | [ v: Vector.t ?A (S ?n) |- _ ] =>
    let hd := fresh "hd" in
    let tl := fresh "tl" in
    rewrite (Vector.eta v) in *;
    set (Vector.hd v: A) as hd; clearbody hd;
    set (Vector.tl v: Vector.t A n) as tl; clearbody tl;
    clear v
  | [ v: Vector.t _ 0 |- _ ] =>
    rewrite (Vector_destruct_O v) in *; clear v
  end.

Lemma pseudoHeader_checksum'_ok :
  forall (srcAddr : ByteBuffer.t 4)
        (destAddr : ByteBuffer.t 4)
         (measure : nat)
         (udpLength : word 16)
         (protoCode : word 8)
         {sz} (packet: ByteBuffer.t sz),
    pseudoHeader_checksum srcAddr destAddr measure udpLength protoCode packet =
    pseudoHeader_checksum' srcAddr destAddr measure udpLength protoCode packet.
Proof.
  unfold pseudoHeader_checksum, pseudoHeader_checksum'.
  intros.
  repeat explode_vector.
  Opaque split1.
  Opaque split2.
  simpl in *.
  unfold ByteBuffer_checksum, InternetChecksum.ByteBuffer_checksum_bound, add_w16_into_checksum,
  add_bytes_into_checksum, ByteBuffer_fold_left16, ByteBuffer_fold_left_pair.
  fold @ByteBuffer_fold_left_pair.
  setoid_rewrite Buffer_fold_left16_acc_oneC_plus.
  rewrite combine_split.
  rewrite !OneC_plus_wzero_r, !OneC_plus_wzero_l, OneC_plus_comm.
  repeat (f_equal; [ ]).
  rewrite <- !OneC_plus_assoc.
  reflexivity.
Qed.

Definition calculate_PseudoChecksum {S} {sz}
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (measure : nat)
           (udpLength : word 16)
           (protoCode : word 8)
           (idx' : nat)
  : AlignedEncodeM (S := S) sz :=
  (fun v idx s =>
     (let checksum := pseudoHeader_checksum' srcAddr destAddr measure udpLength protoCode v in
      (fun v idx s => SetByteAt (n := sz) idx' v 0 (wnot (split2 8 8 checksum)) ) >>
                                                                                (fun v idx s => SetByteAt (n := sz) (1 + idx') v 0 (wnot (split1 8 8 checksum)))) v idx s)%AlignedEncodeM.

Lemma ByteBuffer_to_list_append {sz sz'}
  : forall (v : ByteBuffer.t sz)
           (v' : ByteBuffer.t sz'),
    ByteBuffer.to_list (v ++ v')%vector
    = ((ByteBuffer.to_list v) ++ (ByteBuffer.to_list v'))%list.
Proof.
  induction v.
  - reflexivity.
  - simpl; intros.
    unfold ByteBuffer.to_list at 1; unfold to_list.
    f_equal.
    apply IHv.
Qed.

Import VectorNotations.

Lemma Pseudo_Checksum_Valid_bounded
      {A}
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (udpLength : word 16)
      protoCode
      (predicate : A -> Prop)
      (format_A format_B : FormatM A ByteString)
      (len_format_A : A -> nat)
      (len_format_A_OK : forall a' b ctx ctx',
          computes_to (format_A a' ctx) (b, ctx')
          -> length_ByteString b = len_format_A a')
      (len_format_B : A -> nat)
      (len_format_B_OK : forall a' b ctx ctx',
          computes_to (format_B a' ctx) (b, ctx')
          -> length_ByteString b = len_format_B a')
      (byte_aligned_A : forall a : A, len_format_A a mod 8 = 0)
      (byte_aligned_B : forall a : A, len_format_B a mod 8 = 0)
  : forall (data : A) (x : ByteString) (x0 : CacheFormat) (x1 : ByteString) (x2 : CacheFormat)
           (ext ext' : ByteString) (env : CacheFormat) (c : word 16),
    predicate data ->
    format_A data env ∋ (x, x0) ->
    format_B data (addE x0 16) ∋ (x1, x2) ->
    Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode
                          (bin_measure (mappend x (mappend (format_checksum ByteString monoid ByteString_QueueMonoidOpt 16 c) x1)))
                          (mappend (mappend x (mappend (format_checksum ByteString monoid ByteString_QueueMonoidOpt 16 c) x1)) ext) ->
    Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode
                          (bin_measure (mappend x (mappend (format_checksum ByteString monoid ByteString_QueueMonoidOpt 16 c) x1)))
                          (mappend (mappend x (mappend (format_checksum ByteString monoid ByteString_QueueMonoidOpt 16 c) x1)) ext').
Proof.
  intros.
    unfold Pseudo_Checksum_Valid in *.
    revert H2.
    rewrite !ByteString2ListOfChar_Over; eauto.
    simpl; rewrite padding_eq_mod_8.
    rewrite !length_ByteString_enqueue_ByteString.
    rewrite Nat.add_mod by omega.
    apply len_format_A_OK in H0.
    apply len_format_B_OK in H1.
    unfold format_checksum; rewrite length_encode_word', measure_mempty.
    rewrite H0, byte_aligned_A, plus_O_n, NPeano.Nat.mod_mod, Nat.add_mod by omega.
    rewrite H1, byte_aligned_B, <- plus_n_O, NPeano.Nat.mod_mod by omega.
    reflexivity.
    simpl; rewrite padding_eq_mod_8.
    rewrite !length_ByteString_enqueue_ByteString.
    rewrite Nat.add_mod by omega.
    apply len_format_A_OK in H0.
    apply len_format_B_OK in H1.
    rewrite H0, byte_aligned_A, plus_O_n, NPeano.Nat.mod_mod, Nat.add_mod by omega.
    rewrite H1, byte_aligned_B, <- plus_n_O, NPeano.Nat.mod_mod by omega.
    unfold format_checksum; rewrite length_encode_word'; reflexivity.
Qed.

Lemma compose_PseudoChecksum_format_correct' {A}
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (udpLength : word 16)
      protoCode
      (predicate : A -> Prop)
      (P : CacheDecode -> Prop)
      (P_inv : (CacheDecode -> Prop) -> Prop)
      (P_invM : (CacheDecode -> Prop) -> Prop)
      (format_A format_B : FormatM A ByteString)
      (subformat : FormatM A ByteString)
      (decode_measure : DecodeM (nat * _) _)
      (len_format_A : A -> nat)
      (len_format_A_OK : forall a' b ctx ctx',
          computes_to (format_A a' ctx) (b, ctx')
          -> length_ByteString b = len_format_A a')
      (len_format_B : A -> nat)
      (len_format_B_OK : forall a' b ctx ctx',
          computes_to (format_B a' ctx) (b, ctx')
          -> length_ByteString b = len_format_B a')
      View_Predicate
      format_measure
  : cache_inv_Property P (fun P => P_inv P /\ P_invM P) ->
    (forall a, NPeano.modulo (len_format_A a) 8 = 0)
    -> (forall a, NPeano.modulo (len_format_B a) 8 = 0)
    ->
    forall decodeA : _ -> CacheDecode -> option (A * _ * CacheDecode),
      (cache_inv_Property P P_inv ->
       CorrectDecoder monoid predicate predicate eq (format_A ++ format_unused_word 16 ++ format_B)%format decodeA P (format_A ++ format_unused_word 16 ++ format_B)%format) ->
      (cache_inv_Property P P_invM ->
          CorrectRefinedDecoder monoid predicate View_Predicate
                                (fun a n => len_format_A a + 16 + len_format_B a = n * 8)
                                (format_A ++ format_unused_word 16 ++ format_B)%format
                                subformat
                                decode_measure P
                                format_measure) ->
      (Prefix_Format _ (format_A ++ format_unused_word 16 ++ format_B) subformat)%format->
      CorrectDecoder monoid predicate predicate eq
                     (format_A ThenChecksum (Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode) OfSize 16 ThenCarryOn format_B)
                     (fun (bin : _) (env : CacheDecode) =>
                          `(n, _, _) <- decode_measure bin env;
                            if weqb (onesComplement (wzero 8 :: protoCode ::
                                                       to_list srcAddr ++ to_list destAddr ++ to_list (splitLength udpLength)
                                                       ++(ByteString2ListOfChar (n * 8) bin))%list) (wones 16) then
                              decodeA bin env
                            else None)
                     P
                     (format_A ThenChecksum (Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode) OfSize 16 ThenCarryOn format_B).
Proof.
  intros.
  rename H4 into H4'; rename H3 into H4; rename H2 into H3.
  eapply format_decode_correct_alt.
  Focus 7.
  (*7: {*)
  {eapply (composeChecksum_format_correct'
                 A _ monoid _ 16 (Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode)).
       - eapply H.
       - specialize (H4 (proj2 H)).
         split.
         2: eauto.
         eapply injection_decode_correct with (inj := fun n => mult n 8).
         4: simpl.
         eapply H4.
         + intros.
           instantiate (1 := fun a n => len_format_A a + 16 + len_format_B a = n).
           eapply H6.
         + intros; instantiate (1 := fun v => View_Predicate (Nat.div v 8)).
           cbv beta.
           rewrite Nat.div_mul; eauto.
         + intros; apply unfold_computes; intros.
           split.
           2: rewrite unfold_computes in H5; intuition.
           intros.
           rewrite unfold_computes in H5; intuition.
           instantiate (1 := fun v env t => format_measure (Nat.div v 8) env t).
           cbv beta; rewrite Nat.div_mul; eauto.
       - simpl; intros.
         destruct t1; destruct t2; simpl fst in *; simpl snd in *.
         apply unfold_computes in H7; apply unfold_computes in H6.
         erewrite len_format_A_OK; eauto.
         erewrite (len_format_B_OK _ b0); eauto.
         unfold format_checksum; rewrite length_encode_word', measure_mempty.
         rewrite <- H2; omega.
       - eauto.
       - eapply Pseudo_Checksum_Valid_bounded; eauto. }
  all: try unfold flip, pointwise_relation, impl;
    intuition eauto using EquivFormat_reflexive.
  all: try unfold flip, pointwise_relation, impl;
    intuition eauto using EquivFormat_reflexive.
    instantiate (1 := fun (n : nat) a =>
                    weq
       (onesComplement
          (wzero 8
           :: (protoCode
               :: to_list srcAddr ++
                  to_list destAddr ++ to_list (splitLength udpLength) ++ ByteString2ListOfChar n a)%list))
       (wones 16)).
  unfold Compose_Decode.
  Local Opaque Nat.div.
  destruct (decode_measure a a0) as [ [ [? ?] ? ] | ]; simpl; eauto.
  symmetry.
  find_if_inside.
  eapply weqb_true_iff in e; rewrite e; eauto.
  destruct (weqb
      (add_bytes_into_checksum (wzero 8) protoCode
         (onesComplement
            (to_list srcAddr ++
             to_list destAddr ++ split2 8 8 udpLength :: (split1 8 8 udpLength :: ByteString2ListOfChar (n * 8) a)%list)))
      WO~1~1~1~1~1~1~1~1~1~1~1~1~1~1~1~1) eqn: ? ; eauto.
  eapply weqb_true_iff in Heqb0.
  congruence.
Qed.

Fixpoint aligned_Pseudo_checksum
         (srcAddr : ByteBuffer.t 4)
         (destAddr : ByteBuffer.t 4)
         (pktlength : word 16)
         id
         measure
         {sz}
         (v : t Core.char sz) (idx : nat)
         {struct idx}
  := match idx with
     | 0 =>
       weqb (pseudoHeader_checksum' srcAddr destAddr measure pktlength id v)
            (wones 16)
     | S idx' =>
       match v with
       | Vector.cons _ _ v' => aligned_Pseudo_checksum srcAddr destAddr pktlength id measure v' idx'
       | _ => false
       end
     end.

Lemma Vector_checksum_bound_acc'
  : forall sz'' sz sz' (sz_lt : le sz' sz'') (v : Vector.t _ sz) b1 b2 acc,
    Vector_checksum_bound sz' v (add_bytes_into_checksum b1 b2 acc) =
    add_bytes_into_checksum b1 b2 (Vector_checksum_bound sz' v acc).
Proof.
  induction sz''; intros.
  - inversion sz_lt.
    subst; reflexivity.
  - inversion sz_lt; subst.
    + clear sz_lt.
      destruct sz''; simpl.
      * destruct v; simpl; eauto.
        rewrite add_bytes_into_checksum_swap; eauto.
      * destruct v; simpl; eauto.
        destruct v; simpl; eauto.
        rewrite add_bytes_into_checksum_swap; eauto.
        rewrite !IHsz'' by omega.
        rewrite add_bytes_into_checksum_swap; eauto.
    + eauto.
Qed.

Lemma Vector_checksum_bound_acc
  : forall sz sz' (v : Vector.t _ sz) b1 b2 acc,
    Vector_checksum_bound sz' v (add_bytes_into_checksum b1 b2 acc) =
    add_bytes_into_checksum b1 b2 (Vector_checksum_bound sz' v acc).
Proof.
  intros; eapply Vector_checksum_bound_acc'.
  reflexivity.
Qed.

Lemma dequeue_byte_ByteString2ListOfChar
  : forall m sz (v : Vector.t _ sz) b,
    ByteString2ListOfChar ((S m) * 8) (build_aligned_ByteString (b :: v))
    = cons b (ByteString2ListOfChar (m * 8) (build_aligned_ByteString (v))).
Proof.
  intros; erewrite <- ByteString2ListOfChar_push_char.
  f_equal.
  pose proof (build_aligned_ByteString_append v [b]) as H; simpl in H.
  rewrite H.
  reflexivity.
Qed.

Lemma ByteString2ListOfChar_overflow
  : forall n,
    ByteString2ListOfChar ((S n) * 8) (build_aligned_ByteString [])
    = cons (wzero 8) (ByteString2ListOfChar (n * 8) (build_aligned_ByteString [])).
Proof.
  reflexivity.
Qed.

Lemma InternetChecksum_To_ByteBuffer_Checksum':
  forall (sz m : nat) (v : t Core.char sz),
    checksum (ByteString2ListOfChar (m * 8) (build_aligned_ByteString v)) = ByteBuffer_checksum_bound m v.
Proof.
  intros.
  assert ((exists m', m = 2 * m') \/ (exists m', m = S (2 * m'))).
  { induction m; eauto.
    destruct IHm; destruct_ex; subst; eauto.
    left; exists (S x); omega.
  }
  destruct H as [ [? ?] | [? ?] ]; subst.
  - rewrite (mult_comm 2).
    apply InternetChecksum_To_ByteBuffer_Checksum.
  - revert sz v.
    induction x.
    + intros; destruct v.
      * reflexivity.
      * rewrite dequeue_byte_ByteString2ListOfChar.
        reflexivity.
    + intros; destruct v.
      * replace (S (2 * S x)) with ((S (S (S (2 * x))))) by omega.
        rewrite ByteString2ListOfChar_overflow.
        rewrite ByteString2ListOfChar_overflow.
        unfold checksum; fold checksum.
        rewrite IHx.
        unfold ByteBuffer_checksum_bound, ByteBuffer_fold_left16.
        simpl.
        destruct (2 * x); eauto.
      * rewrite dequeue_byte_ByteString2ListOfChar.
        destruct v.
        replace (2 * S x * 8) with ((S (S (2 * x))) * 8) by omega.
        rewrite ByteString2ListOfChar_overflow.
        unfold checksum; fold checksum.
        rewrite IHx.
        rewrite <- !ByteBuffer_checksum_bound_ok.
        simpl.
        destruct (2 * x); eauto.
        replace (2 * S x * 8) with ((S (S (2 * x))) * 8) by omega.
        rewrite dequeue_byte_ByteString2ListOfChar.
        replace
          (checksum (h :: (h0 :: ByteString2ListOfChar (S (2 * x) * 8) (build_aligned_ByteString v))%list))
          with
            (add_bytes_into_checksum
               h h0
               (checksum (ByteString2ListOfChar (S (2 * x) * 8) (build_aligned_ByteString v))%list))
          by reflexivity.
        rewrite IHx.
        rewrite <- !ByteBuffer_checksum_bound_ok.
        replace (2 * S x) with (S (S ( 2 * x))) by omega.
        rewrite <- Vector_checksum_bound_acc; reflexivity.
Qed.

Lemma aligned_Pseudo_checksum_OK_1
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (pktlength : word 16)
      id
      measure
      {sz}
  : forall (v : t Core.char sz),
    weqb
      (InternetChecksum.add_bytes_into_checksum (wzero 8) id
                                                (onesComplement(to_list srcAddr ++ to_list destAddr ++ split2 8 8  pktlength :: split1 8 8 pktlength :: (ByteString2ListOfChar (measure * 8) (build_aligned_ByteString v)))%list))
      WO~1~1~1~1~1~1~1~1~1~1~1~1~1~1~1~1
    = aligned_Pseudo_checksum srcAddr destAddr pktlength id measure v 0.
Proof.
  simpl; intros.
  unfold onesComplement.
  rewrite <- pseudoHeader_checksum'_ok.
  rewrite checksum_eq_Vector_checksum.
  unfold pseudoHeader_checksum.
  rewrite <- ByteBuffer_checksum_bound_ok.
  unfold ByteBuffer.t in *.
  simpl.
    replace srcAddr with
      (Vector.hd srcAddr :: Vector.hd (Vector.tl srcAddr)
                 :: Vector.hd (Vector.tl (Vector.tl srcAddr))
                 :: Vector.hd (Vector.tl (Vector.tl (Vector.tl srcAddr)))
                 :: (@Vector.nil _))
    by abstract (pattern srcAddr;
              repeat (apply caseS'; let t' := fresh in intros ? t'; pattern t'); apply case0;
              reflexivity).
  replace destAddr with
      (Vector.hd destAddr :: Vector.hd (Vector.tl destAddr)
                 :: Vector.hd (Vector.tl (Vector.tl destAddr))
                 :: Vector.hd (Vector.tl (Vector.tl (Vector.tl destAddr)))
                 :: (@Vector.nil _))
    by abstract (pattern destAddr;
              repeat (apply caseS'; let t' := fresh in intros ? t'; pattern t'); apply case0;
              reflexivity).
  simpl.
  repeat rewrite Vector_checksum_bound_acc.
  rewrite <- checksum_eq_Vector_checksum.
  f_equal.
  rewrite ByteBuffer_checksum_bound_ok.
  repeat rewrite (add_bytes_into_checksum_swap _ id); f_equal.
  repeat rewrite (add_bytes_into_checksum_swap _ (Vector.hd (Vector.tl srcAddr))); f_equal.
  repeat rewrite (add_bytes_into_checksum_swap _ (Vector.hd (Vector.tl (Vector.tl (Vector.tl srcAddr))))); f_equal.
  repeat rewrite (add_bytes_into_checksum_swap _ (Vector.hd (Vector.tl destAddr))); f_equal.
  repeat rewrite (add_bytes_into_checksum_swap _ (Vector.hd (Vector.tl (Vector.tl (Vector.tl destAddr))))); f_equal.
  f_equal.
  apply InternetChecksum_To_ByteBuffer_Checksum'.
Qed.

Lemma aligned_Pseudo_checksum_OK_2
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (pktlength : word 16)
      id
      measure
      {sz}
  : forall (v : ByteBuffer.t (S sz)) (idx : nat),
    aligned_Pseudo_checksum srcAddr destAddr pktlength id measure v (S idx) =
    aligned_Pseudo_checksum srcAddr destAddr pktlength id measure (Vector.tl v) idx.
Proof.
  intros v; pattern sz, v.
  apply Vector.caseS; reflexivity.
Qed.

Fixpoint aligned_IPchecksum
         m
         {sz}
         (v : t Core.char sz) (idx : nat)
         {struct idx}
  := match idx with
     | 0 =>
       weqb (InternetChecksum.ByteBuffer_checksum_bound m v) (wones 16)
     | S idx' =>
       match v with
       | Vector.cons _ _ v' => aligned_IPchecksum m v' idx'
       | _ => false
       end
     end.

Corollary aligned_IPChecksum_OK_1 :
  forall m sz (v : ByteBuffer.t sz),
    IPChecksum_Valid_check (m * 8) (build_aligned_ByteString v)
    = aligned_IPchecksum m v 0.
Proof.
  intros.
  unfold IPChecksum_Valid_check, aligned_IPchecksum.
  rewrite InternetChecksum_To_ByteBuffer_Checksum'.
  higher_order_reflexivity.
Qed.

Lemma aligned_IPChecksum_OK_2
      m
      {sz}
  : forall (v : ByteBuffer.t (S sz)) (idx : nat),
    aligned_IPchecksum m v (S idx) =
    aligned_IPchecksum m (Vector.tl v) idx.
Proof.
  intros v; pattern sz, v.
  apply Vector.caseS; reflexivity.
Qed.

Hint Extern 4 (weqb _ _ = _) =>
rewrite aligned_Pseudo_checksum_OK_1; higher_order_reflexivity.
Hint Extern 4 => eapply aligned_Pseudo_checksum_OK_2.

Hint Extern 4 (IPChecksum_Valid_check _ _ = _ ) =>
rewrite aligned_IPChecksum_OK_1; higher_order_reflexivity.
Hint Extern 4  => eapply aligned_IPChecksum_OK_2.

Ltac destruct_unit :=
  repeat match goal with
         | v : () |- _ => destruct v
         end.

Ltac solve_seq_padding :=
  repeat (intros; match goal with
                  | H : sequence_Encode ?e1 ?e2 _ _ = Some (?t, _) |- padding ?t = 0 =>
                    eapply sequence_Encoding_padding_0 in H; eauto
                  end).

Lemma EncodeMEquivAlignedEncodeMForChecksum
      {S A}
      (enc1' enc2' enc3' : EncodeM S ByteString)
      (encA' : A -> EncodeM S ByteString)
      (f' : ByteString -> A)
      (enc1 enc2 enc3 : forall sz, AlignedEncodeM sz)
      (encA : A -> forall sz, AlignedEncodeM sz)
      (f : nat -> forall {sz}, Vector.t (word 8) sz -> nat -> A)
      (enc1_OK : EncodeMEquivAlignedEncodeM enc1' enc1)
      (enc2_OK : EncodeMEquivAlignedEncodeM enc2' enc2)
      (enc3_OK : EncodeMEquivAlignedEncodeM enc3' enc3)
      (encA_OK : forall a, EncodeMEquivAlignedEncodeM (encA' a) (encA a))
      (enc1'_aligned : forall s ce t ce', enc1' s ce = Some (t, ce') -> padding t = 0)
      (enc2'_aligned : forall s ce t ce', enc2' s ce = Some (t, ce') -> padding t = 0)
      (enc3'_aligned : forall s ce t ce', enc3' s ce = Some (t, ce') -> padding t = 0)
      (encA'_aligned : forall a s ce t ce', encA' a s ce = Some (t, ce') -> padding t = 0)
      (enc'_sz_eq : forall s a ce t1 ce1 t2 ce2,
          enc2' s ce = Some (t1, ce1) ->
          encA' a s ce = Some (t2, ce2) ->
          bin_measure t1 = bin_measure t2)
      (f_OK : forall (b : ByteString)
                idx (v1 : Vector.t Core.char idx)
                m (v2 : Vector.t Core.char m)
                (v : Vector.t Core.char (idx + ((numBytes b) + m))),
          ByteString_enqueue_ByteString (build_aligned_ByteString v1)
                                        (ByteString_enqueue_ByteString b (build_aligned_ByteString v2))
          = build_aligned_ByteString v ->
          f' b = f (numBytes b) v idx)
  : EncodeMEquivAlignedEncodeM
      (fun s ce =>
         `(p, _) <- sequence_Encode enc1' (sequence_Encode enc2' enc3') s ce;
           (fun a => (sequence_Encode enc1' (sequence_Encode (encA' a) enc3'))) (f' p) s ce)
      (fun sz => EncodeAgain (enc1 sz >> enc2 sz >> enc3 sz)
                          (fun idx' v idx s =>
                             (fun a sz => enc1 sz >> encA a sz >> enc3 sz)
                               (f (idx'-idx) v idx) sz v idx s))%AlignedEncodeM.
Proof.
  apply EncodeMEquivAlignedEncodeMDep; eauto; intros;
    repeat (apply Append_EncodeMEquivAlignedEncodeM; eauto);
    solve_seq_padding.
  apply sequence_Encoding_inv in H.
  apply sequence_Encoding_inv in H0. destruct_conjs.
  apply sequence_Encoding_inv in H9.
  apply sequence_Encoding_inv in H4. destruct_conjs.
  subst. rewrite !mappend_measure.
  simpl in *. destruct_unit.
  substss. injections. eauto.
Qed.

Definition f_bit_aligned_free {A} (f : nat -> forall {sz}, Vector.t (word 8) sz -> nat -> A)
  : ByteString -> A :=
  fun bs => f (numBytes bs) (byteString bs) 0.

Lemma CorrectAlignedEncoderForChecksum
      {S}
      checksum_sz (checksum_valid : nat -> ByteString -> Prop)
      (format_A format_B : FormatM S ByteString)
      (encode_A encode_B : forall sz, AlignedEncodeM sz)
      (encode_C' : forall sz, AlignedEncodeM sz)
      (encode_C : word checksum_sz -> forall sz, AlignedEncodeM sz)
      (encoder_A_OK : CorrectAlignedEncoder format_A encode_A)
      (encoder_B_OK : CorrectAlignedEncoder format_B encode_B)
      (f : nat -> forall {sz}, Vector.t (word 8) sz -> nat -> word checksum_sz)
      (f_OK : forall idx n m (t : Vector.t Core.char n)
                (v1 : Vector.t Core.char idx) (v2 : Vector.t Core.char m),
         f n t 0 = f n (v1 ++ t ++ v2) idx)
      (enc2' : EncodeM S ByteString)
      (enc2_OK : EncodeMEquivAlignedEncodeM enc2' encode_C')
      (encode_C_OK : forall a, EncodeMEquivAlignedEncodeM
                            ((fun a _ ce => Some (format_checksum _ _ _ _ a, ce)) a) (encode_C a))
      (enc2'_aligned : forall s ce t ce', enc2' s ce = Some (t, ce') -> padding t = 0)
      n1
      (format_B_sz_eq : forall s ce t ce',
          format_B s ce ∋ (t, ce') -> bin_measure t = n1)
      (checksum_sz_OK : checksum_sz mod 8 = 0)
      (checksum_sz_OK' : forall s ce t ce', enc2' s ce = Some (t, ce') ->
                                       checksum_sz = bin_measure t)
      (checksum_OK : forall s ce b' ce',
          enc2' s ce = Some (b', ce') ->
          forall b1 b3 ext,
            (exists s ce ce', format_B s ce ∋ (b1, ce')) ->
            (exists s ce ce', format_A s ce ∋ (b3, ce')) ->
            let a := (f_bit_aligned_free f) (mappend b1 (mappend b' b3)) in
            let b2 := format_checksum _ _ _ checksum_sz a in
            checksum_valid
              (bin_measure (mappend b1 (mappend b2 b3)))
              (mappend (mappend b1 (mappend b2 b3)) ext))
      (checksum_OK' : forall s ce,
          enc2' s ce = None ->
          forall b1 b2 b3 ext,
            ~checksum_valid
              (bin_measure (mappend b1 (mappend b2 b3)))
              (mappend (mappend b1 (mappend b2 b3)) ext))
  : CorrectAlignedEncoder
      (format_B ThenChecksum checksum_valid OfSize checksum_sz ThenCarryOn format_A)
      (fun sz => EncodeAgain (encode_B sz >> encode_C' sz >> encode_A sz)
                          (fun idx' v idx s =>
                             encode_C (f (idx'-idx) v idx) sz v (idx+n1/8) s))%AlignedEncodeM.
Proof.
  destruct encoder_A_OK as [enc3' [HA1 [HA2 HA3]]].
  destruct encoder_B_OK as [enc1' [HB1 [HB2 HB3]]].
  exists (fun s ce =>
       `(p, _) <- sequence_Encode enc1' (sequence_Encode enc2' enc3') s ce;
       (fun a => (sequence_Encode enc1' (sequence_Encode
                                        ((fun a _ ce => Some (format_checksum _ _ _ _ a, ce)) a)
                                        enc3'))) ((f_bit_aligned_free f) p) s ce).
  split; [| split]; intros.
  - unfold composeChecksum, sequence_Encode. split; intros; simpl in *.
    + intros [? ?]. intros.
      computes_to_inv. injections.
      destruct enc1' eqn:Henc1; [| discriminate]; destruct_conjs; simpl in *.
      destruct enc2' eqn:Henc2; [| discriminate]; destruct_conjs; simpl in *.
      destruct enc3' eqn:Henc3; [| discriminate]; destruct_conjs; simpl in *.
      destruct_unit.
      rewrite Henc3 in H. simpl in *. injections.
      repeat computes_to_econstructor; eauto.
      edestruct HB1. apply H in Henc1. apply Henc1.
      repeat computes_to_econstructor; eauto. simpl.
      edestruct HA1. apply H in Henc3. apply Henc3.
      repeat computes_to_econstructor; eauto.
      intros. eapply checksum_OK; eauto.
      repeat eexists. simpl. apply HB1 in Henc1. eauto.
      repeat eexists. simpl. apply HA1 in Henc3. eauto.
      simpl. apply eq_ret_compute. repeat f_equal.
    + intro. unfold Bind2 in H0.
      computes_to_inv. destruct_conjs. injections. simpl in *. destruct_unit.
      destruct enc1' eqn:Henc1; destruct_conjs; simpl in *; destruct_unit.
      destruct enc2' eqn:Henc2; destruct_conjs; simpl in *; destruct_unit.
      destruct enc3' eqn:Henc3; destruct_conjs; simpl in *; destruct_unit.
      discriminate.
      edestruct HA1. intuition eauto. intuition eauto.
      edestruct HB1. intuition eauto.
  - destruct sequence_Encode; [| discriminate]; destruct_conjs; simpl in *.
    eapply sequence_Encoding_padding_0; try apply H; eauto.
    intros.
    eapply sequence_Encoding_padding_0; try apply H0; eauto.
    intros. simpl in *. injections.
    unfold format_checksum. rewrite encode_word'_padding. eauto.
  - eapply EncodeMEquivAlignedEncodeM_morphism; cycle 1.
    apply EncodeMEquivAlignedEncodeMForChecksum
      with (encA' := (fun a _ ce => Some (format_checksum _ _ _ _ a, ce))); eauto.
    + intros. injections.
      unfold format_checksum. rewrite encode_word'_padding. eauto.
    + intros. injections. unfold format_checksum.
      rewrite length_encode_word'. rewrite measure_mempty.
      apply checksum_sz_OK' in H. omega.
    + unfold f_bit_aligned_free. instantiate (1:=f). intros.
      assert (padding b = 0) as L1. {
        apply (f_equal padding) in H. simpl in H.
        rewrite padding_ByteString_enqueue_aligned_ByteString in H; eauto.
        rewrite ByteString_enqueue_ByteString_padding_eq in H. simpl padding in H.
        rewrite (Nat.mod_add _ 0 8) in H; eauto.
        rewrite Nat.mod_small in H; eauto.
        destruct b. simpl in *. auto.
      }
      assert (b = build_aligned_ByteString (byteString b)) as L2. {
        revert L1. clear. intros.
        destruct b. simpl. unfold build_aligned_ByteString. simpl in *.
        subst. f_equal. eauto using shatter_word_0.
        apply Core.le_uniqueness_proof.
      }
      revert L2 H. revert v.
      generalize (byteString b). generalize (numBytes b). intros. rewrite L2 in H.
      rewrite <- !build_aligned_ByteString_append in H.
      apply build_aligned_ByteString_inj in H. subst.
      eauto.
    + intros. unfold EncodeAgain.
      set (fun sz => encode_C' sz >> encode_A sz)%AlignedEncodeM as enc'.
      set (fun sz => encode_B sz >> enc' sz)%AlignedEncodeM as enc.
      match goal with
      | |- context[Ifopt ?b _ _ _ _ as _ Then _ Else _] => replace b with (enc sz) by reflexivity
      end.
      destruct enc eqn:?; eauto.
      destruct_conjs. simpl. unfold AppendAlignedEncodeM.

      edestruct @AlignedEncoder_inv2 as [n2 [n3 [t1 [t23 [v23 ?]]]]]; try apply Heqa.
      apply Append_EncodeMEquivAlignedEncodeM. 2 : eauto. eauto.
      apply Append_EncodeMEquivAlignedEncodeM. 2 : eauto. eauto. eauto.
      solve_seq_padding.
      destruct_conjs. subst. simpl in *.

      edestruct @AlignedEncoder_append_inv with (enc1:=encode_B)
        as [nB [nC [nC3 [tB [tC3 [vB [vC3 [?ce ?]]]]]]]].
      5 : apply H0.
      3 : apply Append_EncodeMEquivAlignedEncodeM. all : eauto.
      solve_seq_padding. destruct_conjs.
      subst. simpl in *. destruct H. simpl in *.

      assert (encode_B (idx + (nB + (nC + nC3))) (t1 ++ vB ++ vC3) idx w c =
              Some (t1 ++ vB ++ vC3, idx+nB, ce)). {
        assert (idx + nB + (nC + nC3) = idx + (nB + (nC + nC3))) as L by omega.
        rewrite Vector_append_assoc with (H:=L). destruct L. simpl.
        epose proof AlignedEncoder_extr as H'. eapply H'; eauto. clear H'.
        eapply @AlignedEncoder_fixed; eauto.
        eapply AlignedEncoder_extl; eauto.
      } rewrite H. simpl.

      assert (nB = n1/8). {
        erewrite <- format_B_sz_eq; eauto.
        epose proof AlignedEncoder_inv0 as L. eapply L in H1; eauto. clear L.
        instantiate (1:=build_aligned_ByteString vB).
        rewrite length_ByteString_no_padding; eauto.
        rewrite Nat.mul_comm.
        rewrite Nat.div_mul by auto.
        reflexivity.
        eapply HB1; eauto. eauto using AlignedEncoder_inv0.
      } destruct H3.
      destruct_unit. destruct encode_C eqn:?; eauto. destruct_conjs. simpl in *.

      edestruct @AlignedEncoder_append_inv with (enc1:=encode_C')
        as [nC' [nA [nA3 [tC [tA3 [vC [vA3 [?ce ?]]]]]]]]; try apply H2; eauto.
      destruct_conjs.
      subst. simpl in *. destruct H3. simpl in *. rename nC' into nC.

      match goal with
      | H : encode_C ?a _ _ _ _ _ = _ |- _ =>
        edestruct @AlignedEncoder_inv with (enc:=(encode_C a))
          as [nC' [n3' [t1C [tC' [tC3 [vC' ?]]]]]]
      end; eauto.
      intros. simpl in *. injections.
      unfold format_checksum. rewrite encode_word'_padding. eauto.
      destruct_conjs.

      assert (nC = nC'). {
        injections.
        assert (checksum_sz = 8 * nC').
        eapply (f_equal bin_measure) in H12. simpl in *.
        unfold format_checksum in H12.
        rewrite length_encode_word' in H12. rewrite measure_mempty in H12.
        unfold length_ByteString in H12. simpl in H12. omega.
        assert (checksum_sz = 8 * nC).
        epose proof AlignedEncoder_inv0 as L. eapply L in H4; eauto. clear L.
        apply checksum_sz_OK' in H4.
        unfold length_ByteString in H4. simpl in H4. omega.
        omega.
      } subst nC'.

      assert (n3' = nA + nA3) as L by omega. subst n3'.
      assert (encode_A (idx + (nB + (nC + (nA + nA3)))) t n w c =
              Some (t, n+nA, c)). {
        rewrite (Vector_append_assoc _ _ _ H3) in H8.
        clear Heqa0. destruct H3. simpl in *.
        subst. apply Vector_append_inj in H8. destruct_conjs.
        apply Vector_append_inj in H7. destruct_conjs. subst.
        assert (idx + nB + nC + (nA + nA3) = idx + nB + (nC + (nA + nA3))) as L by omega.
        rewrite (Vector_append_assoc _ _ _ L). destruct L. simpl.
        epose proof AlignedEncoder_extl as L. eapply L; eauto. clear L.
        destruct_unit.
        eapply @AlignedEncoder_fixed; eauto.
      } rewrite H11. simpl. reflexivity.
      Unshelve.
      eauto.
Qed.

Definition encode_word_const
           {S : Type}
           {n}
  : word (n*8) -> forall sz, AlignedEncodeM (S:=S) sz :=
  fun w sz v idx _ ce => SetCurrentBytes v idx w ce.

Definition encode_word_16_const {S : Type} := @encode_word_const S 2.

Definition encode_word_16_0 {S : Type} := @encode_word_16_const S (wzero 16).

Lemma format_word_is_encode_word {T}
      {cache : Cache} {cacheAddNat : CacheAdd cache nat}
      {monoid : Monoid T} {monoidUnit : QueueMonoidOpt monoid bool}
      {n} (enc : EncodeM (word n) T)
  : (forall s env,
        (forall t env', enc s env = Some (t, env')
                   -> refine (format_word s env) (ret (t, env')))
        /\ (enc s env = None ->
           forall benv', ~ computes_to (format_word s env) benv')) ->
    forall s env, enc s env = Some (encode_word' _ s mempty, addE env n).
Proof.
  intros. destruct enc eqn:?; destruct_conjs.
  - apply H in Heqe.
    eapply Return_inv in Heqe; eauto. congruence.
  - exfalso. eapply H; eauto.
    computes_to_econstructor; eauto.
Qed.

Lemma EncodeMEquivAlignedEncodeM_const
      {S A} {cache : Cache}
      (enc' : EncodeM A ByteString)
      (enc : forall sz, AlignedEncodeM sz)
      (enc_OK : EncodeMEquivAlignedEncodeM enc' enc)
  : forall a, EncodeMEquivAlignedEncodeM (S:=S)
           (fun _ ce => enc' a ce)
           (fun sz v n _ ce => enc sz v n a ce).
Proof.
  intros. repeat split; simpl; intros;
            edestruct enc_OK as [? [? [? ?]]]; eauto.
Qed.

Lemma EncodeMEquivAlignedEncodeM_word_const
      {S}
      {n}
  : forall (w : word (n*8)),
    EncodeMEquivAlignedEncodeM
      (fun (_ : S) env => Some (encode_word' _ w mempty, addE env (n*8)))
      (encode_word_const w).
Proof.
  intros.
  destruct CorrectAlignedEncoderForFormatNChar with (sz:=n); eauto. destruct_conjs.
  eapply EncodeMEquivAlignedEncodeM_const in H1.
  match goal with
  | H : EncodeMEquivAlignedEncodeM ?a _ |- EncodeMEquivAlignedEncodeM ?b _ =>
    replace b with a
  end. eauto.
  extensionality s. extensionality ce.
  eauto using format_word_is_encode_word.
Qed.

Definition AlignedEncoderForChecksum
      {S}
      n1
      (encode_B encode_A : forall sz, AlignedEncodeM sz)
      (f : nat -> forall {sz}, ByteBuffer.t sz -> nat -> word 16)
  := (fun sz => EncodeAgain (encode_B sz >> encode_word_16_0 sz >> encode_A sz)
                         (fun idx' v idx (s : S) =>
                            encode_word_16_const (f (idx'-idx) v idx) sz v (idx+n1/8) s))%AlignedEncodeM.

Local Opaque encode_word'.
Lemma CorrectAlignedEncoderForChecksum'
      {S}
      (checksum_valid : nat -> ByteString -> Prop)
      (format_A format_B : FormatM S ByteString)
      (encode_A encode_B : forall sz, AlignedEncodeM sz)
      (encoder_A_OK : CorrectAlignedEncoder format_A encode_A)
      (encoder_B_OK : CorrectAlignedEncoder format_B encode_B)
      (f : nat -> forall {sz}, Vector.t (word 8) sz -> nat -> word 16)
      (f_OK : forall idx n m (t : Vector.t Core.char n)
                (v1 : Vector.t Core.char idx) (v2 : Vector.t Core.char m),
         f n t 0 = f n (v1 ++ t ++ v2) idx)
      n1
      (format_B_sz_eq : forall s ce t ce',
          format_B s ce ∋ (t, ce') -> bin_measure t = n1)
      (checksum_OK :
         forall b1 b3 ext,
            (exists s ce ce', format_B s ce ∋ (b1, ce')) ->
            (exists s ce ce', format_A s ce ∋ (b3, ce')) ->
           let a := (f_bit_aligned_free f)
                      (mappend b1 (mappend (format_checksum _ _ _ _ (wzero 16)) b3)) in
           let b2 := format_checksum _ _ _ 16 a in
           checksum_valid
             (bin_measure (mappend b1 (mappend b2 b3)))
             (mappend (mappend b1 (mappend b2 b3)) ext))
  : CorrectAlignedEncoder
      (format_B ThenChecksum checksum_valid OfSize 16 ThenCarryOn format_A)
      (AlignedEncoderForChecksum n1 encode_B encode_A f).
Proof.
  unfold AlignedEncoderForChecksum.
  eapply CorrectAlignedEncoderForChecksum; intros;
    eauto using (EncodeMEquivAlignedEncodeM_word_const (n:=2));
    simpl in H; match goal with
                | H : Some (?a, _) = Some (?b, _) |- _ => replace b with a in * by congruence
                | H : Some _ = None |- _ => discriminate H
                end.
  - apply encode_word'_padding.
  - rewrite length_encode_word'. reflexivity.
  - unfold format_checksum in *. eauto.
Qed.

Lemma ByteBuffer_checksum_bound_tail
      {sz}
  : forall (v : ByteBuffer.t sz) (idx : nat) (v' : ByteBuffer.t idx),
    InternetChecksum.ByteBuffer_checksum_bound sz (v++v') =
    InternetChecksum.ByteBuffer_checksum_bound sz v.
Proof.
  simpl. intros.
  rewrite <- !InternetChecksum_To_ByteBuffer_Checksum'.
  f_equal. rewrite build_aligned_ByteString_append.
  replace (sz * 8) with (bin_measure (build_aligned_ByteString v)).
  rewrite ByteString2ListOfChar_Over; eauto.
  simpl. unfold length_ByteString. simpl. omega.
Qed.

Lemma padding_append_0
      b1 b2
  : padding b1 = 0 -> padding b2 = 0 -> padding (mappend b1 b2) = 0.
Proof.
  intros.
  rewrite padding_ByteString_enqueue_aligned_ByteString; eauto.
Qed.

Lemma refine_ret_ret_eq
      {A} (a b : A)
  : refine (ret a) (ret b) -> a = b.
Proof.
  eauto using Return_inv.
Qed.

Lemma checksum_app_0
  : forall l, checksum (wzero 8 :: wzero 8 :: l)%list = checksum l.
Proof.
  intros. replace (wzero 8 :: (wzero 8 :: l))%list with ([wzero 8; wzero 8] ++ l)%list by reflexivity.
  rewrite checksum_split; simpl; eauto.
  exists 1. reflexivity.
Qed.

Lemma ByteString2ListOfChar_format_checksum
      (w : word 16)
  : let b := format_checksum _ _ _ _ w in
    ByteString2ListOfChar (bin_measure b) b = [hi8 w; lo8 w]%list.
Proof.
  simpl.
  assert (refine (format_word w ()) (ret (build_aligned_ByteString ([hi8 w; lo8 w]), ()))). {
    etransitivity.
    Focus 2.
    (* 2 : { *)
      pose proof AlignedFormat2Char.
      apply H; eauto. higher_order_reflexivity.
    (* } *)
    autorewrite with monad laws.
    higher_order_reflexivity.
  }
  unfold format_word, format_checksum in *. simpl in *.
  apply refine_ret_ret_eq in H. injections. rewrite H0.
  rewrite ByteString2ListOfChar_eq'; eauto.
  Unshelve.
  simpl. exact ().
Qed.

Local Arguments Nat.modulo : simpl never.
Lemma ByteString_mod_16_padding
  : forall b, bin_measure b mod 16 = 0 -> padding b = 0.
Proof.
  intros.
  rewrite padding_eq_mod_8. simpl in *.
  apply Nat.mod_divides; eauto.
  apply Nat.mod_divides in H; eauto.
  destruct H as [x ?]. rewrite H. clear.
  exists (2*x). omega.
Qed.

Fixpoint calculate_aligned_IPchecksum''
         m
         {sz}
         (v : ByteBuffer.t sz) (idx : nat)
         {struct idx}
  := match idx with
     | 0 =>
       InternetChecksum.ByteBuffer_checksum_bound m v
     | S idx' =>
       match v with
       | Vector.cons _ _ v' => calculate_aligned_IPchecksum'' m v' idx'
       | _ => wzero _
       end
     end.

Require AlignedByteBuffer.

Definition calculate_aligned_IPchecksum'
         m
         {sz}
         (v : ByteBuffer.t sz) (idx : nat)
  := match idx with
     | 0 =>
       InternetChecksum.ByteBuffer_checksum_bound m v
     | S idx' =>
       InternetChecksum.ByteBuffer_checksum_bound m (ByteBuffer.drop idx v)
     end.

  (* := match idx with *)
  (*    | 0 => *)
  (*      InternetChecksum.ByteBuffer_checksum_bound m v *)
  (*    | S idx' => *)
  (*      let (_, v') := (AlignedByteBuffer.bytebuffer_of_bytebuffer_range idx (sz - idx) v) in *)
  (*      InternetChecksum.ByteBuffer_checksum_bound m v' *)
  (*    end. *)

Lemma calculate_aligned_IPchecksum_eq
      m
      {sz}
      (v : ByteBuffer.t sz) (idx : nat)
  : calculate_aligned_IPchecksum' m v idx = calculate_aligned_IPchecksum'' m v idx.
Proof.
  unfold calculate_aligned_IPchecksum'.
  revert v m. revert sz.
  induction idx; simpl; intros; try easy.
  destruct v; simpl;
    rewrite <- Eqdep_dec.eq_rect_eq_dec; try apply Nat.eq_dec.
  - destruct m. reflexivity. destruct m; reflexivity.
  - destruct idx; eauto.
    assert (n = n - 0) as L by omega.
    replace (ByteBuffer.drop 0 v) with (eq_rect _ _ v _ L).
    destruct L. simpl. reflexivity.
    simpl.
    f_equal. apply Eqdep_dec.UIP_dec.
    apply Nat.eq_dec.

  (* unfold calculate_aligned_IPchecksum'. *)
  (* revert v. revert sz. revert m. *)
  (* induction idx; simpl; intros; try easy. *)
  (* Local Opaque to_list. *)
  (* destruct v; simpl. *)
  (* - destruct m. reflexivity. *)
  (*   unfold ByteBuffer_checksum_bound, ByteBuffer_fold_left16. simpl. destruct m; reflexivity. *)
  (* - rewrite <- IHidx. destruct idx; simpl. *)
  (*   + f_equal. *)
  (*     replace (to_list (h :: v)) with (h :: (to_list v))%list by reflexivity. *)
  (*     replace (n - 0) with (length (to_list v)). *)
  (*     rewrite firstn_all. *)
  (*   + f_equal. *)
Qed.

Lemma calculate_aligned_IPchecksum_front
      m
      {sz}
  : forall (v : ByteBuffer.t sz) (idx : nat) (v' : ByteBuffer.t idx),
    calculate_aligned_IPchecksum' m (v'++v) idx =
    calculate_aligned_IPchecksum' m v 0.
Proof.
  intros.
  rewrite calculate_aligned_IPchecksum_eq.
  induction v'; eauto.
Qed.

Lemma calculate_aligned_IPchecksum_tail
      {sz}
  : forall (v : ByteBuffer.t sz) (idx : nat) (v' : ByteBuffer.t idx),
    calculate_aligned_IPchecksum' sz (v++v') 0 =
    calculate_aligned_IPchecksum' sz v 0.
Proof.
  intros. simpl.
  rewrite ByteBuffer_checksum_bound_tail. reflexivity.
Qed.

Definition calculate_aligned_IPchecksum
         m
         {sz}
         (v : ByteBuffer.t sz) (idx : nat)
  := wnot (calculate_aligned_IPchecksum' m v idx).

Lemma CorrectAlignedEncoderForIPChecksumThenC
      {S}
      (format_A format_B : FormatM S ByteString)
      (encode_A : forall sz, AlignedEncodeM sz)
      (encode_B : forall sz, AlignedEncodeM sz)
      (encoder_B_OK : CorrectAlignedEncoder format_B encode_B)
      (encoder_A_OK : CorrectAlignedEncoder format_A encode_A)
      n1
      (format_B_sz_OK' : forall (s : S) (b : ByteString) (env env' : CacheFormat),
          format_B s env ∋ (b, env') -> bin_measure b = n1)
      (format_B_sz_OK : n1 mod 16 = 0)
      (len_format_A : S -> nat)
      (format_A_sz_OK' : forall (s : S) (b : ByteString) (env env' : CacheFormat),
          format_A s env ∋ (b, env') -> bin_measure b = len_format_A s)
      (format_A_sz_OK : forall (s : S), len_format_A s mod 8 = 0)
  : CorrectAlignedEncoder
      (format_B ThenChecksum IPChecksum_Valid OfSize 16 ThenCarryOn format_A)
      (AlignedEncoderForChecksum n1 encode_B encode_A calculate_aligned_IPchecksum).
Proof.
  eapply CorrectAlignedEncoderForChecksum';
    unfold calculate_aligned_IPchecksum; eauto; intros.
  - rewrite calculate_aligned_IPchecksum_front.
    rewrite calculate_aligned_IPchecksum_tail. reflexivity.
  - destruct_conjs.
    assert (bin_measure b1 mod 16 = 0) as L1 by (erewrite format_B_sz_OK'; eauto).
    assert (padding b3 = 0) as L2 by (rewrite padding_eq_mod_8; erewrite format_A_sz_OK'; eauto).
    unfold f_bit_aligned_free, IPChecksum_Valid, onesComplement. simpl in *.
    rewrite <- InternetChecksum_To_ByteBuffer_Checksum'.
    rewrite ByteString2ListOfChar_Over.
    rewrite !build_aligned_ByteString_byteString_idem.
    match goal with
    | |- context [ByteString2ListOfChar (numBytes ?a * 8) ?a] =>
      replace (numBytes a * 8) with (bin_measure a)
    end.
    2 : rewrite length_ByteString_no_padding; try omega.
    rewrite !ByteString2ListOfChar_append.
    rewrite !ByteString2ListOfChar_format_checksum.

    assert (forall n, n mod 16 = 0 ->
                 forall b, exists x, |ByteString2ListOfChar n b| = 2 * x) as L. {
      clear. intros.
      apply Nat.mod_divides in H; eauto. destruct H as [x ?].
      replace (16 * x) with (8 * (2*x)) in * by omega.
      subst.
      rewrite ByteString2ListOfChar_len. eauto.
    }
    destruct L with (n:=bin_measure b1) (b:=b1) as [x1 ?]; eauto.
    clear L.
    set (ByteString2ListOfChar (bin_measure b1) b1) as l1 in *.
    set (ByteString2ListOfChar (bin_measure b3) b3) as l3 in *.

    unfold hi8, lo8.
    rewrite split1_wzero, split2_wzero.
    match goal with
    | |- context [wnot (checksum ?a)] =>
      assert (checksum a = checksum (l1 ++ l3))%list as L
    end. {
      rewrite !checksum_split; eauto. exists 1. reflexivity.
    } rewrite L. clear L.
    rewrite checksum_split; eauto.
    rewrite checksum_split; eauto.
    match goal with
    | |- context [?a ^1+ (?b ^1+ ?c)] =>
        assert (a ^1+ (b ^1+ c) = b ^1+ a ^1+ c) as L
    end. {
      rewrite <- !OneC_plus_assoc.
      rewrite OneC_plus_comm.
      rewrite <- !OneC_plus_assoc.
      f_equal.
      rewrite OneC_plus_comm.
      reflexivity.
    } rewrite L. clear L.
    rewrite <- !checksum_split; eauto.
    apply checksum_correct.
    simpl. exists (1+x1). unfold Core.char in *. rewrite H7. omega.
    simpl. exists 1. omega.
    simpl. exists 1. omega.
    all : destruct_conjs; simpl.
    all : repeat match goal with
                 | |- padding (ByteString_enqueue_ByteString _ _) = _ =>
                   rewrite !padding_ByteString_enqueue_aligned_ByteString
                 | |- padding (format_checksum _ _ _ _ _) = _ =>
                   unfold format_checksum; rewrite encode_word'_padding
                 | _ => eauto using ByteString_mod_16_padding
                 end.
Qed.

Definition calculate_aligned_Pseudochecksum'
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (udpLength : word 16)
           (protoCode : word 8)
           (m : nat)
           {sz} (v : ByteBuffer.t sz) (idx : nat) :=
  ByteBuffer_checksum srcAddr ^1+
  ByteBuffer_checksum destAddr ^1+
  zext protoCode 8 ^1+
  udpLength ^1+
  calculate_aligned_IPchecksum' m v idx.

Definition calculate_aligned_Pseudochecksum''
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (udpLength : word 16)
           (protoCode : word 8)
           (m : nat)
           {sz} (v : ByteBuffer.t sz) :=
  calculate_aligned_IPchecksum' (12 + m)
                                ([wzero 8; protoCode] ++ srcAddr ++ destAddr ++
                                         (splitLength udpLength) ++ v) 0.

Lemma calculate_aligned_Pseudochecksum_equiv :
  forall (srcAddr : ByteBuffer.t 4)
    (destAddr : ByteBuffer.t 4)
    (udpLength : word 16)
    (protoCode : word 8)
    (m : nat)
    {sz} (v : ByteBuffer.t sz),
    calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode m v 0 =
    calculate_aligned_Pseudochecksum'' srcAddr destAddr udpLength protoCode m v.
Proof.
  unfold calculate_aligned_Pseudochecksum', calculate_aligned_Pseudochecksum''.
  intros.
  repeat explode_vector.
  Opaque split1.
  Opaque split2.
  simpl in *.
  unfold ByteBuffer_checksum, InternetChecksum.ByteBuffer_checksum_bound, add_w16_into_checksum,
  add_bytes_into_checksum, ByteBuffer_fold_left16, ByteBuffer_fold_left_pair.
  fold @ByteBuffer_fold_left_pair.
  setoid_rewrite Buffer_fold_left16_acc_oneC_plus.
  rewrite combine_split.
  rewrite !OneC_plus_wzero_r, !OneC_plus_wzero_l, OneC_plus_comm.
  repeat (f_equal; [ ]).
  rewrite <- !OneC_plus_assoc.
  match goal with
  | |- _ = ?a ^1+ ?b =>
    assert (a ^1+ b = b ^1+ a) by (apply OneC_plus_comm)
  end.
  simpl in *. rewrite H.
  rewrite <- !OneC_plus_assoc.
  reflexivity.
Qed.

Lemma calculate_aligned_Pseudochecksum_front
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (udpLength : word 16)
      (protoCode : word 8)
      m
      {sz}
  : forall (v : ByteBuffer.t sz) (idx : nat) (v' : ByteBuffer.t idx),
    calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode m (v'++v) idx =
    calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode m v 0.
Proof.
  intros. unfold calculate_aligned_Pseudochecksum'.
  f_equal. apply calculate_aligned_IPchecksum_front.
Qed.

Lemma calculate_aligned_Pseudochecksum_tail
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (udpLength : word 16)
      (protoCode : word 8)
      {sz}
  : forall (v : ByteBuffer.t sz) (idx : nat) (v' : ByteBuffer.t idx),
    calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode sz (v++v') 0 =
    calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode sz v 0.
Proof.
  intros. unfold calculate_aligned_Pseudochecksum'.
  f_equal. apply calculate_aligned_IPchecksum_tail.
Qed.

Definition calculate_aligned_Pseudochecksum
           (srcAddr : ByteBuffer.t 4)
           (destAddr : ByteBuffer.t 4)
           (udpLength : word 16)
           (protoCode : word 8)
           (m : nat)
           {sz} (v : ByteBuffer.t sz) (idx : nat) :=
  wnot (calculate_aligned_Pseudochecksum' srcAddr destAddr udpLength protoCode m v idx).

Lemma CorrectAlignedEncoderForPseudoChecksumThenC
      {S}
      (srcAddr : ByteBuffer.t 4)
      (destAddr : ByteBuffer.t 4)
      (udpLength : word 16)
      (protoCode : word 8)
      (format_A format_B : FormatM S ByteString)
      (encode_A : forall sz, AlignedEncodeM sz)
      (encode_B : forall sz, AlignedEncodeM sz)
      (encoder_B_OK : CorrectAlignedEncoder format_B encode_B)
      (encoder_A_OK : CorrectAlignedEncoder format_A encode_A)
      n1
      (format_B_sz_OK' : forall (s : S) (b : ByteString) (env env' : CacheFormat),
          format_B s env ∋ (b, env') -> bin_measure b = n1)
      (format_B_sz_OK : n1 mod 16 = 0)
      (len_format_A : S -> nat)
      (format_A_sz_OK' : forall (s : S) (b : ByteString) (env env' : CacheFormat),
          format_A s env ∋ (b, env') -> bin_measure b = len_format_A s)
      (format_A_sz_OK : forall (s : S), len_format_A s mod 8 = 0)
  : CorrectAlignedEncoder
      (format_B ThenChecksum (Pseudo_Checksum_Valid srcAddr destAddr udpLength protoCode) OfSize 16
                ThenCarryOn format_A)
      (AlignedEncoderForChecksum n1 encode_B encode_A
                                 (calculate_aligned_Pseudochecksum srcAddr destAddr udpLength protoCode)).
Proof.
  apply CorrectAlignedEncoderForChecksum';
    unfold calculate_aligned_Pseudochecksum; eauto with nocore; intros.
  - rewrite calculate_aligned_Pseudochecksum_front.
    rewrite calculate_aligned_Pseudochecksum_tail. reflexivity.
  - destruct_conjs.
    assert (bin_measure b1 mod 16 = 0) as L1 by (erewrite format_B_sz_OK'; eauto).
    assert (padding b3 = 0) as L2 by (rewrite padding_eq_mod_8; erewrite format_A_sz_OK'; eauto).
    unfold f_bit_aligned_free, Pseudo_Checksum_Valid, onesComplement.
    rewrite calculate_aligned_Pseudochecksum_equiv.
    unfold calculate_aligned_Pseudochecksum''.
    unfold calculate_aligned_IPchecksum'.
    rewrite <- InternetChecksum_To_ByteBuffer_Checksum'.
    rewrite ByteString2ListOfChar_Over.
    rewrite !build_aligned_ByteString_append.
    rewrite !build_aligned_ByteString_byteString_idem.
    rewrite Nat.mul_add_distr_r.
    match goal with
    | |- context [numBytes ?a * 8] =>
      replace (numBytes a * 8) with (bin_measure a)
    end.
    2 : rewrite length_ByteString_no_padding; try omega.

    match goal with
    | |- context [wnot (checksum (ByteString2ListOfChar ?a ?b))] =>
      assert (a = bin_measure b) as L
    end. {
      rewrite !@mappend_measure.
      rewrite !Nat.add_assoc.
      reflexivity.
    } rewrite L. clear L.
    rewrite !ByteString2ListOfChar_append.
    rewrite !ByteString2ListOfChar_format_checksum.

    assert (forall n, n mod 16 = 0 ->
                 forall b, exists x, |ByteString2ListOfChar n b| = 2 * x) as L. {
      clear. intros.
      apply Nat.mod_divides in H; eauto. destruct H as [x ?].
      replace (16 * x) with (8 * (2*x)) in * by omega.
      subst.
      rewrite ByteString2ListOfChar_len. eauto.
    }
    destruct L with (n:=bin_measure b1) (b:=b1) as [x1 ?]; eauto.
    clear L.
    set (ByteString2ListOfChar (bin_measure b1) b1) as l1 in *.
    set (ByteString2ListOfChar (bin_measure b3) b3) as l3 in *.

    rewrite !ByteString2ListOfChar_eq'.
    simpl Core.byteString. unfold ByteBuffer.to_list.

    unfold hi8, lo8.
    rewrite split1_wzero, split2_wzero.
    match goal with
    | |- checksum (?a :: ?b :: ?l)%list = _ =>
      replace (a :: b :: l)%list with (to_list (a :: [b])%vector ++ l)%list by reflexivity
    end.
    rewrite !app_assoc.
    unfold Core.char in *.
    match goal with
    | |- context [(?a ++ l1)%list] => set (a ++ l1)%list as l1'
    end.
    assert (exists x, | l1' | = 2 * x) as L'. {
      subst l1'. simpl.
      rewrite !app_length. rewrite <- !ByteBuffer.to_list_length. rewrite H7.
      exists (6+x1). omega.
    } destruct L' as [x L'].
    rewrite <- !app_assoc.
    match goal with
    | |- context [wnot (checksum ?a)] =>
      assert (checksum a = checksum (l1' ++ l3))%list as L
    end. {
      rewrite !checksum_split; eauto.
      exists 1. reflexivity.
    } rewrite L. clear L.
    rewrite checksum_split; eauto.
    rewrite checksum_split; eauto.
    match goal with
    | |- context [?a ^1+ (?b ^1+ ?c)] =>
        assert (a ^1+ (b ^1+ c) = b ^1+ a ^1+ c) as L
    end. {
      rewrite <- !OneC_plus_assoc.
      rewrite OneC_plus_comm.
      rewrite <- !OneC_plus_assoc.
      f_equal.
      rewrite OneC_plus_comm.
      reflexivity.
    } rewrite L. clear L.
    rewrite <- !checksum_split; eauto.
    apply checksum_correct.
    rewrite !app_length. rewrite L'. simpl. exists (1+x). omega.
    simpl. exists 1. omega.
    simpl. exists 1. omega.

    all : destruct_conjs; simpl.
    all : repeat match goal with
                 | |- padding (ByteString_enqueue_ByteString _ _) = _ =>
                   rewrite !padding_ByteString_enqueue_aligned_ByteString
                 | |- padding (format_checksum _ _ _ _ _) = _ =>
                   unfold format_checksum; rewrite encode_word'_padding
                 | _ => eauto using ByteString_mod_16_padding
                 end.
Qed.


From iris.proofmode Require Import base proofmode classes.
From transfinite.base_logic.lib Require Export fancy_updates.
From melocoton Require Import stdpp_extra language_commons.
From melocoton.language Require Export language weakestpre.
(* FIXME: If we import iris.bi.weakestpre earlier texan triples do not
   get pretty-printed correctly. *)
From iris.bi Require Export weakestpre.
From iris.prelude Require Import options.
Import uPred.

Section wp.
Context `{SI:indexT, !langG val Λ Σ, !invG Σ}.
Implicit Types P : iProp Σ.
Implicit Types Φ : val → iProp Σ.
Implicit Types v : val.
Implicit Types e : expr Λ.
Implicit Types Ψ : protocol val Σ.
Implicit Types prog : lang_prog Λ.
Implicit Types pe : prog_environ Λ Σ.

Definition program_fulfills
  (Ψin : protocol val Σ) (p : mixin_prog (func Λ)) (Ψshould : protocol val Σ) : iProp Σ :=
  ∀ s vv Φ, Ψshould s vv Φ -∗ ⌜p !! s <> None⌝ ∗ WP (of_class _ (ExprCall s vv)) @ ⟨p, Ψin⟩; ⊤ {{ Φ }}.

Definition env_fulfills
  (p:prog_environ Λ Σ) (Ψshould : protocol val Σ) := program_fulfills p.(penv_proto) p.(penv_prog) Ψshould.

Notation "Ψin '|-' p '::' Ψshould" := (program_fulfills Ψin p Ψshould) (at level 25, p, Ψshould at level 26) : bi_scope.

Definition spec_union (p1 p2 : protocol val Σ) : protocol val Σ :=
  λ s vv Φ, (p1 s vv Φ ∨ p2 s vv Φ)%I.
Definition spec_inters (p1 p2 : protocol val Σ) : protocol val Σ :=
  λ s vv Φ, (p1 s vv Φ ∧ p2 s vv Φ)%I.

Class can_link (Ψ1 Ψ2 Ψaxiom Ψres : protocol val Σ) (p1 p2 p3 : lang_prog Λ) : Prop := {
  p1_p2_disjoint : dom p1 ## dom p2;
  Ψres_is_union : ⊢ (∀ s vv Φ, Ψres s vv Φ -∗ (spec_union Ψ1 Ψ2) s vv Φ)%I;
  Ψaxiom_is_axiomatic : ⊢ (∀ s vv Φ, Ψaxiom s vv Φ -∗ ⌜s ∉ dom p3⌝)%I;
  p1_satisfies_Ψ1 : ⊢ (spec_union Ψ2 Ψaxiom |- p1 :: Ψ1)%I;
  p2_satisfies_Ψ2 : ⊢ (spec_union Ψ1 Ψaxiom |- p2 :: Ψ2)%I;
  p3_is_union : p3 = (union_with (λ _ _ : func Λ, None) p1 p2);
}.

Definition pairwise {X} (A:list X) (Ψ : X -> X -> Prop) :=
  forall i j, i <> j -> match (nth_error A i, nth_error A j) with
      (Some a, Some b) => Ψ a b
     | _ => True end.

Definition spec_union_list (P : list (protocol val Σ)) : protocol val Σ :=
  λ s vv Φ, (∃ i, match nth_error P i with Some pp => pp s vv Φ | _ => ⌜False⌝ end)%I.
Definition spec_union_list_except k (P : list (protocol val Σ)) : protocol val Σ :=
  λ s vv Φ, (∃ i, ⌜i <> k⌝ ∗ match nth_error P i with Some pp => pp s vv Φ | _ => ⌜False⌝ end)%I.

Fixpoint union_map_list (P : list (lang_prog Λ)) : lang_prog Λ := match P with
  nil => ∅
 | x::xr => union_with (λ _ _ : func Λ, None) x (union_map_list xr)
end.

Lemma pairwise_subset {X} H (a : list X) x : pairwise (x::a) H -> pairwise a H.
Proof.
  intros HH i j Hne.
  specialize (HH (S i) (S j)). cbn in HH.
  apply HH. congruence.
Qed.

Lemma union_map_list_spec x y (P:list (lang_prog Λ)) :
  (pairwise P (fun p1 p2 => dom p1 ## dom p2)) ->
  ((union_map_list P) !! x = Some y <-> ∃ k, In k P ∧ k !! x = Some y).
Proof.
  intros Hdis.
  revert x y.
  induction P; intros x y; first split; cbn.
  - rewrite lookup_empty. intros H; congruence.
  - intros (? & [] & ?).
  - rewrite lookup_union_with.
    destruct (a !! x) as [va|] eqn:Hva;
    destruct (union_map_list P !! x) as [vP|] eqn:HvP; cbn; split; try congruence.
    + intros _. apply IHP in HvP; last by eapply pairwise_subset.
      destruct HvP as (k & (pos & Hpos)%In_nth_error & Heq).
      specialize (Hdis 0 (S pos)); cbn in Hdis. rewrite Hpos in Hdis.
      exfalso. assert (0 ≠ S pos) as Hne by lia. specialize (Hdis Hne).
      rewrite  elem_of_disjoint in Hdis. eapply Hdis.
      all: by eapply elem_of_dom_2.
    + injection 1; intros ->. exists a; repeat split; by try left.
    + intros (k & [-> | Hk] & Heq); first congruence.
      exfalso. assert (∃ k : lang_prog Λ, In k P ∧ k !! x = Some y) as H by by exists k.
      apply IHP in H; first congruence. by eapply pairwise_subset.
    + apply IHP in HvP; last by eapply pairwise_subset.
      destruct HvP as (k & ? & ?).
      injection 1; intros ->. exists k; repeat split; eauto.
    + intros (k & [-> | Hk] & Heq); first congruence.
      assert (∃ k : lang_prog Λ, In k P ∧ k !! x = Some y) as H by by exists k.
      apply IHP in H; last by eapply pairwise_subset. congruence.
    + intros (k & [-> | Hk] & Heq); first congruence.
      assert (∃ k : lang_prog Λ, In k P ∧ k !! x = Some y) as H by by exists k.
      apply IHP in H; last by eapply pairwise_subset. congruence.
Qed.

Lemma union_map_subset A p :
  (pairwise A (fun p1 p2 => dom p1 ## dom p2)) -> p ∈ A → p ⊆ union_map_list A.
Proof.
  intros Hdis H. apply map_subseteq_spec.
  intros i x Hin.
  apply union_map_list_spec; first done.
  exists p. split; eauto. by apply elem_of_list_In.
Qed.

Class can_link_all
    (Ψaxiom Ψres : protocol val Σ) (pres : lang_prog Λ)
    (A : list (prod (protocol val Σ) (lang_prog Λ))) := {
  all_disjoint : pairwise (map snd A) (fun p1 p2 => dom p1 ## dom p2);
  Ψres_is_big_union : ⊢ (∀ s vv Φ, Ψres s vv Φ -∗ (spec_union_list (map fst A)) s vv Φ)%I;
  Ψaxiom_is_axiomatic_all : ⊢ (∀ s vv Φ, Ψaxiom s vv Φ -∗ ⌜s ∉ dom pres⌝)%I;
  pres_is_union : pres = union_map_list (map snd A);
  one_spec := fun i => spec_union (spec_union_list_except i (map fst A)) Ψaxiom;
  all_satisfy_spec : ∀ i, match (nth_error A i) with None => True |
      Some (Ψi, pi) =>  ⊢ (one_spec i |- pi :: Ψi)%I end
}.

#[global]
Instance can_link_can_link_all Ψaxiom Ψres pres Ψ1 Ψ2 p1 p2 : can_link Ψ1 Ψ2 Ψaxiom Ψres p1 p2 pres
  -> can_link_all Ψaxiom Ψres pres [(Ψ1,p1); (Ψ2,p2)].
Proof.
  intros [H1 H2 H3 H4 H5 H6]; split; cbn.
  - intros [|[|[|i]]] [|[|[|j]]] H; cbn; easy.
  - iIntros (s vv Φ) "Hres". iPoseProof H2 as "H2".
    iDestruct ("H2" $! s vv Φ with "Hres") as "[H2'|H2']".
    + iExists 0; cbn; done.
    + iExists 1; cbn; done.
  - done.
  - rewrite H6. apply map_eq. intros i. rewrite ! lookup_union_with. rewrite lookup_empty.
    destruct (p1 !! i), (p2 !! i); done.
  - intros [|[|[|i]]]; cbn; try done.
    all: iIntros (s vv Φ) "Hres"; cbn.
    1: iPoseProof H4 as "H"; iSpecialize ("H" $! s vv Φ with "Hres").
    2: iPoseProof H5 as "H"; iSpecialize ("H" $! s vv Φ with "Hres").
    all: iDestruct "H" as "[$ HWP]".
    all: iApply wp_proto_mono; last iApply "HWP".
    all: cbn; intros s' vv' Φ'; iIntros "[H|H]".
    2,4: by iRight. all: iLeft.
    1: iExists 1. 2: iExists 0.
    all: iSplitR; first done. all: done.
Qed.

Lemma wp_link_execs Ψaxiom Ψres (pres : gmap string (Λ.(func))) A :
  can_link_all Ψaxiom Ψres pres A
  -> ⊢ ∀ e Φ i, match (nth_error A i) with None => ⌜False⌝ |
          Some (Ψi, pi) => WP e @ ⟨pi, spec_union (spec_union_list_except i (map fst A)) Ψaxiom ⟩; ⊤ {{ Φ }} end
     -∗ WP e @ ⟨pres, Ψaxiom⟩; ⊤ {{ Φ }}.
Proof.
  intros [Hdis HΨres Haxiom -> one_spec' Hsatis].
  iLöb as "IHe". iIntros (e Φ i).
  destruct (nth_error A i) as [[Ψi pi]|] eqn:Heq; last (iIntros "%H"; done).
  rewrite !wp_unfold /wp_pre /=.
  iIntros "H %σ Hσ".
  - iSpecialize ("H" $! σ with "Hσ").
    iMod "H".
    iDestruct "H" as "[(%x & -> & Hσ & H)|[(%s' & %vv' & %K & %HeqK & %H2 & >(%Ξ & Hσ & [HΨ|HΨ] & H3))|(%HH & H3)]]".
    * iModIntro. iLeft. iExists _. iFrame. iPureIntro. done.
    * iDestruct "HΨ" as "(%kidx & %Hknei & HΨ)". rewrite nth_error_map.
      destruct (nth_error A kidx) as [[Ψc pc]|] eqn:Heqk; cbn; last iPure "HΨ" as [].
      specialize (Hsatis kidx). rewrite Heqk in Hsatis.
      iPoseProof (Hsatis) as "Hsatis".
      iDestruct ("Hsatis" $! s' vv' Ξ with "HΨ")as "(%HNone & HΨ)".
      rewrite wp_unfold /wp_pre /=.
      iSpecialize ("HΨ" $! σ with "Hσ").
      iMod "HΨ" as "[(%x & %Heqx & Hσ & H)|[(%s'2 & %vv'2 & %K'2 & %HeqK2 & %H''2 & >(%Ξ' & Hσ & HΨ' & H3'))|(%HH2&H3')]]".
      -- exfalso. apply of_class_inj in Heqx. congruence.
      -- exfalso. assert (K'2 = empty_ectx) as ->.
         2: { rewrite fill_empty in HeqK2. apply of_class_inj in HeqK2. assert (s'2 = s') as -> by congruence; congruence. }
         destruct (fill_class K'2 (of_class Λ (ExprCall s'2 vv'2))) as [H1|[x Hx]].
         { eexists. rewrite <- HeqK2. by rewrite to_of_class. }
        { done. }
          { exfalso. unfold to_val in Hx. rewrite to_of_class in Hx; congruence. }
     -- destruct HH2 as (e' & σ' & (KK & e1' & e2' & Heq2 & -> & Hstep)%prim_step_inv).
          destruct (fill_class KK e1') as [->|[x Hx]].
         { eexists. rewrite <- Heq2. by rewrite to_of_class. }
         2: { apply val_head_stuck in Hstep. congruence. }
         rewrite ! fill_empty in Heq2. subst e1'.
         apply call_head_step in Hstep.
         destruct Hstep as (Fn & Heqp2' & He2' & -> ). iRight. iRight.
         assert (union_map_list (map snd A) !! s' = Some Fn) as Heqs'.
         { apply union_map_list_spec; first done. exists pc; repeat split; try done.
           apply in_map_iff. exists ((Ψc,pc)); repeat split. by eapply nth_error_In. }
         iModIntro. iSplit; first iPureIntro.
         { do 2 eexists. rewrite HeqK. econstructor; first done. 1:done.
           apply call_head_step. exists Fn. repeat split; try done. }
         subst e. iIntros (? ? (e'2 & -> & (Fn3 & Heqp3' & He3' & ->)%call_head_step)%head_reducible_prim_step_ctx).
         2: { do 2 eexists.
              apply call_head_step. exists Fn. repeat split; try done. }
         rewrite Heqs' in Heqp3'. assert (e'2 = e2') as -> by congruence.
         iSpecialize ("H3'" $! σ e2').
         assert (prim_step pc (of_class Λ (ExprCall s' vv')) σ e2' σ) as Hstep2.
         { apply head_prim_step. apply call_head_step. eexists; repeat split; done. }
         iSpecialize ("H3'" $! Hstep2). iMod "H3'". iModIntro. iNext.
         iMod "H3'" as "(? & H3')". iModIntro. iFrame.
         iApply wp_bind. iApply (wp_wand with "[H3']").
         { iApply ("IHe" $! _ _ kidx). rewrite Heqk. iApply "H3'". }
         iIntros (r) "Hr". iSpecialize ("H3" $! r with "Hr"). iApply ("IHe" $! _ _ i). rewrite Heq. done.
    * iRight. iLeft. iModIntro. do 3 iExists _; iSplitR; first done.
      destruct (union_map_list (map snd A) !! s') as [v|] eqn:Heqv.
      { iExFalso. iDestruct (Haxiom $! s' vv' Ξ with "HΨ") as "%Hfalse". exfalso.
        apply Hfalse. eapply elem_of_dom_2. done. }
      cbn. iSplitR; first done. iModIntro. iExists Ξ. iFrame. iNext.
      iIntros (r) "Hr". iSpecialize ("H3" with "Hr"). iApply ("IHe" $! _ _ i). by rewrite Heq.
    * iModIntro. iRight. iRight.
      iSplitR.
      { iPureIntro. eapply reducible_mono; last done.
        apply union_map_subset; first done. apply elem_of_list_In.
        apply in_map_iff. exists ((Ψi, pi)); split; eauto.
        by eapply nth_error_In. }
      iIntros (σ' e' Hstep%prim_step_inv).
      iSpecialize ("H3" $! σ' e').
      destruct Hstep as (K & e1' & e2' & -> & -> & H).
      assert (prim_step pi (fill K e1') σ (fill K e2') σ') as Hstep.
      { econstructor. 1-2:done. destruct (to_class e1') as [[]|] eqn:Heqe1.
        - exfalso. eapply val_head_step. apply of_to_class in Heqe1. erewrite Heqe1. done.
        - apply of_to_class in Heqe1. subst e1'.
          destruct HH as (? & σ'2 & (e2'' & -> & (K2 & eK1 & eK2 & Heq1 & -> & HHstep)%prim_step_inv)%fill_step_inv).
          2: unfold to_val; by rewrite to_of_class.
          assert (K2 = empty_ectx) as ->.
          1: { destruct (fill_class' K2 eK1) as [|[vv Hvv]].
               1: eexists; rewrite <- Heq1; by rewrite to_of_class.
               1: done.
               exfalso. apply of_to_class in Hvv. subst eK1. eapply val_head_step. done. }
          rewrite fill_empty in Heq1. subst eK1. apply call_head_step in H, HHstep.
          destruct H as (FN1 & HFN1 &Happy1 & ->).
          destruct HHstep as (FN2 & HFN2 &Happy2 & ->).
          apply union_map_list_spec in HFN1; last done. destruct HFN1 as (? & [[ΨΨ kk] [<- (kki & Hkki)%In_nth_error]]%in_map_iff & HFN1).
          cbn in HFN1. destruct (decide (kki = i)) as [-> | Hcontr].
          1: apply call_head_step; exists FN2; repeat split; try congruence.
          exfalso. specialize (Hdis kki i Hcontr). rewrite ! nth_error_map in Hdis.
          rewrite Hkki in Hdis. rewrite Heq in Hdis. cbn in Hdis. rewrite elem_of_disjoint in Hdis.
          eapply Hdis; eapply elem_of_dom_2; done.
        - by eapply head_step_no_call. }
      iDestruct ("H3" $! Hstep) as ">H3". iModIntro. iNext.
      iMod "H3" as "(Hσ & HWP)". iModIntro. iFrame. iApply ("IHe" $! _ _ i). rewrite Heq. done.
Qed.

Lemma wp_link_progs Ψaxiom Ψres (pres : gmap string (Λ.(func))) A :
  can_link_all Ψaxiom Ψres pres A
 -> ⊢ Ψaxiom |- pres :: Ψres.
Proof.
  intros [Hdis HΨres Haxiom -> one_spec' Hsatis].
  iIntros (s vv Φ) "H". iPoseProof HΨres as "HΨres".
    iDestruct ("HΨres" with "H") as "[%x Hx]".
    erewrite nth_error_map.
    destruct (nth_error A x) as [[Ψi pi]|] eqn:Heq; last by iExFalso. cbn.
    specialize (Hsatis x) as Hsatis2.
    rewrite Heq in Hsatis2. iDestruct (Hsatis2 $! s vv Φ with "Hx") as "[%Hx1 Hx2]".
    destruct (union_map_list (map snd A) !! s) eqn:Hl2.
  - iSplitR; first done. unshelve iApply (wp_link_execs _ _ _ _ $! _ _ x). 3: by split. rewrite Heq. done.
  - exfalso. destruct (pi !! s) as [f|] eqn:Heqpi; try congruence.
    assert (∃ k, In k ((map snd A)) ∧ k !! s = Some f) as HH.
    + exists pi; repeat split; try done. apply in_map_iff. exists (Ψi,pi).
      repeat split; try done. eapply nth_error_In. done.
    + apply union_map_list_spec in HH; try congruence. done.
Qed.

End wp.


Require Import String.
Require Import Coq.ZArith.ZArith.
Require Import coqutil.Z.Lia.
Require Import Coq.Lists.List. Import ListNotations.
Require Import Kami.Lib.Word.
Require Import Kami.Ex.IsaRv32 riscv.Spec.Decode.
Require Import riscv.Utility.Encode.
Require Import coqutil.Word.LittleEndian.
Require Import coqutil.Word.Properties.
Require Import coqutil.Map.Interface.
Require Import coqutil.Tactics.Tactics.
Require Import coqutil.Tactics.rdelta.
Require Import processor.KamiWord.
Require Import riscv.Utility.Utility.
Require Import riscv.Utility.runsToNonDet.
Require Import riscv.Spec.Primitives.
Require Import riscv.Spec.MetricPrimitives.
Require Import riscv.Spec.Machine.
Require riscv.Platform.Memory.
Require Import riscv.Spec.PseudoInstructions.
Require Import riscv.Proofs.EncodeBound.
Require Import riscv.Proofs.DecodeEncode.
Require Import riscv.Platform.Run.
Require Import riscv.Utility.MkMachineWidth.
Require Import riscv.Utility.Monads. Import MonadNotations.
Require Import coqutil.Datatypes.PropSet.
Require Import riscv.Platform.RiscvMachine.
Require Import riscv.Platform.MetricRiscvMachine.
Require Import riscv.Platform.MinimalMMIO.
Require Import riscv.Platform.MetricMinimalMMIO.
Require Import riscv.Platform.FE310ExtSpec.

Require Import Kami.Syntax Kami.Semantics Kami.Tactics.
Require Import Kami.Ex.MemTypes Kami.Ex.SC Kami.Ex.SCMMInl Kami.Ex.SCMMInv.

Require Export processor.KamiProc.
Require Import processor.Consistency.
Require Import processor.KamiRiscvStep.

Lemma get_of_list_not_In:
  forall (key: Type) (key_dec: forall k1 k2: key, {k1 = k2} + {k1 <> k2})
         (value: Type) (map: map.map key value),
    map.ok map ->
    forall (l: list (key * value)) k,
      ~ In k (List.map fst l) ->
      map.get (map.of_list l) k = None.
Proof.
  induction l as [|[k v] l]; simpl; intros;
    [rewrite map.get_empty; reflexivity|].
  destruct (key_dec k k0).
  - intuition idtac.
  - rewrite map.get_put_diff by auto.
    apply IHl; intuition idtac.
Qed.

Lemma alignedXAddrsRange_zero_bound_in:
  forall n a,
    (wordToN a < N.of_nat n)%N -> In a (alignedXAddrsRange 0 n).
Proof.
  induction n; [blia|].
  intros.
  assert (wordToN a = N.of_nat n \/ wordToN a < N.of_nat n)%N by blia.
  clear H; destruct H0.
  - unfold alignedXAddrsRange; fold alignedXAddrsRange.
    left; apply wordToN_inj.
    rewrite H.
    change (0 + n)%nat with n.
    pose proof (wordToN_bound a); rewrite H in H0.
    rewrite <-wordToN_NToWord_2 with (sz:= Z.to_nat width) (n:= N.of_nat n) by assumption.
    rewrite NToWord_nat, Nnat.Nat2N.id; reflexivity.
  - right; auto.
Qed.

Section Equiv.
  Context {Registers: map.map Z word}
          {mem: map.map word byte}.

  Local Notation M := (free action result).
  Local Notation RiscvMachine := MetricRiscvMachine.
  Local Existing Instance MetricMinimalMMIO.IsRiscvMachine.
  Local Existing Instance MetricMinimalMMIOSatisfiesPrimitives.

  (** * Processor, software machine, and states *)

  Variable (instrMemSizeLg memSizeLg: Z).
  Hypotheses (Hinstr1: 3 <= instrMemSizeLg)
             (Hinstr2: instrMemSizeLg <= width - 2)
             (Hkmem1: 2 + instrMemSizeLg < memSizeLg)
             (Hkmem2: memSizeLg <= width)
             (* 16 used to be disjoint to MMIO addresses.
              * [Hkmem2] is meaningless assuming this [Hkmemdisj]
              * but still having that in context ease some proofs. *)
             (Hkmemdisj: memSizeLg <= 16).
  Local Notation Hinstr := (conj Hinstr1 Hinstr2).

  Variable (memInit: Vec (ConstT (Bit BitsPerByte)) (Z.to_nat memSizeLg)).
  Definition kamiMemInit := ConstVector memInit.
  Local Definition kamiProc :=
    @KamiProc.proc instrMemSizeLg memSizeLg Hinstr kamiMemInit
                   (kami_AbsMMIO (Z.to_N memSizeLg)).

  (* redefine mcomp_sat to simplify for the case where no answer is returned *)
  Local Notation mcomp_sat_unit m initialL post :=
    (mcomp_sat m initialL (fun (_: unit) => post)).

  Context (Registers_ok: map.ok Registers)
          (mem_ok: map.ok mem).

  Local Notation states_related :=
    (@states_related Registers mem instrMemSizeLg memSizeLg Hinstr1 Hinstr2).
  Local Notation kamiStep :=
    (@kamiStep _ _ Hinstr1 Hinstr2 memInit).

  Arguments Z.add: simpl never.

  Lemma KamiLabelR_unique: forall {klbl t t'},
      KamiLabelR klbl t ->
      KamiLabelR klbl t' ->
      t = t'.
  Proof.
    intros. inversion H; inversion H0; clear H H0; subst.
    + reflexivity.
    + rewrite H4 in H1. exfalso. eapply FMap.M.add_empty_neq. eassumption.
    + rewrite H1 in H5. exfalso. eapply FMap.M.add_empty_neq. eassumption.
    + simpl in *. rewrite H1 in H5.
      apply (f_equal (FMap.M.find "mmioExec"%string)) in H5.
      do 2 rewrite FMap.M.find_add_1 in H5.
      apply Option.eq_of_eq_Some in H5.
      apply Eqdep.EqdepTheory.inj_pair2 in H5.
      inversion H5; subst; clear H5.
      reflexivity.
  Qed.

  Inductive KamiLabelSeqR: list LabelT -> list Event -> Prop :=
  | KamiSeqNil: KamiLabelSeqR nil nil
  | KamiSeqCons:
      forall klseq t,
        KamiLabelSeqR klseq t ->
        forall klbl nt,
          KamiLabelR klbl nt ->
          KamiLabelSeqR (klbl :: klseq) (nt ++ t).

  (** * Rephrasing single step soundness in a more readable way *)

  Definition KState: Type := KamiMachine * list LabelT.

  Inductive kstep1: KState -> KState -> Prop :=
  | KStep: forall (m1: KamiMachine) (t1: list LabelT) kupd klbl,
      Step kamiProc m1 kupd klbl ->
      kstep1 (m1, t1) (FMap.M.union kupd m1, klbl :: t1).

  Definition CState: Type := RiscvMachine.

  Definition cstep1(m1: CState)(P: CState -> Prop): Prop := mcomp_sat_unit (run1 iset) m1 P.

  Inductive related: KState -> CState -> Prop :=
  | Related: forall km kt eventTrace rm,
      Kami.Semantics.reachable km kamiProc ->
      states_related (km, eventTrace) rm ->
      KamiLabelSeqR kt eventTrace ->
      related (km, kt) rm.

  Theorem kstep1_sound: forall ks1 ks2 rs1 P,
      related ks1 rs1 ->
      kstep1 ks1 ks2 ->
      cstep1 rs1 P ->
      related ks2 rs1 \/
      exists rs2, related ks2 rs2 /\ P rs2.
  Proof.
    intros.
    pose proof (KamiRiscvStep.kamiStep_sound instrMemSizeLg memSizeLg Hinstr1 Hinstr2
      Hkmem1 Hkmem2 Hkmemdisj memInit Registers_ok mem_ok) as Q.
    unfold cstep1, kamiStep in *.
    inversion H. inversion H0. subst. inversion H8. subst. clear H H0 H8.
    specialize Q with (1 := H2) (3 := H3) (4 := H1).
    edestruct Q as [A | A]; clear Q.
    - eauto.
    - left. destruct A as [Q1 Q2]. econstructor.
      + unfold reachable in *.
        destruct H2 as [sigma H2]. inversion H2. subst; clear H2.
        eexists. constructor. eapply Multi; eassumption.
      + eassumption.
      + inversion H4; subst; clear H4.
        * change (KamiLabelSeqR [klbl] ([] ++ [])).
          eapply KamiSeqCons. 1: exact KamiSeqNil.
          eapply KamiSilent. assumption.
        * inversion Q1. clear Q1. subst.
          change (nt ++ t) with ([] ++ nt ++ t).
          repeat eapply KamiSeqCons; try assumption. constructor. assumption.
    - right. destruct A as [rs2 [tNew' [A [B C]]]].
      exists rs2. split. 2: exact C.
      econstructor. 2: exact B.
      + unfold reachable in *.
        destruct H2 as [sigma H2]. inversion H2. subst; clear H2.
        eexists. constructor. eapply Multi; eassumption.
      + eapply KamiSeqCons; eassumption.
  Qed.

  (** * Multistep and Behavior soundness *)

  Inductive ksteps: KState -> KState -> Prop :=
  | KSteps: forall m1 t1 m2 tNew,
      Multistep kamiProc m1 m2 tNew ->
      ksteps (m1, t1) (m2, tNew ++ t1).

  (* Can't use this one here because we're not doing a simulation from Kami to compiler *)
  Definition csteps: CState -> (CState -> Prop) -> Prop := runsTo cstep1.

  Definition calways{State: Type}(Step: State -> (State -> Prop) -> Prop)(s: State)(P: State -> Prop): Prop :=
    P s /\ forall s', P s' -> Step s' P.

  (* Can't use this one here because we don't really have a `kstep: KState -> KState -> Prop` that
     we lift with *, but Kami requires its own ksteps (its own "star") *)
  Definition kalways0{State: Type}(Step: State -> State -> Prop)(s: State)(P: State -> Prop): Prop :=
    P s /\ forall s1 s2, P s1 -> Step s1 s2 -> P s2.

  Definition kalways{State: Type}(Stepstar: State -> State -> Prop)(s: State)(P: State -> Prop): Prop :=
    forall s', Stepstar s s' -> P s'.

  Theorem ksteps_sound: forall (inv: CState -> Prop) ks1 rs1,
      related ks1 rs1 ->
      calways cstep1 rs1 inv ->
      kalways ksteps ks1 (fun ks2 => exists rs2, related ks2 rs2 /\ inv rs2).
  Proof.
    unfold calways, kalways.
    intros. destruct H0. inversion H1. subst. clear H1. revert H3 rs1 t1 H0 H.
    induction 1; intros.
    - subst. exists rs1. split; assumption.
    - specialize IHMultistep with (1 := H0) (2 := H).
      destruct IHMultistep as [rs2 [A B]].
      edestruct kstep1_sound.
      + exact A.
      + econstructor. eassumption.
      + eapply H2. exact B.
      + eauto.
      + eauto.
  Qed.

  Definition KamiImplMachine: Type := RegsT.

  (* When running the processor on an FPGA, this memory module will be implemented in some
     trusted Verilog code, and will forward requests either to a DRAM or to a source
     from which the program to run can be loaded at startup.
     This source could be a connection to a host computer, an SD card, a ROM, ...
     In any case, we model this in Kami as a huge register file.
     Therefore, a faithful Verilog implementation will have to make sure that all in-range
     addresses behave like memory, including the ones from which the program is loaded.
     One possible implementation would be this:
     For each address which is designated as a "program load address", also have DRAM for it,
     as well as an extra "initialized bit", which is set to false initially.
     Whenever such an address is loaded, if the initialized bit is set, the value from the
     DRAM is returned, otherwise the value from the program source is loaded, stored into
     DRAM, and the bit is set to 1.
     Whenever such an address is stored, the initialized bit is set, and the value is stored
     into DRAM.
     For the proofs, we model this component as a huge register file where the addresses
     designated as "program load adddresses" are initialized to [prog], and the other
     addresses are initialized to zero.
     We'll use the convention that "program load addresses" are from 0 to 4*2^instrMemSizeLg,
     and that the data memory goes from 4*2^instrMemSizeLg to dataMemSize
     because then we don't have to pass the base program load address to this definition,
     and to serve load/store requests in the Kami model, we can just ignore the upper bits
     and use the lower bits to index into the Vector.
   *)

  Definition mm: Modules := Kami.Ex.SC.mm
                              (existT _ rv32DataBytes eq_refl)
                              kamiMemInit (kami_AbsMMIO (Z.to_N memSizeLg)).
  Definition p4mm: Modules := p4mm Hinstr kamiMemInit (kami_AbsMMIO (Z.to_N memSizeLg)).

  Fixpoint setRegsInit (kinits: kword 5 -> kword width) (n: nat): Registers :=
    match n with
    | O => map.put map.empty 0 $0
    | S n' => map.put (setRegsInit kinits n') (Z.of_nat n) (kinits $n)
    end.

  Definition riscvRegsInit: Registers :=
    setRegsInit (evalConstT (rfInit procInit)) 31.
  Lemma regs_related_riscvRegsInit:
    regs_related (evalConstT (rfInit procInit)) riscvRegsInit.
  Proof.
    red; intros.

    clear -Registers_ok.
    pose proof (wordToN_bound w).
    change (NatLib.Npow2 (BinInt.Z.to_nat 5)) with 32%N in H.
    assert (wordToN w = 0 \/ wordToN w = 1 \/ wordToN w = 2 \/ wordToN w = 3 \/
            wordToN w = 4 \/ wordToN w = 5 \/ wordToN w = 6 \/ wordToN w = 7 \/
            wordToN w = 8 \/ wordToN w = 9 \/ wordToN w = 10 \/ wordToN w = 11 \/
            wordToN w = 12 \/ wordToN w = 13 \/ wordToN w = 14 \/ wordToN w = 15 \/
            wordToN w = 16 \/ wordToN w = 17 \/ wordToN w = 18 \/ wordToN w = 19 \/
            wordToN w = 20 \/ wordToN w = 21 \/ wordToN w = 22 \/ wordToN w = 23 \/
            wordToN w = 24 \/ wordToN w = 25 \/ wordToN w = 26 \/ wordToN w = 27 \/
            wordToN w = 28 \/ wordToN w = 29 \/ wordToN w = 30 \/ wordToN w = 31)%N
      by abstract blia.
    clear H.
    repeat match goal with
           | H: _ \/ _ |- _ => destruct H
           end.

    all: match goal with
         | H: wordToN _ = ?n |- _ =>
           change n with (wordToN (sz:= 5) $(N.to_nat n)) in H;
             apply wordToN_inj in H; subst; simpl
         end.
    all: cbv [riscvRegsInit setRegsInit].
    all: repeat rewrite map.get_put_diff by discriminate.
    all: rewrite map.get_put_same.
    all: reflexivity.
  Qed.

  Lemma riscvRegsInit_sound:
    forall reg, 0 < reg < 32 -> map.get riscvRegsInit reg <> None.
  Proof.
    intros.
    assert (reg = 1 \/ reg = 2 \/ reg = 3 \/
            reg = 4 \/ reg = 5 \/ reg = 6 \/ reg = 7 \/
            reg = 8 \/ reg = 9 \/ reg = 10 \/ reg = 11 \/
            reg = 12 \/ reg = 13 \/ reg = 14 \/ reg = 15 \/
            reg = 16 \/ reg = 17 \/ reg = 18 \/ reg = 19 \/
            reg = 20 \/ reg = 21 \/ reg = 22 \/ reg = 23 \/
            reg = 24 \/ reg = 25 \/ reg = 26 \/ reg = 27 \/
            reg = 28 \/ reg = 29 \/ reg = 30 \/ reg = 31)
      by abstract blia.
    clear H.
    repeat match goal with
           | H: _ \/ _ |- _ => destruct H
           end.

    all: subst.
    all: cbv [riscvRegsInit setRegsInit].
    all: repeat rewrite map.get_put_diff by discriminate.
    all: rewrite map.get_put_same.
    all: discriminate.
  Qed.

  Definition riscvMemInit : mem := map.of_list (List.map
    (fun i : nat =>
      (word.of_Z (Z.of_nat i),
       byte.of_Z (uwordToZ (evalConstT kamiMemInit $i))))
    (seq 0 (2 ^ Z.to_nat memSizeLg))).

  Instance kword32: coqutil.Word.Interface.word 32 := KamiWord.word 32.
  Instance kword32_ok: word.ok kword32. eapply KamiWord.ok. reflexivity. Qed.
  Lemma riscvMemInit_get_None:
    forall addr,
      (kunsigned addr <? 2 ^ memSizeLg) = false ->
      map.get riscvMemInit addr = None.
  Proof.
    intros.
    apply get_of_list_not_In; [exact (@weq (Z.to_nat width))|assumption|].

    intro Hx.
    apply in_map_iff in Hx; destruct Hx as [[addr' v] [? Hx]].
    simpl in H0; subst.
    apply in_map_iff in Hx; destruct Hx as [n [? ?]].
    inversion H0; subst; clear H0.
    apply in_seq in H1; destruct H1 as [_ ?]; simpl in H0.

    apply Nat2Z.inj_lt in H0.
    rewrite N_Z_nat_conversions.Nat2Z.inj_pow in H0.
    rewrite Z2Nat.id in H0 by blia.
    simpl in H0.

    match type of H with
    | (?x <? ?y) = false => destruct (Z.ltb_spec x y); [discriminate|clear H]
    end.
    change kunsigned with (word.unsigned (width:= width)) in H1.
    change kofZ with (word.of_Z (width:= width)) in H1.
    rewrite word.unsigned_of_Z in H1.
    cbv [word.wrap] in H1.
    rewrite Z.mod_small in H1
      by (split; [blia|];
          eapply Z.lt_le_trans; [eassumption|];
          apply Z.pow_le_mono_r; blia).
    blia.
  Qed.

  Lemma mem_related_riscvMemInit : mem_related _ (evalConstT kamiMemInit) riscvMemInit.
  Proof.
    cbv [mem_related riscvMemInit].
    intros addr.
    case (kunsigned addr <? 2 ^ memSizeLg) eqn:H.
    2: { apply riscvMemInit_get_None; assumption. }
    assert (#addr < 2 ^ Z.to_nat memSizeLg)%nat.
    { rewrite <-wordToN_to_nat.
      apply Nat2Z.inj_lt.
      rewrite N_nat_Z, N_Z_nat_conversions.Nat2Z.inj_pow.
      rewrite Z2Nat.id by blia.
      apply Z.ltb_lt; assumption.
    }
    erewrite Properties.map.get_of_list_In_NoDup; trivial.
    1: eapply NoDup_nth_error; intros i j ?.
    2: eapply (nth_error_In _ (wordToNat addr)).

    { rewrite map_map; cbn; cbv [kofZ].
      clear dependent addr.
      rewrite !map_length, seq_length in H1.
      rewrite (@map_nth_error _ _ _ _ _ i).
      2: etransitivity; [eapply nth_error_nth'|];
           rewrite ?seq_length, ?seq_nth; trivial.
      destruct (lt_dec j (2^Z.to_nat memSizeLg)).
      { rewrite (@map_nth_error _ _ _ _ _ j).
        2: etransitivity; [eapply nth_error_nth'|];
            rewrite ?seq_length, ?seq_nth; trivial.
        intros HX.
        injection HX; clear HX; intros HX.
        eapply (f_equal (@wordToZ _)) in HX.
        pose proof Z.pow_le_mono_r 2 memSizeLg 31 eq_refl ltac:(blia);
        pose proof N_Z_nat_conversions.Z2Nat.inj_pow 2 memSizeLg ltac:(blia) ltac:(blia);
        change (Z.to_nat 2) with 2%nat in *.
        rewrite 2wordToZ_ZToWord'' in HX; try split;
         change (BinInt.Z.of_nat (Pos.to_nat 32) - 1) with 31;
         blia. }
      { rewrite (proj2 (nth_error_None _ _)); try congruence.
        rewrite map_length, seq_length; blia. } }
    { replace (evalZeroExtendTrunc (BinInt.Z.to_nat memSizeLg) addr)
        with (natToWord (Z.to_nat memSizeLg) (wordToNat addr)).
      2: {
        cbv [evalZeroExtendTrunc].
        destruct (lt_dec _ _); [exfalso; apply Z2Nat.inj_lt in l; blia|].
        apply wordToNat_inj.
        rewrite wordToNat_natToWord_eqn.
        rewrite wordToNat_split1.
        cbv [eq_rec_r eq_rec]; rewrite wordToNat_eq_rect.
        reflexivity.
      }
      rewrite (@map_nth_error _ _ _ _ _ (wordToNat addr)).
      2: {
        etransitivity; [eapply nth_error_nth'|].
        all : rewrite ?seq_length, ?seq_nth; trivial.
      }
      do 2 f_equal.
      eapply word.unsigned_inj.
      rewrite word.unsigned_of_Z.
      cbv [word.wrap]; rewrite <-word.wrap_unsigned; f_equal.
      unfold word.unsigned, word, wordW, KamiWord.word, kword, kunsigned.
      rewrite wordToN_nat, nat_N_Z; reflexivity.
    }
    Unshelve. all: exact O.
  Qed.

  Lemma states_related_init:
    states_related
      (initRegs (getRegInits (proc Hinstr kamiMemInit (kami_AbsMMIO (Z.to_N memSizeLg)))), [])
      {| getMachine :=
           {| RiscvMachine.getRegs := riscvRegsInit;
              RiscvMachine.getPc := word.of_Z 0;
              RiscvMachine.getNextPc := word.of_Z 4;
              RiscvMachine.getMem := riscvMemInit;
              RiscvMachine.getXAddrs := kamiXAddrs instrMemSizeLg;
              RiscvMachine.getLog := nil; (* <-- intended to be nil *) |};
         getMetrics := MetricLogging.EmptyMetricLog; |}.
  Proof.
    econstructor; try reflexivity.
    - econstructor.
    - eapply pRegsToT_init.
    - intros; discriminate.
    - split; reflexivity.
    - apply regs_related_riscvRegsInit.
    - apply mem_related_riscvMemInit.
  Qed.

  Lemma equivalentLabel_preserves_KamiLabelR:
    forall l1 l2,
      equivalentLabel (liftToMap1 (@idElementwise _)) l1 l2 ->
      forall l,
        KamiLabelR l2 l -> KamiLabelR l1 l.
  Proof.
    intros.
    destruct l1 as [ann1 ds1 cs1], l2 as [ann2 ds2 cs2].
    destruct H as [? [? ?]]; simpl in *.
    rewrite SemFacts.liftToMap1_idElementwise_id in H, H1; subst.
    inversion_clear H0; subst.
    - apply KamiSilent; assumption.
    - eapply KamiMMIO; eauto.
  Qed.

  Lemma equivalentLabelSeq_preserves_KamiLabelSeqR:
    forall t1 t2,
      equivalentLabelSeq (liftToMap1 (@idElementwise _)) t1 t2 ->
      forall t,
        KamiLabelSeqR t2 t ->
        KamiLabelSeqR t1 t.
  Proof.
    induction 1; intros; [assumption|].
    inversion_clear H1.
    constructor; auto.
    eapply equivalentLabel_preserves_KamiLabelR; eauto.
  Qed.

  Lemma riscv_init_memory_undef_on_MMIO:
    map.undef_on riscvMemInit isMMIOAddr.
  Proof.
    cbv [map.undef_on map.agree_on]; intros.
    cbv [elem_of] in H.
    pose proof (mmio_mem_disjoint _ Hkmemdisj _ H); clear H.
    rewrite map.get_empty.
    apply riscvMemInit_get_None.
    destruct (Z.ltb_spec (kunsigned k) (2 ^ memSizeLg)); intuition idtac.
  Qed.

  Lemma mmio_init_xaddrs_disjoint:
    disjoint (of_list (kamiXAddrs instrMemSizeLg)) isMMIOAddr.
  Proof.
    cbv [disjoint of_list elem_of]; intros.
    pose proof (mmio_mem_disjoint _ Hkmemdisj x).
    destruct (Z.ltb_spec (kunsigned x) (2 ^ memSizeLg)).
    - right; intro Hx; auto.
    - left; intro Hx.
      apply kamiXAddrs_isXAddr1_bound in Hx.
      apply N2Z.inj_lt in Hx.
      rewrite NatLib.Z_of_N_Npow2 in Hx.
      assert (2 ^ BinInt.Z.of_nat (2 + Z.to_nat instrMemSizeLg) < 2 ^ memSizeLg)
        by (apply Z.pow_lt_mono_r; blia).
      cbv [kunsigned] in *.
      blia.
  Qed.

  Lemma riscv_to_kamiImplProcessor:
    forall (traceProp: list Event -> Prop)
           (* --- hypotheses which will be proven by the compiler --- *)
           (RvInv: RiscvMachine -> Prop)
           (establishRvInv:
              forall (m0RV: RiscvMachine),
                m0RV.(RiscvMachine.getMem) = riscvMemInit ->
                m0RV.(RiscvMachine.getPc) = word.of_Z 0 ->
                m0RV.(RiscvMachine.getNextPc) = word.of_Z 4 ->
                (forall a: word,
                    0 <= word.unsigned a < 2 ^ (2 + instrMemSizeLg) ->
                    In a m0RV.(RiscvMachine.getXAddrs)) ->
                disjoint (of_list m0RV.(RiscvMachine.getXAddrs)) isMMIOAddr ->
                (forall reg, 0 < reg < 32 -> map.get m0RV.(getRegs) reg <> None) ->
                m0RV.(getLog) = nil ->
                RvInv m0RV)
           (preserveRvInv:
              forall (m: RiscvMachine), RvInv m -> mcomp_sat_unit (run1 iset) m RvInv)
           (useRvInv:
              forall (m: RiscvMachine),
                RvInv m -> exists t, traces_related t m.(getLog) /\
                                     traceProp t),
    (* --- final end to end theorem will start here --- *)
    forall (t: list LabelT) (mFinal: KamiImplMachine),
      Behavior p4mm mFinal t ->
      (* --- conclusion ---
         The trace produced by the kami implementation can be mapped to an MMIO trace
         (this guarantees that the only external behavior of the kami implementation is MMIO)
         and moreover, this MMIO trace satisfies some desirable property. *)
      exists (t': list Event), KamiLabelSeqR t t' /\ traceProp t'.
  Proof.
    intros.
    pose proof (@proc_correct instrMemSizeLg memSizeLg Hinstr kamiMemInit) as P.
    unfold traceRefines in P.
    specialize P with (1 := H).
    destruct P as (mFinal' & t' & B & E).
    inversion_clear B.
    edestruct ksteps_sound as (rs2 & Rel & Inv). 2: unfold calways; split.
    - econstructor.
      + eapply Kami.SemFacts.reachable_init.
      + eapply states_related_init.
      + eapply KamiSeqNil.
    - eapply establishRvInv; try reflexivity.
      all: cbv [getXAddrs getMachine]; intros.
      + apply alignedXAddrsRange_zero_bound_in.
        apply N2Z.inj_lt.
        rewrite nat_N_Z.
        cbv [instrMemSize].
        rewrite N_Z_nat_conversions.Nat2Z.inj_pow.
        rewrite Nat2Z.inj_add, Z2Nat.id by blia.
        apply H0.
      + apply mmio_init_xaddrs_disjoint.
      + apply riscvRegsInit_sound; assumption.
    - exact preserveRvInv.
    - econstructor. exact HMultistepBeh.
    - specialize (useRvInv _ Inv).
      inversion Rel. subst. clear Rel.
      simpl in useRvInv.
      destruct useRvInv as (t''' & R' & p).
      eexists. split; [|exact p].
      rewrite app_nil_r in H5.
      inversion H3. subst. clear H3.
      eapply equivalentLabelSeq_preserves_KamiLabelSeqR.
      1: eassumption.
      pose proof (traces_related_unique R' H4). subst t'''.
      assumption.
  Qed.

End Equiv.


(* Middle-end optimizer *)
(* Inerting Speculation in CoreIR programs *)

Require Import RTL.
Require Import IR.
Require Import monad.
Require Import common.
Require Import Coqlib.
Require Import backend.
Require Import primitives.
Require Import anchor_insertion.
Require Import assume_insertion.
Require Import anchor_removal.
Require Import profiler_types.
Require Import Errors.
Require Import monad.

(* Inserting Anchors in all functions according to the profiler wish *)
Fixpoint process_anc_list (p:program) (l:list (fun_id * list label)): program :=
  match l with
  | nil => p
  | (fid, anc_lbl)::l' => process_anc_list (safe_insert_anchor p fid anc_lbl) l'
  end.

(* Processing each middle_end optimization sugeested by the profiler *)
(* Using the safe optimizations: if one fails, it is just ignored by the optimizer *)
Fixpoint process_optim_list (p:program) (l:list (fun_id * middle_wish)): program :=
  match l with
  | nil => p
  | (fid, AS_INS guard anc_lbl)::l' => process_optim_list (safe_insert_assume p fid guard anc_lbl) l'
  end.

Definition middle_end (ps:profiler_state) (p:program): res program :=
    do optims <- OK (middle_end_suggestion ps);
    do fs_list <- OK (anchors_to_insert ps);
    do pfs <- OK (process_anc_list p fs_list);
    do newp <- OK (process_optim_list pfs optims);
    OK (lowering newp).

(* An error in optimization should not stop the execution *)
Definition safe_middle_end (ps:profiler_state) (p:program): program :=
  safe_res (middle_end ps) p.