(** * Combi.LRrule.stdplact : Plactic congruences and standardization *)
(******************************************************************************)
(*      Copyright (C) 2014-2018 Florent Hivert <florent.hivert@lri.fr>        *)
(*                                                                            *)
(*  Distributed under the terms of the GNU General Public License (GPL)       *)
(*                                                                            *)
(*    This code is distributed in the hope that it will be useful,            *)
(*    but WITHOUT ANY WARRANTY; without even the implied warranty of          *)
(*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU       *)
(*    General Public License for more details.                                *)
(*                                                                            *)
(*  The full text of the GPL is available at:                                 *)
(*                                                                            *)
(*                  http://www.gnu.org/licenses/                              *)
(******************************************************************************)
(** * Plactic congruences and standardization

The goal of this file is to show a few important properties of plactic
congruence related to standardisation.

The first one is easy: it says that standardization preserve plactic congruence.
This is Theorem [std_plact] and is stated as [u =Pl v -> std u =Pl std v.]

The next series of results relates the result of the RS map applied on
[u] and [std u]:
- the insertion tableau have the same shape: Theorem [shape_RS_std];
- the recording tableau are equal: Theorem [RSmap_std] and [RStabmap_std].

And finally, inverting a standard word amount to exchange the insertion and
the recording tableaux: Theorem [invseqRSPQE], and Corollaries [invseqRSE],
[invstdRSE], [RSTabmapstdE] and [RSinvstdE].
 *************)
Require Import mathcomp.ssreflect.ssreflect.
From mathcomp Require Import ssrbool ssrfun ssrnat eqtype finfun fintype choice.
From mathcomp Require Import seq tuple finset perm fingroup path order.

Require Import tools combclass ordcast partition Yamanouchi ordtype std tableau.
Require Import stdtab sorted Schensted congr plactic Greene Greene_inv.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Import Order.TTheory.

(** ** Plactic congruence and standardization *)
Section StdRS.

Context {disp : unit} {Alph : inhOrderType disp}.
Implicit Type s u v w : seq Alph.
Implicit Type p : seq nat.
Implicit Type t : seq (seq Alph).

Lemma std_plact1 (u v1 w v2 : seq Alph) :
  v2 \in plact1 v1 -> std (u ++ v1 ++ w) =Pl std (u ++ v2 ++ w).
Proof using.
move/plact1P => [a] [b] [c] [Habc -> ->].
have:= std_cutabc u w a c b => [] [U] [V] [A] [C] [B] [Hsz Hstd].
have Hac : (a < c)%O by move: Habc => /andP []; apply: le_lt_trans.
rewrite Hstd (std_transp Hac Hsz Hstd) -[[:: C; A] ++ B :: V]/([:: C; A; B] ++ V).
apply: plact_catr; apply: plact_catl.
apply: rule_gencongr => /=.
suff -> : (A <= B < C)%O by rewrite inE eq_refl.
have:= eq_inv_std (u ++ [:: a; c; b] ++ w).
rewrite Hstd {Hstd} => /eq_invP [Hsize Hinv].
have:= Hinv (size u) (size u).+2.
set Hyp := (X in (X -> _ ) -> _) => Hinvab.
have {Hyp} /Hinvab : Hyp.
  rewrite /Hyp (leq_trans (leqnSn _) (leqnSn _)) /= size_cat /=.
  by rewrite -2!addSnnS -{1}[(size u).+2]addn0 ltn_add2l.
rewrite {3 4}Hsz !nth_sizeu !nth_sizeu2 => <-.
move/(_ (size u).+1 (size u).+2): Hinv.
set Hyp := (X in (X -> _ ) -> _) => Hinvbc.
have {Hyp}/Hinvbc : Hyp.
  rewrite /Hyp (ltn_trans (ltnSn _) (ltnSn _)) /= size_cat /=.
  by rewrite -2!addSnnS -{1}[(size u).+2]addn0 ltn_add2l.
rewrite {3 4}Hsz !nth_sizeu1 !nth_sizeu2 ltNge => <-.
by rewrite -ltNge.
Qed.

Lemma std_plact2 (u v1 w v2 : seq Alph) :
  v2 \in plact2 v1 -> std (u ++ v1 ++ w) =Pl std (u ++ v2 ++ w).
Proof using.
have reorg3 (T : eqType) (U W : seq T) b a c :
  U ++ [:: b; a; c] ++ W = (U ++ [:: b]) ++ [:: a; c] ++ W by rewrite -catA.
move/plact2P => [a] [b] [c] [Habc -> ->].
have:= std_cutabc u w b a c => [] [U] [W] [B] [A] [C] [Hsz].
have Hac : (a < c)%O by move: Habc => /andP []; apply: lt_le_trans.
rewrite !reorg3.
have Hsz1 : size (u ++ [:: b]) = size (U ++ [:: B]) by rewrite !size_cat Hsz.
move=> Hstd; rewrite Hstd (std_transp Hac Hsz1 Hstd).
rewrite -!reorg3.
apply: plact_catr; apply: plact_catl.
apply: rule_gencongr => /=.
suff -> : (A < B <= C)%O by rewrite !mem_cat inE eq_refl /= !orbT.
rewrite -!reorg3  in Hstd.
have:= eq_inv_std (u ++ [:: b; a; c] ++ w).
rewrite Hstd {Hstd} => /eq_invP [Hsize Hinv].
have:= Hinv (size u) (size u).+2.
set Hyp := (X in (X -> _ ) -> _) => Hinvab.
have {Hyp}/Hinvab : Hyp.
  rewrite /Hyp (leq_trans (leqnSn _) (leqnSn _)) /= size_cat /=.
  by rewrite -2!addSnnS -{1}[(size u).+2]addn0 ltn_add2l.
rewrite {3 4}Hsz !nth_sizeu !nth_sizeu2 => <-.
move/(_ (size u) (size u).+1): Hinv.
set Hyp := (X in (X -> _ ) -> _) => Hinvbc.
have {Hyp}/Hinvbc : Hyp.
  rewrite /Hyp (leqnSn _) /= size_cat /=.
  by rewrite -addSnnS -{1}[(size u).+1]addn0 ltn_add2l.
rewrite {3 4}Hsz !nth_sizeu !nth_sizeu1 ltNge => <-.
by rewrite -ltNge.
Qed.

Theorem std_plact u v : u =Pl v -> std u =Pl std v.
Proof using.
move: v; apply: gencongr_ind; first exact: plact_refl.
move=> a b1 c b2 H Hplact.
rewrite (plact_ltrans H).
move: Hplact {H} => /plactruleP [].
- exact: std_plact1.
- rewrite -plact1I => /std_plact1-/(_ a c)/plact_ltrans <-; exact: plact_refl.
- exact: std_plact2.
- rewrite -plact2I => /std_plact2-/(_ a c)/plact_ltrans <-; exact: plact_refl.
Qed.

(** ** Robinson-Schensted correspondence and standardization *)
Lemma cast_enum u (S : {set 'I_(size u)}) :
  enum (mem (cast_set (esym (size_std u)) S)) =
  map (cast_ord (esym (size_std u))) (enum (mem S)).
Proof using.
rewrite {1}/enum_mem -enumT /=.
rewrite -[filter _ _]map_id (cast_map_cond _ _ (esym (size_std u))).
congr (map _ _).
rewrite {2}/enum_mem -enumT /=.
apply: eq_filter => i /=.
exact/mem_imset/cast_ord_inj.
Qed.

Lemma sorted_std_extract u (S : {set 'I_(size u)}) :
  sorted <=%O (extractpred (in_tuple u) (mem S)) =
  sorted <=%O (extractpred (in_tuple (std u))
                           (mem (cast_set (esym (size_std u)) S))).
Proof using.
rewrite /extractpred cast_enum /= /sorted.
set leqI := (fun i j : 'I_(size u) => i <= j).
have leqI_trans : transitive leqI.
  move=> i j k; rewrite /leqI; exact: leq_trans.
have: sorted leqI (enum (mem S)).
  rewrite {1}/enum_mem -enumT /=.
  apply: (sorted_filter leqI_trans).
  exact: enum_ord_sorted.
case: (enum (mem S)) => [//= | i0 l] {S} /=.
elim: l i0 => [//= | i1 l IHl] i0 /= /andP [Hleqi Hpath].
rewrite -(IHl i1 Hpath) {IHl Hpath}; congr (_ && _).
rewrite !(tnth_nth inh) /=.
have:= eq_inv_std u => /eq_invP [Hsz]; apply.
move: Hleqi; rewrite /leqI => -> /=.
exact: ltn_ord.
Qed.

Lemma ksupp_inj_std u k : ksupp_inj <=%O <=%O k u (std u).
Proof using.
rewrite /ksupp_inj /ksupp => ks /and3P [Hsz Htriv /forallP Hall].
exists (cast_set (esym (size_std u)) @: ks).
apply/and4P; split.
- rewrite cover_cast /cast_set /=.
  by rewrite card_imset; last exact: cast_ord_inj.
- apply: (@leq_trans #|ks|); last exact: Hsz.
  exact: leq_imset_card.
- apply: imset_trivIset; last exact: Htriv.
  exact: cast_ord_inj.
- apply/forallP => Sstd; apply/implyP => /imsetP [S HS -> {Sstd}].
  move/(_ S): Hall; rewrite HS /=.
  by rewrite sorted_std_extract.
Qed.

Lemma ksupp_inj_stdI u k : ksupp_inj <=%O <=%O k (std u) u.
Proof using.
rewrite /ksupp_inj /ksupp => ks /and3P [Hsz Htriv /forallP Hall].
exists (cast_set (size_std u) @: ks).
apply/and4P; split.
- rewrite cover_cast /cast_set /=.
  by rewrite card_imset; last exact: cast_ord_inj.
- apply: (@leq_trans #|ks|); last exact: Hsz.
  exact: leq_imset_card.
- apply: imset_trivIset; last exact: Htriv.
  exact: cast_ord_inj.
- apply/forallP => Sstd; apply/implyP => /imsetP [S HS -> {Sstd}].
  move/(_ S): Hall; rewrite HS /=.
  rewrite sorted_std_extract; congr (path.sorted _ ).
  rewrite /extractpred; congr (map _ _).
  apply: eq_enum => i.
  rewrite inE /cast_set /= -imset_comp.
  rewrite (eq_imset (g := id)); first last.
    by move=> j; rewrite /= (cast_ordK (size_std u)).
  by rewrite imset_id inE.
Qed.

Lemma Greene_std u k : Greene_row (std u) k = Greene_row u k.
Proof using.
apply/eqP; rewrite eqn_leq; apply/andP; split; apply: leq_Greene.
+ exact: ksupp_inj_stdI.
+ exact: ksupp_inj_std.
Qed.

Theorem shape_RS_std u : shape (RS (std u)) = shape (RS u).
Proof using. by apply: Greene_row_eq_shape_RS; exact: Greene_std. Qed.

End StdRS.

Theorem RSmap_std (d : unit) (T : inhOrderType d) (w : seq T) :
  (RSmap (std w)).2 = (RSmap w).2.
Proof.
move Hn : (size w) => n.
elim: n d T w Hn => [/= | n IHn] d T w; first by move/eqP/nilP => ->.
case/lastP Hw : w => [//= | w' wn] Hn.
have:= shape_RS_std (rcons w' wn).
rewrite -!RSmapE !shape_RSmap_eq.
case/lastP H : (std (rcons w' wn)) => [/= | st stn].
  exfalso; move: H => /(congr1 size).
  by rewrite size_std_rec size_rcons.
case HRS : ((RSmap (rcons w' wn)).2) => [/= | iw yamw].
  exfalso; move: HRS => /(congr1 size).
  by rewrite size_RSmap2 size_rcons.
have Hyamw : yamw = (RSmap w').2.
  move: HRS; rewrite /RSmap rev_rcons /=.
  case: (RSmap_rev (rev w')) => [t rows] /=.
  by case: (instabnrow t wn) => [tr nrow] /= [_ ->].
have Hsize : size w' = n by move: Hn; rewrite size_rcons => /succn_inj.
have /std_rconsK Hst : std (rcons w' wn) = std (rcons st stn).
  by rewrite -H std_stdE.
rewrite {wn Hw Hn H HRS} Hyamw /= -(IHn _ _ _ Hsize).
have Hsizest : size st = n.
  move: Hst; rewrite /std -{2}(size_std_rec (size st) st) => <-.
  by rewrite size_std_rec.
rewrite Hst (IHn _ _ _ Hsizest) {Hst IHn Hsize Hsizest}.
rewrite /RSmap rev_rcons /= -[RSmap_rev (rev st)]/(RSmap st).
case HRS : (RSmap st) => [t0 rows].
by case Hins : (instabnrow t0 stn) => [tr irow] /= /incr_nth_inj ->.
Qed.

Corollary RStabmap_std (d : unit) (T : inhOrderType d) (w : seq T) :
  (RStabmap (std w)).2 = (RStabmap w).2.
Proof.
rewrite /RStabmap.
move H : (RSmap w) => [P Q].
move Hs : (RSmap (std w)) => [Ps Qs] /=.
by rewrite -[Q]/(P, Q).2 -H -[Qs]/(Ps, Qs).2 -Hs RSmap_std.
Qed.

(** ** Robinson-Schensted correspondence and inversion *)
Section KsuppInj.

Variable s t : seq nat.
Hypothesis Hinv : invseq s t.

Let Hinvst : linvseq s t.
Proof using Hinv. by have:= Hinv; rewrite /invseq => /andP []. Qed.
Let Hinvts : linvseq t s.
Proof using Hinv. by have:= Hinv; rewrite /invseq => /andP []. Qed.

Local Definition val2pos :=
  fun (i : 'I_(size s)) => Ordinal (linvseq_ltn_szt Hinvst (ltn_ord i)).

Lemma val2posE : val \o val2pos =1 nth (size t) s.
Proof using. by []. Qed.
Lemma val2pos_inj : injective val2pos.
Proof using.
move: Hinvst => /linvseqP Hv.
move=> i j; rewrite /val2pos; set posi := Ordinal _ => /(congr1 val) /= Heq.
by apply/val_inj; rewrite /= -(Hv i (ltn_ord i)) -(Hv j (ltn_ord j)) Heq.
Qed.

Lemma val2pos_enum (p : {set 'I_(size s)}) :
  (* Hypothesis : val2pos sorted in p *)
  sorted <=%O [seq tnth (in_tuple s) i | i <- enum (mem p)] ->
  enum (mem [set val2pos x | x in p]) = [seq val2pos x | x <- enum p].
Proof using.
rewrite /enum_mem (eq_filter (a2 := mem p)) // -!enumT /= => H.
apply: (inj_map val_inj).
rewrite -map_comp (eq_map val2posE).
rewrite (eq_filter (a2 := mem [set val2pos x | x in p])) //.
move: H; rewrite (eq_map (f2 := fun i : 'I_(_) => nth (size t) s i)); first last.
  by move=> i /=; exact: tnth_nth.
set l := (X in sorted _ X); rewrite -[RHS]/l.
move=> H; apply: (sorted_eq (leT := leq)).
- exact: leq_trans.
- by move=> i j H1; apply/eqP; rewrite eqn_leq.
- apply: (subseq_sorted leq_trans (s2 := map val (enum 'I_(size t)))).
  + by apply: map_subseq; exact: filter_subseq.
  + by rewrite val_enum_ord; exact: iota_sorted.
- move: H; rewrite /sorted; case: l => [//= | l0 l].
  by rewrite (eq_path (e' := leq)); last by move=> i j /=; rewrite leEnat.
- apply: uniq_perm.
  + rewrite map_inj_in_uniq.
    * by apply: filter_uniq; exact: enum_uniq.
    * by move=> i j Hi _ /=; exact: val_inj.
  + rewrite map_inj_in_uniq.
    * by apply: filter_uniq; exact: enum_uniq.
    * move=> i j Hij _ /= HH; apply: val_inj.
      have:= Hinvst => /linvseqP Heq.
      by rewrite /= -(Heq _ (ltn_ord i)) -(Heq _ (ltn_ord j)) HH.
  + rewrite /l {l H} => i /=; apply/idP/idP => /mapP [j].
    * rewrite mem_filter => /andP [/= /imsetP [i' Hi' -> _] -> {i}].
      apply/mapP; exists i'; last by [].
      by rewrite mem_filter /= Hi' /=; rewrite mem_enum.
    * rewrite mem_filter => /andP [/= Hj _ -> {i}].
      apply/mapP; exists (Ordinal (linvseq_ltn_szt Hinvst (ltn_ord j))); last by [].
      rewrite mem_filter; apply/andP; split => /=; last by rewrite mem_enum.
      by rewrite /val2pos; rewrite imset_f.
Qed.

Lemma ksupp_inj_invseq k : ksupp_inj <=%O <=%O k s t.
Proof using Hinvst.
rewrite /ksupp_inj /ksupp => ks /and3P [Hsz Htriv /forallP Hall].
exists [set val2pos @: (p : {set 'I_(size s)}) | p in ks].
apply/and4P; split.
- by rewrite /=; apply/eqP; apply: size_cover_inj; exact: val2pos_inj.
- exact: (leq_trans (leq_imset_card _ _) Hsz).
- exact: (imset_trivIset val2pos_inj).
- apply/forallP => ptmp; apply/implyP => /imsetP [p Hp -> {ptmp}].
  move/(_ p): Hall; rewrite Hp /= /extractpred.
  move/val2pos_enum ->; rewrite -map_comp /=.
  rewrite (eq_map (f2 := nat_of_ord)); first last.
    move=> i /=; rewrite (tnth_nth (size s)) /=.
    by have:= Hinvst => /linvseqP ->.
  set l := map _ _; have : subseq l [seq val x | x  <- enum 'I_(size s)].
    apply: map_subseq; rewrite /enum_mem.
    rewrite [X in subseq _ X](eq_filter (a2 := xpredT)) //.
    by rewrite filter_predT; exact: filter_subseq.
  move/subseq_sorted; apply; first by move=> a b c H1 H2; exact: le_trans H1 H2.
  by rewrite val_enum_ord; exact: iota_sorted.
Qed.

End KsuppInj.

Lemma Greene_invseq s t k : invseq s t -> Greene_row s k = Greene_row t k.
Proof.
move=> Hst; have Hts := invseq_sym Hst.
by apply/eqP; rewrite eqn_leq; apply/andP; split;
  apply: leq_Greene; exact: ksupp_inj_invseq.
Qed.

Lemma shape_invseq s t : invseq s t -> shape (RS s) = shape (RS t).
Proof.
move=> Hinv; apply: Greene_row_tab_eq_shape; try apply: is_tableau_RS.
move=> k; rewrite -!(Greene_row_invar_plactic (congr_RS _) k).
exact: Greene_invseq.
Qed.

Lemma std_rcons_shiftinv t tn :
  is_std (rcons t tn) -> std t = map (shiftinv_pos tn) t.
Proof.
move=> H; apply/eqP/stdP; constructor.
- rewrite /is_std size_map perm_sym.
  apply: perm_size_uniq.
  + exact: iota_uniq.
  + move=> i; move: H; rewrite /is_std => /perm_mem Hperm.
    case (ltnP i tn) => Hi.
    * have {2}-> : i = shiftinv_pos tn i by rewrite /shiftinv_pos Hi.
      move=> Hiota; apply: map_f; have:= Hperm i; move: Hiota.
      rewrite !mem_iota /= !add0n size_rcons ltnS => /ltnW ->.
      by rewrite mem_rcons inE (ltn_eqF Hi) /= => ->.
    * have {2}-> : i = shiftinv_pos tn i.+1.
        by rewrite /shiftinv_pos ltnNge (leq_trans Hi (leqnSn _)).
      move=> Hiota; apply: map_f; have:= Hperm i.+1; move: Hiota.
      rewrite !mem_iota /= !add0n size_rcons ltnS => ->.
      rewrite mem_rcons inE.
      by move: Hi; rewrite -ltnS => /ltn_eqF; rewrite eq_sym => -> /= ->.
  + by rewrite size_map size_iota.
- apply/eq_invP; split; first by rewrite size_map.
  move=> i j /andP [Hij Hj].
  rewrite (nth_map inh); last exact (leq_ltn_trans Hij Hj).
  rewrite (nth_map inh); last exact Hj.
  rewrite !leEnat.
  apply/idP/idP; first exact: shiftinv_pos_incr.
  apply: contraLR.
  rewrite -!ltnNge !ltn_neqAle => /andP [Hneq /shiftinv_pos_incr ->].
  rewrite andbT.
  have:= std_uniq H; rewrite rcons_uniq => /andP [Htn Huniqt].
  move: Hneq; set xi := nth _ _ i; set xj := nth _ _ j.
  rewrite /shiftinv_pos.
  case (ltnP xi tn) => Hxi; case (ltnP xj tn) => Hxj.
  + by apply.
  + apply: contra => /eqP H1; move: Hxj Hxi; rewrite -H1 {H1 i xi Hij} => H1 H2.
    exfalso; have Hneq : xj != tn.
      by move: Htn; apply: contra => /eqP <-; rewrite /xj mem_nth.
    move: H2; have {H1 Hneq} : tn < xj by rewrite ltn_neqAle eq_sym Hneq H1.
    case: xj => [//= | x] /=.
    rewrite ltnS => H1 H2; have:= leq_ltn_trans H1 H2.
    by rewrite ltnn.
  + apply: contra => /eqP H1; move: Hxj Hxi.
    have {Hij Hj} Hi := leq_ltn_trans Hij Hj; rewrite H1 {H1 j xj} => H1 H2.
    exfalso.
    have Hneq : xi != tn.
      by move: Htn; apply: contra => /eqP <-; rewrite /xi mem_nth.
    have {H2 Hneq} : tn < xi by rewrite ltn_neqAle eq_sym Hneq H2.
    case: xi H1 => [//= | x] /=.
    rewrite ltnS => H1 H2; have := leq_ltn_trans H2 H1.
    by rewrite ltnn.
  + apply: contra => /eqP H1.
    have Hi := leq_ltn_trans Hij Hj.
    have {Hxi} : tn < xi.
      rewrite ltn_neqAle Hxi andbT.
      by move: Htn; apply: contra => /eqP ->; rewrite /xi mem_nth.
    have {Hxj} : tn < xj.
      rewrite ltn_neqAle Hxj andbT.
      by move: Htn; apply: contra => /eqP ->; rewrite /xj mem_nth.
    case: xi H1 => [//= | xi] /=.
    by case: xj => [//= | xj] /= ->.
Qed.

Lemma posbig_invseq s0 s t tn :
  invseq (s0 :: s) (rcons t tn) -> posbig (s0 :: s) = tn.
Proof.
move => Hinv; have:= Hinv; rewrite /invseq => /andP [Hst Hts].
have Hpos := posbig_size_cons s0 s.
have Hszt : size t < size (rcons t tn) by rewrite size_rcons.
have:= linvseq_ltn_szt Hts Hszt.
rewrite nth_rcons ltnn eq_refl (size_invseq Hinv) => Htn.
have:= invseq_nthE Hinv Hpos Hszt.
rewrite (set_nth_default inh _ Hpos) nth_posbig.
rewrite nth_rcons ltnn eq_refl => ->.
rewrite (std_max (invseq_is_std Hinv)); apply/eqP.
by move: Hinv => /size_invseq; rewrite size_rcons /= => /eqP.
Qed.

Lemma nth_std_pos s i x :
  is_std s -> i < size s -> i != posbig s -> nth x s i < (size s).-1.
Proof.
case: s => [//= | s0 s] Hstd Hi Hipos.
rewrite [nth _ _ _ < _]ltn_neqAle -ltnS; apply/andP; split.
- rewrite /= -(std_max Hstd) -(nth_posbig inh s0 s).
  rewrite (set_nth_default inh x Hi).
  by rewrite (nth_uniq _ Hi (posbig_size_cons s0 s) (std_uniq Hstd)).
- rewrite -[(size (s0 :: s)).-1.+1]/(size (s0 :: s)).
  by rewrite -(mem_std _ Hstd); exact: mem_nth.
Qed.

Lemma linvseqK s0 s t tn :
  invseq (s0 :: s) (rcons t tn) -> linvseq (rembig (s0 :: s)) (std t).
Proof.
move=> Hinv.
rewrite (std_rcons_shiftinv (invseq_is_std (invseq_sym Hinv))).
have := Hinv; rewrite /invseq => /andP [/linvseqP Hst /linvseqP Hts].
apply/linvseqP => i Hi.
rewrite (set_nth_default inh _ Hi) -nth_rembig.
rewrite (nth_map (size (s0 :: s))); first last.
  set si := shift_pos _ _; have Hsipos : si != posbig (s0 :: s).
    rewrite /si /shift_pos; case: (ltnP i (posbig (s0 :: s))).
    + by move /ltn_eqF ->.
    + by rewrite -ltnS => /gtn_eqF ->.
  have Hsi : si < (size (s0 :: s)).
    rewrite /si /shift_pos; case: (ltnP i (posbig (s0 :: s))) => H.
    + by apply: (ltn_trans Hi); rewrite size_rembig /=.
    + by apply: (leq_trans Hi); rewrite size_rembig /=.
  apply: (leq_trans (nth_std_pos _ (invseq_is_std Hinv) Hsi Hsipos)).
  by move: Hinv => /size_invseq /=; rewrite size_rcons => /succn_inj ->.
rewrite (posbig_invseq Hinv).
have Hpos : shift_pos tn i < size (s0 :: s).
  move: Hi; rewrite /shift_pos size_rembig /=.
  by case (i < tn); first by move/leq_trans; apply.
rewrite (set_nth_default (size (rcons t tn)) inh Hpos).
move: (Hst _ Hpos); rewrite nth_rcons.
rewrite (_ : nth _ _ _ < size t); first by move ->; rewrite shift_posK.
have -> : size t = (size (s0 :: s)).-1 by rewrite (size_invseq Hinv) size_rcons.
apply: (nth_std_pos (size (rcons t tn)) (invseq_is_std Hinv) Hpos).
rewrite (posbig_invseq Hinv) /shift_pos.
case (ltnP i tn); first by move /ltn_eqF ->.
by rewrite -ltnS => /gtn_eqF ->.
Qed.

Lemma invseqK s0 s t tn :
  invseq (s0 :: s) (rcons t tn) -> invseq (rembig (s0 :: s)) (std t).
Proof.
move=> Hinv; apply: linvseq_sizeP; first exact: (linvseqK Hinv).
rewrite size_rembig /= size_std.
by move: Hinv => /size_invseq; rewrite size_rcons /= => /succn_inj ->.
Qed.

Theorem invseqRSPQE s t :
  invseq s t -> (RStabmap s).1 = (RStabmap t).2.
Proof.
rewrite /RStabmap /= => Hinv.
case HRSs : (RSmap s) => [Ps Qs] /=.
have -> : Ps = RS s by rewrite -RSmapE HRSs.
case HRSt : (RSmap t) => [Pt Qt] /=.
have -> : Qt = (RSmap t).2 by rewrite HRSt.
move Hn: (size s) => n {Ps Qs HRSs Pt Qt HRSt}.
have Hsize := size_invseq Hinv.
elim: n s t Hn Hinv Hsize => [/= | n IHn] s t Hn Hinv Hsize.
  by move: Hn Hsize => /eqP/nilP Hs; subst s => /= /esym/eqP/nilP -> /=.
have:= shape_invseq Hinv.
rewrite -!RSmapE [shape (RSmap t).1]shape_RSmap_eq !RSmapE.
case Hs : s Hn Hinv Hsize => [//= | s0 s']; rewrite -Hs => Hn Hinv Hsize.
rewrite Hs (rembig_RS_last_big s0 s') shape_append_nth -Hs.
set irow := last_big _ _.
move=> Hshape.
have Hszrem : size (rembig (s0 :: s')) = n.
  by move: Hn; rewrite Hs size_rembig /= => /succn_inj.
case/lastP Ht : t => [//= | t' tn].
  by exfalso; move: Hsize; rewrite Hn Ht.
have:= Hinv; rewrite Hs Ht => /invseqK Hrec.
have Hsize' : size (rembig (s0 :: s')) = size (std t').
  by rewrite size_rembig -Hs Hsize Ht size_rcons /std size_std_rec.
move/(_ _ _ Hszrem Hrec Hsize'): IHn => IHn.
move: Hshape; rewrite -Hs Ht /RSmap rev_rcons /= -/(RSmap t').
case HRSt : (RSmap t') => [t0 rowt'].
case Hins : (instabnrow t0 tn) => [tr irowt'] /=.
move: IHn; rewrite -Hs RSmap_std HRSt /= => ->.
rewrite shape_stdtab_of_yam => /incr_nth_inj ->.
congr (append_nth _ _ _).

have:= (congr1 (fun t => sumn (shape t)) (RSmapE t')).
rewrite (shape_RSmap_eq t') evalseq_eq_size HRSt => ->.
rewrite -/(size_tab (RS t')) size_RS.
move: Hsize; rewrite Ht Hn size_rcons => /succn_inj <-.
move: Hinv => /invseq_is_std; rewrite Hs => /std_max ->.
by apply succn_inj; rewrite -Hn Hs.
Qed.

Corollary invseqRSE s t :
  invseq s t -> RStabmap s = ((RStabmap t).2, (RStabmap t).1).
Proof.
move=> Hinv.
rewrite -(invseqRSPQE Hinv) (invseqRSPQE (invseq_sym Hinv)).
by case (RStabmap s).
Qed.

Corollary invstdRSE s :
  is_std s -> let (P, Q) := RStabmap (invstd s) in RStabmap s = (Q, P).
Proof.
move=> Hstd.
case H : (RStabmap (invstd s)) => [P Q].
by rewrite (invseqRSE (invseq_invstd Hstd)) H.
Qed.

Corollary RSTabmapstdE (disp : unit) (T : inhOrderType disp) (w : seq T) :
  (RStabmap (invstd (std w))).1 = (RStabmap (std w)).2.
Proof.
have := invstdRSE (std_is_std w).
by case (RStabmap (invstd (std w))) => [P Q] /= ->.
Qed.

Corollary RSinvstdE (disp : unit) (T : inhOrderType disp) (w : seq T) :
  RS (invstd (std w)) = (RStabmap w).2.
Proof.
rewrite -RStabmapE RSTabmapstdE /RStabmap.
rewrite [LHS](_ : _ = stdtab_of_yam (RSmap (std w)).2); last by case: RSmap.
rewrite [RHS](_ : _ = stdtab_of_yam (RSmap w).2); last by case: RSmap.
by congr (stdtab_of_yam _); exact: RSmap_std.
Qed.


From stdpp Require Import tactics.
From stdpp Require Import options.
Local Set Universe Polymorphism.

(* Not using [list Type] in order to avoid universe inconsistencies *)
Inductive tlist := tnil : tlist | tcons : Type → tlist → tlist.

Inductive hlist : tlist → Type :=
  | hnil : hlist tnil
  | hcons {A As} : A → hlist As → hlist (tcons A As).

Fixpoint tapp (As Bs : tlist) : tlist :=
  match As with tnil => Bs | tcons A As => tcons A (tapp As Bs) end.
Fixpoint happ {As Bs} (xs : hlist As) (ys : hlist Bs) : hlist (tapp As Bs) :=
  match xs with hnil => ys | hcons x xs => hcons x (happ xs ys) end.

Definition hhead {A As} (xs : hlist (tcons A As)) : A :=
  match xs with hnil => () | hcons x _ => x end.
Definition htail {A As} (xs : hlist (tcons A As)) : hlist As :=
  match xs with hnil => () | hcons _ xs => xs end.

Fixpoint hheads {As Bs} : hlist (tapp As Bs) → hlist As :=
  match As with
  | tnil => λ _, hnil
  | tcons _ _ => λ xs, hcons (hhead xs) (hheads (htail xs))
  end.
Fixpoint htails {As Bs} : hlist (tapp As Bs) → hlist Bs :=
  match As with
  | tnil => id
  | tcons _ _ => λ xs, htails (htail xs)
  end.

Fixpoint himpl (As : tlist) (B : Type) : Type :=
  match As with tnil => B | tcons A As => A → himpl As B end.

Definition hinit {B} (y : B) : himpl tnil B := y.
Definition hlam {A As B} (f : A → himpl As B) : himpl (tcons A As) B := f.
Global Arguments hlam _ _ _ _ _ / : assert.

Definition huncurry {As B} (f : himpl As B) (xs : hlist As) : B :=
  (fix go {As} xs :=
    match xs in hlist As return himpl As B → B with
    | hnil => λ f, f
    | hcons x xs => λ f, go xs (f x)
    end) _ xs f.
Coercion huncurry : himpl >-> Funclass.

Fixpoint hcurry {As B} : (hlist As → B) → himpl As B :=
  match As with
  | tnil => λ f, f hnil
  | tcons x xs => λ f, hlam (λ x, hcurry (f ∘ hcons x))
  end.

Lemma huncurry_curry {As B} (f : hlist As → B) xs :
  huncurry (hcurry f) xs = f xs.
Proof. by induction xs as [|A As x xs IH]; simpl; rewrite ?IH. Qed.

Fixpoint hcompose {As B C} (f : B → C) {struct As} : himpl As B → himpl As C :=
  match As with
  | tnil => f
  | tcons A As => λ g, hlam (λ x, hcompose f (g x))
  end.


From iris.proofmode Require Import tactics.
From iris.program_logic Require Export weakestpre.
From iris.heap_lang Require Export lang.
From iris.heap_lang Require Import proofmode notation.
From iris.prelude Require Import options.

(** A library defining binary [minimum] and [maximum] functions, together with
their expected specs. These operations come up often when working manipulating
array indices (checking for bounds). *)

Definition minimum : val :=
  λ: "m" "n", if: "m" < "n" then "m" else "n".

Definition maximum : val :=
  λ: "m" "n", if: "m" < "n" then "n" else "m".

Lemma minimum_spec `{!heapG Σ} s E (Φ : val → iProp Σ) (m n : Z) :
  ▷ Φ #(m `min` n) -∗
  WP minimum #m #n @ s;E {{ Φ }}.
Proof.
  iIntros "HΦ". wp_lam. wp_pures. case_bool_decide; wp_pures.
  - rewrite Z.min_l; [ done | by lia ].
  - rewrite Z.min_r; [ done | by lia ].
Qed.

Lemma minimum_spec_nat `{!heapG Σ} s E (Φ : val → iProp Σ) (m n : nat) :
  ▷ Φ #(m `min` n)%nat -∗
  WP minimum #m #n @ s;E {{ Φ }}.
Proof. iIntros "HΦ". iApply minimum_spec. by rewrite Nat2Z.inj_min. Qed.

Lemma maximum_spec `{!heapG Σ} s E (Φ : val → iProp Σ) (m n : Z) :
  ▷ Φ #(m `max` n) -∗
  WP maximum #m #n @ s;E {{ Φ }}.
Proof.
  iIntros "HΦ". wp_lam. wp_pures. case_bool_decide; wp_pures.
  - rewrite Z.max_r; [ done | by lia ].
  - rewrite Z.max_l; [ done | by lia ].
Qed.

Lemma maximum_spec_nat `{!heapG Σ} s E (Φ : val → iProp Σ) (m n : nat) :
  ▷ Φ #(m `max` n)%nat -∗
  WP maximum #m #n @ s;E {{ Φ }}.
Proof. iIntros "HΦ". iApply maximum_spec. by rewrite Nat2Z.inj_max. Qed.


Require Import Kami.AllNotations StdLibKami.RegStruct.
Require Import Kami.Utila.
Require Import List.
Import ListNotations.

Section Granule.
  Variable lgGranuleSize : nat.
  Notation n := (Nat.pow 2 lgGranuleSize).

  Notation divCeil x y := (Nat.div (x + (y - 1)) y).
  Notation div_packn k := (divCeil (size k) n).
  Notation lg_packn k := (Nat.log2_up (divCeil (size k) n)).
  Notation pow2_packn k := (Nat.pow 2 (lg_packn k)).

  Notation getStartGranule addr := (wordToNat (wsplitr _ _ addr)).
  Notation getFinishGranule addr k := (getStartGranule addr + div_packn k).
  Notation getFinishPacket addr k maskSz :=
    (divCeil (getFinishGranule addr k) maskSz).
  Notation getFinishPacketGranule addr k maskSz :=
    (getFinishPacket addr k maskSz * maskSz).

  Definition putRightPosition ty k (val: k @# ty) start finish :=
    {< $$ (natToWord (finish - (start + size k)) 0), pack val, $$ (natToWord start 0)>}%kami_expr.

  Lemma divCeil_ge x y: y <> 0 -> divCeil x y * y >= x.
  Proof.
    intros.
    pose proof (Nat.div_mod (x + (y-1)) y ltac:(Omega.omega)) as sth.
    rewrite Nat.mul_comm in sth.
    pose proof (Nat.mod_le (x + (y-1)) _ H) as sth2.
    assert (sth3: divCeil x y * y = x + (y-1) - ((x + (y-1))mod y)) by Omega.omega.
    Opaque Nat.div.
    simpl.
    rewrite sth3.
    Transparent Nat.div.
    pose proof (Nat.mod_bound_pos (x + (y-1)) y ltac:(Omega.omega) ltac:(Omega.omega)).
    Omega.omega.
  Qed.

  Definition byteAlign ty k (e: k @# ty): (Bit (div_packn k * n) @# ty).
    refine (castBits _ (ZeroExtend (div_packn k * n - size k) (pack e))).
    abstract (pose proof (@divCeil_ge (size k) n (pow2_ne_zero _)); Omega.omega).
  Defined.

  Section RegMapper.
    Variable ty: Kind -> Type.
    Variable lgMaskSz realAddrSz: nat.

    Local Notation maskSz := (Nat.pow 2 lgMaskSz).
    Local Notation addrSz := (realAddrSz + lgMaskSz).
    Local Notation dataSz := (maskSz * n).

    (* For tile-link, addr, mask and size should all be compatible,
       which is why maskSz, dataSz are powers of 2 *)

    (* Each entry represents a single data entry of size dataSz,
       indexed by an addr of size realAddrSz.
       The mask represents the granuleEnable for granules of size n that make up the data *)
    Definition RegMapT :=
      STRUCT_TYPE
        { "addr" :: Bit realAddrSz ;
          "mask" :: Bit maskSz ;
          "data" :: Bit dataSz }.

    Definition FullRegMapT :=
      STRUCT_TYPE
        { "isRd" :: Bool ;
          "info" :: RegMapT }.

    Record GenRegField :=
      { grf_addr  : word realAddrSz ;
        grf_mask  : Bit maskSz @# ty;
        grf_read  : RegMapT @# ty -> ActionT ty (Bit dataSz) ;
        grf_write : RegMapT @# ty -> ActionT ty Void }.

    Local Open Scope kami_action.
    Local Open Scope kami_expr.

    (* Given an addr of size realAddrSz and a mask, if the addr matches any entry's address
       and if any granule in the data is enabled by the mask, then do the corresponding
       read or write, sending the whole mask, data and the address itself (address is
       redundant for the final reader/writer) *)
    Definition createRegMap (rq: Maybe FullRegMapT @# ty) (ls: list GenRegField): ActionT ty (Bit dataSz) :=
      If rq @% "valid"
    then (If rq @% "data" @% "isRd"
          then GatherActions (map (fun x =>
                                     If (rq @% "data" @% "info" @% "addr") ==
                                     ($$ (grf_addr x))
                                       && ((rq @% "data" @% "info" @% "mask" .& (grf_mask x)) != $ 0)
                                   then (LETA retVal <- grf_read x (rq @% "data" @% "info");
                                           Ret #retVal)
                                   else Ret ($$ (natToWord dataSz 0))
                                    as retVal;
                                     Ret #retVal) ls) as listVals;
                 Ret (Kor listVals)
          else GatherActions (map (fun x =>
                                     If (rq @% "data" @% "info" @% "addr") ==
                                     ($$ (grf_addr x))
                                       && ((rq @% "data" @% "info" @% "mask" .& (grf_mask x)) != $ 0)
                                   then grf_write x (rq @% "data" @% "info")
                                   else Retv; Retv) ls) as _;
            Ret ($$ (natToWord dataSz 0))
            as retVal;
            Ret #retVal)
    else Ret $0
    as retVal;
      Ret #retVal.
    Local Close Scope kami_expr.
    Local Close Scope kami_action.

    Definition makeSplitBits (addr: word addrSz) (k: Kind) (e: k @# ty): Array (getFinishPacket addr k maskSz) (Bit dataSz) @# ty.
      refine
        (unpack (Array (getFinishPacket addr k maskSz) (Bit dataSz))
                (castBits _ (putRightPosition (byteAlign e) (getStartGranule addr * n) (getFinishPacketGranule addr k maskSz * n)))).
      Opaque Nat.div.
      abstract (
          simpl;
          assert (divCeil (getStartGranule addr + (size k + (n - 1)) / n) maskSz * maskSz * n >= getStartGranule addr * n + (size k + (n-1)) / n * n) by
            (pose proof (divCeil_ge (getStartGranule addr + div_packn k) (pow2_ne_zero lgMaskSz)); simpl in *;
             nia);
          rewrite Nat.mul_assoc;
          lia).
      Transparent Nat.div.
    Defined.

    Definition makeSplitMask (addr: word addrSz) (k: Kind): Array (getFinishPacket addr k maskSz) (Bit maskSz) @# ty.
      refine
        (unpack (Array (getFinishPacket addr k maskSz) (Bit maskSz))
                (castBits _ (putRightPosition ($$ (wones (getFinishGranule addr k - getStartGranule addr)))%kami_expr
                                              (getStartGranule addr) (getFinishPacketGranule addr k maskSz)))).
      Opaque Nat.div.
      abstract (
          simpl;
          assert (divCeil (getStartGranule addr + (size k + (n-1)) / n) maskSz * maskSz >= getStartGranule addr + (size k + (n-1)) / n) by
              (pose proof (divCeil_ge (getStartGranule addr + div_packn k) (pow2_ne_zero lgMaskSz)); simpl in *;
               lia);
          lia).
      Transparent Nat.div.
    Defined.

    Definition makeJoinBits (addr: word addrSz) (k: Kind) (e: Array (getFinishPacket addr k maskSz) (Bit dataSz) @# ty): k @# ty.
      refine
        (unpack k (UniBit
                     (TruncLsb _ (div_packn k * n - size k))
                     (castBits
                        _
                        (ConstExtract (getStartGranule addr * n)
                                      (div_packn k * n)
                                      (getFinishPacket addr k maskSz * dataSz - getFinishGranule addr k * n)
                                      (castBits _ (pack e)))))).
      abstract (pose proof (@divCeil_ge (size k) n (pow2_ne_zero _)); Omega.omega).
      Opaque Nat.div.
      abstract (
          simpl;
          pose proof (@divCeil_ge (getStartGranule addr + (size k + (n-1)) / n) maskSz (pow2_ne_zero lgMaskSz));
          nia).
      Transparent Nat.div.
    Defined.

    (* Represents a register group starting at address represented by srg_addr.
       Note that srg_addr is of size addrSz, which means it indexes into a single
       granule, unlike grf_addr, which indexes into the entire data.
       IMPORTANT: The size of srg_kind can be well beyond dataSz of RegMapT.
       The value is just loaded contiguously starting from srg_addr and padded to fill entire data boundaries *)
    Record SimpleRegGroup :=
      { srg_addr  : word addrSz ;
        srg_kind  : Kind ;
        srg_read  : ActionT ty srg_kind ;
        srg_write : (srg_kind @# ty) -> ActionT ty Void;
        srg_name  : option string
      }.

    Definition expandRqMask (m: Bit maskSz @# ty): Bit dataSz @# ty.
      refine (castBits _ (pack (BuildArray (fun i => replicate (ReadArrayConst (unpack (Array maskSz (Bit 1)) (castBits _ m)) i) n))));
        abstract (auto; simpl; lia).
    Defined.

    Local Open Scope kami_action.
    Local Open Scope kami_expr.
    Definition readWriteGranules_Gen (x: SimpleRegGroup): list GenRegField :=
      map (fun y => {| grf_addr  := (wsplitl realAddrSz _ (srg_addr x) ^+ $(proj1_sig (Fin.to_nat y)))%word ;
                       grf_mask  := ReadArrayConst (makeSplitMask (srg_addr x) (srg_kind x)) y ;
                       grf_read  :=
                         fun rm =>
                           (LETA readK: srg_kind x <- srg_read x ;
                              LET readVal: Bit dataSz <-
                                               ReadArrayConst (makeSplitBits (srg_addr x) #readK)
                                               y;
                              LET maskVal: Bit dataSz <-
                                               ReadArrayConst
                                               (makeSplitBits
                                                  (srg_addr x)
                                                  (Const ty (wones (size (srg_kind x)))))
                                               y;
                              LET finalVal: Bit dataSz <- expandRqMask (rm @% "mask") .& #readVal .& #maskVal;
                              Ret #finalVal) ;
                       grf_write :=
                         fun rm =>
                           (LETA readK: srg_kind x <- srg_read x ;
                              LET readVal <- makeSplitBits (srg_addr x) #readK;
                              LET maskVal <- makeSplitBits (srg_addr x)
                                  (Const ty (wones (size (srg_kind x))));
                              LET maskVali <- ReadArrayConst #maskVal y;
                              LET t1Val <- (expandRqMask (rm @% "mask") .& #maskVali) .& rm @% "data";
                              LET t2Val <- (~(expandRqMask (rm @% "mask") .& #maskVali)) .& (ReadArrayConst # readVal y);
                              LET t3Val <- (#t1Val .| #t2Val);
                              LET finalVal <- UpdateArrayConst #readVal y #t3Val;
                              srg_write x (makeJoinBits (srg_addr x) (srg_kind x) #finalVal)
                           )
                    |})
          (getFins (getFinishPacket (srg_addr x) (srg_kind x) maskSz)).

    Definition readWriteGranules rq ls :=
      createRegMap rq (concat (map readWriteGranules_Gen ls)).

    Record SingleReg :=
      { sr_addr : word addrSz ;
        sr_kind : Kind ;
        sr_name : sum (string * bool) (ConstT sr_kind)}.

    Definition SingleReg_Gen (x: SingleReg) := {| srg_addr  := sr_addr x ;
                                                  srg_kind  := sr_kind x ;
                                                  srg_read  :=
                                                    match sr_name x with
                                                    | inl (name, _) =>
                                                      (Read val : (sr_kind x) <- name ;
                                                         Ret #val)
                                                    | inr uval =>
                                                      Ret ($$ uval)
                                                    end ;
                                                  srg_write :=
                                                    fun val =>
                                                      match sr_name x with
                                                      | inl (name, true) =>
                                                        (Write name : (sr_kind x) <- val ;
                                                           Retv)
                                                      | _ =>
                                                        Retv
                                                      end ;
                                                  srg_name := match sr_name x with
                                                              | inl (name, _) => Some name
                                                              | _ => None
                                                              end
                                               |}.

    Definition readWriteGranules_SingleReg rq ls := createRegMap rq (concat (map (fun x => readWriteGranules_Gen (SingleReg_Gen x)) ls)).

    (* Record SingleNonReg := *)
    (*   { snr_addr : word addrSz ; *)
    (*     snr_kind : Kind ; *)
    (*     snr_val  : ConstT snr_kind }. *)

    (* Definition readWriteGranules_SingleNonReg_Gen (x: SingleNonReg) := *)
    (*   readWriteGranules_Gen {| srg_addr  := snr_addr x ; *)
    (*                         srg_kind  := snr_kind x ; *)
    (*                         srg_read  := Ret (Const ty (snr_val x)) ; *)
    (*                         srg_write := fun val => Retv |}. *)

    (* Definition readWriteGranules_SingleNonReg rq ls := createRegMap rq (concat (map readWriteGranules_SingleNonReg_Gen ls)). *)

    Record GroupReg :=
      { gr_addr : word addrSz ;
        gr_kind : Kind ;
        gr_name : string
      }.

    Definition GroupReg_Gen (x: GroupReg) := {| srg_addr  := gr_addr x ;
                                                srg_kind  := gr_kind x ;
                                                srg_read  := Struct_RegReads ty (gr_kind x) ;
                                                srg_write := fun val => Struct_RegWrites val ;
                                                srg_name  := Some (gr_name x)
                                             |}.

    Definition readWriteGranules_GroupReg rq ls := createRegMap rq (concat (map (fun x => readWriteGranules_Gen (GroupReg_Gen x)) ls)).

    Record MayGroupReg :=
      { mgr_addr : word addrSz ;
        mgr_size : nat ;
        mgr_kind : MayStruct mgr_size ;
        mgr_name : string
      }.

    Definition MayStruct_Struct n (x: MayStruct n) := Struct (fun i => projT1 (vals x i)) (names x).

    Definition MayGroupReg_Gen (x: MayGroupReg) := {| srg_addr  := mgr_addr x ;
                                                      srg_kind  := MayStruct_Struct (mgr_kind x) ;
                                                      srg_read  := MayStruct_RegReads ty (mgr_kind x) ;
                                                      srg_write := fun val => MayStruct_RegWrites (mgr_kind x) val ;
                                                      srg_name  := Some (mgr_name x)
                                                   |}.

    Definition readWriteGranules_MayGroupReg rq ls := createRegMap rq (concat (map (fun x => readWriteGranules_Gen (MayGroupReg_Gen x)) ls)).

    Definition mayStructKind (n : nat) (x : MayStruct n)
      :  Kind
      := Struct (fun i : Fin.t n => projT1 (vals x i)) (names x).

    Variable ContextCodeWidth : nat.
    Definition ContextCodeT := Bit ContextCodeWidth.

    Definition LocationReadWriteInputT
      (k : Kind)
      := STRUCT_TYPE {
           "isRd"        :: Bool;
           "addr"        :: Bit addrSz;
           "contextCode" :: ContextCodeT;
           "data"        :: k
         }.

    Record View :=
      {
        view_context : ContextCodeT @# ty;
        view_size : nat ;
        view_kind : MayStruct view_size
      }.

    (* Represents a memory location supporting context-dependent views. *)
    Record Location :=
      {
        loc_name : string ;
        loc_addr : word addrSz ;
        loc_views : list View
      }.

    Definition viewReadWrite
      (n : nat)
      (k : Kind)
      (req : LocationReadWriteInputT k @# ty)
      (view : MayStruct n)
      :  ActionT ty k
      := LETA read_result
           :  mayStructKind view
           <- MayStruct_RegReads ty view;
         System [
           DispString _ "[viewReadWrite]\n";
           DispString _ "  read context code: ";
           DispDecimal (req @% "contextCode");
           DispString _ "\n";
           DispString _ "  read address: ";
           DispHex (req @% "addr");
           DispString _ "\n";
           DispString _ "  read result: ";
           DispBinary (pack #read_result);
           DispString _ "\n"
         ];
         If !(req @% "isRd")
           then
             LET write_value
               :  mayStructKind view
               <- unpack (mayStructKind view)
                    (ZeroExtendTruncLsb
                      (size (mayStructKind view))
                      (pack (req @% "data")));
             LETA write_result
               :  Void
               <- MayStruct_RegWrites view #write_value;
             System [
               DispString _ "  is write request.\n";
               DispString _ "  write value: ";
               DispBinary (pack #write_value);
               DispString _ "\n"
             ];
             Retv;
         Ret
           (unpack k
             (ZeroExtendTruncLsb (size k)
               (pack (#read_result)))).

    (*
      This function reads and writes to memory. Every memory
      location has an associated set of "views." A view is a set of
      fields. When this function reads a memory location, it selects
      one of the location's views, reads the values stored within
      the view fields, concatenates these field values together,
      and returns them.

      When this function writes to a memory location, it selects a
      view associated with the location. It then parses the write
      value into bit ranges corresponding to the view fields. It
      then writes the values stored within these bit ranges to the
      corresponding view fields.

      This function accepts three arguments: [k], [req] and [entries].

      [req] represents a memory request. Every memory request has
      an associated address and context code.

      [entries] represents a list of memory locations. Each entry
      contains a vector of views - one for each context code
      value. Each view contains a MayStruct that lists the fields
      that comprise the view.

      When the [isRd] field in [req] is true, this function reads
      the word at location [req @% "addr"] and converts it into a
      packet of type [k].

      When the [isRd] fields is false, this function writes the
      value given in [req @% "data"] to memory location [req @%
      "addr"].

    *)
    Definition locationReadWrite
      (k : Kind)
      (req : LocationReadWriteInputT k @# ty)
      (entries : list Location)
      :  ActionT ty (Maybe k)
      := System [
           DispString _ "[locationReadWrite]\n";
           DispString _ "[locationReadWrite] request:\n";
           DispHex req;
           DispString _ "\n"
         ];
         utila_acts_find_pkt
           (map
             (fun addr_entry : Location
               => utila_acts_find_pkt
                    (map
                      (fun view_entry : View
                        => LET entry_match
                             :  Bool
                             <- ((req @% "addr") == $$(loc_addr addr_entry) &&
                                 (req @% "contextCode") == view_context view_entry);
                           If #entry_match
                             then
                               System [
                                 DispString _ "[locationReadWrite]\n";
                                 DispString _ "  location name: ";
                                 DispString _ (loc_name addr_entry);
                                 DispString _ "\n"
                               ];
                               viewReadWrite req (view_kind view_entry)
                             else
                               Ret (unpack k $0)
                             as result;
                           (utila_acts_opt_pkt #result #entry_match))
                      (loc_views addr_entry)))
             entries).

    Local Close Scope kami_expr.
    Local Close Scope kami_action.

  End RegMapper.
End Granule.

  (* Lemma helper_pow2_packn k: (pow2_packn k * n >= size k)%nat. *)
  (* Proof. *)
  (*   remember (size k) as x; clear Heqx. *)
  (*   pose proof (@divCeil_ge x n ltac:(Omega.omega)) as sth. *)
  (*   pose proof (log2_up_pow2 (divCeil x n)). *)
  (*   Omega.omega. *)
  (* Qed. *)

  (* Fixpoint wordSplitter' n t: word (t * n) -> list (word n) := *)
  (*   match t return word (t * n) -> list (word n) with *)
  (*   | 0 => fun _ => nil *)
  (*   | S m => fun w => split1 n (m * n) w :: @wordSplitter' _ m (split2 n (m * n) w) *)
  (*   end. *)

  (* Definition wordSplitter n (pf: n <> 0) sz (w: word sz): list (word n). *)
  (*   refine *)
  (*     (wordSplitter' n (divCeil sz n) (nat_cast word _ ({< natToWord (divCeil sz n * n - sz) 0, w>})%word)). *)
  (*   abstract (pose proof (divCeil_ge sz pf); *)
  (*             Omega.omega). *)
  (* Defined. *)

  (* Fixpoint exprSplitter' ty n t: (Bit (t * n) @# ty) -> list (Bit n @# ty) := *)
  (*     match t return Bit(t * n) @# ty -> list (Bit n @# ty) with *)
  (*     | 0 => fun _ => nil *)
  (*     | S m => fun w => UniBit (TruncLsb n (m * n)) w :: @exprSplitter' _ _ m (UniBit (TruncMsb n (m * n)) w) *)
  (*     end. *)

  (* Definition exprSplitter ty n (pf: n <> 0) sz (w: Bit sz @# ty): list (Bit n @# ty). *)
  (*   refine *)
  (*     (exprSplitter' n (divCeil sz n) (castBits _ ({< Const ty (natToWord (divCeil sz n * n - sz) 0), w>})%kami_expr)). *)
  (*   abstract (pose proof (divCeil_ge sz pf); *)
  (*             Omega.omega). *)
  (* Defined. *)

  (* Fixpoint convertBoolsToWord ls: word (length ls) := *)
  (*   match ls return word (length ls) with *)
  (*   | nil => WO *)
  (*   | x :: xs => WS x (convertBoolsToWord xs) *)
  (*   end. *)

  (* Definition byteAlignMask k: word (div_packn k * n). *)
  (*   refine (nat_cast word _ (combine (wones (size k)) (natToWord (div_packn k * n - size k) 0))). *)
  (*   abstract (pose proof (@divCeil_ge (size k) n ltac:(Omega.omega)); Omega.omega). *)
  (* Defined. *)

  (* Definition putRightPositionWord n (w: word n) start finish := *)
  (*   {< (natToWord (finish - (start + n)) 0), w, (natToWord start 0)>}%word. *)

  (* Eval compute in (map (@evalExpr _) (@exprSplitter type 3 ltac:(Omega.omega) 4 (Const type WO~0~1~0~1))). *)
  (* Eval compute in (map (@evalExpr _) (@exprSplitter type 3 ltac:(Omega.omega) _ (Const type WO~1~1~0~1~1~0~1~0))). *)
  (* Eval compute in ((@wordSplitter 3 ltac:(Omega.omega) _ (WO~1~1~0~1~1~0~1~0))). *)

  (* Goal True. *)
  (*   pose (evalExpr (putRightPosition (Const type (WO~1~1~1)%word) 4 16)). *)
  (*   simpl in f. *)
  (*   assert (f = ($0)~1~1~1~0~0~0~0). *)
  (*   unfold f. *)
  (*   auto. *)
  (*   auto. *)
  (* Defined. *)


Require Import Coq.Bool.Sumbool.
Require Import Coq.FSets.FMapInterface.
Require Import Coq.FSets.FMapFacts.
Require Import Crypto.Compilers.Named.Context.
Require Import Crypto.Compilers.Named.ContextDefinitions.
Require Import Crypto.Util.Tactics.BreakMatch.
Require Import Crypto.Util.Equality.

Module FMapContextFun (E : DecidableType) (W : WSfun E).
  Module Import Properties := WProperties_fun E W.
  Import F.
  Section ctx.
    Context (E_eq_l : forall x y, E.eq x y -> x = y)
            base_type_code (var : base_type_code -> Type)
            (base_type_code_beq : base_type_code -> base_type_code -> bool)
            (base_type_code_bl_transparent : forall x y, base_type_code_beq x y = true -> x = y)
            (base_type_code_lb : forall x y, x = y -> base_type_code_beq x y = true).

    Definition var_cast {a b} (x : var a) : option (var b)
      := match Sumbool.sumbool_of_bool (base_type_code_beq a b), Sumbool.sumbool_of_bool (base_type_code_beq b b) with
         | left pf, left pf' => match eq_trans (base_type_code_bl_transparent _ _ pf) (eq_sym (base_type_code_bl_transparent _ _ pf')) with
                                | eq_refl => Some x
                                end
         | right _, _ | _, right _ => None
         end.
    Definition FMapContext : @Context base_type_code W.key var
      := {| ContextT := W.t { t : _ & var t };
            lookupb t ctx n
            := match W.find n ctx with
               | Some (existT t' v)
                 => var_cast v
               | None => None
               end;
            extendb t ctx n v
            := W.add n (existT _ t v) ctx;
            removeb t ctx n
            := W.remove n ctx;
            empty := W.empty _ |}.
    Lemma FMapContextOk : @ContextOk base_type_code W.key var FMapContext.
    Proof using E_eq_l base_type_code_lb.
      split;
        repeat first [ reflexivity
                     | progress simpl in *
                     | progress intros
                     | rewrite add_eq_o by reflexivity
                     | progress rewrite ?add_neq_o, ?remove_neq_o, ?remove_eq_o in * by intuition
                     | progress rewrite empty_o in *
                     | progress unfold var_cast
                     | progress break_innermost_match_step
                     | rewrite concat_pV
                     | congruence
                     | rewrite base_type_code_lb in * by reflexivity
                     | break_innermost_match_hyps_step
                     | progress unfold var_cast in * ].
    Qed.
  End ctx.

  Section ctx_nd.
    Context {base_type_code var : Type}.

    Definition FMapContext_nd : @Context base_type_code W.key (fun _ => var)
      := {| ContextT := W.t var;
            lookupb t ctx n := W.find n ctx;
            extendb t ctx n v := W.add n v ctx;
            removeb t ctx n := W.remove n ctx;
            empty := W.empty _ |}.
  End ctx_nd.
End FMapContextFun.

Module FMapContext (W : WS) := FMapContextFun W.E W.


Require Import Util Get ListUpdateAt PartialOrder AllInRel.

Set Implicit Arguments.

Inductive terminates X (R: X -> X -> Prop) : X -> Prop :=
| terminatesI x : (forall y, R x y -> terminates R y) -> terminates R x.

Class Terminating X (R: X -> X -> Prop) : Prop :=
  terminating : forall x, terminates R x.

Arguments Terminating X R : clear implicits.

Instance terminating_list Dom `{PO:PartialOrder Dom}
  : Terminating Dom poLt
    -> Terminating (list Dom) poLt.
Proof.
  intros. hnf; intros.
  assert (LE:poLe x x) by reflexivity.
  revert LE. generalize x at 2 3.
  induction x; intros; clear_trivial_eqs.
  - econstructor. intros ? [A B].
    inv A. exfalso. eauto.
  - specialize (H a).
    revert y pf YL H3.
    induction H; intros.
    assert (poLe x x) by reflexivity.
    revert y pf YL H3. specialize (IHx _ H1). clear H1.
    induction IHx; intros.
    econstructor. intros ? [A B]; inv A.
    decide (poEq YL YL0).
    + decide (poEq y y1).
      * exfalso; apply B.
        rewrite p, p0. reflexivity.
      * eapply (H0 y1); eauto.
    + eapply (H2 YL0); eauto.
      split; eauto.
      rewrite H3, H7. reflexivity.
      intro. eapply n.
      eapply poLe_antisymmetric; eauto.
Qed.

Lemma terminates_get_list Dom `{PO:PartialOrder Dom} L
  : terminates poLt L
    -> forall n x, get L n x -> terminates poLt x.
Proof.
  intro Trm.
  induction Trm; intros.
  - econstructor; intros.
    eapply H0. instantiate (1:=list_update_at x n y).
    + revert H1 H2. clear_all.
      general induction x; simpl; isabsurd.
      inv H1.
      * split; eauto with po.
        intro. eapply H2. inv H; eauto.
      * exploit IHx; eauto. split; eauto.
        intro. eapply H. inv H0; eauto.
    + eapply list_update_at_get_3; eauto.
Qed.

Lemma terminates_list_get Dom `{PO:PartialOrder Dom} L
  : (forall n x, get L n x -> terminates poLt x)
    -> terminates poLt L.
Proof.
  intros.
  assert (LE:poLe L L) by reflexivity.
  revert LE. generalize L at 2 3.
  induction L; intros.
  - inv LE. econstructor. intros ? [A B].
    inv A. exfalso. eapply B. reflexivity.
  - invc LE.
    pose proof (H 0 a (getLB _ _)).
    revert y pf YL H3.
    induction H0; intros.
    exploit IHL; [ eauto using get | reflexivity |].
    revert y pf YL H3.
    induction H2; intros.
    econstructor. intros ? [A B]; inv A.
    decide (poEq YL YL0).
    + decide (poEq y y1).
      exfalso; eapply B.
      rewrite p, p0; eauto.
      eapply (H1 y1); eauto.
      intros ? ? Get; inv Get; eauto using get.
    + assert (poLe x0 YL0) by (etransitivity; eauto).
      assert (poLt x0 YL0). {
        split; eauto. intro. eapply n.
        eapply poLe_antisymmetric; eauto.
      }
      eapply H3; eauto.
      * intros ? ? Get. inv Get; eauto using get.
        exploit H2 as Trm; eauto using get.
        eapply terminates_get_list in Trm; eauto.
      * intros. eapply H1; eauto.
        intros ? ? Get; inv Get; eauto using get.
Qed.

Instance terminating_sig Dom `{PO:PartialOrder Dom}
  : Terminating Dom poLt
    -> forall P, Terminating { x : Dom | P x } poLt.
Proof.
  intros Trm P [x Px].
  specialize (Trm x).
  induction Trm.
  econstructor.
  intros [y Py] [LE NEQ]; simpl in *.
  eapply H0. split; eauto.
Qed.

Instance terminating_pair Dom `{PO:PartialOrder Dom} Dom' `{PO':PartialOrder Dom'}
  : Terminating Dom poLt
    -> Terminating Dom' poLt
    -> Terminating (Dom * Dom') poLt.
Proof.
  intros Trm1 Trm2 [x y].
  specialize (Trm1 x).
  specialize (Trm2 y).
  assert (H:poLe y y) by reflexivity; revert H.
  generalize y at 2 3.
  induction Trm1.
  assert (H':poLe x x) by reflexivity; revert H'.
  generalize x at 2 3.
  induction Trm2.
  econstructor.
  intros [z z'] [[LE1 LE2] NEQ]; simpl in *.
  decide (poEq x1 z).
  + decide (poEq y z').
    * exfalso; eapply NEQ; split; simpl; eauto.
    * eapply (H2 z'); eauto.
  + eapply H0; eauto.
Qed.

Instance terminating_bool
  : Terminating bool poLt.
Proof.
  intros x.
  econstructor. intros y [A B].
  destruct x, y; simpl in *; isabsurd.
  econstructor. intros [] [A' B']; isabsurd.
Qed.