Require Import bedrock2.NotationsCustomEntry.

Import Syntax BinInt String List.ListNotations.
Local Open Scope string_scope. Local Open Scope Z_scope. Local Open Scope list_scope.
From coqutil.Word Require Import Interface Properties.

From bedrock2 Require Import WeakestPrecondition WeakestPreconditionProperties
  Syntax Semantics BasicC64Semantics ProgramLogic Scalars Array Loops ZnWords.
From bedrock2.Map Require Import Separation SeparationLogic.
Require Import coqutil.Map.Interface.

From coqutil.Tactics Require Import Tactics letexists eabstract.
Require Import Coq.ZArith.ZArith.
From coqutil.Z Require Import div_mod_to_equations Lia.
Require Import Coq.Program.Tactics.

Require Import Crypto.Arithmetic.WordByWordMontgomery.
Import Markers.

Require Import Crypto.Util.ZUtil.Modulo.PullPush.
Require Import Crypto.Util.ZUtil.Tactics.PullPush.Modulo.

Section WithParameters.

  Import List coqutil.Datatypes.List.
  Import WordByWordMontgomery.

  Context {prime: Z} (r := 64) {ri : Z}.
  Context {ri_correct: (ri * 2^r) mod prime = 1} {prime_reasonable: 1 < prime}.
  (* prime is the modulus; r is the word size; ri is the inverse of 2^r mod prime *)

  Local Notation "m =* P" := ((P%sep) m) (at level 70, only parsing) (* experiment*).

  (* redc_alt ought to take in small arrays A and B, and output an array S *)
  (* S should be small, and should evaluate mod the prime to the same thing as
     A * B * ri *)

  Instance spec_of_redc_alt : spec_of "redc_alt" :=
    fnspec! "redc_alt" Astart Bstart Sstart len / A (aval: Z) B (bval: Z) S R,
    { requires t m :=
        m =* array scalar (word.of_Z 8) Astart A *
                  array scalar (word.of_Z 8) Bstart B *
                  array scalar (word.of_Z 8) Sstart S * R /\
        word.unsigned len = Z.of_nat (List.length A)  /\
        word.unsigned len = Z.of_nat (List.length B)  /\
        word.unsigned len = Z.of_nat (List.length S) /\
        @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned A) = aval /\
        @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned B) = bval;
      ensures t' m' :=  t=t' /\ exists S',
          m' =*
             array scalar (word.of_Z 8) Astart A *
             array scalar (word.of_Z 8) Bstart B *
            array scalar (word.of_Z 8) Sstart S' * R /\
          ( aval * bval * ri^(word.unsigned len) ) mod prime =
            @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned S') mod prime
    }.

  (* redc_step ought to take in small arrays B and S, and value a, and output an array S' *)
  (* S' should be small, and should eval to the same as (a * B + S) * ri modulo the prime *)

  Instance spec_of_redc_step : spec_of "redc_step" :=
    fnspec! "redc_step" a Bstart Sstart len / B (bval: Z) S (sval: Z) R,
      { requires t m :=
          m =* array scalar (word.of_Z 8) Bstart B *
                    array scalar (word.of_Z 8) Sstart S * R /\
          word.unsigned len = Z.of_nat (List.length B) /\
          word.unsigned len = Z.of_nat (List.length S) /\
          @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned B) = bval /\
          @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned S) = sval;
        ensures t' m' := t=t' /\ exists S',
            m' =* array scalar (word.of_Z 8) Bstart B *
              array scalar (word.of_Z 8) Sstart S' * R /\
              word.unsigned len = Z.of_nat (List.length S') /\
              ((word.unsigned a) * bval + sval) * ri mod prime =
                @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned S') mod prime
      }.

  Definition redc_alt :=
    func! (Astart, Bstart, Sstart, len) {
    i = $0;
    while (i < len) {
         store(Sstart + $8*i, $0);
         i = i + $1
      };
    i = $0;
    while (i < len) {
         redc_step ( load(Astart + $8*i), Bstart, Sstart, len );
          i = i + $1
      }
    }.

  Import Coq.Lists.List.

  Let zeros (n: Z) :=
        repeat (@word.of_Z _ word 0) (Z.to_nat n).

  Theorem eval_firstn:
    forall modulus depth number,
      @Core.Positional.eval (UniformWeight.uweight modulus) depth number = @Core.Positional.eval (UniformWeight.uweight modulus) depth (firstn depth number) .
    Proof.
      intros. cbv [Core.Positional.eval Core.Positional.to_associational Core.Associational.eval].
      rewrite combine_firstn_l. rewrite map_length. rewrite seq_length. reflexivity.
    Qed.

    Theorem eval_single:
      forall modulus x,
        Core.Positional.eval (UniformWeight.uweight modulus) (length [x]) [x] = x.
    Proof.
      intros. cbv [length].
      cbv [Core.Positional.eval Core.Positional.to_associational Core.Associational.eval].
      cbv [seq map UniformWeight.uweight ModOps.weight]. simpl.
      rewrite Z.mul_0_r. rewrite Zdiv_0_l. replace (-0) with 0 by Lia.lia.
      rewrite Z.pow_0_r. Lia.lia.
    Qed.

  Theorem eval_one_further:
    forall modulus depth number,  Z.of_nat (length number) > Z.of_nat depth -> 0 <= modulus ->
                                  @eval modulus (depth + 1) number =
                                  @eval modulus depth number + ((2 ^ modulus) ^ Z.of_nat depth) * hd 0 (skipn depth number).
  Proof.
    intros. cbv [eval].
    assert ((firstn depth number ++ skipn depth number) = number) by (apply firstn_skipn).
    destruct (skipn depth number) as [|new junk].
    - assert (length (firstn depth number ++ []) = length number) by (rewrite H1; reflexivity).
      rewrite app_length in H2. simpl in H2. assert ((length (firstn depth number)) <= depth)%nat by apply firstn_le_length. Lia.lia.
    - cbv [hd]. replace (Core.Positional.eval (UniformWeight.uweight modulus) (depth + 1) number) with (Core.Positional.eval (UniformWeight.uweight modulus) (depth + 1) (firstn depth number ++ [new])).
      2: {
        symmetry. rewrite eval_firstn. rewrite <- H1.
        replace (firstn (depth + 1) (firstn depth number ++ new :: junk)) with (firstn depth (firstn depth number ++ new :: junk) ++ [new]).
        1: { reflexivity. }
        replace (depth + 1)%nat with (S depth) by Lia.lia.
        rewrite <- firstn_nth with (d := 0).
        2: { rewrite app_length. replace (length (firstn depth number)) with depth.
             - assert (length (new::junk) = S (length junk)) by apply ListUtil.cons_length. Lia.lia.
             - rewrite ListUtil.List.firstn_length_le; Lia.lia.
        }
        f_equal. assert (depth = length (firstn depth number)) by (rewrite ListUtil.List.firstn_length_le; Lia.lia).
        remember (firstn depth number) as plz_dont_rewrite_me.
        rewrite H2. rewrite nth_middle. reflexivity.
      }
      rewrite UniformWeight.uweight_eval_app with (n := depth).
      + rewrite <- eval_firstn. f_equal. rewrite eval_single. f_equal.
        cbv [UniformWeight.uweight ModOps.weight]. rewrite Z.div_1_r.
        rewrite <- Z.pow_mul_r; try Lia.lia.
        f_equal. Lia.lia.
    + Lia.lia.
    + rewrite ListUtil.List.firstn_length_le; Lia.lia.
    + reflexivity.
 Qed.

  Theorem array_small:
  forall start arr R m,
    m =* array scalar (word.of_Z 8) start arr * R  ->
   Z.of_nat (@length (@word.rep 64 word) arr) * 8 < 2 ^ 64
  .
  Proof.
    Admitted.

 Theorem redc_alt_ok :
      program_logic_goal_for_function! redc_alt.
 Proof.
   repeat straightline.
      (*after the first loop, our output array is full of zeros*)
      refine ( tailrec (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ HList.polymorphic_list.nil))))))))
               ("Astart":: "Bstart" :: "Sstart" :: "len" :: "i" :: nil)
               (fun l A aval B bval S Ra Rb R t m Astart Bstart Sstart len i => PrimitivePair.pair.mk
                                    (m =* array scalar (word.of_Z 8) (word.add Sstart (word.mul (word.of_Z 8) i)) S * R /\
                                       word.unsigned len - word.unsigned i = Z.of_nat (List.length S) /\

                                    l = List.length S )
                                    (fun t' m' Astart' Bstart' Sstart' len' i' =>
                                       (
                                     t = t' /\ Astart = Astart' /\ Bstart = Bstart' /\ Sstart = Sstart' /\ len = len' /\
                                     m' =* array scalar (word.of_Z 8) (word.add Sstart (word.mul (word.of_Z 8) i)) (zeros (word.unsigned len - word.unsigned i)) * R
                                     )
                                    )
               )
               lt _ _ _ _ _ _ _ _ _ _ _ _ _);
        cbn [reconstruct map.putmany_of_list HList.tuple.to_list
         HList.hlist.foralls HList.tuple.foralls
         HList.hlist.existss HList.tuple.existss
         HList.hlist.apply  HList.tuple.apply
         HList.hlist
         List.repeat Datatypes.length
         HList.polymorphic_list.repeat HList.polymorphic_list.length
         PrimitivePair.pair._1 PrimitivePair.pair._2] in *.

      { repeat straightline. }
      { exact Wf_nat.lt_wf. }
      { repeat straightline.
        subst i.
        replace (word.add Sstart (word.mul (word.of_Z 8) (word.of_Z 0))) with (Sstart) by ring.
        repeat split; try eauto.
        - ecancel_assumption.
        - rewrite word.unsigned_of_Z_0. Lia.lia. }

      { repeat straightline. eexists. split. { repeat straightline. } repeat straightline. split.

        (*loop exits properly*)
        2: {
          repeat straightline; repeat split.
          rename x3 into S'; rename x9 into Sstart'; rename x10 into len'; rename x11 into i'; rename x6 into R'.

          destruct (word.ltu i' len') eqn: Hbreak.
          - rewrite word.unsigned_of_Z_1 in H8; try inversion H8; clear H8.
          - destruct (length S') eqn: HS.
          (*length cannot be nonzero*)
          2: { rewrite word.unsigned_ltu in Hbreak. Lia.lia. }
          (*if length is zero, all arrays are the same*)
          cbv [Z.of_nat] in H7. rewrite H7; clear H7 Hbreak.
          apply ListUtil.length0_nil in HS.
          subst S'. apply H6.
        }

        (*loop body is good*)
        repeat straightline.
        rename x3 into S'; rename x9 into Sstart'; rename x10 into len'; rename x11 into i'; rename x6 into R'.
        destruct S'.
        - cbv [length Z.of_nat] in H7.
          destruct (word.ltu i' len') eqn: Hbreak; try (rewrite word.unsigned_of_Z_0 in H8; exfalso; apply H8; trivial).
          rewrite word.unsigned_ltu in Hbreak. Lia.lia.
        - cbn [array] in H6. repeat straightline.
          repeat split; try trivial. exists (S'). repeat split; try trivial. exists ( (scalar (word.add Sstart' (word.mul (word.of_Z 8) i')) (word.of_Z 0)) * R')%sep. exists (length S').
          repeat split; subst v0 a.
          all: try (repeat (destruct H10 as [solver H10]; try assumption; clear solver)); subst i.
          + replace (word.mul (word.of_Z 8) (word.add i' (word.of_Z 1))) with (word.add (word.mul (word.of_Z 8) i') (word.of_Z 8)) by ring. rewrite word.add_assoc. apply sep_comm. apply sep_assoc. ecancel_assumption.
          + destruct (word.unsigned i' + 1 <? 2^64) eqn: Hisize.
            2: { assert (2^64 - 1 <= word.unsigned i') by Lia.lia; clear Hisize.
                 assert (word.unsigned i' < 2^64) by apply word.unsigned_range.
                 assert (Hi': word.unsigned i' = 2^64 - 1) by Lia.lia; clear H10 H11.
                 assert (Hlen': word.unsigned len' < 2^64) by apply word.unsigned_range.
                 assert (Hneg: word.unsigned len' - word.unsigned i' < 1) by Lia.lia; clear Hi' Hlen'.
                 simpl in H7; simpl in Hneg. rewrite Zpos_P_of_succ_nat in H7.
                 assert (0 <= (Z.of_nat(length S'))) by apply Zle_0_nat.
                 Lia.lia.
            }

            rewrite word.unsigned_add. rewrite word.unsigned_of_Z_1. cbv [word.wrap].
            assert (Hsmall: (word.unsigned i' + 1) mod 2^64 = word.unsigned i' + 1).
            {  apply Z.mod_small. assert (0 <= word.unsigned i') by apply word.unsigned_range. Lia.lia. }
            rewrite Hsmall; clear Hisize Hsmall.

            simpl in H7; simpl. rewrite Zpos_P_of_succ_nat in H7.
            assert (0 <= (Z.of_nat(length S'))) by apply Zle_0_nat.
            Lia.lia.
          + subst v.  auto.
          + assert (Hhead: zeros (word.unsigned len' - word.unsigned i') = word.of_Z 0 :: zeros (word.unsigned len' - word.unsigned (word.add i' (word.of_Z 1))) ).
            {
            destruct (word.unsigned i' + 1 <? 2^64) eqn: Hisize.
            2: { assert (2^64 - 1 <= word.unsigned i') by Lia.lia; clear Hisize.
                 assert (word.unsigned i' < 2^64) by apply word.unsigned_range.
                 assert (Hi': word.unsigned i' = 2^64 - 1) by Lia.lia; clear H10 H11.
                 assert (Hlen': word.unsigned len' < 2^64) by apply word.unsigned_range.
                 assert (Hneg: word.unsigned len' - word.unsigned i' < 1) by Lia.lia; clear Hi' Hlen'.
                 simpl in H7; simpl in Hneg. rewrite Zpos_P_of_succ_nat in H7.
                 assert (0 <= (Z.of_nat(length S'))) by apply Zle_0_nat.
                 Lia.lia.
            }

            rewrite word.unsigned_add. rewrite word.unsigned_of_Z_1. cbv [word.wrap].
            assert (Hsmall: (word.unsigned i' + 1) mod 2^64 = word.unsigned i' + 1).
            {  apply Z.mod_small. assert (0 <= word.unsigned i') by apply word.unsigned_range. Lia.lia. }
            rewrite Hsmall; clear Hisize Hsmall.
            assert (Hpos: 0 <= word.unsigned len' - (word.unsigned i' + 1)). {
              simpl in H7; simpl. rewrite Zpos_P_of_succ_nat in H7.
              assert (0 <= (Z.of_nat(length S'))) by apply Zle_0_nat.
              Lia.lia.
            }
            replace (word.unsigned len' - word.unsigned i') with (Z.succ (word.unsigned len' - (word.unsigned i' + 1)) ) by Lia.lia.
            cbv [zeros]. rewrite Z2Nat.inj_succ; try assumption. cbn [repeat]. trivial.
            }
            rewrite Hhead; clear Hhead.
            cbn [array].
            replace (word.add (word.add Sstart' (word.mul (word.of_Z 8) i')) (word.of_Z 8)) with  (word.add Sstart' (word.mul (word.of_Z 8) (word.add i' (word.of_Z 1)))) by ring. ecancel_assumption.
        }

        {
          repeat straightline.
          (*then, the second loop does the multiplication properly*)
          refine ( tailrec (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ (HList.polymorphic_list.cons _ HList.polymorphic_list.nil))))))))
               ("Astart":: "Bstart" :: "Sstart" :: "len" :: "i" :: nil)
               (fun l A aval B bval S Ra Rb R t m Astart Bstart Sstart len i => PrimitivePair.pair.mk
                                           (
                                            m =* array scalar (word.of_Z 8) Astart A *
                                              array scalar (word.of_Z 8) Bstart B *
                                              array scalar (word.of_Z 8) Sstart S * R /\
                                            word.unsigned len = Z.of_nat (List.length A)  /\
                                            word.unsigned len = Z.of_nat (List.length B)  /\
                                            word.unsigned len = Z.of_nat (List.length S) /\
                                            @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned A) = aval /\
                                            @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned B) = bval /\
                                            @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned S) mod prime =
                                            @eval r (Z.to_nat (word.unsigned i)) (List.map word.unsigned A)
                                            * bval * ri^(word.unsigned i) mod prime /\
                                             0 <= word.unsigned i <= word.unsigned len /\
                                             l = Z.to_nat (word.unsigned len - word.unsigned i)
                                           )
                                    (fun t' m' Astart' Bstart' Sstart' len' i' =>
                                     (
                                     t = t' /\ Astart = Astart' /\ Bstart = Bstart' /\ Sstart = Sstart' /\ len = len' /\
                                     exists S',
                                       m' =* array scalar (word.of_Z 8) Astart A *
                                              array scalar (word.of_Z 8) Bstart B *
                                              array scalar (word.of_Z 8) Sstart S' * R /\
                                       @eval r (Z.to_nat (word.unsigned len)) (List.map word.unsigned S') mod prime =
                                       aval * bval * ri^(word.unsigned len) mod prime
                                     )
                                    )
               )
               lt _ _ _ _ _ _ _ _ _ _ _ _ _);
        cbn [reconstruct map.putmany_of_list HList.tuple.to_list
         HList.hlist.foralls HList.tuple.foralls
         HList.hlist.existss HList.tuple.existss
         HList.hlist.apply  HList.tuple.apply
         HList.hlist
         List.repeat Datatypes.length
         HList.polymorphic_list.repeat HList.polymorphic_list.length
         PrimitivePair.pair._1 PrimitivePair.pair._2] in *.

          { repeat straightline. }

          { exact Wf_nat.lt_wf. }

          { repeat straightline; subst_words.
            replace (word.add x1 (word.mul (word.of_Z 8) (word.of_Z 0))) with (x1) in H11 by ring.
            rewrite word.unsigned_of_Z_0 in H11. replace (word.unsigned x2 - 0) with (word.unsigned x2) in H11 by ring.
            repeat split.
            - ecancel_assumption.
            - assumption.
            - assumption.
            - cbv [zeros]. rewrite repeat_length. rewrite H3. rewrite Nat2Z.id. trivial.
            - rewrite word.unsigned_of_Z_0.

              assert (eval_zero: forall beep boop,  @eval r beep (map word.unsigned (zeros boop)) = 0 ).
              { clear; intros.
                cbv [zeros eval Core.Positional.eval Core.Associational.eval Core.Positional.to_associational].
                generalize dependent (map (UniformWeight.uweight r) (seq 0 beep)).
                induction (Z.to_nat boop).
                - simpl. destruct l; trivial.
                - intros l'. cbn [repeat map]. rewrite word.unsigned_of_Z_0. destruct l'; trivial.
                  cbn [combine map fst snd fold_right]. rewrite IHn. ring.
              }

              rewrite eval_zero.

              assert (eval_in_name_only: forall xs, @eval r 0 xs = 0).
              {  clear; intros.
                cbv [eval Core.Positional.eval Core.Associational.eval Core.Positional.to_associational].
                reflexivity.
              }
              rewrite eval_in_name_only. eexists.

            - rewrite word.unsigned_of_Z_0. reflexivity.

            - rewrite word.unsigned_of_Z_0. rewrite H3. apply Nat2Z.is_nonneg.

          }

          { repeat straightline. eexists. split. { repeat straightline. } repeat straightline. split.

            (*loop exits properly*)
            2: {
              repeat straightline; repeat split; eauto. eexists. split.
              1: { ecancel_assumption. }
              destruct (word.ltu x16 x15) eqn: Hbreak.
              - rewrite word.unsigned_of_Z_1 in H15; inversion H15.
              - subst x5.  rewrite word.unsigned_ltu in Hbreak.
                assert (word.unsigned x15 = word.unsigned x16) by Lia.lia.
                rewrite H10 in *. assumption.

            }

            (*loop body good*)
            repeat (repeat eexists; straightline).

            - eapply load_word_of_sep. seprewrite_in @array_address_inbounds H6.

              4: { ecancel_assumption. }

               { destruct (word.ltu x16 x15) eqn: Hbreak; try (rewrite word.unsigned_of_Z_0 in H15; contradiction);
              rewrite word.unsigned_ltu in Hbreak; assert (Hiupper : word.unsigned x16 < Z.of_nat (length x4)) by Lia.lia;
                clear H14 H15 Hbreak.
                 assert (Hilower: 0 <= word.unsigned x16) by apply word.unsigned_range.
                 ZnWords. }

               { clear. ZnWords. }

               {
                 replace (word.sub (word.add x12 (word.mul (word.of_Z 8) x16)) x12)
                   with (word.mul (word.of_Z 8) x16) by ring.
                 eexists. }

            - straightline_call.
              { repeat split.
                - ecancel_assumption.
                - assumption.
                - assumption.
              }
              repeat straightline.

              rename x4 into A'; rename x6 into B'; rename x8 into S'.
              rename x12 into Astart'; rename x13 into Bstart'; rename x14 into Sstart'.
              rename x5 into aval'; rename x7 into bval'; rename x15 into len; rename x16 into i'.
              rename x11 into R'. rename x17 into Snew.

              exists A'. exists aval'. exists B'. exists bval'. exists Snew.
              exists "duopus".
              exists "voltaire".
              eexists.
              exists (Z.to_nat (word.unsigned len - word.unsigned i)).
              split.
              { destruct (word.ltu i' len) eqn: Hbreak; try (rewrite word.unsigned_of_Z_0 in H15; contradiction);
                     rewrite word.unsigned_ltu in Hbreak; assert (Hiupper : word.unsigned i' < word.unsigned len) by Lia.lia;
                  subst i.
                split. {  ecancel_assumption. } repeat split; eauto.
                2: { clear - Hiupper. ZnWords. }
                {
                  rewrite <- H21.

                  assert ( no_overflow: (Z.of_nat (length Snew))*8 < 2^64).
                  2: {

                  rename H7 into HAlen; rename H8 into HBlen; rename H9 into HSoldlen;
                  rename H19 into HSnewlen; rename H13 into Holdeval.

                  assert (Hivalid: word.unsigned (word.add i' (word.of_Z 1)) = word.unsigned i' + 1) by
                    (clear - Hiupper; ZnWords).

                  rewrite Hivalid. assert (8*(word.unsigned i') < 2^64) by Lia.lia.

                  replace (@word.unsigned 64 word (@word.mul 64 word (@word.of_Z 64 word 8) i') /
                             @word.unsigned 64 word (@word.of_Z 64 word 8)) with (@word.unsigned 64 word i').

                  2:{
                    rewrite word.unsigned_mul.
                    rewrite word.unsigned_of_Z. cbv [word.wrap].
                    replace (8 mod 2^64) with 8 by reflexivity.
                    replace ((8 * word.unsigned i') mod 2 ^ 64) with (8 * word.unsigned i') by
                        (symmetry; apply Z.mod_small; Lia.lia).
                    rewrite Z.mul_comm.
                    symmetry; apply Z_div_mult. reflexivity.
                  }

                  rewrite H12.

                  instantiate (1 := (word.of_Z 0)).

                  rewrite Z2Nat.inj_add; try Lia.lia.

                  rewrite eval_one_further.

                  2: { rewrite map_length. rewrite Z2Nat.id; Lia.lia. }
                  2: { Lia.lia. }

                  rewrite Z.mul_mod_l. rewrite Z.add_mod_r. rewrite Holdeval.
                  rewrite <- Z.add_mod_r. rewrite <- Z.mul_mod_l.
                  rewrite Z2Nat.id; try Lia.lia.
                  rewrite Z.mul_add_distr_r.
                  repeat rewrite <- Z.mul_assoc.
                  replace (ri ^ word.unsigned i' * ri) with (ri * ri ^ word.unsigned i') by apply Z.mul_comm.
                  rewrite Pow.Z.pow_mul_base; try Lia.lia.
                  repeat rewrite Z.mul_assoc.
                  rewrite ListUtil.skipn_map. rewrite <- hd_map. rewrite word.unsigned_of_Z_0.

                  pose (Abig := hd 0 (map word.unsigned (skipn (Z.to_nat (word.unsigned i')) A')) ).
                  replace (hd 0 (map word.unsigned (skipn (Z.to_nat (word.unsigned i')) A'))) with Abig; try trivial.
                  pose (Arest := @eval r (Z.to_nat (@word.unsigned 64 word i')) (@map (@word.rep 64 word) Z (@word.unsigned 64 word) A')).
                  replace (@eval r (Z.to_nat (@word.unsigned 64 word i')) (@map (@word.rep 64 word) Z (@word.unsigned 64 word) A')) with Arest; try trivial.

                  rewrite Z.mul_add_distr_r. rewrite Z.mul_add_distr_r. rewrite Z.add_comm.

                  repeat rewrite <- Z.mul_assoc.
                  replace ((2^r) ^ word.unsigned i' * (Abig * (bval' * ri ^ (word.unsigned i' + 1)))) with ((Abig * (bval' * ri ^ (word.unsigned i' + 1)))*(2^r)^word.unsigned i') by apply Z.mul_comm.
                  repeat rewrite <- Z.mul_assoc.

                  replace ((ri ^ (word.unsigned i' + 1) * (2^r) ^ word.unsigned i')) with ( ri * (ri*(2^r)) ^ (word.unsigned i')).
                  2: {
                    rewrite <- Pow.Z.pow_mul_base; try rewrite Z.pow_mul_l; Lia.lia.
                  }

                  rewrite Z.add_mod_r. symmetry. rewrite Z.add_mod_r.

                  replace ((Abig * (bval' * (ri * (ri * (2^r)) ^ word.unsigned i'))) mod prime) with ((Abig * (bval' * ri)) mod prime).
                  1: { trivial. }

                    rewrite Z.mul_mod_r. symmetry. rewrite Z.mul_mod_r.

                  replace (((bval' * (ri * (ri * (2^r)) ^ word.unsigned i')) mod prime)) with ((bval' * ri) mod prime).
                  1: { trivial. }

                  rewrite Z.mul_mod_r. symmetry. rewrite Z.mul_mod_r.

                  replace ((ri * (ri * (2^r)) ^ word.unsigned i') mod prime) with (ri mod prime).
                  1: {trivial. }

                  rewrite Z.mul_mod_r. replace ((ri * (2^r)) ^ word.unsigned i' mod prime) with 1.
                  1: { rewrite Z.mul_1_r. reflexivity. }

                  rewrite Z.mod_pow_full. rewrite ri_correct. rewrite Z.pow_1_l; try Lia.lia. rewrite Zmod_1_l; try Lia.lia.

                  }
                  assert (exists Rsnew, a0 =* array scalar (word.of_Z 8) Sstart' Snew * Rsnew).
                  2: {
                    destruct H17 as [Rsnew H17]. apply array_small with (start := Sstart') (m := a0) (R := Rsnew). apply H17.
                  }
                  eexists. ecancel_assumption.
              }

              {
                clear -Hbreak.
                assert (word.unsigned i' < word.unsigned len) by Lia.lia.
                ZnWords.
              }

              }

              {
                 destruct (word.ltu i' len) eqn: Hbreak; try (rewrite word.unsigned_of_Z_0 in H15; contradiction);
                     rewrite word.unsigned_ltu in Hbreak; assert (Hiupper : word.unsigned i' < word.unsigned len) by Lia.lia;
                  subst i.
                repeat split.
                 - clear -Hiupper H14. ZnWords.
                 - repeat (destruct H17 as [maybewin H17]; try trivial; clear maybewin).
                 - repeat (destruct H17 as [maybewin H17]; try trivial; clear maybewin).
                 - repeat (destruct H17 as [maybewin H17]; try trivial; clear maybewin).
                 - repeat (destruct H17 as [maybewin H17]; try trivial; clear maybewin).
                 - repeat (destruct H17 as [maybewin H17]; try trivial; clear maybewin).
                 - destruct H17 as [_ H17]. destruct H17 as [_ H17]. destruct H17 as [_ H17].
                   destruct H17 as [_ H17]. destruct H17 as [_ H17]. destruct H17 as [Ssss H17].
                   destruct H17 as [Hmem Heval]. exists Ssss. split.
                   + ecancel_assumption.
                   + trivial.

              }
              }
              repeat straightline. repeat split; trivial. exists x9. split.
          - ecancel_assumption.
          - symmetry. subst aval. subst bval. trivial.
        }
        Unshelve.
   - eauto.
   - eauto.
   - eauto.
   - eauto.
   - eauto.
   - eauto.
   - exact String.HelloWorld.
   - exact " :) ".

    Qed.

End WithParameters.


(*vars_to_clique attaches each v_i and v_i' to x_i.  If in the boolean formula,
**u_i is set to true in a satisfiable environment, then v_i gets assigned the color u_i
**and v_i' gets assigned the color C*)
Require Export ThreeSatReduction.

(*attaching v_i to x_j for i <> j preserves coloring*)
Theorem connectXColorable : forall Gamma Delta G C eta eta' u,
                              ~In u Delta -> setVertices Gamma C 0 eta' eta -> In (u,3*u,3*u+1,3*u+2) Gamma ->
                              connectX Gamma Delta (3*u) G -> coloring eta' G C.
Proof.
  intros. genDeps {{ eta; C; eta'}}. remember (3*u). induction H2; intros.
  {constructor. }
  {subst. destruct(eq_nat_dec u u0).
   {subst. exfalso. apply H. simpl. auto. }
   {copy H3. copy H1. eapply V_V'MapToUOrC in H1; eauto. eapply XMapsToU in H0; eauto.
    inv H1. invertHyp. econstructor; eauto. omega. omega. eapply IHconnectX; eauto.
    intros c. apply H. simpl; auto. invertHyp. econstructor; eauto. omega.
    omega. eapply IHconnectX; eauto. intros c. apply H. simpl; auto. }
  }
Qed.

(*attaching v_i' to x_j for i <> j preserves coloring*)
Theorem connectXColorable' : forall Gamma Delta G C eta eta' u,
                              ~In u Delta -> setVertices Gamma C 0 eta' eta -> In (u,3*u,3*u+1,3*u+2) Gamma ->
                              connectX Gamma Delta (3*u+1) G -> coloring eta' G C.
Proof.
  intros. genDeps {{ eta; C; eta'}}. remember (3*u+1). induction H2; intros.
  {constructor. }
  {subst. destruct(eq_nat_dec u u0).
   {subst. exfalso. apply H. simpl. auto. }
   {copy H3. copy H1. eapply V_V'MapToUOrC in H1; eauto. eapply XMapsToU in H0; eauto.
    inv H1. invertHyp. econstructor; eauto. omega. omega. eapply IHconnectX; eauto.
    intros c. apply H. simpl; auto. invertHyp. econstructor; eauto. omega.
    omega. eapply IHconnectX; eauto. intros c. apply H. simpl; auto. }
  }
Qed.

(*attaching x_i to v_i and v_i' for i <> j preserves coloring*)
Theorem connectVColorable : forall Gamma Delta G C eta eta' u,
                              ~In u Delta -> setVertices Gamma C 0 eta' eta -> In (u,3*u,3*u+1,3*u+2) Gamma ->
                              connectV Gamma Delta (3*u+2) G -> coloring eta' G C.
Proof.
  intros. genDeps {{ eta; C; eta'}}. remember (3*u+2). induction H2; intros.
  {constructor. }
  {subst. destruct(eq_nat_dec u u0).
   {subst. exfalso. apply H. simpl. auto. }
   {copy H3. copy H1. eapply XMapsToU in H1; eauto. eapply V_V'MapToUOrC in H0; eauto.
    inv H0.
    {invertHyp. econstructor; eauto; try omega. econstructor; eauto; try omega.
     eapply IHconnectV; eauto. intros contra. apply H. simpl. auto. }
    {invertHyp. econstructor; eauto; try omega. econstructor; eauto; try omega.
     eapply IHconnectV; eauto. intros contra. apply H. simpl; auto. }
   }
  }
Qed.

Theorem varsToCliqueColorable : forall Gamma eta Delta C G eta' U,
                            setVertices Gamma C 0 eta' eta ->
                            unique U Delta ->
                            vars_to_clique Gamma Delta G -> coloring eta' G C.
Proof.
  intros. genDeps {{ U; eta'; eta; C }}. induction H1; intros.
  {constructor. }
  {inv H5. constructor. eapply IHvars_to_clique; eauto. constructor.
   eapply connectXColorable; eauto. eapply uniqueNotIn; eauto. apply Union_intror.
   constructor. constructor. eapply connectXColorable'; eauto.
   eapply uniqueNotIn; eauto. apply Union_intror. constructor.
   eapply connectVColorable; eauto. eapply uniqueNotIn; eauto.
   apply Union_intror. constructor. }
Qed.


From mathcomp Require Import all_ssreflect all_algebra.
Require Import Psatz Poly_complements.

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Import GRing.Theory.
Local Open Scope ring_scope.

Section RBase.

Variable R : ringType.
Variable F : nat -> {poly R}.

Lemma polybase_widen (p : {poly R}) n :
    (size p <= n)%N -> \sum_(i < n) p`_i *: F i = \sum_(i < size p) p`_i *: F i.
Proof.
move=> Hs.
rewrite -(big_mkord xpredT (fun i => p`_i *: F i)).
rewrite (big_cat_nat _ _ _ _ Hs) //= big_mkord.
rewrite addrC big_nat_cond big1 ?add0r // => i.
rewrite andbT => /andP[/(@leq_sizeP R) // H1 H2].
by rewrite H1 ?scale0r.
Qed.

Lemma sumrA (p q: {poly R}):
	\sum_(i < size (p + q)) (p + q)`_i *: F i = \sum_(i < size p) p`_i *: F i + \sum_(i < size q) q`_i *: F i.
Proof.
rewrite -(@polybase_widen (p + q) (maxn (size p) (size q))); last exact: size_add.
rewrite -(@polybase_widen p (maxn (size p) (size q))); last by rewrite leq_maxl.
rewrite -(@polybase_widen q (maxn (size p) (size q))); last by rewrite leq_maxr.
elim: (maxn (size p) (size q)) => [ | n ih]; first by rewrite !big_ord0 rm0.
rewrite !big_ord_recr/= ih coefD scalerDl.
rewrite -!addrA; congr (_ + _).
by rewrite addrCA.
Qed.

Lemma sumrZ a (p: {poly R}):
	\sum_(i < size (a *: p)) (a *: p)`_i *: F i = a*: (\sum_(i < size p) p`_i *: F i).
Proof.
rewrite -(@polybase_widen (a *: p) (size p)); last exact: size_scale_leq.
elim: (size p) => [ | n ih]; first by rewrite !big_ord0 rm0.
by rewrite !big_ord_recr/= ih -!mul_polyC mulrDr !mul_polyC scalerA -!mul_polyC coefCM.
Qed.
End RBase.

Section Base.
Variable R : unitRingType.
Variable F : nat -> {poly R}.
Hypothesis F_size : forall n, size (F n) = n.+1.
Hypothesis F_lead: forall n, GRing.rreg (lead_coef (F n)).

Lemma size_seqbase n l :
   size (\sum_(i < n) l`_ i *: F i) = \max_(i < n | l`_i != 0) i.+1.
Proof.
elim: n => [|n IH]; first by rewrite !big_ord0 size_poly0.
rewrite [RHS]big_mkcond /= !big_ord_recr /= -big_mkcond /=.
have [/eqP Hp| /eqP Hnp /=] := boolP (l`_n == 0).
  by rewrite Hp scale0r addr0 IH /= maxn0.
have lp : l`_ n * (F n)`_n != 0.
	apply /eqP => eq; apply Hnp.
	apply (@F_lead n).
	by rewrite lead_coefE F_size rm0.
have Hs: size (l`_n *: F n) = n.+1.
	suff/leP: (n.+1 <= size (l`_n *: F n))%nat.
		have/leP:= @size_scale_leq _ (l`_n) (F n).
		rewrite F_size; lia.
	by apply gtn_size; rewrite coefZ.
case: (n) Hs IH => [Hs _|n1 Hs IH].
  by rewrite !big_ord0 add0r Hs.
rewrite addrC size_addl Hs.
  rewrite (maxn_idPr _) //.
  apply/bigmax_leqP=> i _.
  by apply: leq_trans (ltn_ord i) _; rewrite ltnW.
rewrite ltnS IH.
by apply/bigmax_leqP => i _.
Qed.

Lemma size_polybase (p: {poly R}):
	size (\sum_(i < size p) p`_ i *: F i) = size p.
Proof.
rewrite size_seqbase.
case E: (size p) => [ | n]; first by rewrite big_ord0.
rewrite big_mkcond/=.
rewrite big_ord_recr /=.
rewrite -big_mkcond/=.
case: ifP => [_ | ].
apply/ maxn_idPr.
by apply/ bigmax_leqP => /= i _; rewrite ltnS ltnW.
case/ idP.
have ->: n = (size p).-1 by rewrite E.
by rewrite -lead_coefE lead_coef_eq0 -size_poly_eq0 E.
Qed.

Lemma Fneq0: forall n, F n != 0.
Proof.
by move => n; rewrite -size_poly_gt0 F_size.
Qed.

Lemma seqbase_coef_eq0 n (l: seq R) :
   \sum_(i < n) l`_ i *: F i = 0 -> forall i, (i < n)%N -> l`_i = 0.
Proof.
move=> H i Hi.
suff P j : (0 < j < n)%N -> l`_j = 0.
  have := H.
  case: (n) Hi P => // n1 Hi P /eqP.
  rewrite big_ord_recl big1 => [|j _]; last first.
    by rewrite P ?scale0r // lift0 /= ltnS.
  rewrite addr0.
  move => /eqP eq; move: eq.
  rewrite -polyP => prp.
  specialize (prp 0%nat).
  rewrite coef0 coefZ in prp.
  move: prp.
  rewrite [0%nat]pred_Sn.
  rewrite -(F_size 0).
  rewrite -lead_coefE => prp.
  case: i Hi => [_ |i Hi]; first by apply (@F_lead 0); rewrite rm0.
  by apply: P.
move=> /andP[HP1 HP2].
have /eqP := (H).
rewrite -size_poly_eq0 size_seqbase -leqn0 => /bigmax_leqP/(_ (Ordinal HP2)).
by rewrite /=; case: eqP => //= _ /(_ isT).
Qed.

Lemma polybase_eq0 (p: {poly R}) :
   \sum_(i < size p) p`_ i *: F i = 0 <-> p = 0.
Proof.
split => [prp | ->]; last by rewrite big1 => // i _; rewrite coef0 rm0.
rewrite [LHS]coef_Xn_eq.
rewrite big1 => // i _.
by rewrite (seqbase_coef_eq0 prp); first by rewrite !rm0.
Qed.

End Base.


Require Import Frap Pset9Sig.
(* Authors: Adam Chlipala (adamc@csail.mit.edu),
 * Peng Wang (wangpeng@csail.mit.edu) *)

(* The Forgetful Loop Rule
 *
 * In this pset, we explore a different proof rule for loops, which combines
 * some of the nifty small-footprint reasoning of the frame rule.  Specifically,
 * we consider loops where we traverse linked data structures, *forgetting*
 * about nodes as we pass through them, narrowing our focus to just the subsets
 * of nodes that future loop iterations might touch.  Recall how, to prove
 * linked-list length, we needed to do some grunt work with a predicate for
 * linked-list segments, even though the function will never again access the
 * segments described in the loop invariant.  The forgetful loop rule will allow
 * us to skip the segments and write a loop invariant the matches the overall
 * function specification. *)

(* From the Sig file:
Inductive hoare_triple : forall {result}, assertion -> cmd result -> (result -> assertion) -> Prop :=
(* First, some basic rules that look exactly the same as before *)
| HtReturn : forall P {result : Set} (v : result),
    hoare_triple P (Return v) (fun r => P * [| r = v |])%sep
| HtBind : forall P {result' result} (c1 : cmd result') (c2 : result' -> cmd result) Q R,
    hoare_triple P c1 Q
    -> (forall r, hoare_triple (Q r) (c2 r) R)
    -> hoare_triple P (Bind c1 c2) R

(* THIS RULE IS DIFFERENT. *)
| HtLoop : forall {acc res : Set} (init : acc) (body : acc -> cmd (loop_outcome acc res)) P Q,
    (* As before, the premise forces us to consider any accumulator at the start
     * of a loop iteration, proving a Hoare triple for each case. *)
    (forall acc,
        (* Important difference: now the rule is parameterized over both a
         * precondition [P] and a postcondition [Q], each of which takes, as an
         * extra argument, the latest accumulator value. *)
        hoare_triple (P acc) (body acc)
                     (fun r =>
                        match r with
                        | Done res =>
                          Q acc res
                          (* The loop is done?  Then the postcondition had
                           * better be satisfied directly.  Note that it takes
                           * the "before" and "after" accumulators as arguments.
                           * We'll see shortly why that pays off.... *)
                        | Again acc' =>
                          (* It's time for more iterations?  Then we'd better
                           * satisfy [P] w.r.t. the "after" accumulator, but
                           * with a twist.  We are allowed to *forget* some
                           * state, captured by the arbitrary frame predicate
                           * [R].  The idea is that the state we shunt into [R]
                           * will not be touched again until the loop finishes
                           * running. *)
                          exists R, P acc' * R
                                    (* There is another important requirement on
                                     * [R]: Assume that the loop finishes, so
                                     * that the postcondition [Q] is satisfied
                                     * w.r.t. the new accumulator [acc'].  If we
                                     * *put back* [R], we should then arrive at
                                     * a state where the postcondition is
                                     * satisfied w.r.t. the "before" accumulator
                                     * [acc]! *)
                                    * [| forall r, Q acc' r * R ===> Q acc r |]
                        end%sep))
    -> hoare_triple (P init) (Loop init body) (Q init)
(* All that may be a bit abstract, but we will show an example
 * verification below, to illustrate. *)

| HtFail : forall {result},
    hoare_triple (fun _ => False) (Fail (result := result)) (fun _ _ => False)

| HtRead : forall a R,
    hoare_triple (exists v, a |-> v * R v)%sep (Read a) (fun r => a |-> r * R r)%sep
| HtWrite : forall a v v',
    hoare_triple (a |-> v)%sep (Write a v') (fun _ => a |-> v')%sep
| HtAlloc : forall numWords,
    hoare_triple emp%sep (Alloc numWords) (fun r => [| r <> 0 |] * r |--> zeroes numWords)%sep
(* ----------------------------------------------------^^^^^^
 * DIFFERENCE FROM CLASS: Now we record that a freshly allocated object has
 * a nonnull address, so that we are free to use null (0) for a special purpose
 * in linked data structures. *)

| HtFree : forall a numWords,
    hoare_triple (a |->? numWords)%sep (Free a numWords) (fun _ => emp)%sep

| HtConsequence : forall {result} (c : cmd result) P Q (P' : assertion) (Q' : _ -> assertion),
    hoare_triple P c Q
    -> P' ===> P
    -> (forall r, Q r ===> Q' r)
    -> hoare_triple P' c Q'
| HtFrame : forall {result} (c : cmd result) P Q R,
    hoare_triple P c Q
    -> hoare_triple (P * R)%sep c (fun r => Q r * R)%sep.

Notation "{{ P }} c {{ r ~> Q }}" :=
  (hoare_triple P%sep c (fun r => Q%sep)) (at level 90, c at next level).
*)

(** * EXAMPLE VERIFICATION: linked-list length revisited *)

(* First, here's essentially the same list-predicate definition from class. *)

Fixpoint llist' (ls : list nat) (p : nat) : hprop :=
  match ls with
  | nil => [| p = 0 |]
  | x :: ls' => [| p <> 0 |] * exists p', p |--> [x; p'] * llist' ls' p'
  end%sep.
(* [m] for memory! *)
(* Let's define a less precise version, which forgets exactly which data a list
 * stores, only remembering that there is indeed a list rooted at [p]. *)
Definition llist (p : nat) :=
  (exists ls, llist' ls p)%sep.
(* In general with this pset, we'll work with less precise predicates like this
 * [llist], to give you a bit of a break! *)

(* We can prove some logical equivalences on our predicates. *)

Lemma llist'_null : forall {ls p}, p = 0
  -> llist' ls p === [| ls = nil |].
Proof.
  heq; cases ls; cancel.
Qed.

Theorem llist_null : forall p, p = 0
  -> llist p === emp.
Proof.
  unfold llist; simplify.
  setoid_rewrite (llist'_null H).
  (* setoid_rewrite does not support "with", just positional arguments *)
  heq; cancel.
Qed.

Lemma llist'_nonnull : forall {ls p}, p <> 0
  -> llist' ls p === exists ls' x p', [| ls = x :: ls' |] * p |--> [x; p'] * llist' ls' p'.
Proof.
  heq; cases ls; cancel.
  equality.
  invert H0; cancel.
Qed.

Theorem llist_nonnull : forall {p}, p <> 0
  -> llist p === exists x p', p |--> [x; p'] * llist p'.
Proof.
  unfold llist; simplify.
  setoid_rewrite (llist'_nonnull H).
  heq; cancel.
Qed.

Opaque llist.
(* It's important that we mark [llist] as opaque after we've finished proving
 * the lemmas, so that its definition is never again unfolded.  Rather, we
 * reason about it only with the two lemmas we proved for it. *)

(* Now here's linked-list length again. *)
Definition llength (p : nat) :=
  for a := (p, 0) loop
    if fst a ==n 0 then
      Return (Done (snd a))
    else
      y <- Read (fst a + 1);
      Return (Again (y, snd a + 1))
  done.

(* And here's the simpler proof. However, this time around, we
 * don't prove any functional correctness. We only confirm
 * the absence of memory errors and that if a [llength] call finishes,
 * [p] still points to some linked list.
 *)
Theorem llength_ok : forall p,
  {{llist p}}
    llength p
  {{_ ~> llist p}}.
Proof.
  unfold llength.
  simp.
  (* We have reached the loop, and it's time to pick an invariant.  The
   * forgetful loop rule asks for both a precondition and a postcondition, so
   * the [loop_inv] tactic takes both as separate arguments. *)
  loop_inv (fun a : nat * nat => llist (fst a))
           (fun (a : nat * nat) (_ : nat) => llist (fst a)).
    (* We can use the most natural invariant: there is a list rooted at the first
     * component of the accumulator [a]. *)
  -
    cases (a ==n 0).
    + step.
      cancel.
    + rewrite llist_nonnull by assumption.
      step.
      step.
      simp.
      step.
      (* Here's where we encounter the extra quantified [R] from the forgetful loop
       * rule.  The automation isn't quite smart enough to pick a good [R] for us,
       * and anyway we might prefer to be in control of what we forget!  We use the
       * lemma [exis_right] to manually instantiate an existential quantifier
       * immediately to the right of [===>]. *)
      apply exis_right with (x := ((a+1) |-> r * exists n0, a |-> n0)%sep).
      (* The right choice in this case: forget the list cell that [a] points to.  We
       * are done with this cell and can continue the loop using only the cells that
       * follow it. *)
      cancel.
      rewrite (llist_nonnull n).
      (* We specify the hypothesis [n] of [llist_nonnull] so that Coq chooses to
       * rewrite the correct occurrence of [llist] in the goal.  Try without that
       * detail and watch Coq make the wrong choice! *)
      cancel.
  - cancel.
  - cancel.
Qed.

(** * Binary trees *)

(* Now we define binary trees and ask you to verify two of their classic
 * operations. This verification task only concerns memory safety, not
   functional correctness -- which you already tackled in pset 4! *)

(*
Inductive tree :=
| Leaf
| Node (l : tree) (x : nat) (r : tree).
*)

(* [m] for memory! *)
Fixpoint mtree' (t : tree) (p : nat) : hprop :=
  match t with
  | Leaf => [| p = 0 |]
  | Node l x r => [| p <> 0 |]
                  * exists p1 p2, p |--> [p1; x; p2]
                                  * mtree' l p1
                                  * mtree' r p2
  end%sep.

(* Here's the version that forgets exactly which tree it is. *)
Definition mtree (p : nat) : hprop :=
  (exists t, mtree' t p)%sep.

(* And here's an extra layer of indirection: a mutable pointer to a tree, which
 * comes in handy for operations that modify the tree. *)
Definition mtreep (p : nat) : hprop :=
  (exists p', [| p <> 0 |] * p |-> p' * mtree p')%sep.

(* Your task: verify the lookup and insertion methods below.
 *
 * Before diving into the proof hacking, it might be a good idea to review the
 * relevant material from the lecture. To help you do that, we suggest that you
 * briefly answer each of the questions below. This exercise is not graded, but
 * we hope it will help you understand the material better.
 * We will also reference this list in office hours to see where you
 * might be stuck.
 *
 * - What does A * B mean?
 *
 * - What does h === g mean?
 *
 * - What does [| P |] do? (When is it necessary?)
 *
 * - What does p |--> [x;y] mean?
 *
 * - What does emp mean?
 *
 * - What memory does an empty llist use?
 *
 * - How can code detect that an llist is empty?
 *
 * - What memory does an llist cons cell use?
 *
 * - How can code detect that an llist starts with a cons cell?
 *
 * - If an llist starts with a cons cell, which lemma/theorem can we use to learn something about its tail, and what does it tell us?
 *
 * - What is the difference between llist' and llist?
 *
 * - How do proofs of lemmas about llist use lemmas about llist'?
 *
 * (extra pedantry:)
 *
 * - Have you read the proof of llength_ok?
 *
 * - Have you stepped through the proof of llength_ok?
 *
 * - Have you read the goal before and after every commented tactic invocation in llength_ok?
 *
 * Now proving correctness of binary-tree manipulation shouldn't be too bad. Cheers!
*)

(* IMPORTANT NOTE:
 * The difficulty of this problem set is scoped assuming you will use the
 * tactics included in the Sig file, which are very similar to the ones from
 * lecture!  Trying to attack these proofs from first principles will likely
 * lead to proof size spiraling out of control. The main ingredients you should
 * expect to reuse are:
 *  - [step]: A tactic to choose the right Hoare-logic rule to apply next, when
 *    your goal is a Hoare triple
 *  - [cancel]: A tactic to prove an implication between two separation-logic
 *    assertions, or reduce it to a simpler implication by *cancel*ing matching
 *    subformulas; before calling [cancel], be sure you have marked [Opaque]
 *    all of the predicates standing for data structures!  There is an [Opaque
 *    mtree] command right below for precisely this reason; please put any lemmas
 *    about [mtree] before it and all program logic proofs after.
 *  - [heq]: A tactic to reduce a separation-logic assertion implication (using
 *    [===]) to two implications (using [===>])
 *  - [loop_inv P0 Q0]: A tactic to apply the loop rule, given an invariant
 *    split in a peculiar way; see above for an example
 *  - [setoid_rewrite H]: A tactic to rewrite within a separation-logic
 *    implication, using an equivalence that applies to one of its subterms
 *
 * Our solution also has a few direct uses of rules [HtConsequence] and
 * [HtFrame], sometimes using the [with] syntax of [apply] to specify values
 * for some variables that appear in these rules' statements. In particular,
 * you can use [HtConsequence] before [HtFrame] to ensure that the predicate
 * that you want to frame out is in the desired syntactic position so that
 * [HtFrame] will apply. The [exis_right] lemma will also be handy, as used
 * in the example above. *)

(* Space is provided here for additional lemmas about [mtree] and [mtree']. *)

Lemma mtree'_null : forall {t p}, p = 0
  -> mtree' t p === [| t = Leaf |].
Proof.
  heq; cases t; cancel.
Qed.

Theorem mtree_null : forall p, p = 0
  -> mtree p === emp.
Proof.
  unfold mtree; simplify.
  setoid_rewrite (mtree'_null H).
  (* setoid_rewrite does not support "with", just positional arguments *)
  heq; cancel.
Qed.

Lemma mtree'_nonnull : forall {t p}, p <> 0
  -> mtree' t p === exists l r lp v rp, [| t = Node l v r |] * p |--> [lp ; v; rp] * mtree' l lp * mtree' r rp.
Proof.
  heq; cases t; cancel.
  equality.
  invert H0; cancel.
Qed.

Theorem mtree_nonnull : forall {p}, p <> 0
  -> mtree p === exists v lp rp, p |--> [lp; v; rp] * mtree lp * mtree rp.
Proof.
  unfold mtree; simplify.
  setoid_rewrite (mtree'_nonnull H).
  heq; cancel.
Qed.

Opaque mtree.
(* ^-- Keep predicates opaque after you've finished proving all the key
* algebraic properties about them, in order for them to work well with
* the [cancel] tactic. *)

(* Here's the usual lookup operation. *)
Definition lookup (x p : nat) :=
  t <- Read p; (* First peel away the initial layer of indirection.
                * You will want to use the regular old frame rule to forget
                * about some of the state that you won't need after this
                * point! *)
  for a := t loop
    (* The accumulator tells us: the node of the tree we have reached
     * (for [Again]) or whether the key [x] has been found (for [Done]). *)
    if a ==n 0 then
      (* Oh, the pointer is null.  Sorry, didn't find [x]. *)
      Return (Done false)
    else
      (* Read the data value of the current node (which must be nonnull). *)
      y <- Read (a + 1);
      if x ==n y then
        (* Found it! *)
        Return (Done true)
      else if x <=? y then
        (* The key must be earlier in the tree.  Read the left-child pointer and
         * continue looping with it. *)
        l <- Read a;
        Return (Again l)
      else
        (* The key must be later in the tree.  Read the right-child pointer and
         * continue looping with it. *)
        r <- Read (a + 1 + 1);
        (* Why [+ 1 + 1] instead of [+ 2]?  It happens to work better with the
         * automation we're using. ;) *)
        Return (Again r)
  done.

Theorem lookup_ok : forall x p,
  {{mtreep p}}
    lookup x p
  {{_ ~> mtreep p}}.
Proof.
  intros.
  unfold mtreep in *.
  repeat step.
  simp.
  apply HtWeaken with (P := (mtree r * p |-> r * [|p = 0 -> False|])%sep); cancel.
  apply HtStrengthen with (Q := fun _ => (mtree r * p |-> r * [| p <> 0 |])%sep); cancel.
  repeat eapply HtFrame.
  loop_inv (fun r : nat => mtree r)
           (fun (r : nat) (_ : bool) => mtree r); cancel.
  cases (acc ==n 0).
  + step.
    cancel.
  + step.
    setoid_rewrite mtree_nonnull; try assumption.
    step.
    cases (x ==n r0).
    ++
      step.
      cancel.
      rewrite (mtree_nonnull n).
      cancel.
    ++
      cases (x <=? r0); (
        repeat step;
        eapply exis_right;
        repeat cancel;
        rewrite (mtree_nonnull n);
        cancel).
Qed.

(* And here's the operation to add a new key to a tree. *)
Definition insert (x p : nat) :=
  for a := p loop
    (* Note that now the accumulator is not the latest tree root, but instead
     * *a pointer to it*, so that we may overwrite that pointer if necessary.
     * We start by reading the actual root out of the pointer [p]. *)
    q <- Read a;
    if q ==n 0 then
      (* It's a null pointer?  Perfect.  This is the spot to insert a new
       * node. *)
      node <- Alloc 3;
      (* Initialize its data field with [x]. *)
      _ <- Write (node + 1) x;
      (* Redirect the pointer [p] to the new node. *)
      _ <- Write a node;
      Return (Done tt)
    else
      (* Nonnull?  Read the data field into [y]. *)
      y <- Read (q + 1);
      if x <=? y then
        (* The right spot to insert must be to the left.  Recurse thataway. *)
        Return (Again q)
      else
        (* The right spot to insert must be to the right.  Recurse thataway. *)
        Return (Again (q + 1 + 1))
  done.

(* Something very subtle happened in that loop: we iterated using a pointer into
 * *the interior of a struct*, in each branch of the last [if]!  This is a fun
 * example of the kinds of tricks that can be played in a low-level language,
 * and the verification techniques are up to the challenge. *)

Theorem insert_ok : forall x p,
  {{mtreep p}}
    insert x p
  {{_ ~> mtreep p}}.
Proof.
  simp.
  loop_inv (fun p : nat => mtreep p)
           (fun (p : nat) (_  : unit) => mtreep p).
  all: unfold mtreep in *; cancel.
  cases (r ==n 0).
  ++
    repeat step.
    cancel.
    rewrite (mtree_nonnull H).
    cancel.
    assert (mtree 0 === emp) by (apply mtree_null; linear_arithmetic).
    rewrite H1; cancel.
  ++
    step.
    apply mtree_nonnull in n.
    apply HtWeaken with (P := (mtree r * acc |-> r * [|acc = 0 -> False|])%sep); cancel.
    repeat eapply HtFrame.
    rewrite n.
    step.
    cases (x <=? r0).
    +++
      step.
      eapply exis_right.
      repeat cancel.
      rewrite H0.
      cancel.
    +++
      step.
      eapply mtree_nonnull in n.
      eapply exis_right.
      repeat cancel.
      rewrite n.
      cancel.
      linear_arithmetic.
Qed.

(* Our solution also includes a proof that the Hoare triples in this pset
 * correspond to the usual operational semantics... which you do not need
 * to prove. *)


Require Import List.
Import ListNotations.

Example inversion_ex_3 : forall (X : Type) (x y z : X) (l j : list X),
  x :: y :: l = z :: j ->
  y :: l = x :: j ->
  x = y.
Proof.
  intros X x y z l j eq1 eq2.
  inversion eq1.
  inversion eq2.
  symmetry.
  exact H0.
Qed.


(* This file is part of the Linear Logic  formalization in Coq: https://github.com/meta-logic/coq-ll *)

(** ** Cut Elimination
Here we prove cut-elimination for linear logic and, as a corollary, we show the consistency of the logic
 *)

(*Add LoadPath "../". *)

Require Import Coq.Relations.Relations.
Require Import Coq.Arith.EqNat.
Require Import Coq.Classes.Morphisms.
Require Import Coq.Setoids.Setoid.
Require Export Coq.Sorting.PermutSetoid.
Require Export Coq.Sorting.PermutEq.
Require Import Coq.Program.Equality.
(* Require Export StructuralRules. *)
Require Export LL.SequentCalculi.
Require Export LL.Eqset.
Require Import LL.CutCases.

Module CElimination (DT : Eqset_dec_pol).
  Module Export Sys :=  CCases DT.

  Ltac simpl_duals H :=
    try rewrite one_bot in H;
    try rewrite bot_one in H;
    try rewrite top_zero in H;
    try rewrite zero_top in H;
    try rewrite atom_perp in H;
    try rewrite perp_atom in H;
    try rewrite tensor_par in H;
    try rewrite par_tensor in H;
    try rewrite with_plus in H;
    try rewrite plus_with in H;
    try rewrite bang_quest in H;
    try rewrite quest_bang in H.

  Theorem cut_aux B L M1 M2 F :
    L =mul= M1 ++ M2 ->
    0 |~> 0 ; B ; F :: M1 ->
                  0 |~> 0 ; B ; F° :: M2 ->
                                exists m, m |~> 0 ; B ; L.
  Proof.
    intros P Hn1 Hn2.
    inversion Hn1; subst.
    *
      apply axPair in H.
      destruct H;
        firstorder; subst.
    +
      simpl_duals Hn2.
      inversion Hn2; subst.
    - (* init rule *)
      apply axPair in H.
      destruct H;
        firstorder; [lexp_contr H |].
      apply PerpEq in H; subst.
      eexists; eapply sig3_init.
      eauto.
    - (* contradiction *)
      eapply resolvers2 in H; firstorder; lexp_contr H.
    - (* top rule *)
      simpl_cases2.
      top_commutes.
      + (* symmetric case *)
        simpl_duals Hn2.
        inversion Hn2; subst.
    - (* init rule *)
      apply axPair in H.
      destruct H;
        firstorder; [|lexp_contr H].
      apply AtomEq in H; subst.
      eexists; eapply sig3_init.
      eauto.
    - (* contradiction *)
      eapply resolvers2 in H; firstorder; lexp_contr H.
    - (* top rule *)
      simpl_cases2.
      top_commutes.
      *
        eapply resolvers2 in H; firstorder.
        subst.
        simpl_duals Hn2.
        inversion Hn2; subst.
      + (* contradiction *)
        apply resolvers in H; firstorder; lexp_contr H.
      + (* contradiction *)
        eapply resolvers2 in H; firstorder; lexp_contr H.
      + (* top rule *)
        simpl_cases2.
        top_commutes.
        *
          destruct (FEqDec F Top).
      +
        subst. simpl_cases4.
        simpl_duals Hn2.
        inversion Hn2; subst.
    - (* contradiction *)
      apply resolvers in H0; firstorder; lexp_contr H0.
    - (* contradiction *)
      eapply resolvers2 in H0; firstorder; lexp_contr H0.
    - (* top rule *)
      simpl_cases2.
      top_commutes.
      + (* top rule *)
        simpl_cases2.
        top_commutes.
  Qed.
  Arguments cut_aux[B L M1 M2 F].

  Theorem ccut_aux  B L M1 M2 F w :
    S w = Lexp_weight (! F) -> L =mul= M1 ++ M2 ->
    0 |~> 0 ; B ; ! F :: M1 -> 0 |~> 0 ; F° :: B ; M2 ->
                                                   exists m, m |~> 0 ; B ; L.
  Proof.
    intros Hw P Hn1 Hn2.
    inversion Hn1; subst;
      [ apply resolvers in H; firstorder; lexp_contr H |
        apply resolvers2 in H; firstorder; lexp_contr H |
        simpl_cases2; top_commutes].
  Qed.
  Arguments ccut_aux [B L M1 M2 F w].

  Lemma cut_elimination_base: forall n B L,
      n |~> 1 ; B ; L -> exists m, m |~> 0 ; B ; L.
  Proof.
    intros.
    revert dependent B.
    revert dependent L.
    induction n using strongind; intros.
    - inversion H.
    - inversion H0; subst.
      + assert (exists m, m |~> 0; B; M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_bot _ Ht); eauto.
      + assert (exists m, m |~> 0; B; F :: G :: M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_par _ Ht); resolve_rewrite.
      + apply Nat.eq_add_1 in H2.
        destruct H2; destruct H1; subst.
        assert (exists m, m |~> 0; B; F :: M) as Hyp.
        refine (H _ _ _ _ H4); resolve_max.
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_tensor _ Ht H5); resolve_rewrite; auto.
        assert (exists m, m |~> 0; B; G :: N) as Hyp.
        refine (H _ _ _ _ H5); resolve_max.
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_tensor _ H4 Ht); resolve_rewrite.
      + assert (exists m, m |~> 0; B; F :: M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_plus1 _ Ht); resolve_rewrite.
      + assert (exists m, m |~> 0; B; G :: M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_plus2 _ Ht); resolve_rewrite.
      + apply Nat.eq_add_1 in H2.
        destruct H2; destruct H1; subst.
        assert (exists m, m |~> 0; B; F :: M) as Hyp.
        refine (H _ _ _ _ H4); resolve_max.
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_with _ Ht H5); resolve_rewrite. auto.
        assert (exists m, m |~> 0; B; G :: M) as Hyp.
        refine (H _ _ _ _ H5); resolve_max.
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_with _ H4 Ht); resolve_rewrite.
      + rewrite H3 in H4.
        assert (exists m, m |~> 0; B; F:: L) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_copy _ _ Ht); auto.
      + assert (exists m, m |~> 0; F:: B; M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_quest _ Ht); auto.
      + assert (exists m, m |~> 0; B; [F]) as Hyp by solve [eapply H; auto].
        destruct Hyp as [t Ht]; eexists;
          refine (sig3_bang _ Ht); auto.
      + assert (exists m, m |~> 0; B; Subst FX t :: M) as Hyp by solve [eapply H; auto].
        destruct Hyp as [p Hp]; eexists; eapply sig3_ex; eauto.
      +
        assert (forall y : Term, exists m, m |~> 0; B; Subst FX y :: M) as Hyp.
        intro.
        assert (n |~> 1; B; Subst FX y :: M) by auto.
        eapply H; auto.
        apply fx_swap_sig3h in Hyp.
        destruct Hyp as [p Hp];
          eexists; eapply sig3_fx; eauto.
      + (* caso CUT *)
        clear H0. clear H.
        revert dependent B.
        revert dependent L.
        revert dependent n.

        dependent induction w using strongind;

          dependent induction h using strongind; intros.

        * (* w=h=0 *)
          inversion H3; subst.
          **
            cut_free.
            simpl in H0, H1.
            cut_free.
            eapply cut_aux; eauto.
          **
            cut_free.
            simpl in H1, H0.
            inversion H0.
        *
          inversion H3; subst.
          **

            cut_free.
            simpl in H1.
            rename H1 into Hw, H2 into Hh.
            rename H4 into P, H5 into Hn1, H9 into Hn2.
            rename m into n1, n0 into n2.
            rename M into M1, N into M2.
            caseLexp F; subst; simpl; try inversion Hw;
              simpl_duals Hn1; simpl_duals Hn2.
            refine (aux_CUT_ap _ Hh _ P Hn1 Hn2); eauto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_ap _ Hh _ _ Hn2 Hn1); eauto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_bo _ Hh _ _ Hn2 Hn1); eauto.
            refine (aux_CUT_bo _ Hh _ _ Hn1 Hn2); eauto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_tz _ Hh _ _ Hn2 Hn1); eauto.
            refine (aux_CUT_tz _ Hh _ _ Hn1 Hn2); eauto.
          **
            cut_free.
            simpl in H1.
            inversion H1.
        * (* h=0 e S w *)
          inversion H3; subst.
          **
            cut_free.
            simpl in H2.
            cut_free.
            eapply cut_aux; eauto.
          **
            cut_free.
            simpl in H2.
            cut_free.
            eapply ccut_aux; eauto.
        * inversion H3; subst; try cut_free.
          -- (** cut *)
            rename H2 into Hw, H4 into Hh.
            rename H5 into P, H6 into Hn1, H10 into Hn2.
            rename m into n1, n0 into n2.
            rename M into M1, N into M2.
            clear H3.
            caseLexp F; subst; simpl; try inversion Hw;
              simpl_duals Hn1; simpl_duals Hn2.
            refine (aux_CUT_tp _ _ _ _ P Hn1 Hn2); eauto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_tp _ _ _ _ P Hn2 _); eauto.
            rewrite <- !lweight_dual; auto.
            rewrite <- !Ng_involutive; auto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_pw _ _ _ _ P Hn2 _); eauto.
            rewrite <- !lweight_dual; auto.
            rewrite <- !Ng_involutive; auto.
            refine (aux_CUT_pw _ _ _ _ P Hn1 Hn2); eauto.
            refine (aux_CUT_bq _ _ _ _ P Hn1 Hn2); eauto.
            rewrite union_comm in P.
            rewrite Nat.add_comm in Hh.
            refine (aux_CUT_bq _ _ _ _ P Hn2 _); eauto.
            rewrite <- !lweight_dual; auto.
            rewrite <- !Ng_involutive; auto.
            refine (aux_CUT_ef _ _ Hw _ P Hn1 Hn2 ); eauto.
            refine (aux_CUT_fe _ _ Hw _ P Hn1 Hn2 ); eauto.

          -- (** ccut *)
            rename H2 into Hw, H4 into Hh.
            rename H5 into P, H6 into Hn1, H10 into Hn2.
            rename m into n1, n0 into n2.
            rename M into M1, N into M2.

            inversion Hn1;subst.
            eapply resolvers in H1; intuition;
              lexp_contr H2.
            eapply resolvers2 in H1; intuition;
              lexp_contr H2.

            (* Top Commutes *)
            simpl_cases2.
            eexists.
            rewrite P.
            rewrite H2.
            eapply sig3_top. eauto.

            (* Bot Commutes *)
            simpl_cases2.
            assert (exists m, m |~> 0 ; B; x ++ M2) as Hyp.
            rewrite H1 in H2.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H2 Hn2);
                  auto; try resolve_rewrite]; resolve_max.

            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_bot _ Ht); auto; resolve_rewrite.

            (* Par Commutes *)
            simpl_cases2.
            assert (exists m, m |~> 0 ; B; F0 :: G :: (x ++ M2)) as Hyp.
            rewrite H1 in H2;
              rewrite union_rotate_cons in H2.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H2 Hn2);
                  auto; try resolve_rewrite]; resolve_max.

            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_par _ Ht); auto; resolve_rewrite.

            (* Tensor Commutes *)
            simpl_cases2.
            cut_free.
            simpl_cases_tns.

            assert (exists m, m |~> 0 ; B; F0 :: (L1 ++ M2)) as Hyp.
            rewrite HL1 in H4;
              rewrite meq_swap_cons in H4.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H4 Hn2);
                  auto; try resolve_rewrite].
            inversion Hh.  resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_tensor _ Ht H6); auto; resolve_rewrite.

            assert (exists m, m |~> 0 ; B; G :: (L2 ++ M2)) as Hyp.

            rewrite HL2 in H6;
              rewrite meq_swap_cons in H6.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H6 Hn2);
                  auto; try resolve_rewrite].
            inversion Hh.  resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_tensor _ H4 Ht); auto; resolve_rewrite.

            (* Plus1 Commutes *)
            simpl_cases2.
            assert (exists m : nat, m |~> 0 ; B; F0 :: x ++ M2) as Hyp.
            rewrite H1 in H2;
              rewrite meq_swap_cons in H2.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H2 Hn2);
                  auto; try resolve_rewrite]; resolve_max.

            destruct Hyp as [t Ht];
              try rewrite union_assoc in Ht;
              eexists;
              refine (sig3_plus1 _ Ht); auto; resolve_rewrite.

            (* Plus2 Commutes *)
            simpl_cases2.
            assert (exists m : nat, m |~> 0 ; B; G :: x ++ M2) as Hyp.
            rewrite H1 in H2;
              rewrite meq_swap_cons in H2.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H2 Hn2);
                  auto; try resolve_rewrite]; resolve_max.

            destruct Hyp as [t Ht];
              try rewrite union_assoc in Ht;
              eexists;
              refine (sig3_plus2 _ Ht); auto; resolve_rewrite.

            (* With Commutes *)
            simpl_cases2.
            cut_free.
            assert (exists m, m |~> 0 ; B; F0 :: (x ++ M2)) as Hyp1.
            rewrite H2 in H4;
              rewrite meq_swap_cons in H4;
              refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H4 Hn2);
                  auto; try resolve_rewrite];resolve_max.
            inversion Hh;  resolve_max.

            assert (exists m, m |~> 0 ; B; G :: (x ++ M2)) as Hyp2.
            rewrite H2 in H6; rewrite meq_swap_cons in H6;
              refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H6 Hn2);
                  auto; try resolve_rewrite];resolve_max.
            inversion Hh;  resolve_max.
            destruct Hyp1 as [t1 Ht1];
              destruct Hyp2 as [t2 Ht2];

              eexists;
              refine (sig3_with _ Ht1 Ht2); resolve_rewrite.

            (* Copy Commutes *)
            assert (exists m, m |~> 0 ; B ; F0 :: L) as Hyp.
            rewrite meq_swap_cons in H2.
            rewrite H2 in H4.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H4 Hn2);
                  auto; try resolve_rewrite];resolve_max.

            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_copy _ _ Ht); auto; resolve_rewrite.

            (* Quest Commutes *)
            simpl_cases2.

            assert (exists m, m |~> 0 ; F0 :: B; x ++ M2) as Hyp.
            rewrite H1 in H2.
            eapply height_preserving_weakning_sig3 with (D:=[F0]) in Hn2.
            rewrite <- app_comm_cons in Hn2.
            rewrite union_comm in Hn2.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ H2 Hn2);
                  auto; try resolve_rewrite]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_quest _ Ht); auto; resolve_rewrite.

            (* Bang Principal *)
            eapply resolvers2 in H1; intuition.
            rewrite H5 in P, Hn1; clear H5.
            rewrite <- P in *. clear P.
            LexpSubst.

            inversion Hn2; subst; inversion Hh; try cut_free.
            eexists; eapply sig3_init; eassumption.
            eexists; eapply sig3_one; eassumption.
            eexists; eapply sig3_top; eassumption.
            assert (exists m : nat, m |~> 0; B; M) as Hyp.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H4);
                auto]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_bot _ Ht); auto; resolve_rewrite.
            assert (exists m, m |~> 0; B; F0 :: G :: M) as Hyp.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H4);
                auto]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_par _ Ht); auto; resolve_rewrite.
            assert (exists m, m |~> 0; B; F0 :: M) as Hyp1.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H5);
                auto]; resolve_max.
            assert (exists m, m |~> 0; B; G :: N) as Hyp2.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H7);
                auto]; resolve_max.
            destruct Hyp1 as [t1 Ht1];
              destruct Hyp2 as [t2 Ht2];
              eexists;
              refine (sig3_tensor _ Ht1 Ht2); auto; resolve_rewrite.
            assert (exists m, m |~> 0; B; F0 :: M) as Hyp.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H4);
                auto]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_plus1 _ Ht); auto; resolve_rewrite.
            assert (exists m, m |~> 0; B; G :: M) as Hyp.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H4);
                auto]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_plus2 _ Ht); auto; resolve_rewrite.
            assert (exists m, m |~> 0; B; F0 :: M) as Hyp1.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H5);
                auto]; resolve_max.
            assert (exists m, m |~> 0; B; G :: M) as Hyp2.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H7);
                auto]; resolve_max.
            destruct Hyp1 as [t1 Ht1];
              destruct Hyp2 as [t2 Ht2];
              eexists;
              refine (sig3_with _ Ht1 Ht2); auto; resolve_rewrite.

            2:{
              assert (exists m, m |~> 0; F0 :: B; M) as Hyp.
              rewrite meq_swap_cons in H4.
              eapply height_preserving_weakning_sig3 with (D:=[F0]) in Hn1.
              rewrite union_comm in Hn1.

              refine (H0 _ _ _ _ _ _);
                [ |
                  change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ Hn1 H4);
                  auto]; resolve_max.
              destruct Hyp as [t Ht];
                eexists;
                refine (sig3_quest _ Ht); auto; resolve_rewrite.
            }

            2:{
              assert (exists m, m |~> 0; B; [F0]) as Hyp.
              refine (H0 _ _ _ _ _ _);
                [ |
                  change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ Hn1 H4);
                  auto]; resolve_max.
              destruct Hyp as [t Ht];
                eexists;
                refine (sig3_bang _ Ht); auto; resolve_rewrite.
            }

            destruct (FEqDec F0 F°); subst.
            assert (exists m, m |~> 0; B; F° :: L) as Hyp.
            rewrite H4 in H5.
            refine (H0 _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0);
                  refine (sig3_ccut _ _ _ Hn1 H5);
                  auto; resolve_rewrite]; resolve_max.

            destruct Hyp as [t Ht].

            refine (H _ _ _ _ _ _ _);
              [ | change (0%nat) with (plus 0 0)].
            2:{  refine (sig3_cut _ _ _ H2 Ht); auto. }
            inversion Hw.  unfold Lexp_weight; auto.
            apply not_eqLExp_sym in n1.
            simpl_cases2.
            assert (exists m, m |~> 0; B; F0 :: L) as Hyp.
            rewrite H4 in H5.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H5);
                auto]; resolve_max.
            destruct Hyp as [t Ht];
              eexists;
              refine (sig3_copy _ _ Ht); auto.

            assert (exists m, m |~> 0; B; Subst FX t :: M) as Hyp.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H4);
                auto]; resolve_max.
            destruct Hyp as [p Hp];
              eexists;
              refine (sig3_ex _ Hp); auto; resolve_rewrite.

            assert (forall y : Term, exists m, m |~> 0; B; Subst FX y :: M) as Hyp.
            intro.
            assert (n0 |~> 0; F° :: B; Subst FX y :: M) by auto.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ Hn1 H5);
                auto]; resolve_max.
            apply fx_swap_sig3h in Hyp.
            destruct Hyp as [p Hp];
              eexists;
              refine (sig3_fx _ Hp); auto; resolve_rewrite.

            simpl_cases2.
            assert (exists m, m |~> 0; B; Subst FX t :: (x++M2)) as Hyp.
            rewrite H1 in H2.
            rewrite meq_swap_cons in H2.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ H2 Hn2) ;
                auto]; resolve_max.
            destruct Hyp as [p Hp];
              eexists;
              refine (sig3_ex _ Hp); auto; resolve_rewrite.

            simpl_cases2.
            assert (forall y : Term, exists m, m |~> 0; B; Subst FX y :: (x++M2)) as Hyp.
            intro.
            assert (n |~> 0; B; Subst FX y :: M) by auto.
            rewrite H1 in H5.
            rewrite meq_swap_cons in H5.
            refine (H0 _ _ _ _ _ _);
              [ |
                change (0%nat) with (plus 0 0);
                refine (sig3_ccut _ _ _ H5 Hn2) ;
                auto]; resolve_max.
            apply fx_swap_sig3h in Hyp.
            destruct Hyp as [p Hp];
              eexists;
              refine (sig3_fx _ Hp); auto; resolve_rewrite.
  Qed.

  Theorem cut_elimination : forall B L n c, n |~> c ; B ; L ->
                                                          exists m, m |~> 0 ; B ; L.
  Proof.
    intros.
    dependent induction n generalizing B L using strongind; induction c using strongind.
    + (* Case <0,0> *)
      eexists; eauto.
    + (* Case <0, S c> *)
      inversion H.
    + (* case <Sn, 0> *)
      inversion H0;subst;try (rewrite H2).
      ++ (* BOTTOM*) eexists. eapply sig3_bot;eauto.
      ++ (* PAR *)  eexists. eapply sig3_par;eauto.
      ++ (* TENSOR *)
        cut_free.
        eexists.
        change (0%nat) with (plus 0 0).
        eapply sig3_tensor;eauto.
      ++ (* OPLUS F *)  eexists. eapply sig3_plus1; eauto.
      ++ (* OPLUS G *)  eexists. eapply sig3_plus2; eauto.
      ++ (* WITH *)
        cut_free.
        eexists.
        change (0%nat) with (plus 0 0).
        eapply sig3_with;eauto.
      ++ (* COPY *) eexists. eapply sig3_copy; eauto.
      ++ (* QUEST *) eexists. eapply sig3_quest; eauto.
      ++ (* BANG *) eexists. eapply sig3_bang; eauto.
      ++ (* EX *) eexists. eapply sig3_ex; eauto.
      ++ (* FX *) eexists. eapply sig3_fx; eauto.
    + (* Case <Sn, Sc> *)
      rename H1 into IHn.

      inversion H0;subst.
      ++   (* bottom *)
        apply H in H3; auto. destruct H3.
        eexists.
        eapply sig3_bot; eauto.

      ++   (* PAR *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_par; eauto.

      ++
        (* TENSOR *)
        apply H in H4;auto. destruct H4.
        apply H in H5;auto. destruct H5.
        change (0%nat) with (plus 0 0).
        eexists.
        eapply sig3_tensor;eauto.
      ++ (* PLUS1 *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_plus1; eauto.

      ++ (* PLUS2 *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_plus2; eauto.

      ++ (* WITH *)
        apply H in H4;auto.  destruct H4.
        apply H in H5;auto.  destruct H5.

        change (0%nat) with (plus 0 0).
        eexists.
        eapply sig3_with;eauto.
      ++ (* copy *)
        apply H in H4;auto.  destruct H4.
        eexists.
        eapply sig3_copy; eauto.
      ++  (* quest *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_quest;eauto.
      ++ (* bang *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_bang; eauto.
      ++ (* ex *)
        apply H in H3;auto.  destruct H3.
        eexists.
        eapply sig3_ex; eauto.
      ++ (* fx *)
        assert (forall y : Term, exists m, m |~> 0; B; Subst FX y :: M) as Hyp.
        intro.
        assert (n |~> S c; B; Subst FX y :: M) by auto.
        eapply H; auto.
        apply fx_swap_sig3h in Hyp.
        destruct Hyp as [p Hp].
        eexists.
        eapply sig3_fx;eauto.
      ++ (* CUT *)
        destruct H3.
      --
        apply H in H4;auto. destruct H4.
        apply H in H5;auto. destruct H5.

        assert(Hc: S(Init.Nat.max x0 x) |~> S(0 + 0) ; B ; L).
        eapply sig3_CUT.
        eapply  sig3_cut;eauto.

        apply cut_elimination_base in Hc.
        destruct Hc.
        eexists. eauto.
      --
        (* CCUT *)
        apply H in H4;auto. destruct H4.
        apply H in H5;auto. destruct H5.

        assert(Hc: S(Init.Nat.max x0 x) |~> S(0 + 0) ; B ; L).
        eapply sig3_CUT.
        eapply  sig3_ccut;eauto.

        apply cut_elimination_base in Hc.
        destruct Hc.
        eexists; eauto.
  Qed.

  Theorem weak_consistency : forall n c, ~ sig3 n c [] [0].
  Proof.
    intros.
    intro.
    apply cut_elimination in H.
    destruct H.
    inversion H; subst;
      try simpl_cases0.
    eapply DestructMulFalse; eauto.
  Qed.

  Lemma strong_cons : forall n c, ~ sig3 n c [] [].
  Proof.
    intros.
    intro.
    apply cut_elimination in H.
    destruct H.
    inversion H; subst;
      try solve [eapply DestructMulFalse; eauto].
  Qed.

  Theorem strong_consistency : forall n c, ~ sig3 n c [] [Bot].
  Proof.
    intros.
    intro.
    apply cut_elimination in H.
    destruct H.
    inversion H; subst;
      try simpl_cases0.
    rewrite Hm in H1.
    eapply strong_cons; eauto.
    eapply DestructMulFalse; eauto.
  Qed.
End CElimination.


(** * Natural transformations

Authors: Benedikt Ahrens, Chris Kapulkin, Mike Shulman (January 2013)

*)

(** ** Contents
  - Definition of natural transformations
  - Equality is pointwise equality
  - Identity natural transformation
  - Composition of natural transformations
  - Natural isomorphisms
 *)

Require Import UniMath.Foundations.Propositions.
Require Import UniMath.MoreFoundations.Tactics.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Isos.
Require Import UniMath.CategoryTheory.Core.TransportMorphisms.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.Core.Univalence.

Local Open Scope cat.
Section nat_trans.

  (** ** Definition of natural transformations *)

Definition nat_trans_data {C C' : precategory_ob_mor} (F F' : functor_data C C'): UU :=
  ∏ x : ob C, F x -->  F' x.

Definition is_nat_trans {C C' : precategory_data}
  (F F' : functor_data C C') (t : nat_trans_data F F') :=
  ∏ (x x' : ob C)(f : x --> x'), # F f · t x' = t x · #F' f.

Lemma isaprop_is_nat_trans (C C' : precategory_data) (hs: has_homsets C')
  (F F' : functor_data C C') (t : nat_trans_data F F'):
  isaprop (is_nat_trans F F' t).
Proof.
  repeat (apply impred; intro).
  apply hs.
Qed.

Definition nat_trans {C C' : precategory_data} (F F' : functor_data C C') : UU :=
  total2 (fun t : nat_trans_data F F' => is_nat_trans F F' t).

Notation "F ⟹ G" := (nat_trans F G) (at level 39) : cat.
(* to input: type "\==>" with Agda input method *)

Definition make_nat_trans {C C' : precategory_data} (F F' : functor_data C C')
           (t : nat_trans_data F F') (H : is_nat_trans F F' t) :
           nat_trans F F'.
Proof.
exists t.
exact H.
Defined.

Lemma isaset_nat_trans {C C' : precategory_data} (hs: has_homsets C')
  (F F' : functor_data C C') : isaset (nat_trans F F').
Proof.
  apply (isofhleveltotal2 2).
  + apply impred; intro t; apply hs.
  + intro x; apply isasetaprop, isaprop_is_nat_trans, hs.
Qed.

Definition nat_trans_data_from_nat_trans {C C' : precategory_data}
  {F F' : functor_data C C'}(a : nat_trans F F') :
  nat_trans_data F F' := pr1 a.

Definition nat_trans_data_from_nat_trans_funclass {C C' : precategory_data}
  {F F' : functor_data C C'}(a : nat_trans F F') :
  ∏ x : ob C, F x -->  F' x := pr1 a.
Coercion nat_trans_data_from_nat_trans_funclass : nat_trans >-> Funclass.

Definition nat_trans_ax {C C' : precategory_data}
  {F F' : functor_data C C'} (a : nat_trans F F') :
  ∏ (x x' : ob C)(f : x --> x'),
    #F f · a x' = a x · #F' f := pr2 a.

(** Equality between two natural transformations *)

Lemma nat_trans_eq {C C' : precategory_data} (hs: has_homsets C')
  (F F' : functor_data C C')(a a' : nat_trans F F'):
  (∏ x, a x = a' x) -> a = a'.
Proof.
  intro H.
  assert (H' : pr1 a = pr1 a').
  { now apply funextsec. }
  apply (total2_paths_f H'), proofirrelevance, isaprop_is_nat_trans, hs.
Qed.

Lemma nat_trans_eq_alt {C C' : category} (F F' : functor C C') (a a' : nat_trans F F'):
  (∏ x, a x = a' x) -> a = a'.
Proof.
  apply nat_trans_eq.
  apply homset_property.
Qed.

Section nat_trans_eq.

  Context {C D : precategory}.
  Variable hsD : has_homsets D.
  Context {F G : functor C D}.
  Variables alpha beta : nat_trans F G.

  Definition nat_trans_eq_weq : (alpha = beta) ≃ (∏ c, alpha c = beta c).
  Proof.
    eapply weqcomp.
    - apply subtypeInjectivity.
      intro x. apply isaprop_is_nat_trans. apply hsD.
    - apply weqtoforallpaths.
  Defined.

End nat_trans_eq.

(* Can be given as an instance of general equality lemmas,
but useful to have specifically defined. *)
Definition nat_trans_eq_pointwise {C C' : precategory_data}
   {F F' : functor_data C C'} {a a' : nat_trans F F'}:
      a = a' -> ∏ x, a x = a' x.
Proof.
  intro. apply toforallpaths, maponpaths. assumption.
Qed.

(** a more intuitive variant of [functor_data_eq] *)
Lemma functor_data_eq_from_nat_trans (C C': precategory) (F F' : functor_data C C')
      (H : F ~ F') (H1 : is_nat_trans F F' (fun c:C => idtomor _ _ (H c))) :
      F = F'.
Proof.
  apply (functor_data_eq _ _ _ _ H).
  intros c1 c2 f.
  rewrite double_transport_idtoiso.
  rewrite <- assoc.
  apply z_iso_inv_on_right.
(* make the coercion visible: *)
  unfold z_iso_mor.
  do 2 rewrite eq_idtoiso_idtomor.
  apply H1.
Qed.

(** ** Identity natural transformation *)

Lemma is_nat_trans_id {C : precategory_data}{C' : precategory}
  (F : functor_data C C') : is_nat_trans F F
     (λ c : ob C, identity (F c)).
Proof.
  intros ? ? ? .
  now rewrite id_left, id_right.
Qed.

Definition nat_trans_id {C:precategory_data}{C' : precategory}
  (F : functor_data C C') : nat_trans F F :=
    tpair _ _ (is_nat_trans_id F).

(** ** Composition of natural transformations *)

Lemma is_nat_trans_comp {C : precategory_data}{C' : precategory}
  {F G H : functor_data C C'}
  (a : nat_trans F G)
  (b : nat_trans G H): is_nat_trans F H
     (λ x : ob C, a x · b x).
Proof.
  intros ? ? ?.
  now rewrite assoc, nat_trans_ax, <- assoc, nat_trans_ax, assoc.
Qed.

Definition nat_trans_comp {C:precategory_data}{C' : precategory}
  (F G H: functor_data C C')
  (a : nat_trans F G)
  (b : nat_trans G H): nat_trans F H :=
    tpair _ _ (is_nat_trans_comp a b).

(** Natural transformations for reasoning about various compositions of functors *)
Section nat_trans_functor.

Context {A B C D : precategory}.

Definition nat_trans_functor_id_right (F : functor A B) :
  nat_trans (functor_composite F (functor_identity B)) F.
Proof.
  apply nat_trans_id.
Defined.

Definition nat_trans_functor_id_right_inv (F : functor A B) :
  nat_trans F (functor_composite F (functor_identity B)) :=
    nat_trans_functor_id_right F.

Definition nat_trans_functor_id_left (F : functor A B) :
  nat_trans (functor_composite (functor_identity A) F) F :=
    nat_trans_functor_id_right F.

Definition nat_trans_functor_id_left_inv (F : functor A B) :
  nat_trans F (functor_composite (functor_identity A) F) :=
    nat_trans_functor_id_right F.

Definition nat_trans_functor_assoc (F1 : functor A B) (F2 : functor B C) (F3 : functor C D) :
  nat_trans (functor_composite (functor_composite F1 F2) F3)
            (functor_composite F1 (functor_composite F2 F3)).
Proof.
  apply nat_trans_id.
Defined.

Definition nat_trans_functor_assoc_inv (F1 : functor A B) (F2 : functor B C) (F3 : functor C D) :
  nat_trans (functor_composite F1 (functor_composite F2 F3))
            (functor_composite (functor_composite F1 F2) F3) :=
    nat_trans_functor_assoc F1 F2 F3.

End nat_trans_functor.

(** Reasoning about composition of natural transformations *)
Section nat_trans_comp_laws.

  Context {A B: precategory} (hs: has_homsets B).

Definition nat_trans_comp_id_right (F G : functor A B) (α: nat_trans F G):
  nat_trans_comp _ _ _ α (nat_trans_id G) = α.
Proof.
  apply nat_trans_eq; try exact hs.
  intro a.
  simpl.
  apply id_right.
Qed.

Definition nat_trans_comp_id_left (F G : functor A B) (α: nat_trans F G):
  nat_trans_comp _ _ _ (nat_trans_id F) α = α.
Proof.
  apply nat_trans_eq; try exact hs.
  intro a.
  simpl.
  apply id_left.
Qed.

Definition nat_trans_comp_assoc (F1 F2 F3 F4 : functor A B)
           (α: nat_trans F1 F2) (β: nat_trans F2 F3) (γ: nat_trans F3 F4):
  nat_trans_comp _ _ _ α (nat_trans_comp _ _ _ β γ) = nat_trans_comp _ _ _ (nat_trans_comp _ _ _ α β) γ.
Proof.
  apply nat_trans_eq; try exact hs.
  intro a.
  simpl.
  apply assoc.
Qed.

(** analogously to [assoc'], for convenience *)
Definition nat_trans_comp_assoc' (F1 F2 F3 F4 : functor A B)
           (α: nat_trans F1 F2) (β: nat_trans F2 F3) (γ: nat_trans F3 F4):
  nat_trans_comp _ _ _ (nat_trans_comp _ _ _ α β) γ = nat_trans_comp _ _ _ α (nat_trans_comp _ _ _ β γ).
Proof.
  apply pathsinv0, nat_trans_comp_assoc.
Qed.

End nat_trans_comp_laws.

(** ** Natural isomorphisms *)

Definition is_nat_iso {C D : precategory_data} {F G : functor_data C D} (μ : F ⟹ G) : UU :=
∏ (c : C), is_iso (μ c).

Definition isaprop_is_nat_iso
           {C D : category}
           {F G : C ⟶ D}
           (α : F ⟹ G)
  : isaprop (is_nat_iso α).
Proof.
  apply impred.
  intro.
  apply isaprop_is_iso.
Defined.

Definition is_nat_id {C D : precategory} {F : C ⟶ D} (μ : F ⟹ F) : UU :=
∏ (c : C), μ c = identity (F c).

Definition nat_iso {C D : precategory} (F G : C ⟶ D) : UU
:= ∑ (μ : F ⟹ G), is_nat_iso μ.

Definition make_nat_iso {C D : precategory} (F G : C ⟶ D) (μ : F ⟹ G) (is_iso : is_nat_iso μ) : nat_iso F G.
Proof.
  exists μ.
  exact is_iso.
Defined.

Definition nat_iso_pointwise_iso {C D : precategory} {F G : C ⟶ D} (μ : nat_iso F G) (c: C): iso (F c) (G c) := (pr1 μ c,,pr2 μ c).

Definition iso_inv_after_iso' {C : precategory} {a b : C} (f : a --> b) (f' : iso a b) (deref : pr1 f' = f) : f · inv_from_iso f' = identity _.
Proof.
  rewrite <- deref.
  exact (iso_inv_after_iso f').
Defined.

Definition iso_after_iso_inv' {C : precategory} {a b : C} (f : a --> b) (f' : iso a b) (deref : pr1 f' = f) : inv_from_iso f' · f = identity _.
Proof.
  rewrite <- deref.
  exact (iso_after_iso_inv f').
Defined.

Definition nat_iso_inv_trans
           {C D : precategory}
           {F G : C ⟶ D}
           (μ : nat_iso F G)
  : G ⟹ F.
Proof.
  use make_nat_trans.
  - exact (λ x, inv_from_iso (make_iso _ (pr2 μ x))).
  - abstract
      (intros x y f ; cbn ;
       refine (!_) ;
       use iso_inv_on_right ; cbn ;
       rewrite !assoc ;
       use iso_inv_on_left ; cbn ;
       exact (!(nat_trans_ax (pr1 μ) _ _ f))).
Defined.

Definition nat_iso_inv {C D : precategory} {F G : C ⟶ D} (μ : nat_iso F G) : nat_iso G F.
Proof.
  use make_nat_iso.
  - exact (nat_iso_inv_trans μ).
  - intro x.
    apply is_iso_inv_from_iso.
Defined.

Definition nat_iso_to_trans {C D : precategory} {F G : C ⟶ D} (ν : nat_iso F G) : F ⟹ G :=
  pr1 ν.

Coercion nat_iso_to_trans : nat_iso >-> nat_trans.

(* ⁻¹ *)
Definition nat_iso_to_trans_inv {C D : precategory} {F G : C ⟶ D} (ν : nat_iso F G) : G ⟹ F :=
  pr1 (nat_iso_inv ν).

Definition nat_comp_to_endo {C D : precategory} {F G : C ⟶ D} (eq : F = G) {c : C} (f : F c --> G c) : F c --> F c.
Proof.
  rewrite <- eq in f.
  assumption.
Defined.

Definition is_nat_iso_id {C D : precategory} {F G : C ⟶ D} (eq : F = G) (ν : nat_iso F G) : UU :=
  ∏ (c : C), nat_comp_to_endo eq (nat_iso_to_trans ν c) = identity (F c).

Definition induced_precategory_incl {M : precategory} {X:Type} (j : X -> ob M) :
  induced_precategory M j ⟶ M.
Proof.
  use make_functor.
  - use make_functor_data.
    + exact j.
    + intros a b f. exact f.
  - repeat split.
Defined.

(* ** analogous development for [z_iso] *)

Definition is_nat_z_iso {C D : precategory_data} {F G : functor_data C D} (μ : nat_trans_data F G) : UU :=
∏ (c : C), is_z_isomorphism (μ c).

Definition isaprop_is_nat_z_iso
           {C D : category}
           {F G : C ⟶ D}
           (α : nat_trans_data F G)
  : isaprop (is_nat_z_iso α).
Proof.
  apply impred.
  intro.
  apply isaprop_is_z_isomorphism.
Defined.

Definition nat_z_iso {C D : precategory_data} (F G : C ⟶ D) : UU
:= ∑ (μ : F ⟹ G), is_nat_z_iso μ.

Definition make_nat_z_iso {C D : precategory_data} (F G : C ⟶ D) (μ : F ⟹ G) (is_z_iso : is_nat_z_iso μ) : nat_z_iso F G.
Proof.
  exists μ.
  exact is_z_iso.
Defined.

Definition nat_z_iso_to_trans {C D : precategory_data} {F G : C ⟶ D} (μ : nat_z_iso F G) : F ⟹ G :=
  pr1 μ.

Coercion nat_z_iso_to_trans : nat_z_iso >-> nat_trans.

Definition pr2_nat_z_iso {C D : precategory_data} {F G : C ⟶ D} (μ : nat_z_iso F G) : is_nat_z_iso μ :=
  pr2 μ.

Definition nat_z_iso_pointwise_z_iso {C D : precategory_data} {F G : C ⟶ D} (μ : nat_z_iso F G) (c: C): z_iso (F c) (G c) := (pr1 μ c,,pr2 μ c).

(* ⁻¹ *)
Definition nat_z_iso_to_trans_inv {C : precategory_data} {D : precategory} {F G : C ⟶ D} (μ : nat_z_iso F G) : G ⟹ F.
Proof.
  apply (make_nat_trans G F (fun c => is_z_isomorphism_mor (pr2 μ c))).
  red.
  intros c c' f.
  set (h := μ c,,pr2 μ c : z_iso (F c) (G c)).
  set (h' := μ c',,pr2 μ c' : z_iso (F c') (G c')).
  change (# G f · inv_from_z_iso h' = inv_from_z_iso h · # F f).
  apply pathsinv0.
  apply z_iso_inv_on_right.
  rewrite assoc.
  apply z_iso_inv_on_left.
  apply pathsinv0.
  apply (nat_trans_ax μ).
Defined.

Definition nat_z_iso_inv {C : precategory_data} {D : precategory} {F G : C ⟶ D} (μ : nat_z_iso F G) : nat_z_iso G F.
Proof.
  exists (nat_z_iso_to_trans_inv μ).
  intro c.
  red.
  exists (μ c).
  red.
  split.
  - apply (pr2 (is_z_isomorphism_is_inverse_in_precat (pr2 μ c))).
  - apply (pr1 (is_z_isomorphism_is_inverse_in_precat (pr2 μ c))).
Defined.

Definition is_nat_z_iso_comp {C : precategory_data} {D : precategory} {F G H: C ⟶ D} {μ : F ⟹ G} {ν : G ⟹ H}
           (isμ: is_nat_z_iso μ) (isν: is_nat_z_iso ν) : is_nat_z_iso (nat_trans_comp F G H μ ν).
Proof.
  intro c.
  use make_is_z_isomorphism.
  - exact (is_z_isomorphism_mor (isν c) · is_z_isomorphism_mor (isμ c)).
  - exact (is_inverse_in_precat_comp (pr2 (isμ c)) (pr2 (isν c))).
Defined.

Definition nat_z_iso_comp {C : precategory_data} {D : precategory} {F G H: C ⟶ D}
           (μ: nat_z_iso F G) (ν: nat_z_iso  G H) : nat_z_iso F H.
Proof.
  use make_nat_z_iso.
  - exact (nat_trans_comp F G H μ ν).
  - exact (is_nat_z_iso_comp (pr2 μ) (pr2 ν)).
Defined.

Definition is_nat_z_iso_id {C D : precategory} {F G : C ⟶ D} (eq : F = G) (ν : nat_z_iso F G) : UU :=
  ∏ (c : C), nat_comp_to_endo eq (nat_z_iso_to_trans ν c) = identity (F c).

Lemma is_nat_z_iso_nat_trans_id {C D : precategory} (F : C ⟶ D): is_nat_z_iso (nat_trans_id F).
Proof.
  intro c.
  exists (identity (F c)).
  split; apply id_right.
Defined.

End nat_trans.

Definition constant_nat_trans
           (C₁ : category)
           {C₂ : category}
           {x y : C₂}
           (f : x --> y)
  : nat_trans
      (constant_functor C₁ C₂ x)
      (constant_functor C₁ C₂ y).
Proof.
  use make_nat_trans.
  - exact (λ _, f).
  - abstract
      (intros ? ? ? ;
       cbn ;
       rewrite id_left, id_right ;
       apply idpath).
Defined.

Notation "F ⟹ G" := (nat_trans F G) (at level 39) : cat.
(* to input: type "\==>" with Agda input method *)

(* Commutativity of functors expressed as natural isomorphisms *)

Lemma nat_z_iso_functor_comp_assoc
      {C1 C2 C3 C4 : category}
      (F1 : functor C1 C2)
      (F2 : functor C2 C3)
      (F3 : functor C3 C4)
  : nat_z_iso (F1 ∙ (F2 ∙ F3)) ((F1 ∙ F2) ∙ F3).
Proof.
  use make_nat_z_iso.
  - exists (λ _, identity _).
    abstract (intro ; intros ; exact (id_right _ @ ! id_left _)).
  - intro.
    exists (identity _).
    abstract (split ; apply id_right).
Defined.

Lemma functor_commutes_with_id
      {C D : category} (F : functor C D)
  : nat_z_iso (F ∙ functor_identity D) (functor_identity C ∙ F).
Proof.
  use make_nat_z_iso.
  - exists (λ _, identity _).
    abstract (intro ; intros ; exact (id_right _ @ ! id_left _)).
  - intro.
    exists (identity _).
    abstract (split ; apply id_right).
Defined.

Lemma nat_z_iso_comp_assoc
      {C D : category} {F1 F2 F3 F4 : functor C D}
      (α1 : nat_z_iso F1 F2)
      (α2 : nat_z_iso F2 F3)
      (α3 : nat_z_iso F3 F4)
  : nat_z_iso_comp α1 (nat_z_iso_comp α2 α3)
    = nat_z_iso_comp (nat_z_iso_comp α1 α2) α3.
Proof.
  use total2_paths_f.
  2: { apply isaprop_is_nat_z_iso. }

  use nat_trans_eq.
  { apply homset_property. }
  intro.
  apply assoc.
Qed.