(* Copyright (c) 2012. Ben Ferrell, Kenneth Miller,
 *  Matthew Pettersson, Justin Sahs, and Brett Webster.
 *
 *  This file is part of REINS.
 *
 *  This file is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU General Public License as
 *  published by the Free Software Foundation; either version 2 of
 *  the License, or (at your option) any later version.*)

(* TODO need to find out the size of LPBYTE in order to correctly observe
 * Image Thunk Data structures and others... *)

Require Import Coq.Lists.List.
Require Import PEFormat.
Require Import Bits.
Require Import Coq.ZArith.BinInt.
Require Import Coq.NArith.BinPos.
Require Import Coq.ZArith.Zdiv.

Open Scope vector_scope.

Definition block_size : Z := Z_of_nat 3072.

Definition Z_to_nat (z : Z) : nat :=
    match z with
    | Z0 => 0
    | Zpos p => nat_of_P p
    | Zneg p => 0
    end.

Definition word_to_nat (w : WORD) : nat :=
    Z_to_nat (Word.unsigned w).

Definition dword_to_nat (w : DWORD) : nat :=
    Z_to_nat (Word.unsigned w).

Definition ptr_to_nat {A} (p : Ptr A) : nat :=
    match p with
    | ptr d _ => dword_to_nat d
    end.

Definition word_to_Z (w : WORD) : Z :=
    Word.unsigned w.

Definition dword_to_Z (w : DWORD) : Z :=
    Word.unsigned w.

Definition ptr_to_Z {A} (p : Ptr A) : Z :=
    match p with
    | ptr d _ => dword_to_Z d
    end.

Definition getChunk (data : list (list int8)) (z : Z) : list int8 :=
   nth (Z_to_nat (Zdiv z block_size)) data (cons (Word.repr 0) nil).

(* use to cast int8 as int16 *)
Definition int8_to_int16 (x : int8) : int16 :=
   Word.repr(Word.unsigned x)
.

(* takes two int16 (cast from int8) and turns into WORD *)
Definition bytes_to_word (msb : int16) (lsb : int16) : WORD :=
   Word.add (Word.mul msb (Word.repr 256)) lsb
.

Definition parseByte (data : list(list BYTE)) (z : Z) : BYTE :=
   let chunk := getChunk data z in
   let n := Z_to_nat (Zmod z block_size) in
   nth n chunk (Word.repr 0)
.

(* takes list of bytes and idx and returns a WORD *)
Definition parseWord (data : list (list BYTE)) (z : Z) : WORD :=
   let chunk := getChunk data z in
   let n := Z_to_nat (Zmod z block_size) in
   bytes_to_word (int8_to_int16 (nth (S n) chunk (Word.repr 0)))
                 (int8_to_int16 (nth n chunk (Word.repr 0)))
.

(* use to cast int8 as int32 *)
Definition int8_to_int32 (x : int8) : int32 :=
   Word.repr(Word.unsigned x)
.

(* takes four int32 (cast from int8) and turns into a DWORD *)
Definition bytes_to_dword (w : int32) (x : int32)
                          (y : int32) (z : int32)
                          : DWORD :=
   Word.add
      (Word.add (Word.mul w (Word.repr 16777216))
                (Word.mul x (Word.repr 65536)))
      (Word.add (Word.mul y (Word.repr 256)) z)
.

(* returns DWORD at idx *)
Definition parseDoubleWord (data : list (list BYTE)) (z : Z) : DWORD :=
   let chunk := getChunk data z in
   let n := Z_to_nat (Zmod z block_size) in
   bytes_to_dword (int8_to_int32 (nth (S (S (S n))) chunk (Word.repr 0)))
                  (int8_to_int32 (nth (S (S n)) chunk (Word.repr 0)))
                  (int8_to_int32 (nth (S n) chunk (Word.repr 0)))
                  (int8_to_int32 (nth n chunk (Word.repr 0)))
.

Definition parsePtr (data : list (list BYTE)) (z : Z) {A : Type} : Ptr A :=
   ptr (parseDoubleWord data z) A
.

Fixpoint parseVector {A : Type} (parse : list (list BYTE) -> Z -> A)
                                (data : list(list BYTE))
                                (size : Z)
                                (n : Z)
                                (count : nat)
                                 : vector count A :=
    match count with
    | 0 => []
    | S count' => (parse data n) :: (parseVector parse data size (n + size) count')
    end.

Fixpoint copySection (data : list (list BYTE)) (start : Z) (size : Z) (n : nat) : list (list BYTE) :=
   match n with
   | O => cons (vtolist (parseVector parseByte data 1 start (Z_to_nat size))) nil
   | S n' => cons (vtolist (parseVector parseByte data 1 start (Z_to_nat block_size)))
             (copySection data (start+block_size) (size - block_size) n')
end.

Definition parseImageDosHeader (data : list (list BYTE)) : _IMAGE_DOS_HEADER :=
 mkImageDosHeader
   (parseWord data 0)   (* e_magic *)
   (parseWord data 2)   (* e_cblp *)
   (parseWord data 4)   (* e_cp *)
   (parseWord data 6)   (* e_crlc *)
   (parseWord data 8)   (* e_cparhdr *)
   (parseWord data 10)  (* e_minalloc *)
   (parseWord data 12)  (* e_maxalloc *)
   (parseWord data 14)  (* e_ss *)
   (parseWord data 16)  (* e_sp *)
   (parseWord data 18)  (* e_csum *)
   (parseWord data 20)  (* e_ip *)
   (parseWord data 22)  (* e_cs *)
   (parseWord data 24)  (* e_lfarlc *)
   (parseWord data 26)  (* e_ovno *)
   ((parseWord data 28)::
   (parseWord data 30)::
   (parseWord data 32)::
   (parseWord data 34)::[]) (* e_res *)
   (parseWord data 36)  (* e_oemid *)
   (parseWord data 38)  (* e_oeminfo *)
   ((parseWord data 40)::
   (parseWord data 42)::
   (parseWord data 44)::
   (parseWord data 46)::
   (parseWord data 48)::
   (parseWord data 50)::
   (parseWord data 52)::
   (parseWord data 54)::
   (parseWord data 56)::
   (parseWord data 58)::[]) (* e_res2 *)
   (parsePtr data 60) (* e_lfanew *)
.

Definition parseImageFileHeader (data : list (list BYTE)) (n : Z) : _IMAGE_FILE_HEADER :=
    mkImageFileHeader
      (parseWord data n) 		(* Machine *)
      (parseWord data (n+2)) 		(* NumberOfSections *)
      (parseDoubleWord data (n+4)) 	(* TimeDateStamp_IFH *)
      (parseDoubleWord data (n+8)) 	(* PointerToSymbolTable *)
      (parseDoubleWord data (n+12)) 	(* NumberOfSymbols *)
      (parseWord data (n+16)) 		(* SizeOfOptionalHeader *)
      (parseWord data (n+18)) 		(* Characteristics_IFH *)
.

Definition parseImageDataDirectory (data : list (list BYTE)) (n : Z) : _IMAGE_DATA_DIRECTORY :=
    mkImageDataDirectory
      (parseDoubleWord data n)		(* VirtualAddress_IDD *)
      (parseDoubleWord data (n + 4)) 	(* Size *)
.

Definition parseImageOptionalHeader (data : list (list BYTE)) (n : Z) : _IMAGE_OPTIONAL_HEADER :=
    mkImageOptionalHeader
      (parseWord data n)			(*Magic *)
      (parseByte data (n + 2))			(*MajorLinkerVersion *)
      (parseByte data (n + 3))			(*MinorLinkerVersion *)
      (parseDoubleWord data (n + 4))		(*SizeOfCode *)
      (parseDoubleWord data (n + 8))		(*SizeOfInitializedData *)
      (parseDoubleWord data (n + 12))		(*SizeOfUninitializedData *)
      (parseDoubleWord data (n + 16))		(*AddressOfEntryPoint *)
      (parseDoubleWord data (n + 20))		(*BaseOfCode *)
      (parseDoubleWord data (n + 24))		(*BaseOfData *)
      (parseDoubleWord data (n + 28))		(*ImageBase *)
      (parseDoubleWord data (n + 32))		(*SectionAlignment *)
      (parseDoubleWord data (n + 36))		(*FileAlignment *)
      (parseWord data (n + 40))			(*MajorOperatingSystemVersion *)
      (parseWord data (n + 42))			(*MinorOperatingSystemVersion *)
      (parseWord data (n + 44))			(*MajorImageVersion *)
      (parseWord data (n + 46))			(*MinorImageVersion *)
      (parseWord data (n + 48))			(*MajorSubsystemVersion *)
      (parseWord data (n + 50))			(*MinorSubsystemVersion *)
      (parseDoubleWord data (n + 52))		(*Win32VersionValue *)
      (parseDoubleWord data (n + 56))		(*SizeOfImage *)
      (parseDoubleWord data (n + 60))		(*SizeOfHeaders *)
      (parseDoubleWord data (n + 64))		(*CheckSum *)
      (parseWord data (n + 68))			(*Subsystem *)
      (parseWord data (n + 70))			(*DllCharacteristics *)
      (parseDoubleWord data (n + 72))		(*SizeOfStackReserve *)
      (parseDoubleWord data (n + 76))		(*SizeOfStackCommit *)
      (parseDoubleWord data (n + 80))		(*SizeOfHeapReserve *)
      (parseDoubleWord data (n + 84))		(*SizeOfHeapCommit *)
      (parseDoubleWord data (n + 88))		(*LoaderFlags *)
      (parseDoubleWord data (n + 92))		(*NumberOfRvaAndSizes *)
      (parseVector parseImageDataDirectory data 8 (n + 96) 16)(*DataDirectory *)
.

Definition parseImageNtHeader (data : list (list BYTE)) (n : Z) : _IMAGE_NT_HEADER :=
    mkImageNtHeader
      (parseDoubleWord data n) 			(*Signature  *)
      (parseImageFileHeader data (n + 4)) 	(*FileHeader  *)
      (parseImageOptionalHeader data (n + 24)) 	(*OptionalHeader  *)
.

Definition ISSN : Z := Z_of_nat IMAGE_SIZEOF_SHORT_NAME.

Definition parseImageSectionHeader (data : list (list BYTE)) (n : Z) : _IMAGE_SECTION_HEADER :=
    mkImageSectionHeader
      (parseVector (parseByte) data 1 n IMAGE_SIZEOF_SHORT_NAME)(*Name_ISH *)
      (parseDoubleWord data (n + ISSN))	(*PhysicalAddressORVirtualSize*)
      (parseDoubleWord data (n + ISSN + 4))	(*VirtualAddress_ISH *)
      (parseDoubleWord data (n + ISSN + 8))	(*SizeOfRawData *)
      (parseDoubleWord data (n + ISSN + 12))	(*PointerToRawData *)
      (parseDoubleWord data (n + ISSN + 16))	(*PointerToRelocations *)
      (parseDoubleWord data (n + ISSN + 20))	(*PointerToLinenumbers *)
      (parseWord data (n + ISSN + 24))	(*NumberOfRelocations *)
      (parseWord data (n + ISSN + 26))	(*NumberOfLinenumbers *)
      (parseDoubleWord data (n + ISSN + 28))	(*Characteristics_ISH *)
.

Definition parseImageExportDirectory (data : list (list BYTE)) (n : Z) : _IMAGE_EXPORT_DIRECTORY :=
    mkImageExportDirectory
      (parseDoubleWord data n)		(*Characteristics_IED *)
      (parseDoubleWord data (n+4))	(*TimeDateStamp *)
      (parseWord data (n+8))		(*MajorVersion *)
      (parseWord data (n+10))		(*MinorVersion *)
      (parseDoubleWord data (n+12))	(*Name_IED *)
      (parseDoubleWord data (n+16))	(*Base *)
      (parseDoubleWord data (n+20))	(*NumberOfFunctions *)
      (parseDoubleWord data (n+24))	(*NumberOfNames *)
      (parseDoubleWord data (n+28))	(*AddressOfFunctions *)
      (parseDoubleWord data (n+32))	(*AddressOfNames *)
      (parseDoubleWord data (n+36))	(*AddressOfNameOrdinals *)
.

Definition parseImageThunkData (data : list (list BYTE)) (n : Z) : _IMAGE_THUNK_DATA :=
    mkImageThunkData
    	(parseByte data n)		(*ForwarderString *)
	(parsePtr data (n+1)) 	(*Function *)
	(parseDoubleWord data (n+5))	(*Ordinal *)
	(parsePtr data (n+9))	(*AddressOfData *)
.

Fixpoint findSection (data : list (list BYTE)) (rva : DWORD) (n : Z) (num_sec : nat) : option _IMAGE_SECTION_HEADER :=
    match num_sec with
    | 0 => None
    | S n' => let curSection := parseImageSectionHeader data n in
              let v_start := VirtualAddress_ISH curSection in
              let v_end := Word.add v_start (SizeOfRawData curSection) in
              if andb (Word.lequ v_start rva) (Word.ltu rva v_end) then
                 Some curSection
              else
                 findSection data rva (n + 40) n'
    end.

Definition vAddr_to_offset (vaddr : DWORD) (header : _IMAGE_SECTION_HEADER) : DWORD :=
    Word.add (Word.sub vaddr (VirtualAddress_ISH header)) (PointerToRawData header).

Definition derefImageNtHeader (data : list (list BYTE)) (p : Ptr _IMAGE_NT_HEADER) : _IMAGE_NT_HEADER :=
        match p with
        | ptr d _ => parseImageNtHeader data (dword_to_Z d)
        end.

Definition getExports (data : list (list BYTE)) : list DWORD :=
    let dosHeader := parseImageDosHeader data in
    let ntHeader := derefImageNtHeader data (e_lfanew dosHeader) in
    let rva := (VirtualAddress_IDD (nth 0 (vtolist (DataDirectory (OptionalHeader ntHeader)))
                                           {| VirtualAddress_IDD := Word.repr 0; Size := Word.repr 0 |})) in
    match Word.unsigned rva with
    | 0%Z => nil
    | _ =>
           let sectionHeader := findSection data rva
                               ((ptr_to_Z (e_lfanew dosHeader)) + 248)
                               (word_to_nat (NumberOfSections (FileHeader ntHeader))) in
           match sectionHeader with
           | None => nil
           | Some header =>
               let exportDir := parseImageExportDirectory data (dword_to_Z (vAddr_to_offset rva header)) in
               vtolist (parseVector (parseDoubleWord) data 4
                           (dword_to_Z (vAddr_to_offset (AddressOfFunctions exportDir) header))
                           (dword_to_nat (NumberOfFunctions exportDir)))
           end
    end.

(* given a file, check that all exported symbols target
 * low memory chunk boundaries *)
Definition checkExports (data : list (list BYTE)) (mask : DWORD) : bool :=
    let exports := getExports data in
    let check (addr : DWORD) : bool :=
        Word.eq addr (Word.and addr mask)
    in
    List.fold_left (andb) (List.map check exports) true
.

(*----------------------------------------------------------- *)

Definition derefDataDirectoryIAT optionalHeader : _IMAGE_DATA_DIRECTORY :=
	(nth IMAGE_DIRECTORY_ENTRY_IAT (vtolist (DataDirectory optionalHeader))
		{| VirtualAddress_IDD := Word.repr 0; Size := Word.repr 0 |})
.

Definition derefImageOptionalHeader data : _IMAGE_OPTIONAL_HEADER :=
	let dosHeader := parseImageDosHeader data in
	let ntHeader := derefImageNtHeader data (e_lfanew dosHeader) in
	OptionalHeader ntHeader
.

(* Unused, has a bug
Definition parseImageImportDescriptor (data : list BYTE) (n : nat) : _IMAGE_IMPORT_DESCRIPTOR:=
    mkImageImportDescriptor
    	(parseDoubleWord data n) 	(* OriginalFirstThunk *)
	(parseDoubleWord data (n+4))	(* TimeDataStamp_IDD *)
	(parseDoubleWord data (n+8))	(* Ordinal *)
	(parseDoubleWord data (n+12))	(* AddressOfData *)
.

unused, has a bug
Definition derefImageImportDescriptor (data : list DWORD)
  (dataDirectory : _IMAGE_DATA_DIRECTORY) : _IMAGE_IMPORT_DESCRIPTOR :=
	(parseImageImportDescriptor data
		(dword_to_nat dataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress_IDD))
.
*)

Inductive IATBounds : Type :=
| iatbounds : (DWORD * DWORD) -> IATBounds.

(* this returns the RVA start address and the size as two DWORD elements *)
Definition getIATBounds (data : list (list BYTE)) : IATBounds :=
	let optionalHeader := derefImageOptionalHeader data in
	let IAT := derefDataDirectoryIAT optionalHeader in
	iatbounds ((VirtualAddress_IDD IAT), (Size IAT))
.

(* Grabs executable sections defined in the section headers *)
Fixpoint getExecutableBounds (data : list (list BYTE)) (start : Z) (n : nat)
                             : list (DWORD * DWORD * list (list BYTE)) :=
   match n with
   | O => nil
   | S n' =>
             let sectionHeader := parseImageSectionHeader data start in
             let containsExec := Word.eq (Word.and IMAGE_SCN_CNT_CODE (Characteristics_ISH sectionHeader)) IMAGE_SCN_CNT_CODE in
             match containsExec with
             | false => getExecutableBounds data (start+40) n'
             | true =>
                            let rstart := PointerToRawData sectionHeader in
                            let rsize := SizeOfRawData sectionHeader in
                            let numSections := Z_to_nat (Zdiv (Word.unsigned rsize) block_size) in
                             cons (rstart,rsize,copySection data (Word.unsigned rstart) (Word.unsigned rsize) numSections)
                            (getExecutableBounds data (start+40) n')
             end
end.

(* Finds location of section headers, then calls getExecutableBounds to get executable sections *)
Definition getExecutableSections (data : list (list BYTE))
                                 : list (DWORD * DWORD * list (list BYTE)) :=
   let dosHeader := parseImageDosHeader data in
   let ntHeader := derefImageNtHeader data (e_lfanew dosHeader) in
   getExecutableBounds data ((ptr_to_Z (e_lfanew dosHeader)) + 248)
                                          (word_to_nat (NumberOfSections (FileHeader ntHeader)))
.

Close Scope vector_scope.


Require Import Reals.
From iris.program_logic Require Export weakestpre.
From iris.algebra Require Import gmap auth agree gset coPset.
From iris.base_logic Require Import big_op soundness.
From iris.base_logic.lib Require Import wsat.
From iris.proofmode Require Import tactics.
From iris.program_logic Require Import prob_language lifting.
Set Default Proof Using "Type".
Import uPred.

From discprob.idxval Require Import pival pival_dist pidist_singleton idist_pidist_pair ival_dist irrel_equiv ival_pair.
From discprob.basic Require Import monad.

Definition prob_state := {X : Type & pidist X}.
Definition ps_eq : relation prob_state.
  intros (X1&Ix1) (X2&Ix2).
  exact ( ∃ (Hpf: X1 = X2), eq_pidist (eq_rect _ _ Ix1 _ Hpf) (Ix2)).
Defined.

Global Instance pidist_Equiv {X} : @Equiv (@pidist X) := eq_pidist.
Global Instance pidist_equiv {X} : @Equivalence (@pidist X) eq_pidist.
Proof.
  split.
  - intros ?; reflexivity.
  - intros ???; by symmetry.
  - intros ????; etransitivity; eauto.
Qed.

Global Instance prob_state_Equiv : @Equiv prob_state := ps_eq.
Global Instance prob_state_equiv : @Equivalence prob_state ps_eq.
Proof.
  rewrite /ps_eq.
  split.
  - intros (X&Ix). exists (Logic.eq_refl) => //=.
  - intros (X&Ix) (Y&Iy) (HeqTy&Heqelt&HeqIs).
    exists (Logic.eq_sym HeqTy); split; destruct HeqTy => //=.
  - intros (X&Ix) (Y&Iy) (Z&Iz) (HeqTy1&Heqelt1&HeqIs1) (HeqTy2&Heqelt2&HeqIs2).
    destruct HeqTy1, HeqTy2; exists (Logic.eq_refl) => //=.
    rewrite //= in Heqelt1 Heqelt2 HeqIs1 HeqIs2.
    split; etransitivity; eauto.
Qed.

Class probG' (Σ : gFunctors) := ProbG {
  probG_invG : invG Σ;
  probG_name : gname;
  probG_inG :> inG Σ (authR (optionUR (exclR (discreteC prob_state))));
}.

Class stateG' (Λ: probLanguage) (Σ : gFunctors) := StateG {
  stateG_interp : (language.state Λ) → iProp Σ
}.

Notation probG Σ := (probG' Σ).

Definition ownProbRaw `{@probG' Σ} {X} (Is: pidist X) : iProp Σ :=
  own probG_name (◯ (Excl' (existT X Is : discreteC prob_state))).
Typeclasses Opaque ownProbRaw.
Instance: Params (@ownProbRaw) 9.
Global Instance ownProbRaw_timeless `{@probG' Σ} {X} (I: pidist X)  : Timeless (ownProbRaw I).
Proof. rewrite /ownProbRaw. apply _. Qed.

Definition ownProb `{@probG' Σ} {X} (Is: pidist X) : iProp Σ :=
  (∃ Is', ⌜ irrel_pidist Is Is' ⌝ ∗ ownProbRaw Is')%I.
Typeclasses Opaque ownProb.
Instance: Params (@ownProb) 9.
Global Instance ownProb_timeless `{@probG' Σ} {X} (I: pidist X)  : Timeless (ownProb I).
Proof. rewrite /ownProb. apply _. Qed.

Global Instance ownProbRaw_proper `{@probG' Σ} {X}:
  Proper ((≡) ==> (⊣⊢)) (@ownProbRaw _ _ X).
Proof.
  intros Is Is' Hequiv.
  rewrite /ownProbRaw.
  assert (existT X Is ≡ existT X Is') as Hequiv'.
  { rewrite /equiv/ps_eq//=. exists Logic.eq_refl => //=. }
  setoid_rewrite Hequiv'.
  done.
Qed.

Global Instance ownProb_proper `{@probG' Σ} {X}:
  Proper ((≡) ==> (⊣⊢)) (@ownProb _ _ X).
Proof.
  intros Is Is' Hequiv.
  rewrite /ownProb. setoid_rewrite Hequiv. done.
Qed.

Global Instance ownProb_anti `{@probG' Σ} {X}:
  Proper ((le_pidist) ==> flip (⊢)) (@ownProb _ _ X).
Proof.
  intros Is Is' Hequiv.
  rewrite /ownProb.
  assert (irrel_pidist Is Is') as Hequiv'.
  {
    eapply irrel_pidist_proper; first eassumption.
    { reflexivity.  }
    { reflexivity.  }
  }
  setoid_rewrite Hequiv'. done.
Qed.

Global Instance ownProb_anti_irrel `{@probG' Σ} {X}:
  Proper ((irrel_pidist) ==> flip (⊢)) (@ownProb _ _ X).
Proof.
  intros Is Is' Hequiv.
  rewrite /ownProb.
  setoid_rewrite Hequiv. done.
Qed.

Lemma ownProb_anti_irrel' `{@probG' Σ} {X} (Is1 Is2: pidist X):
  irrel_pidist Is2 Is1 →
  ownProb Is1 ⊢ ownProb Is2.
Proof.
  intros. by apply ownProb_anti_irrel.
Qed.

Record ctype {Λ} `{@probG' Σ} e1 σ1 (c: prob_state) :=
  { choice_hd_type : Type;
     choice_hd_pidist : pidist choice_hd_type;
     choice_cont : choice_hd_type →  pidist (projT1 c);
     choice_equiv : irrel_pidist  (mbind (choice_cont) (choice_hd_pidist)) (projT2 c);
     choice_post : prim_step_result → choice_hd_type → Prop;
     choice_coupling :> irrel_couplingP (@ivdist_prim_step Λ e1 σ1) choice_hd_pidist choice_post}.

Arguments choice_cont {_ _ _ _ _ _} _ _ .

Lemma probG_irisG_response_inhabited1 `{@probG' Σ} Λ:
  ∀ (e1 : language.expr (prob_lang Λ)) (σ1 : language.state (prob_lang Λ))
  (c1 : prob_state) (c2 : ctype e1 σ1 c1) (e2 : language.expr (prob_lang Λ))
  (σ2 : language.state (prob_lang Λ)) (efs : list (language.expr (prob_lang Λ))),
  language.prim_step e1 σ1 e2 σ2 efs
  → rsupport (choice_coupling e1 σ1 c1 c2) (prim_res_step e2 σ2 efs).
Proof.
  intros e1 σ1 c1 c2 e2 σ2 efs Hprim.
  destruct c1 as (X&x&Ix).
  destruct c2 as [Y m f Hequiv P [Istep' Is' Hirrel1 ? Ic]].
  specialize (ivdist_step_idx e1 _ _ _ _ Hprim); eauto.
  intros (istep&Hind&Hgt)%ClassicalEpsilon.constructive_indefinite_description.
  rewrite /response_type//=/rsupport//=.
  rewrite //=/rsupport.
  destruct Ic as [? ? Ic].
  destruct Ic as [Ic Hp1 Hp2].
  destruct (irrel_ivd_support_coerce _ _ Hirrel1 (prim_res_step e2 σ2 efs)) as (_&Hco2).
  eapply (ClassicalEpsilon.constructive_indefinite_description) in Hco2; last first.
  { rewrite //=. exists istep; split; auto. }
  edestruct Hco2 as (istep'&Hind'&Hgt').
  rewrite //=.

  apply (ClassicalEpsilon.constructive_indefinite_description) in Hp1.
  destruct Hp1 as (h1&Hp1).
  unshelve eexists.
  { refine (snd (proj1_sig (ival.ind Ic _))).
    destruct (h1 (ival.coerce_supp _ _ Hgt')) as ((iic&?)&?).
    exact iic.
  }
  destruct Hp1 as (?&?&?&Hind_equiv&Hval_equiv).
  specialize (Hval_equiv (ival.coerce_supp _ _ Hgt')).
  specialize (Hind_equiv (ival.coerce_supp _ _ Hgt')).
  destruct (h1 _) as ((iic&?)&?) => //=.
  case_eq (ival.ind Ic iic). intros (?&?) ? Heq.
  rewrite //= Heq //= in Hind_equiv.
  exists iic; unshelve eexists.
  {
  rewrite //= in Hind'. rewrite -Hind' -Hind_equiv //=.
  }
  rewrite //= in Hind_equiv *; split.
  * clear -Heq. apply sval.sval_inj_pi => //=. rewrite Heq => //=.
    f_equal. by rewrite Hind_equiv Hind'.
  * rewrite //= in Hval_equiv.
    rewrite Rmult_1_r in Hval_equiv; auto.
    rewrite Hval_equiv. auto with *.
Qed.

Lemma probG_irisG_response_inhabited2 `{@probG' Σ} Λ:
  ∀ (e1 : language.expr (prob_lang Λ)) (σ1 : language.state (prob_lang Λ))
  (c1 : prob_state) (c2 : ctype e1 σ1 c1) (e2 : language.expr (prob_lang Λ))
  (σ2 : language.state (prob_lang Λ)) (efs : list (language.expr (prob_lang Λ))),
  language.prim_step e1 σ1 e2 σ2 efs
  → { x : choice_hd_type _ _ _ c2 | choice_post _ _ _ c2 (prim_res_step e2 σ2 efs) x}.
Proof.
  intros.
  edestruct (probG_irisG_response_inhabited1 _ e1 σ1 c1 c2 e2 σ2 efs) as (x&Hpf); first eassumption.
  exists x. destruct Hpf as (i&Hpf&?); eauto.
Qed.

Instance probG_irisG  `{@probG' Σ} `{stateG' Λ Σ} : irisG' Λ Σ.
Proof.
refine {|
  iris_invG := probG_invG;
  state_interp σ := stateG_interp σ;
  aux_state := prob_state;
  aux_interp c := own probG_name (● (Excl' (c:discreteC prob_state)));
  choice_type := λ e1 σ1 c1, ctype e1 σ1 c1;
  response_type := λ e1 σ1 c1 c2 e2 σ2 efs, { x : choice_hd_type _ _ _ c2 | choice_post _ _ _ c2 (prim_res_step e2 σ2 efs) x};
  step_interpR := λ e1 σ1 c1 c2 e2 σ2 efs c3,
                  own probG_name (● (Excl' (existT (projT1 c1)
                                           (@choice_cont _ _ _ _ _ _ c2 (proj1_sig c3) : pidist (projT1 c1)) :
                                         discreteC prob_state)))
|}.
- abstract (iIntros (e1 σ1 c1 c2 e2 σ2 efs c3) "H";
  destruct c1 as (X&(x&Ix));
  destruct c2 as [Y m f Hequiv P Icoupling];
  destruct c3 as (y&Hgt);
  rewrite //=;
  iExists (existT X (f y) : discreteC prob_state);
  done).
- apply probG_irisG_response_inhabited2.
Defined.

Section lifting.
Context {Λ : probLanguage}.
Context `{stateG' Λ Σ} `{probG Σ}.
Hint Extern 2 (language.to_val _ = None) => apply (reducible_not_val _ inhabitant).

Lemma wp_pure_step_fupd `{Inhabited (state Λ)} s E E' e1 e2 φ Φ :
  PureExec φ e1 e2 →
  φ →
  (|={E,E'}▷=> WP e2 @ s; E {{ Φ }}) ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros ([??] Hφ) "HWP".
  iApply (wp_lift_pure_det_step with "[HWP]"); [eauto|naive_solver|].
  { intros.  destruct s; auto. }
  iIntros (σ1 c1) "Ha". iMod "HWP". iModIntro. iNext.
  destruct c1 as (X&Ix).
  unshelve(
  iExists {| choice_hd_type := unit;
             choice_hd_pidist := mret tt;
             choice_cont := (λ x, _);
             choice_equiv := _;
             choice_post := λ x y, True;
             choice_coupling := irrel_coupling_trivial _ _|}).
  { exact Ix. }
  { rewrite //=. setoid_rewrite pidist_left_id. reflexivity. }
  iMod "HWP". iModIntro. iIntros (c3).
  iFrame "HWP".
  rewrite //= /response_type//= in c3 *.
  iFrame.
Qed.

Lemma wp_pure_step_later `{Inhabited (state Λ)} s E e1 e2 φ Φ :
  PureExec φ e1 e2 →
  φ →
  ▷ WP e2 @ s; E {{ Φ }} ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  intros ??. rewrite -wp_pure_step_fupd //. rewrite -step_fupd_intro //.
Qed.

Lemma prob_update {X Y} (Is: pidist X) (Is': pidist Y) :
  ◯ Excl' (existT X Is : discreteC prob_state)
    ⋅ ● Excl' (existT X Is : discreteC prob_state) ~~>
    ◯ Excl' (existT Y Is' : discreteC prob_state)
    ⋅ ● Excl' (existT Y Is' : discreteC prob_state).
Proof.
  apply (auth_update (Excl' (existT X Is : discreteC prob_state))).
  apply option_local_update, exclusive_local_update.
  econstructor.
Qed.

Lemma wp_lift_step_couple {X Y} s E Φ e1 (Is: pidist X) (g: X → pidist Y) R:
  to_val e1 = None →
   ▷ ownProb (x ← Is; g x) ∗ (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
   ▷ ∃ (ic: irrel_couplingP (ivdist_prim_step e1 σ1) Is R), ∀ e2 σ2 efs (x: X),
      (⌜prim_step e1 σ1 e2 σ2 efs⌝ ∗ ⌜R (prim_res_step e2 σ2 efs) x⌝) ∗ ownProb (g x)  ={∅,E}=∗
      state_interp σ2 ∗
      WP e2 @ s; E {{ λ v, (Φ v) }} ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
           ⊢ WP e1 @ s; E {{ λ v, (Φ v) }}.
Proof.
  iIntros (Htoval) "(Hown&Hwand)". iApply wp_lift_step; first done.
  iIntros (σ1 c1) "(Hσ&Ha)".
  iMod ("Hwand" $! σ1 with "Hσ") as "($&H)".
  iModIntro.
  rewrite /ownProb/ownProbRaw.
  iNext.
  iDestruct "Hown" as (Is'') "(%&Hprob)".
  iDestruct (own_valid_2 with "Ha Hprob") as %Hval.
  apply auth_valid_discrete_2 in Hval as (Hequiv%Excl_included&?).
  symmetry in Hequiv.
  setoid_rewrite <-Hequiv at 1.
  symmetry in Hequiv.
  iDestruct "H" as (ic) "H".
  destruct c1 as (X'&Is').
  rewrite /equiv/ofe_equiv/ps_eq//= in Hequiv.
  destruct Hequiv as (Heq&Heq_spidist).
  rewrite //= in Heq_spidist.
  unshelve (iExists _).
  { unshelve (refine ({| choice_hd_type := X;
             choice_hd_pidist := Is;
             choice_cont := _;
             choice_equiv := _;
             choice_post := _;
             choice_coupling := _ |})); try eapply Ic'.
    { subst. apply g. }
    { eapply R. }
    { abstract (rewrite //=; subst; rewrite /internal_eq_rew_r_dep//=; rewrite -Heq_spidist //=). }
    { apply ic. }
  }
  iIntros (e2' σ2' ef' c3).
  iIntros (Hstep).
  destruct c3 as (x&Hx).
  iSpecialize ("H" $! e2' σ2' ef' x).
  unshelve (iMod (own_update_2  with "Hprob Ha") as "[Hσ Hσf]").
  { exact ( ◯ Excl' (existT _ ((g x)) : discreteC prob_state)). }
  { exact ( ● Excl' (existT _ ((g x)) : discreteC prob_state)). }
  { eapply prob_update. }
  iMod ("H" with "[Hσf]") as "($&Hval&$)".
  { iSplitR.
    * abstract (iSplit; iPureIntro;  eauto; destruct Hx as (i&Hpf&Hval&?); eauto).
    * iExists _; iFrame. iPureIntro; auto; reflexivity.
  }
  iModIntro. iFrame.
  rewrite //=. clear Hx.
  rewrite /choice_hd_type in x *.
  iApply own.own_proper; last auto.
  apply Auth_proper; auto.
  apply Some_proper.
  apply Excl_proper; auto.
  apply Some_proper.
  apply Excl_proper; auto.
  symmetry. unshelve (esplit); eauto.
  subst => //=.
Qed.

Lemma wp_lift_step_nocouple s E Φ e1 :
  to_val e1 = None →
  (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
    ▷ ∀ e2 σ2 efs,
      ⌜prim_step e1 σ1 e2 σ2 efs⌝  ={∅,E}=∗
      state_interp σ2 ∗ WP e2 @ s; E {{ Φ }} ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (?) "H". iApply (wp_lift_step s E _ e1)=>//; iIntros (σ1 c1) "(Hσ1&Ha)".
  iMod ("H" $! σ1 with "Hσ1") as "[$ H]".
  iModIntro; iNext.
  destruct c1 as (X&Ix).
  unshelve(
  iExists {| choice_hd_type := unit;
             choice_hd_pidist := mret tt;
             choice_cont := (λ x, _);
             choice_equiv := _;
             choice_post := λ x y, True;
             choice_coupling := irrel_coupling_trivial _ _|}).
  { exact Ix. }
  { rewrite //=. setoid_rewrite pidist_left_id. reflexivity. }
  iIntros (e2 σ2 efs c3) "%".
  iMod ("H" $! e2 σ2 efs with "[#]") as "($ & HΦ & $)"; first by eauto.
  iModIntro. iSplitR "Ha"; auto.
Qed.

Lemma wp_lift_atomic_step_couple {X Y} s E Φ e1 (Is: pidist X) (g: X → pidist Y) R:
  to_val e1 = None →
   ▷ ownProb (x ← Is; g x) ∗ (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
   ▷ ∃ (ic: irrel_couplingP (ivdist_prim_step e1 σ1) Is R), ∀ e2 σ2 efs (x: X),
      (⌜prim_step e1 σ1 e2 σ2 efs⌝ ∗ ⌜R (prim_res_step e2 σ2 efs) x⌝) ∗ ownProb (g x) ={∅,E}=∗
      state_interp σ2 ∗
      (default False (to_val e2) Φ) ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
           ⊢ WP e1 @ s; E {{ λ v, (Φ v) }}.
Proof.
  iIntros (Htoval) "(Hown&Hwand)".
  iApply wp_lift_step_couple; auto.
  iFrame "Hown". iIntros (σ1) "Hσ".
  iMod ("Hwand" $! σ1 with "Hσ") as "($&H)".
  iModIntro.
  iNext. iDestruct "H" as (ic) "H". iExists ic. iIntros (e2 ???) "Hp".
  iMod ("H" with "Hp") as "(Hσ&Hphi&?)"; first eauto.
  iFrame. iModIntro.
  destruct (to_val e2) eqn:Heq_val; last auto.
  iApply wp_value; eauto.
Qed.

(*
Lemma wp_lift_atomic_step_couple' {X Y} s E Φ e1 (Is: pidist X) (g: X → pidist Y) R:
  to_val e1 = None →
   ▷ ownProb (x ← Is; g x) ∗ (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
   ▷ ∃ (ic: irrel_couplingP (ivdist_prim_step e1 σ1) Is R), ∀ e2 σ2 efs (x: X),
      (⌜prim_step e1 σ1 e2 σ2 efs⌝ ∗ ⌜R (prim_res_step e2 σ2 efs) x⌝)  ={∅,E}=∗
      state_interp σ2 ∗
      default False (to_val e2) Φ ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
           ⊢ WP e1 @ s; E {{ λ v, (Φ v) ∗ ∃ σ2 efs x, ⌜ R (prim_res_step (of_val v) σ2 efs) x ⌝
                                                         ∗ ownProb (g x)}}.
Proof.
  iIntros (Htoval) "(Hown&Hwand)".
  iApply wp_lift_step_couple; auto.
  iFrame "Hown". iIntros (σ1) "Hσ".
  iMod ("Hwand" $! σ1 with "Hσ") as "($&H)".
  iModIntro.
  iNext. iDestruct "H" as (ic) "H". iExists ic. iIntros (e2 ??? Hp).
  iMod ("H" with "[%]") as "(Hσ&?&?)"; first eauto.
  iFrame. iModIntro. iIntros "Hown".
  destruct (to_val e2) eqn:Heq_val; last auto.
  iApply wp_value; eauto.
  iFrame. iExists _, _, _; iFrame; iPureIntro.
  rewrite (of_to_val e2); auto. intuition eauto.
Qed.

Lemma wp_lift_atomic_step_couple_prop {X Y} s E Φ e1 (Is: pidist X) (g: X → pidist Y) R:
  to_val e1 = None →
   ▷ ownProb (x ← Is; g x) ∗ (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
   ▷ ⌜ irrel_coupling_propP (ivdist_prim_step e1 σ1) Is R ⌝ ∗ ∀ e2 σ2 efs (x: X),
      (⌜prim_step e1 σ1 e2 σ2 efs⌝ ∗ ⌜R (prim_res_step e2 σ2 efs) x⌝)  ={∅,E}=∗
      state_interp σ2 ∗
      default False (to_val e2) Φ ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
           ⊢ WP e1 @ s; E {{ λ v, (Φ v) ∗ ∃ σ2 efs x, ⌜ R (prim_res_step (of_val v) σ2 efs) x ⌝
                                                         ∗ ownProb (g x)}}.
Proof.
  iIntros (?) "(Hown&Hwand)". iApply wp_lift_atomic_step_couple; eauto.
  iFrame "Hown". iIntros. iMod ("Hwand" $! _ with "[$]") as "H".
  iModIntro. iDestruct "H" as "($&Hprop)".
  iNext. iDestruct "Hprop" as (Hp) "Hc".
  unshelve (iExists _).
  { apply ic_prop_to_wit; eauto. }
  auto.
Qed.
*)

Lemma wp_lift_pure_step_nocouple `{Inhabited (state Λ)} s E E' Φ e1 :
  (∀ σ1, reducible e1 σ1) →
  (∀ σ1 e2 σ2 efs, prim_step e1 σ1 e2 σ2 efs → σ1 = σ2) →
  (|={E,E'}▷=> ∀ e2 efs σ, ⌜prim_step e1 σ e2 σ efs⌝ →
    WP e2 @ s; E {{ Φ }} ∗ [∗ list] ef ∈ efs, WP ef @ s ; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (? Hpuredet) "H". iApply (wp_lift_pure_step s E); try done.
  { destruct s; auto. }
  iIntros (σ1 c1) "Ha". iMod "H".
  iMod fupd_intro_mask' as "Hclose"; last iModIntro; first set_solver.
  iNext.
  destruct c1 as (X&Ix).
  unshelve(
  iExists {| choice_hd_type := unit;
             choice_hd_pidist := mret tt;
             choice_cont := (λ x, _);
             choice_equiv := _;
             choice_post := λ x y, True;
             choice_coupling := irrel_coupling_trivial _ _|}).
  { exact Ix. }
  { rewrite //=. setoid_rewrite pidist_left_id. reflexivity. }
  iMod "H". iModIntro. iIntros (e2 efs c3) "Hstep".
  rewrite //=. iFrame; iApply "H". auto.
Qed.

Lemma wp_lift_atomic_step_nocouple {s E Φ} e1 :
  to_val e1 = None →
  (∀ σ1, state_interp σ1 ={E}=∗
    ⌜if s is NotStuck then reducible e1 σ1 else True⌝ ∗
    ▷ ∀ e2 σ2 efs, ⌜prim_step e1 σ1 e2 σ2 efs⌝ ={E}=∗
      state_interp σ2 ∗
      default False (to_val e2) Φ ∗ [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (?) "H". iApply (wp_lift_step s E _ e1)=>//; iIntros (σ1 c1) "(Hσ1&Ha)".
  iMod ("H" $! σ1 with "Hσ1") as "[$ H]".
  iMod (fupd_intro_mask' E ∅) as "Hclose"; first set_solver.
  iModIntro; iNext.
  destruct c1 as (X&Ix).
  unshelve(
  iExists {| choice_hd_type := unit;
             choice_hd_pidist := mret tt;
             choice_cont := (λ x, _);
             choice_equiv := _;
             choice_post := λ x y, True;
             choice_coupling := irrel_coupling_trivial _ _|}).
  { exact Ix. }
  { rewrite //=. setoid_rewrite pidist_left_id. reflexivity. }
  iIntros (e2 σ2 efs c3) "%". iMod "Hclose" as "_".
  iMod ("H" $! e2 σ2 efs with "[#]") as "($ & HΦ & $)"; first by eauto.
  destruct (to_val e2) eqn:?; last by iExFalso.
  iModIntro. iSplitR "Ha"; auto.
  by iApply wp_value.
Qed.

Lemma wp_lift_pure_det_step_nocouple `{Inhabited (state Λ)} {s E E' Φ} e1 e2 efs :
  (∀ σ1, reducible e1 σ1) →
  (∀ σ1 e2' σ2 efs', prim_step e1 σ1 e2' σ2 efs' → σ1 = σ2 ∧ e2 = e2' ∧ efs = efs')→
  (|={E,E'}▷=> WP e2 @ s; E {{ Φ }} ∗ [∗ list] ef ∈ efs, WP ef @ s ; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (? Hpuredet) "H". iApply (wp_lift_pure_step_nocouple s E); try done.
  { by intros; eapply Hpuredet. }
  iApply (step_fupd_wand with "H"); iIntros "H".
  by iIntros (e' efs' σ (_&->&->)%Hpuredet).
Qed.

Lemma wp_prob_bind K `{!ProbLanguageCtx K} s E e Φ :
  WP e @ s; E {{ v, WP K (of_val v) @ s; E {{ Φ }} }} ⊢ WP K e @ s; E {{ Φ }}.
Proof.
  apply: wp_bind.
  intros e1 σ1 c1 c2.
  destruct c1 as (X&x&Ix).
  destruct c2 as [Y m f Hequiv P Ic].
  unshelve (eexists).
  {
    refine {| choice_hd_type := Y;
              choice_hd_pidist := m;
              choice_cont := f;
              choice_equiv := Hequiv;
              choice_post := _;
              choice_coupling := _ |}.
    unshelve (apply ivdist_bind_coupling; eauto; destruct s; auto); last eapply Ic.
  }
  rewrite /response_type//=.
  intros e2 σ2 efs (y&e2'&Heq).
  rewrite //= in Heq *.
  destruct Heq as (Hpf&Hind&?).
  assert (e2 = e2').
  { apply: (@fill_inj _ K). auto. }
  subst.
  unshelve (eexists).
  { exists y. eapply Hpf. }
  auto.
Qed.

Lemma wp_prob_bind_inv K `{!ProbLanguageCtx K} s E e Φ :
  WP K e @ s; E {{ Φ }} ⊢ WP e @ s; E {{ v, WP K (of_val v) @ s; E {{ Φ }} }}.
Proof.
  apply: weakestpre.wp_bind_inv.
  intros e1 σ1 c1 c2.
  destruct c1 as (X&x&Ix).
  destruct c2 as [Y m f Hequiv P Ic].
  unshelve (eexists).
  {
    refine {| choice_hd_type := Y;
              choice_hd_pidist := m;
              choice_cont := f;
              choice_equiv := Hequiv;
              choice_post := _;
              choice_coupling := _ |}.
    unshelve (apply ivdist_bind_inv_coupling; eauto; destruct s; eauto); last eapply Ic.
  }
  rewrite /response_type//=.
  intros e2 σ2 efs (y&e2'&Heq).
  rewrite //= in Heq *.
  destruct Heq as (Hpf&Hind&?).
  unshelve (eexists).
  { exists y; eauto. }
  auto.
Qed.

End lifting.

From iris.program_logic Require Import ectx_language ectx_lifting.
Section ectx_lifting.
Context {Λ : ectxLanguage}.
Context  `{probG Σ} `{stateG' Λ Σ}.

Context {Hinh : Inhabited (state Λ)}.
Implicit Types P : iProp Σ.
Implicit Types Φ : val Λ → iProp Σ.
Implicit Types v : val Λ.
Implicit Types e : expr Λ.
Hint Extern 2 (reducible ?e1 ?σ1) => apply head_prim_reducible.
Hint Extern 2 (head_step ?e1 ?σ1 ?e2 ?σ2 ?efs) => apply head_reducible_prim_step.

Lemma wp_lift_atomic_head_step_couple {X Y} {s E Φ} e1 (Is: pidist X) (g: X → pidist Y) R:
  to_val e1 = None →
   ▷ ownProb (x ← Is; g x) ∗ (∀ σ1, state_interp σ1 ={E,∅}=∗
    ⌜head_reducible e1 σ1⌝ ∗
   ▷ ∃ (ic: irrel_couplingP (ivdist_prim_step e1 σ1) Is R), ∀ e2 σ2 efs (x: X),
      (⌜head_step e1 σ1 e2 σ2 efs⌝ ∗ ⌜R (prim_res_step e2 σ2 efs) x⌝) ∗ ownProb (g x) ={∅,E}=∗
      state_interp σ2 ∗
      default False (to_val e2) Φ ∗
      [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
           ⊢ WP e1 @ s; E {{ λ v, (Φ v) }}.
Proof using Hinh.
  iIntros (?) "(Hown&H)". iApply wp_lift_atomic_step_couple.
  { eauto. }
  iFrame "Hown".
  iIntros (σ1) "Hσ1". iMod ("H" $! σ1 with "Hσ1") as "[% H]"; iModIntro.
  iSplit; first by destruct s; auto. iNext.
  iDestruct "H" as (c) "H".
  iExists c. iIntros (e2 σ2 efs c3) "(%&Hp)". iApply "H"; auto.
  iFrame.
  iSplit; iPureIntro; intuition; eauto.
Qed.

Lemma wp_lift_pure_head_step_nocouple {s E E' Φ} e1 :
  (∀ σ1, head_reducible e1 σ1) →
  (∀ σ1 e2 σ2 efs, head_step e1 σ1 e2 σ2 efs → σ1 = σ2) →
  (|={E,E'}▷=> ∀ e2 efs σ, ⌜head_step e1 σ e2 σ efs⌝ →
    WP e2 @ s; E {{ Φ }} ∗ [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof using Hinh.
  iIntros (??) "H". iApply wp_lift_pure_step_nocouple; eauto.
  iApply (step_fupd_wand with "H"); iIntros "H".
  iIntros (????). iApply "H"; eauto.
Qed.

Lemma wp_lift_atomic_head_step_nocouple {s E Φ} e1 :
  to_val e1 = None →
  (∀ σ1, state_interp σ1 ={E}=∗
    ⌜head_reducible e1 σ1⌝ ∗
    ▷ ∀ e2 σ2 efs, ⌜head_step e1 σ1 e2 σ2 efs⌝ ={E}=∗
      state_interp σ2 ∗
      default False (to_val e2) Φ ∗ [∗ list] ef ∈ efs, WP ef @ s; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (?) "H". iApply wp_lift_atomic_step_nocouple; eauto.
  iIntros (σ1) "Hσ1". iMod ("H" $! σ1 with "Hσ1") as "[% H]"; iModIntro.
  iSplit; first by (destruct s; eauto). iNext. iIntros (e2 σ2 efs) "%". iApply "H"; auto.
Qed.

Lemma wp_lift_atomic_head_step_no_fork_nocouple {s E Φ} e1 :
  to_val e1 = None →
  (∀ σ1, state_interp σ1 ={E}=∗
    ⌜head_reducible e1 σ1⌝ ∗
    ▷ ∀ e2 σ2 efs, ⌜head_step e1 σ1 e2 σ2 efs⌝ ={E}=∗
      ⌜efs = []⌝ ∗ state_interp σ2 ∗ default False (to_val e2) Φ)
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof.
  iIntros (?) "H". iApply wp_lift_atomic_head_step_nocouple; eauto.
  iIntros (σ1) "Hσ1". iMod ("H" $! σ1 with "Hσ1") as "[$ H]"; iModIntro.
  iNext; iIntros (v2 σ2 efs) "%".
  iMod ("H" $! v2 σ2 efs with "[# //]") as "(% & $ & $)"; subst; auto.
Qed.

Lemma wp_lift_pure_det_head_step_nocouple {s E E' Φ} e1 e2 efs :
  (∀ σ1, head_reducible e1 σ1) →
  (∀ σ1 e2' σ2 efs',
    head_step e1 σ1 e2' σ2 efs' → σ1 = σ2 ∧ e2 = e2' ∧ efs = efs') →
  (|={E,E'}▷=> WP e2 @ s; E {{ Φ }} ∗ [∗ list] ef ∈ efs, WP ef @ s ; ⊤ {{ _, True }})
  ⊢ WP e1 @ s; E {{ Φ }}.
Proof using Hinh. eauto using wp_lift_pure_det_step_nocouple. Qed.

Lemma wp_prob_ectx_bind {s E e} K Φ :
  WP e @ s; E {{ v, WP fill K (of_val v) @ s; E {{ Φ }} }} ⊢ WP fill K e @ s; E {{ Φ }}.
Proof. apply: wp_prob_bind. Qed.

Lemma wp_prob_ectx_bind_inv {s E Φ} K e :
  WP fill K e @ s; E {{ Φ }} ⊢ WP e @ s; E {{ v, WP fill K (of_val v) @ s; E {{ Φ }} }}.
Proof. apply: wp_prob_bind_inv. Qed.

End ectx_lifting.


(* This code is copyrighted by its authors; it is distributed under  *)
(* the terms of the LGPL license (see LICENSE and description files) *)

(* Contribution to the Coq Library   V6.3 (July 1999)                    *)
(****************************************************************************)
(*                                                                          *)
(*                                                                          *)
(*                   Solange Coupet-Grimal & Line Jakubiec                  *)
(*                                                                          *)
(*                                                                          *)
(*              Laboratoire d'Informatique de Marseille                     *)
(*               CMI-Technopole de Chateau-Gombert                          *)
(*                   39, Rue F. Joliot Curie                                *)
(*                   13453 MARSEILLE Cedex 13                               *)
(*           e-mail:{Solange.Coupet,Line.Jakubiec}@lim.univ-mrs.fr          *)
(*                                                                          *)
(*                                                                          *)
(*                                Coq V5.10                                 *)
(*                              May 30th 1996                               *)
(*                                                                          *)
(****************************************************************************)
(*                       Factorization_Prog.v                               *)
(****************************************************************************)

Require Export Factorization.

Section Factorization_for_Realizer.

  Variable A : Set.
  Variable BASE : BT.
  Let b := base BASE.
  Let Num := num BASE.
  Let Val_bound := val_bound BASE.
  Let Digit := digit BASE.
  Let Tl := tl Digit.

  Variable FR : forall n : nat, A -> inf n -> inf n -> A.

  Let R (n : nat) (a : A) (x y : inf n) (a' : A) : Prop := a' = FR n a x y.

  Notation Factorizable := (factorizable _) (only parsing).
  Notation Proper := (proper _) (only parsing).

  Theorem factorization_for_realizer :
   factorizable _ R ->
   proper _ BASE R ->
   forall (n : nat) (a : A) (X Y : Num n),
   {a' : A | R (exp b n) a (Val_bound n X) (Val_bound n Y) a'}.

  Fixpoint Impl (n : nat) : A -> Num n -> Num n -> A :=
    fun a : A =>
    match n return (Num n -> Num n -> A) with
    | O => fun X Y : Num 0 => a
    | S p =>
        fun X Y : Num (S p) =>
        Impl p (FR b a (Hd Digit p X) (Hd Digit p Y))
          (Tl (S p) X) (Tl (S p) Y)
    end.
  refine (fun H H' n a X Y => exist _ (Impl n a X Y) _).
  generalize X Y a; clear X Y a; elim n.
  intros X Y a.
  simpl in |- *; unfold Val_bound in |- *; unfold R in |- *; apply prop_Rel;
   auto.    (*is_proper*)
  unfold R in |- *; unfold Val_bound in |- *; unfold b in |- *;
   intros m Hr a X Y.
  apply fact_Rel; simpl in |- *; auto.  (*is_factorizable*)
  Qed.

End Factorization_for_Realizer.

(*************************************************************************************)


(** * The Imp language  *)

(** We now demonstrate how to use ITrees in the context of verified compilation.
    To this end, we will consider a simple compiler from an imperative language
    to a control-flow-graph language.  The meaning of both languages will be
    given in terms of ITrees, so that the proof of correctness is expressed by
    proving a bisimulation between ITrees.

    We shall emphasize two main satisfying aspects of the resulting
    formalization.

    - Despite the correctness being termination-sensitive, we do not write any
      cofixpoints. All reasoning is purely equational, and the underlying
      coinductive reasoning is hidden on the library side.

    - We split the correctness in two steps. First, a linking theory of the CFG
      language is proved correct. Then, this theory is leveraged to prove the
      functional correctness of the compiler. The first piece is fairly generic,
      and hence reusable.
 *)

(** This tutorial is composed of the following files:
    - Utils_tutorial.v     : Utilities
    - Fin.v                : Finite types as a categorical embedding
    - KTreeFin.v           : Subcategory of ktrees over finite types
    - Imp.v                : Imp language, syntax and semantics
    - Asm.v                : Asm language, syntax and semantics
    - AsmCombinators.v     : Linking theory for Asm
    - Imp2Asm.v            : Compiler from Imp to Asm
    - Imp2AsmCorrectness.v : Proof of correctness of the compiler
    - AsmOptimizations.v   : (Optional) optimization passes for the Asm language
    The intended entry point for reading is Imp.v.
 *)

(** We start by introducing a simplified variant of Software
    Foundations' [Imp] language.  The language's semantics is first expressed in
    terms of [itree]s.  The semantics of the program can then be obtained by
    interpreting the events contained in the trees.
*)

(* begin hide *)
From Coq Require Import
     Arith.PeanoNat
     Lists.List
     Strings.String
     Morphisms
     Setoid
     RelationClasses.

From ExtLib Require Import
     Data.String
     Structures.Monad
     Structures.Traversable
     Data.List.

From ITree Require Import
     ITree
     ITreeFacts
     Events.MapDefault
     Events.StateFacts.

Import Monads.
Import MonadNotation.
Local Open Scope monad_scope.
Local Open Scope string_scope.
(* end hide *)

(* ========================================================================== *)
(** ** Syntax *)

(** Imp manipulates a countable set of variables represented as [string]s: *)
Definition var : Set := string.

(** For simplicity, the language manipulates [nat]s as values. *)
Definition value : Type := nat.

(** Expressions are made of variables, constant literals, and arithmetic operations. *)
Inductive expr : Type :=
| Var (_ : var)
| Lit (_ : value)
| Plus  (_ _ : expr)
| Minus (_ _ : expr)
| Mult  (_ _ : expr).

(** The statements are straightforward. The [While] statement is the only
 potentially diverging one. *)

Inductive stmt : Type :=
| Assign (x : var) (e : expr)    (* x = e *)
| Seq    (a b : stmt)            (* a ; b *)
| If     (i : expr) (t e : stmt) (* if (i) then { t } else { e } *)
| While  (t : expr) (b : stmt)   (* while (t) { b } *)
| Skip                           (* ; *)
.

(* ========================================================================== *)
(** ** Notations *)

Module ImpNotations.

  (** A few notations for convenience.  *)
  Definition Var_coerce: string -> expr := Var.
  Definition Lit_coerce: nat -> expr := Lit.
  Coercion Var_coerce: string >-> expr.
  Coercion Lit_coerce: nat >-> expr.

  Declare Scope expr_scope.
  Bind Scope expr_scope with expr.

  Infix "+" := Plus : expr_scope.
  Infix "-" := Minus : expr_scope.
  Infix "*" := Mult : expr_scope.

  Declare Scope stmt_scope.
  Bind Scope stmt_scope with stmt.

  Notation "x '←' e" :=
    (Assign x e) (at level 60, e at level 50): stmt_scope.

  Notation "a ';;;' b" :=
    (Seq a b)
      (at level 100, right associativity,
       format
         "'[v' a  ';;;' '/' '[' b ']' ']'"
      ): stmt_scope.

  Notation "'IF' i 'THEN' t 'ELSE' e 'FI'" :=
    (If i t e)
      (at level 100,
       right associativity,
       format
         "'[v ' 'IF'  i '/' '[' 'THEN'  t  ']' '/' '[' 'ELSE'  e ']' 'FI' ']'").

  Notation "'WHILE' t 'DO' b" :=
    (While t b)
      (at level 100,
       right associativity,
       format
         "'[v  ' 'WHILE'  t  'DO' '/' '[v' b  ']' ']'").

End ImpNotations.

Import ImpNotations.

(* ========================================================================== *)
(** ** Semantics *)

(** _Imp_ produces effects by manipulating its variables.  To account for this,
    we define a type of _external interactions_ [ImpState] modeling reads and
    writes to global variables.

    A read, [GetVar], takes a variable as an argument and expects the
    environment to answer with a value, hence defining an event of type
    [ImpState value].

    Similarly, [SetVar] is a write event parameterized by both a variable and a
    value to be written, and defines an event of type [ImpState unit], no
    informative answer being expected from the environment.  *)
Variant ImpState : Type -> Type :=
| GetVar (x : var) : ImpState value
| SetVar (x : var) (v : value) : ImpState unit.

Section Denote.

  (** We now proceed to denote _Imp_ expressions and statements.
      We could simply fix in stone the universe of events to be considered,
      taking as a semantic domain for _Imp_ [itree ImpState X]. That would be
      sufficient to give meaning to _Imp_, but would prohibit us from relating this
      meaning to [itree]s stemmed from other entities. Therefore, we
      parameterize the denotation of _Imp_ by a larger universe of events [eff],
      of which [ImpState] is assumed to be a subevent. *)

  Context {eff : Type -> Type}.
  Context {HasImpState : ImpState -< eff}.

  (** _Imp_ expressions are denoted as [itree eff value], where the returned
      value in the tree is the value computed by the expression.
      In the [Var] case, the [trigger] operator smoothly lifts a single event to
      an [itree] by performing the corresponding [Vis] event and returning the
      environment's answer immediately.
      A constant (literal) is simply returned.
      Usual monadic notations are used in the other cases: we can [bind]
      recursive computations in the case of operators as one would expect. *)

  Fixpoint denote_expr (e : expr) : itree eff value :=
    match e with
    | Var v     => trigger (GetVar v)
    | Lit n     => ret n
    | Plus a b  => l <- denote_expr a ;; r <- denote_expr b ;; ret (l + r)
    | Minus a b => l <- denote_expr a ;; r <- denote_expr b ;; ret (l - r)
    | Mult a b  => l <- denote_expr a ;; r <- denote_expr b ;; ret (l * r)
    end.

  (** We turn to the denotation of statements. As opposed to expressions,
      statements do not return any value: their semantic domain is therefore
      [itree eff unit]. The most interesting construct is, naturally, [while].

      To define its meaning, we make use of the [iter] combinator provided by
      the [itree] library:

      [iter : (A -> itree E (A + B)) -> A -> itree E B].

      The combinator takes as argument the body of the loop, i.e. a function
      that maps inputs of type [A], the accumulator, to an [itree] computing
      either a new [A] that can be fed back to the loop, or a return value of
      type [B]. The combinator builds the fixpoint of the body, hiding away the
      [A] argument from the return type.

      Compared to the [mrec] and [rec] combinators introduced in
      [Introduction.v], [iter] is more restricted in that it naturally
      represents tail recursive functions.  It, however, enjoys a rich equational
      theory: its addition grants the type of _continuation trees_ (named
      [ktree]s in the library), a structure of _traced monoidal category_.

      We use [loop] to first build a new combinator [while].
      The accumulator degenerates to a single [unit] value indicating
      whether we entered the body of the while loop or not. Since the
      the operation does not return any value, the return type is also
      taken to be [unit].
      That is, the right tag [inr tt] says to exit the loop,
      while the [inl tt] says to continue. *)

  Definition while (step : itree eff (unit + unit)) : itree eff unit :=
    iter (C := Kleisli _) (fun _ => step) tt.

  (** Casting values into [bool]:  [0] corresponds to [false] and any nonzero
      value corresponds to [true].  *)
  Definition is_true (v : value) : bool := if (v =? 0)%nat then false else true.

  (** The meaning of Imp statements is now easy to define.  They are all
      straightforward, except for [While], which uses our new [while] combinator
      over the computation that evaluates the conditional, and then the body if
      the former was true.  *)
  Fixpoint denote_imp (s : stmt) : itree eff unit :=
    match s with
    | Assign x e =>  v <- denote_expr e ;; trigger (SetVar x v)
    | Seq a b    =>  denote_imp a ;; denote_imp b
    | If i t e   =>
      v <- denote_expr i ;;
      if is_true v then denote_imp t else denote_imp e

    | While t b =>
      while (v <- denote_expr t ;;
	           if is_true v
             then denote_imp b ;; ret (inl tt)
             else ret (inr tt))

    | Skip => ret tt
    end.

End Denote.

(* ========================================================================== *)
(** ** EXAMPLE: Factorial *)

Section Example_Fact.

  (** We briefly illustrate the language by writing the traditional factorial.
      example.  *)

  Open Scope expr_scope.
  Open Scope stmt_scope.
  Variable input: var.
  Variable output: var.

  Definition fact (n:nat): stmt :=
    input ← n;;;
    output ← 1;;;
    WHILE input
    DO output ← output * input;;;
       input  ← input - 1.

  (** We have given _a_ notion of denotation to [fact 6] via [denote_imp].
      However, this is naturally not actually runnable yet, since it contains
      uninterpreted [ImpState] events.  We therefore now need to _handle_ the
      events contained in the trees, i.e. give a concrete interpretation of the
      environment.  *)

End Example_Fact.

(* ========================================================================== *)
(** ** Interpretation *)

(* begin hide *)
From ITree Require Import
     Events.MapDefault.

From ExtLib Require Import
     Data.String
     Core.RelDec
     Structures.Maps
     Data.Map.FMapAList.

(** We provide an _ITree event handler_ to interpret away [ImpState] events.  We
    use an _environment event_ to do so, modeling the environment as a
    0-initialized map from strings to values.  Recall from [Introduction.v] that
    a _handler_ for the events [ImpState] is a function of type [forall R, ImpState R
    -> M R] for some monad [M].  Here we take for our monad the special case of
    [M = itree E] for some universe of events [E] required to contain the
    environment events [mapE] provided by the library. It comes with an event
    interpreter [interp_map] that yields a computation in the state monad.  *)
Definition handle_ImpState {E: Type -> Type} `{mapE var 0 -< E}: ImpState ~> itree E :=
  fun _ e =>
    match e with
    | GetVar x => lookup_def x
    | SetVar x v => insert x v
    end.

(** We now concretely implement this environment using ExtLib's finite maps. *)
Definition env := alist var value.

(** Finally, we can define an evaluator for our statements.
   To do so, we first denote them, leading to an [itree ImpState unit].
   We then [interp]ret [ImpState] into [mapE] using [handle_ImpState], leading to
   an [itree (mapE var value) unit].
   Finally, [interp_map] interprets the latter [itree] into the state monad,
   resulting in an [itree] free of any event, but returning the final
   _Imp_ env.
 *)
(* SAZ: would rather write something like the following:
 h : E ~> M A
 h' : F[void1] ~> M A
forall eff, {pf:E -< eff == F[E]} (t : itree eff A)
        interp pf h h' t : M A
*)

Definition interp_imp  {E A} (t : itree (ImpState +' E) A) : stateT env (itree E) A :=
  let t' := interp (bimap handle_ImpState (id_ E)) t in
  interp_map t'.

Definition eval_imp (s: stmt) : itree void1 (env * unit) :=
  interp_imp (denote_imp s) empty.

(** Equipped with this evaluator, we can now compute.
    Naturally since Coq is total, we cannot do it directly inside of it.
    We can either rely on extraction, or use some fuel.
 *)
Compute (burn 200 (eval_imp (fact "input" "output" 6))).

(* ========================================================================== *)
Section InterpImpProperties.
  (** We can lift the underlying equational theory on [itree]s to include new
      equations for working with [interp_imp].

      In particular, we have:
         - [interp_imp] respects [≈]
         - [interp_imp] commutes with [bind].

      We could justify more equations than just the ones below.  For instance,
      _Imp_ programs also respect a coarser notation of equivalence for the
      [env] state.  We exploit this possibility to implement optimzations
      at the _Asm_ level (see AsmOptimizations.v).
   *)

  Context {E': Type -> Type}.
  Notation E := (ImpState +' E').

  (** This interpreter is compatible with the equivalence-up-to-tau. *)
  Global Instance eutt_interp_imp {R}:
    Proper (@eutt E R R eq ==> eq ==> @eutt E' (prod (env) R) (prod _ R) eq)
           interp_imp.
  Proof.
    repeat intro.
    unfold interp_imp.
    unfold interp_map.
    rewrite H0. eapply eutt_interp_state_eq; auto.
    rewrite H. reflexivity.
  Qed.

  (** [interp_imp] commutes with [bind]. *)
  Lemma interp_imp_bind: forall {R S} (t: itree E R) (k: R -> itree E S) (g : env),
      (interp_imp (ITree.bind t k) g)
    ≅ (ITree.bind (interp_imp t g) (fun '(g',  x) => interp_imp (k x) g')).
  Proof.
    intros.
    unfold interp_imp.
    unfold interp_map.
    repeat rewrite interp_bind.
    repeat rewrite interp_state_bind.
    apply eqit_bind; [ reflexivity | ].
    red. intros.
    destruct a as [g'  x].
    simpl.
    reflexivity.
  Qed.

End InterpImpProperties.

(** We now turn to our target language, in file [Asm].v *)


From mathcomp Require Import ssreflect ssrfun seq choice.
From rlzrs Require Import all_rlzrs.
Require Import all_cs_base func classical_func classical_cont.
Require Import Morphisms FunctionalExtensionality ClassicalChoice.

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

Local Open Scope cs_scope.
Section discreteness.
  Definition discrete (X: cs) := forall (Y: cs) (f: X -> Y), f \is_continuous.

  Lemma dscrt_prpr: Proper (isomorphic ==> iff) discrete.
  Proof.
    move => X Y [[g ass] [[h ass'] [/=/sec_cncl cncl /sec_cncl cncl']]].
    split => dscrt Z f.
    - rewrite /continuous -(comp_id_r (F2MF f)) /mf_id -cncl' -comp_assoc.
      apply/comp_hcr; first by apply/ass_cont.
      by rewrite F2MF_comp_F2MF; apply/dscrt.
    rewrite /continuous -(comp_id_r (F2MF f)) /mf_id -cncl -comp_assoc.
    apply/comp_hcr; first by apply/ass_cont.
    by rewrite F2MF_comp_F2MF; apply/dscrt.
  Qed.
End discreteness.
Notation "X \is_discrete" := (discrete X) (at level 40).

Section cs_id.
  Context (S: Type).
  Definition id_rep:= make_mf (fun phi (s: S) => phi tt = s).

  Lemma id_rep_sur: id_rep \is_cototal.
  Proof. by move => s; exists (fun str => s). Qed.

  Lemma id_rep_sing: id_rep \is_singlevalued.
  Proof. by move => s t t' <- <-. Qed.

  Context (s: S) (S_count: S \is_countable).
  Definition cs_id_class:= @continuity_space.Class _ _ _
    (interview.Mixin id_rep_sur) (dictionary.Mixin id_rep_sing)
    (continuity_space.Mixin tt s unit_count S_count).

  Definition cs_id:= continuity_space.Pack cs_id_class.

  Lemma cs_id_dscrt:
    discrete cs_id.
  Proof.
    move => Y f.
    pose R:= make_mf (fun phi psi => psi \describes (f (phi tt)) \wrt Y).
    have Rtot: R \is_total by move => phi; apply/rep_sur.
    have [F icf]:= choice _ Rtot.
    exists (F2MF F); split; first by rewrite F2MF_rlzr_F2MF => fn n <-/=; apply/icf.
    rewrite cont_F2MF F2MF_cont_choice => phi q'; exists [::tt] => psi [eq _].
    by have ->: phi = psi by apply functional_extensionality => str; elim str.
  Qed.
End cs_id.

Lemma dscrt_id (X: cs) (x: X) (Xcount: X \is_countable):
  X \is_discrete -> X ~=~ (cs_id x Xcount).
Proof.
  move => dscrt.
  exists (exist_c (dscrt (cs_id x Xcount) id)).
  by exists (exist_c (@cs_id_dscrt X x Xcount X id)).
Qed.

Section TERMINAL.
Canonical cs_unit_class:= cs_id_class tt unit_count.
Canonical cs_unit:= cs_id tt unit_count.

Lemma unit_dscrt: discrete cs_unit.
Proof. exact/cs_id_dscrt. Qed.

Definition unit_fun (X: cs) (x: X) := tt: cs_unit.

Lemma trmnl_uprp_fun (X: cs): exists! f: X -> unit, True.
Proof.
by exists (@unit_fun X); split => // f _; apply functional_extensionality => x; elim (f x).
Qed.

Definition unit_fun_rlzr (X: cs): questions X ->> questions cs_unit
  := (F2MF (fun (_: names X) (_: queries cs_unit) => tt: answers cs_unit)).

Lemma unit_fun_rlzr_spec (X: cs) : (@unit_fun_rlzr X) \realizes (F2MF (@unit_fun X)).
Proof. by rewrite F2MF_rlzr_F2MF. Qed.

Lemma unit_fun_rlzr_cntop (X: cs): (@unit_fun_rlzr X) \is_continuous_operator.
Proof. by rewrite cont_F2MF; exists (fun _ => nil). Qed.

Lemma unit_fun_cont (X: cs): (@unit_fun X) \is_continuous.
Proof.
by exists (@unit_fun_rlzr X); split; [exact/unit_fun_rlzr_spec | exact/unit_fun_rlzr_cntop].
Qed.

Lemma unit_fun_hcr (X: cs): (F2MF (@unit_fun X): X ->> cs_unit) \has_continuous_realizer.
Proof. exact/unit_fun_cont. Qed.

Definition unit_fun_ass (X: cs) (Lq: seq (answers X) * queries cs_unit) :=
  inr tt : seq (queries X) + answers cs_unit.

Lemma unit_fun_ass_eval (X: cs): (U (@unit_fun_ass X)) \evaluates_to (@unit_fun_rlzr X).
Proof.
apply/mon_eval; first exact/U_mon; first exact/F2MF_sing.
by move => phi _ <- q'; exists 2; rewrite /U/=.
Qed.

Lemma unit_fun_ass_spec (X: cs): associate X cs_unit (@unit_fun_ass X) (@unit_fun X).
Proof. exact/ntrvw.tight_rlzr/unit_fun_ass_eval/unit_fun_rlzr_spec. Qed.

Lemma trmnl_uprp_cont (X: cs): exists! f: X c-> cs_unit, True.
Proof.
have cdom: (@unit_fun X) \from codom (associate X cs_unit) by exists (@unit_fun_ass X); exact/unit_fun_ass_spec.
exists (exist (fun p => p \from codom (associate X cs_unit)) _ cdom).
split => // f' _.
apply/eq_sub/functional_extensionality => x /=.
by case: (sval f' x); case: (unit_fun x).
Qed.
End TERMINAL.

Section BOOL.
  Canonical cs_bool_class:= cs_id_class true bool_count.
  Canonical cs_bool:= cs_id true bool_count.

  Lemma bool_dscrt: discrete cs_bool.
  Proof. exact/cs_id_dscrt. Qed.
End BOOL.

Section NATURALS.
  Canonical cs_nat_class := cs_id_class 0 nat_count.
  Canonical cs_nat := cs_id 0 nat_count.

  Lemma S_cont: (S: cs_nat -> cs_nat) \is_continuous.
  Proof.
    exists (F2MF (fun phi q =>S (phi q))).
    split; first by rewrite F2MF_rlzr => /= n phi -> [m]; by exists m.
    by rewrite cont_F2MF => phi; exists (fun _ => [:: tt]) => str psi []; elim: str => ->.
  Qed.

  Lemma nat_dscrt (X: cs) (f: cs_nat -> X): f \is_continuous.
  Proof. exact/cs_id_dscrt. Qed.

  Lemma nat_nat_cont (f: nat -> nat -> nat):
    (fun (p: cs_nat \*_cs cs_nat) => f p.1 p.2: cs_nat) \is_continuous.
  Proof.
    exists (F2MF (fun phi q => f (phi (inl tt)).1 (phi (inr tt)).2)).
    split; first by rewrite F2MF_rlzr_F2MF => phi [n m] [/= <- <-].
    by rewrite cont_F2MF => phi; exists (fun _ => [:: inl tt; inr tt]) => psi str [-> [->]].
  Qed.
End NATURALS.


Require Utf8.

Require Import S00_setoid_basics S01_Wty S02_PolyFun S10_Wstd_Obj
         S30_IdType S31_DiscStd S32_IdWty S33_PtwEq S34_Eqvr.

(* The setoid of extensional trees is isomorphic to
   the subsetoid of pointwise equal trees on the equivariant ones. *)

Section Extensional_as_Equivariant.
  Context {X : Set}{B : setoidfamily (dscrStd X)}(isdscrsf : isDscrStdFam B).

  (* Lemma 5.6(i) *)

  Lemma WPER2EqWty_aux : ∀ z, WstdPER_aux B z → EqWty (|B|) (fst z) (snd z).
  Proof.
    intros z r. induction r as [[w w'] p br IH]. simpl in p. simpl in br. simpl in IH. simpl.
    apply supd with (i := (w , w'))
                    (a := p).
    intros [[b b'] q]. simpl. simpl in *|-.
    set (l := projT1 (trspDscrSF isdscrsf p b)).
    set (qq := projT2 (trspDscrSF isdscrsf p b) ⊙ᵐ q).
    set (qqq := snd (isdscrsf _ (B • p b) b') (existT _ l qq)).
    set (H := existT _ (b , b') qqq : TotStdfamEq B p).
    apply (IH H).
  Qed.

  Lemma WPER2EqWty {w w' : Wty (|B|)}
    : WstdPER B w w' → EqWty (|B|) w w'.
  Proof.
    intro r. apply (WPER2EqWty_aux (w , w') r).
  Qed.

  (* Lemma 5.6(ii) *)

  Lemma ExtIsEqvr {w : Wty (|B|)}
    : WstdPER B w w → isEquivariant (|B|) w.
  Proof.
    induction w as [a f IH]. set (w := sup (|B|) a f).
    intro wext. set (u := (existT (λ x, WstdPER B x x) w wext) : Wstd B).
    refine (supd _ _ _ _ _).
    intros l b. cbn. cbn in *|-.
    apply WPER2EqWty.
    apply (WstdPER_ist wext). apply setoidfamilyrefgeneral_tactical.
    apply isdscrsf.
    exists (l ⁻¹). apply (transpIsEqv2 (|B|) l).
    intro b. specialize (IH b (ist u ↱ setoidrefl (B a) b)).
    apply IH.
  Qed.

  Definition jmap_fnc : Wstd B → EqvrTrees (|B|)
    := λ u, existT _ (projT1 u) (ExtIsEqvr (projT2 u)).

  Definition jmap : Wstd B ⇒ EqvrTrees (|B|).

    apply (Build_setoidmap (Wstd B) (EqvrTrees (|B|)) jmap_fnc).
    intros u u' r. apply (WPER2EqWty r).
  Defined.

  (* Lemma 5.7 *)

  Lemma eqvrEqWty2WPER {w : Wty (|B|)}(E : isEquivariant (|B|) w)
    : ∀ w', isEquivariant (|B|) w'
            → EqWty (|B|) w w' → WstdPER B w w'.
  Proof.
    induction E as [w eqvr br IH]. intros w' E' eq.
    set (p := nodeEqWty eq).
    apply (isMatch (A := dscrStd X) p).
    intros b b' q. cbn in *|-.
    specialize (IH b (Wty_ist w' b')).
    apply IH. apply (DWty_ist _ _ w' E' b').
    set (l := projT1 (trspDscrSF isdscrsf p b)).
    set (rr := Id2EqWty (|B|) (ap (f := Wty_ist w') (projT2 (trspDscrSF isdscrsf p b)) ⁻¹)).
    (*   : EqWty (|B|) (Wty_ist (|B|) w' (|B|•ᵐp b)) (Wty_ist (|B|) w' (|B|•ᵐl (B•p b)))*)
    set (l' := projT1 (fst (isdscrsf _ (B•p b) b') q)).
    set (rr' := Id2EqWty (|B|) (ap (f := Wty_ist w') (projT2 (fst (isdscrsf _ (B•p b) b') q)))).
  (*   : Wty_ist (|B|) w' (|B|•ᵐl' (B•p b)) == Wty_ist (|B|) w' b'*)
    cbn in *|-.
    refine (EqWty_trans _ _ _ (subtEqWty eq rfl) _).
    refine (EqWty_trans _ _ _ rr _).
    refine (EqWty_trans _ _ _ (eqvrBr E' l (B•p b)) _).
    apply (EqWty_trans _ _ _ (EqWty_symm _ _ (eqvrBr E' l' (B•p b))) rr').
  Qed.

  Definition jinv_fnc : EqvrTrees (|B|) → Wstd B
    := λ ew, existT (λ w, WstdPER B w w)
                    (projT1 ew)
                    (eqvrEqWty2WPER (projT2 ew) (projT1 ew) (projT2 ew)
                                    (setoidrefl (PtwTrees (|B|)) (projT1 ew))).

  Definition jinv : EqvrTrees (|B|) ⇒ Wstd B.

    apply (Build_setoidmap (EqvrTrees (|B|)) (Wstd B) jinv_fnc).
    intros ew ew' ptw. apply (eqvrEqWty2WPER (projT2 ew) (projT1 ew') (projT2 ew') ptw).
  Defined.

  (* Theorem 5.8 *)

  Theorem jiso : jmap ∘ jinv ≈ idmap ∧ jinv ∘ jmap ≈ idmap.
  Proof.
    split.
    intro ew. apply (setoidrefl (EqvrTrees (|B|)) ew).
    intro u. apply (setoidrefl (Wstd B) u).
  Qed.

End Extensional_as_Equivariant.


(* the only way to compute zero is for p to become canonical.
The only canonical member of A=A is eq_refl.
However, it is impossible to that p is propositionally equal to eq_refl.
In particular the univalence axiom allows for non-refl proofs. *)
Fixpoint zero (A : Set)  (p : A = A) {struct p} : nat := 0.

(* although this axiom breaks canonicity, it is believed to be consistent *)
Axiom strong_exm : Set -> nat.
Axiom strong_exm_true : strong_exm True = 0.
Axiom strong_exm_false : strong_exm False = 1.

(* same type as [zero] above, but provably non parametric *)
Definition nonParam (A : Set)  (p : A = A) : nat := strong_exm A.

(* because zero cannot be unfolded, it seems safe to assume the following *)
Axiom zeroOpaque :(forall x, zero x = nonParam x).

Inductive eq_R (A₁ A₂ : Type) (A_R : A₁ -> A₂ -> Type) (x₁ : A₁) (x₂ : A₂)
(x_R : A_R x₁ x₂)
  : forall (H : A₁) (H0 : A₂), A_R H H0 -> x₁ = H -> x₂ = H0 -> Prop :=
    eq_refl_R : eq_R A₁ A₂ A_R x₁ x₂ x_R x₁ x₂ x_R eq_refl eq_refl.

Lemma zero_not_parametric :
(forall (A₁ A₂ : Set) (A_R : A₁ -> A₂ -> Set) (p₁ : A₁ = A₁) (p₂ : A₂ = A₂),
eq_R Set Set (fun H1 H2 : Set => H1 -> H2 -> Set) A₁ A₂ A_R A₁ A₂ A_R p₁
  p₂ -> (zero A₁ p₁) = (zero A₂ p₂)) -> False.
Proof using.
  intros Hc.
  specialize (Hc True False (fun _ _ => True) eq_refl eq_refl).
  do 2 rewrite zeroOpaque in Hc.
  unfold nonParam in Hc. simpl in Hc.
  rewrite strong_exm_true in Hc.
  rewrite strong_exm_false in Hc.
  specialize (Hc (@eq_refl_R _ _ _ _ _ _)).
  discriminate.
Qed.


Require Export Iron.Language.SystemF2Effect.Type.Relation.KindT.
Require Export Iron.Language.SystemF2Effect.Type.Operator.LiftTT.
Require Export Iron.Language.SystemF2Effect.Type.Operator.SubstTT.Base.
Require Export Iron.Language.SystemF2Effect.Type.Operator.SubstTT.LiftTT.
Require Export Iron.Language.SystemF2Effect.Type.Operator.SubstTT.LowerTT.
Require Export Iron.Language.SystemF2Effect.Type.Operator.SubstTT.EquivT.

(********************************************************************)
(* Commuting substitutions. *)
Lemma substTT_substTT
 :  forall n m t1 t2 t3
 ,  substTT (n + m) t3 (substTT n t2 t1)
 =  substTT n (substTT (n + m) t3 t2)
              (substTT (1 + n + m) (liftTT 1 n t3) t1).
Proof.
 intros. gen n m t2 t3.
 induction t1; intros;
  try (solve [simpl; f_equal; rewritess; norm]).

 Case "TVar".
  repeat (simpl; split_match; repeat norm_nat_compare);
   first [omega | burn].

 Case "TForall".
  simpl.
  rewrite (IHt1 (S n) m).
  rewrite (liftTT_substTT_1 0 (n + m)).
  rewrite (liftTT_liftTT_11 0 n).
  burn.
Qed.

(********************************************************************)
(* Substitution of types in types preserves kinding.
   Must also subst new new type into types in env higher than ix
   otherwise indices that reference subst type are broken, and
   the resulting type env would not be well formed *)
Theorem subst_type_type_ix
 :  forall ix ke sp t1 k1 t2 k2
 ,  get ix ke = Some k2
 -> KindT ke sp t1 k1
 -> KindT (delete ix ke) sp t2 k2
 -> KindT (delete ix ke) sp (substTT ix t2 t1) k1.
Proof.
 intros. gen ix ke sp t2 k1 k2.
 induction t1; intros; simpl; inverts_kind; eauto.

 Case "TVar".
  fbreak_nat_compare.
  SCase "n = ix".
   rewrite H in H5.
   inverts H5. auto.

  SCase "n < ix".
   apply KiVar. rewrite <- H5.
   apply get_delete_above; auto.

  SCase "n > ix".
   apply KiVar. rewrite <- H5.
   destruct n.
    burn.
    simpl. nnat. apply get_delete_below. omega.

 Case "TForall".
  apply KiForall.
  rewrite delete_rewind.
  eapply IHt1; eauto.
   apply kind_kienv_weaken; auto.

 Case "TCon2".
  eapply KiCon2.
  destruct tc. simpl in *. inverts H4.
  destruct t.  simpl in *. eauto.
  eauto.
  destruct tc. simpl in *. inverts H4.
  eauto.
Qed.

Theorem subst_type_type
 :  forall ke sp t1 k1 t2 k2
 ,  KindT (ke :> k2) sp t1 k1
 -> KindT ke         sp t2 k2
 -> KindT ke sp (substTT 0 t2 t1) k1.
Proof.
 intros.
 unfold substTT.
 rrwrite (ke = delete 0 (ke :> k2)).
 eapply subst_type_type_ix; burn.
Qed.

(* If two types are equivalent, and we substitute some third type
   into both, then the result is also equivalent. *)
Lemma substTT_EquivT
 :  forall  ke sp t1 t2 t3 k3 k
 ,  KindT    ke sp t3 k3
 -> EquivT (ke :> k3) sp t1 t2 k
 -> EquivT  ke sp (substTT 0 t3 t1) (substTT 0 t3 t2) k.
Proof.
 intros ke sp t1 t2 t3 k3 k HK HE.
 have HK1: (KindT (ke :> k3) sp t1 k).
 have HK2: (KindT (ke :> k3) sp t2 k).
 remember (ke :> k3) as X.
 induction HE; subst.

 - Case "EqRefl".
   eapply EqRefl;
    eauto using subst_type_type.

 - Case "EqSym".
   eapply EqSym;
    eauto using subst_type_type.

 - Case "EqTrans".
   eapply EqTrans with (t2 := substTT 0 t3 t2); eauto.

 - Case "EqSumCong".
   eapply EqSumCong; fold substTT. auto.
   eauto using subst_type_type.
   eauto using subst_type_type.

 - Case "EqSumBot".
   apply EqSumBot; fold substTT. auto.
   eauto using subst_type_type.

 - Case "EqSumIdemp".
   eapply EqSumIdemp. auto.
   eauto using subst_type_type.

 - Case "EqSumComm".
   eapply EqSumComm; fold substTT. auto.
   eauto using subst_type_type.
   eauto using subst_type_type.

 - Case "EqSumAssoc".
   eapply EqSumAssoc; fold substTT. auto.
   eauto using subst_type_type.
   eauto using subst_type_type.
   eauto using subst_type_type.
Qed.


Require Import Coq.Program.Syntax.
Require Export ZArith.
Require Import Coq.Program.Basics.
Require Import SetoidTactics.
Require Import SetoidClass.
Require Import Coq.Strings.String.
Require Import Coq.Strings.Ascii.
Require Import Coq.Lists.List.

Require Import HDI.Util.
Require Import HDI.Syntax.
Require Import HDI.Heap.
Require Import HDI.OpSem.
Require Import HDI.Bisimulation.

Local Open Scope stmt.

CoInductive safe (h: heap) : stmt -> Prop :=
  | safe_ret: forall v,
    safe h (ret v)
  | safe_exists_step: forall s,
    (exists h' s',
      step h s h' s' /\ safe h' s') ->
    safe h s.
Hint Constructors safe : safe.

  Lemma safe_bisim: forall h s1 s2,
    bisimilar s1 s2 ->
    safe h s1 ->
    safe h s2.
  Proof.
    cofix; intros.
    inversion H0; subst; clear H0.
    * apply bisim_ret1 in H; subst; auto with safe.
    * destruct H1 as [h'[s'[Hstep Hsafe']]].
    specialize (safe_bisim h').
    constructor.
    edestruct bisim_step1 as [s2'[??]]; eauto.
  Qed.

  Lemma safe_heap_equiv: forall h1 h2 s,
    heap_equiv h1 h2 ->
    safe h1 s ->
    safe h2 s.
  Proof.
    cofix; intros.
    inversion H0; subst; clear H0.
    * constructor.
    * destruct H1 as [h'[s'[Hstep Hsafe']]].
    specialize (safe_heap_equiv h').
    edestruct step_heap_equiv1 as [h2'[??]]; eauto.
    constructor; eauto.
  Qed.

  Add Parametric Morphism : safe
    with signature equiv ==> bisimilar ==> iff
    as safe_m.
  Proof.
    split; intros;
      [ | symmetry in H, H0];
      eauto using safe_bisim, safe_heap_equiv.
  Qed.

  Lemma safe_bind_ret: forall h v k,
    safe h (k v) ->
    safe h (x <- ret v; k x).
  Proof.
    constructor.
    exists h (k v); split; auto.
    apply s_bind_ret.
  Qed.
  Hint Resolve safe_bind_ret : safe.

  Lemma safe_read: forall h a v k,
    safe h (k v) ->
    h a = Some v ->
    safe h (x <- read a; k x).
  Proof.
    constructor.
    exists h (ret v >>= k); split;
      auto with safe.
  Qed.
  Hint Resolve safe_read : safe.

  Lemma safe_write: forall h a v v0 k,
    safe (upd_heap h a v) (ret 1 >>= k) ->
    h a = Some v0 ->
    safe h (write a v >>= k).
  Proof.
    constructor.
    exists (upd_heap h a v) (ret 1 >>= k); split;
      eauto with safe.
  Qed.
  Hint Resolve safe_write : safe.

  Lemma safe_step: forall h s h' s',
    safe h s ->
    step h s h' s' ->
    safe h' s'.
  Proof.
    intros.
    inversion H; subst.
    * inversion H0.
    * destruct H1 as [h''[s''[??]]].
    edestruct step_deterministic
      with (h1':=h'') (h2':=h') as [??];
      eauto; subst h'' s''.
    assumption.
  Qed.
  Hint Resolve safe_step : safe.

  Lemma safe_bind_step: forall h s k h' s',
    step h s h' s' ->
    safe h' (s' >>= k) ->
    safe h (s >>= k).
  Proof.
    constructor.
    exists h' (s' >>= k); split; auto with safe.
  Qed.
  Hint Resolve safe_bind_step : safe.

  Lemma unsafe_abort: forall h,
    safe h abort ->
    False.
  Proof.
    intros.
    inversion H; subst.
    destruct H0 as [?[?[Hstep ?]]].
    inversion Hstep.
  Qed.
  Hint Resolve unsafe_abort : safe.


(******************************************************************************)
(*        A program that performs a perfect play for 4 in a row               *)
(*        This is directly inspired by a program by John Tromp                *)
(******************************************************************************)

Require Import ssreflect ZArith Ascii List String Int63.

Open Scope int63_scope.

Require Import PArray.

(*
Parameter array : Type -> Type.
Parameter set : forall {T : Type}, array T -> int -> T -> array T.
Parameter get : forall {T : Type}, array T -> int -> T.
Parameter foldi : forall {A : Type}, (int -> A -> A) -> int -> int -> A -> A.
*)

(* Naive implementation of foldi *)
Fixpoint nfoldi {A : Type} (f : int -> A -> A) n v r :=
  if n is (S n1) then nfoldi f n1 (v - 1) (f v r) else f v r.
Definition foldi {A : Type} (f : int -> A -> A) v1 v2 r :=
  if v1 <=? v2 then nfoldi f (Z.to_nat (to_Z (v2 - v1))) v2 r
  else r.

(*
Notation " a '.[' b ']' " := (get a b) (at level 32,
  format " a '.[' b ] ").
Notation " a '.[' b '<-' c ']' " := (set a b c) (at level 32,
  format " a '.[' b  '<-'  c ] ").
*)

(*
Module PArray.

Parameter make : forall {T : Type}, int -> T -> array T.

End PArray.
*)

Fixpoint init_matrix (A : Type) n nn a (v : A) m {struct n} :=
  match n with
  | O => a
  | S n1 => init_matrix A n1 (nn - 1) a.[nn - 1 <- make m v] v m
  end.
Arguments init_matrix[A].

Definition to_nat n := Z.to_nat (to_Z n).

Definition make_matrix (A: Type) n m (v : A) :=
  let a := make n (make 0 v) in
  init_matrix (to_nat n) n a v m.

Arguments make_matrix[A].

(* 63 Arithmetic *)

Definition one := of_Z 1.
Definition zero := of_Z 0.

Definition decr s := s - 1.
Definition incr s := s + 1.
Definition is_zero s := eqb s zero.
Definition is_nzero s := negb (eqb s zero).

(* Width of the board *)
Definition width := 7.
Definition nwidth := 7%nat.
(* Height of the board *)
Definition height := 6.
Definition nheight := 6%nat.
(* Shift for moving horizontally *)
Definition horizontal := Eval compute in height + 1.
Definition horizontal2 := Eval compute in 2 * horizontal.
(* Shift for moving vertically *)
Definition vertical := 1.
Definition vertical2 := Eval compute in 2 * vertical.
(* Shift for moving up right *)
Definition up_right := Eval compute in horizontal + 1.
Definition up_right2 := Eval compute in 2 * up_right.
(* Shift for moving up left *)
Definition up_left := Eval compute in horizontal - 1.
Definition up_left2 := Eval compute in 2 * up_left.

Register horizontal as Inline.
Register horizontal2 as Inline.
Register vertical as Inline.
Register vertical2 as Inline.
Register up_right as Inline.
Register up_right2 as Inline.
Register up_left as Inline.
Register up_left2 as Inline.

(* Size of the board *)
Definition size := width * height.
(* Real size of the board *)
Definition number_of_cells := width * horizontal.
(* All the cells are set to one *)
Definition all_set := decr (one << number_of_cells).
(* Cells of the first column *)
Definition first_column := decr (one << height).
(* Cells of the first column + border *)
Definition full_first_column := decr (one << horizontal).
(* Cells at the bottom of the board *)
Definition bottom := div all_set (decr (one << horizontal)).

(* Cells at the top *)
Definition top := bottom << height.
(* Unknown valuation *)
Definition unknown := 0.
(* Loss valuation *)
Definition loss := 1.
(* Loss or draw valuation *)
Definition lossdraw := 2.
(* Draw valuation *)
Definition draw := 3.
(* Draw or win valuation *)
Definition drawwin := 4.
(* Win valuation *)
Definition win := 5.
(* Loss + win valuation (to reverse) *)
Definition losswin := loss + win.
(* Size of the lock 25 bits *)
Definition locksize := 25.
(* Size of the remain bits that are not in the lock *)
Definition slocksize := number_of_cells - locksize.
(* Size of hash table should have HPRIME > 2^SLOCKSIZE *)
(* Definition hprime := 8388609. *)
Definition hprime := 16777121.
(* Mask for the lock  *)
Definition lockmask := (1 << locksize) - 1.
(* Size of the score 3 *)
Definition scoresize := 3.
(* Mask for the size *)
Definition scoremask := (1 << scoresize) - 1.
(* Size of the lock + score *)
Definition scorelocksize := locksize + scoresize.
(* Size of the lock + score *)
Definition scorelockmask := (1 << scorelocksize) - 1.
(* Log of the number of hash tables *)
Definition lhash := 3.
(* Number of hash tables *)
Definition nhash := 1 << lhash.
(* Mask for hash tables *)
Definition mhash := nhash - 1.
(* Symmetry level *)
Definition sym_level := 10.

(* Hash table because of size limitation we create a matrix *)
Definition make_hash (u : unit) :=
  make_matrix nhash (2 * (hprime/nhash) + 1) 0.

Definition min m n := match m ?= n with Lt => m | _ => n end.
Definition max m n := match m ?= n with Lt => n | _ => m end.

(* Score of the different cells *)
Definition values :=
  let t := PArray.make number_of_cells 0 in
  foldi (fun i t =>
    foldi (fun j t =>
      let v := 3 + min 3 i
                 + min 3 (min i j) + min 3 j
                 + min 3 (min ((width + 1) / 2 - i) j) in
      let t := (t.[horizontal * i + j <- v]) in
      let t := (t.[horizontal * (width - 1 - i) + height - 1 - j <- v]) in
      let t := (t.[horizontal * i + height - 1 - j <- v]) in
      (t.[horizontal * (width - 1 - i) + j <- v]))
      0 ((horizontal / 2) - 1) t)
    0 ((width + 1) / 2 - 1) t.

Definition logand2 s dir dir2 :=
  let s1 := s land (s >> dir) in
  let s2 := s1 land (s1 >> dir2) in is_zero s2.

Register logand2 as Inline.

(* Check if the position is won *)
Definition is_won state :=
   (* horizontal win *)
  if logand2 state horizontal horizontal2 then
  if logand2 state vertical vertical2 then
  if logand2 state up_right up_right2 then
  if logand2 state up_left up_left2 then false
  else true
  else true
  else true
  else true.

(* Get the border *)
Definition get_border (wstate bstate : int) :=
  bottom + (wstate lor bstate).

(* Perform a move *)
Definition  make_move move state := move lor state.

(* Get the log 2 of a number *)
Definition get_log2 (v : int) : int :=
   62 - head0 v.

(* List of possible moves, no move = draw *)
Inductive moves := EmptyMove | Move (m : int) (v : int) (l : moves).

(* Moves are ordered by their values *)
Fixpoint insert_fmove m (v : int) l :=
match l with
| EmptyMove => Move m v EmptyMove
| Move m1 v1 l1 =>
  match v ?= v1 with
  |Lt => Move m1 v1 (insert_fmove m v l1)
  | _ => Move m v l
  end
end.

Inductive fmove :=
 | Win
 | Draw
 | Forced (_ : int)
 | Moves (_: moves).

Definition make_moves l :=
  match l with EmptyMove => Draw | _ => Moves l end.

Section FindMoves.

Variables (wstate bstate border: int).

Fixpoint make_colums i column :=
  match i with
      O => nil
  | S i => column :: (make_colums i (column << horizontal))
  end.

Definition columns := Eval compute in make_colums nwidth first_column.

(* Check for a direct win after a threat *)
Fixpoint fmt columns res :=
  match columns with
  | nil => res
  | column :: columns =>
      let move := border land column in
      if is_zero move then fmt columns res
      else
      if is_won (make_move move wstate) then Win
      else fmt columns res
  end.

Fixpoint fms columns res :=
  match columns with
  | nil => make_moves res
  | column :: columns =>
      let move := border land column in
      if is_zero move then fms columns res
      else
      if is_won (make_move move wstate) then Win
      else
      if is_won (make_move move bstate) then
        fmt columns (Forced move)
      else
        let v := (values.[get_log2 move]) in
        fms columns (insert_fmove move v res)
   end.

Lemma fmsE columns res :
fms columns res =
  match columns with
  | nil => make_moves res
  | column :: columns =>
      let move := border land column in
      if is_zero move then fms columns res
      else
      if is_won (make_move move wstate) then Win
      else
      if is_won (make_move move bstate) then
        fmt columns (Forced move)
      else
        let v := (values.[get_log2 move]) in
        fms columns (insert_fmove move v res)
   end.
Proof.
by case: columns.
Qed.

End FindMoves.

(* Find possible moves *)
Definition find_moves wstate bstate :=
  let border := get_border wstate bstate in
  fms wstate bstate border columns EmptyMove.

(* Auxillary parsing function from string to states *)
Fixpoint parsei s i j wstate bstate (turn : bool) :=
  match
    match width ?= j with
    |Eq =>
       match i ?= 0 with Eq => None | _ => Some (i-1,0) end
    | _ => Some (i,j) end
  with None => (wstate,bstate,turn)
  | Some (i,j) =>
    match s with
    | EmptyString => (wstate,bstate,turn)
    | String "X"%char s1 =>
       let move := one << (j * horizontal + i) in
       parsei s1 i (j + 1) (make_move wstate move) bstate (negb turn)
    | String "O"%char s1 =>
       let move := one << (j * horizontal + i) in
       parsei s1 i (j + 1) wstate (make_move bstate move) (negb turn)
    | String "_"%char s1 =>
       parsei s1 i (j + 1) wstate bstate turn
    | String _ s1 => parsei s1 i j wstate bstate turn
    end
  end.

(* Parsing function from string to states *)
Definition parse_string s :=
  parsei s height width zero zero true.

(* Newline String *)
Definition nl := String "013" EmptyString.

(* Auxillary function that turns states into into a string *)
Fixpoint to_stringi m i j wstate bstate :=
  match m with O => ""%string | (S m1) =>
  match
    match width ?= j with
    |Eq =>
       match i ?= 0 with Eq => None | _ => Some (i-1,0,nl) end
    | _ => Some (i,j,""%string) end
  with
  | None => nl
  | Some (i,j,ts) =>
    (ts ++
   (let move := one << (j * horizontal + i) in
    if is_nzero (move land wstate) then "X"%string ++ (to_stringi m1 i (j + 1) wstate bstate)
    else if is_nzero (move land bstate) then "O"%string ++ (to_stringi m1 i (j + 1) wstate bstate)
    else "_"%string ++ (to_stringi m1 i (j + 1) wstate bstate)))%string
  end
  end.

(* Turn states into a string *)
Definition to_string wstate bstate :=
 (to_stringi (nheight * nwidth) height width wstate bstate)%string.

(* Turn the score in a string *)
Definition string_of_score (score : int) :=
  if eqb score unknown then "UNKNOWN"%string else
  if eqb score loss then "LOSS"%string else
  if eqb score draw then "DRAW"%string else
  if eqb score win then "WIN"%string else
  if eqb score drawwin then "DRAWWIN"%string else
  if eqb score lossdraw then "LOSSDRAW"%string else
  "????"%string.

(* Reverse the valuation *)
Definition rev_val value := losswin - value.

Fixpoint sym_code i sres res :=
  match i with
  | O => sres
  | S i =>
      let sres :=  (sres << horizontal) lor
                       (res land full_first_column) in
      let res := res >> horizontal in
      sym_code i sres res
  end.

(* Get the unique code of a position *)
Definition get_code wstate bstate turn height :=
  let res := (match turn with true => wstate | false => bstate end) lor
        (get_border wstate bstate) in
  if height <=? sym_level
  then
     let sres := sym_code nwidth zero res in
     min sres res
  else res.

(* Put an element in the hash-table
    The layout of the two-entry hash-table
      at key : high bits = work first entry, low bits = lock first entry
      at key + 1 : high bits = score first entry then score second entry
                   low bits = lock second entry
 *)

Definition hput wstate bstate turn work score hash_table height :=
   let code := get_code wstate bstate turn height in
   let fkey := code mod hprime in
   let key := 2 * (fkey >> lhash) in
   let r :=  fkey land mhash in
   let lock := (code >> slocksize) in
   let ht := (hash_table.[r]) in
   let val1 := (ht.[key]) in
   let val2 := (ht.[key + 1]) in
   if orb ((val1 land lockmask) =? lock) ((val1 >> locksize) <=? work) then
       let ht := (ht.[key <- (work << locksize) lor lock]) in
       let ht := (ht.[key + 1 <-
                   ((score << scorelocksize) lor (val2 land scorelockmask))]) in
        (hash_table.[r <- ht])
   else
      let ht := (ht.[key + 1 <-
        ((((val2 >> scorelocksize) << scoresize) lor score) << locksize)
              lor lock]) in
        (hash_table.[r <- ht]).

(* Get an element in the hash-table *)
Definition hget (wstate bstate : int) (turn : bool)
         (hash_table : array (array int)) height :=
   let code := get_code wstate bstate turn height in
   let fkey := code mod hprime in
   let key := 2 * (fkey >> lhash) in
   let r :=  fkey land mhash in
   let lock := (code >> slocksize) in
   let ht := (hash_table.[r]) in
   let val1 := (ht.[key]) in
   let val2 := (ht.[key + 1]) in
   if ((val1 land lockmask) =? lock) then
       val2 >> scorelocksize
   else if ((val2 land lockmask) =? lock) then
       (val2 >> locksize) land scoremask
   else unknown.

Definition is_nempty_move m :=
  match m with EmptyMove => false | Move _ _ _ => true end.

(* Process result *)
Inductive pres := PRes (s : int) (v : int) (t : array (array int)).

Section Process.

Variables (wstate bstate : int) (turn : bool) (beta : int) (lvisited : int)
          (height hscore :  int)
          (alpha_beta : int -> int -> bool -> int -> int -> int ->
                         array (array int) -> pres).
Fixpoint process ms alpha score visited hash_table :=
  match ms with
  | EmptyMove =>
      let score := if (score =? losswin - hscore) then draw else score in
      let work := get_log2 (sub visited lvisited) in
      let hash_table := hput wstate bstate turn work score hash_table height in
      PRes score (incr visited) hash_table
  | Move move _ ms1 =>
    let (nscore,visited,hash_table) :=
      alpha_beta bstate (make_move move wstate) (negb turn)
           (rev_val beta) (rev_val alpha) visited hash_table in
    let nscore := rev_val nscore in
    if nscore <=? score then process ms1 alpha score visited hash_table
    else
    let score := nscore in
    if score <=? alpha then process ms1 alpha score visited hash_table
    else
    let alpha := score in
    if alpha <? beta  then process ms1 alpha score visited hash_table
    else
      let score :=
        if (andb (score =? draw) (is_nempty_move ms1)) then drawwin
        else score in
      let score := if (score =? losswin - hscore) then draw else score in
      let work := get_log2 (sub visited  lvisited) in
      let hash_table := hput wstate bstate turn work score hash_table height in
      PRes score (incr visited) hash_table
    end.

End Process.

(* alpha-beta result *)
Inductive ares := ARes (a : int) (b : int) (c : bool).

Section Alpha.

(* alpha beta pruning search *)
Fixpoint alpha_beta nstruct height wstate bstate turn alpha beta
                    visited hash_table :=
  let hscore := hget wstate bstate turn hash_table height in
  let (alpha,beta,flag) :=
    (if (hscore =? unknown) then ARes alpha beta false else
    if negb ((hscore land 1) =? 0) then ARes alpha beta true else
    if (hscore =? drawwin) then
      if (beta =? draw) then ARes alpha beta true else ARes draw beta false
    else
      if (alpha =? draw) then ARes alpha beta true
      else ARes alpha draw false) in
  if flag then PRes hscore visited hash_table else
  match find_moves wstate bstate with
  | Win => PRes win visited hash_table
  | Draw => PRes draw visited hash_table
  | Forced move =>
      match nstruct with
     | 0%nat => PRes unknown visited hash_table
     | S nstruct =>
      let (score,visited,hash_table) :=
        alpha_beta nstruct (height + 1) bstate (make_move move wstate)
                   (negb turn) (rev_val beta) (rev_val alpha)
                    visited hash_table : pres in
      PRes (rev_val score) visited hash_table
     end
  | Moves ms =>
     match nstruct with
     | 0%nat => PRes unknown visited hash_table
     | S nstruct =>
     process wstate bstate turn beta visited height hscore
            (alpha_beta nstruct (height + 1))
             ms alpha loss visited hash_table
     end
  end.

Lemma alpha_betaE nstruct height wstate bstate turn alpha beta
                    visited hash_table  :
 alpha_beta nstruct height wstate bstate turn alpha beta
                    visited hash_table =
  let hscore := hget wstate bstate turn hash_table height in
  let (alpha,beta,flag) :=
    (if (hscore =? unknown) then ARes alpha beta false else
    if negb ((hscore land 1) =? 0) then ARes alpha beta true else
    if (hscore =? drawwin) then
      if (beta =? draw) then ARes alpha beta true else ARes draw beta false
    else
      if (alpha =? draw) then ARes alpha beta true
      else ARes alpha draw false) in
  if flag then PRes hscore visited hash_table else
  match find_moves wstate bstate with
  | Win => PRes win visited hash_table
  | Draw => PRes draw visited hash_table
  | Forced move =>
      match nstruct with
     | 0%nat => PRes unknown visited hash_table
     | S nstruct =>
      let (score,visited,hash_table) :=
        alpha_beta nstruct (height + 1) bstate (make_move move wstate)
                   (negb turn) (rev_val beta) (rev_val alpha)
                    visited hash_table : pres in
      PRes (rev_val score) visited hash_table
     end
  | Moves ms =>
     match nstruct with
     | 0%nat => PRes unknown visited hash_table
     | S nstruct =>
     process wstate bstate turn beta visited height hscore
            (alpha_beta nstruct (height + 1))
             ms alpha loss visited hash_table
     end
  end.
Proof. by case: nstruct. Qed.

Definition eval_position s :=
   match parse_string s with
   (wstate,bstate,turn) =>
   let (wstate,bstate) := if turn then (wstate,bstate) else (bstate,wstate) in
   let (score, _, _) :=
     alpha_beta (1 + nheight * nwidth)%nat 0 wstate bstate turn loss win zero
                (make_hash tt) in
   score
   end.

End Alpha.

Definition ex1 := (
                 "___O___"
              ++ "___X___"
              ++ "___O___"
              ++ "___X___"
              ++ "__OO___"
              ++ "__XX___")%string.

Definition ex2 := (
                 "___X___"
              ++ "__OX___"
              ++ "__XO___"
              ++ "__OX___"
              ++ "__XO___"
              ++ "__OX__O")%string.

Definition ex3 := (
                 "___O___"
              ++ "___X___"
              ++ "___O___"
              ++ "___X___"
              ++ "___O___"
              ++ "XO_X___")%string.

Definition ex4 := ("______" ++ "______" ++ "______" ++
                   "______" ++ "______" ++ "______" ++ "______")%string.

Time Eval native_compute in string_of_score (eval_position ex1).
Time Eval native_compute in string_of_score (eval_position ex2).
Time Eval native_compute in string_of_score (eval_position ex3).
Time Eval native_compute in string_of_score (eval_position ex4).


Load LFindLoad.
From lfind Require Import LFind.
From QuickChick Require Import QuickChick.
From adtind Require Import goal33.

Derive Show for natural.

              Derive Arbitrary for natural.

              Instance Dec_Eq_natural : Dec_Eq natural.

              Proof. dec_eq. Qed.

Lemma conj25eqsynthconj2 : forall  (lv0 : natural) (lv1 : natural), (@eq natural (plus lv0 lv1) (plus Zero (plus lv1 lv0))).
Admitted.
QuickChick conj25eqsynthconj2.