(** * Typechecking: A Typechecker for STLC *)

(** The [has_type] relation of the STLC defines what it means for a
    term to belong to a type (in some context).  But it doesn't, by
    itself, give us an algorithm for _checking_ whether or not a term
    is well typed.

    Fortunately, the rules defining [has_type] are _syntax directed_
    -- that is, for every syntactic form of the language, there is
    just one rule that can be used to give a type to terms of that
    form.  This makes it straightforward to translate the typing rules
    into clauses of a typechecking _function_ that takes a term and a
    context and either returns the term's type or else signals that
    the term is not typable.  *)

(** This short chapter constructs such a function and proves it
    correct. *)

Set Warnings "-notation-overridden,-parsing".
From Coq Require Import Bool.Bool.
From PLF Require Import Maps.
From PLF Require Import Smallstep.
From PLF Require Import Stlc.
From PLF Require MoreStlc.

Module STLCTypes.
Export STLC.

(* ################################################################# *)
(** * Comparing Types *)

(** First, we need a function to compare two types for equality... *)

Locate "Bool".

Fixpoint eqb_ty (T1 T2:ty) : bool :=
  match T1,T2 with
  | <{ Bool }> , <{ Bool }> =>
      true
  | <{ T11->T12 }>, <{ T21->T22 }> =>
      andb (eqb_ty T11 T21) (eqb_ty T12 T22)
  | _,_ =>
      false
  end.

(** ... and we need to establish the usual two-way connection between
    the boolean result returned by [eqb_ty] and the logical
    proposition that its inputs are equal. *)

Lemma eqb_ty_refl : forall T,
  eqb_ty T T = true.
Proof.
  intros T. induction T; simpl.
    reflexivity.
    rewrite IHT1. rewrite IHT2. reflexivity.  Qed.

Lemma eqb_ty__eq : forall T1 T2,
  eqb_ty T1 T2 = true -> T1 = T2.
Proof with auto.
  intros T1. induction T1; intros T2 Hbeq; destruct T2; inversion Hbeq.
  - (* T1=Bool *)
    reflexivity.
  - (* T1 = T1_1->T1_2 *)
    rewrite andb_true_iff in H0. inversion H0 as [Hbeq1 Hbeq2].
    apply IHT1_1 in Hbeq1. apply IHT1_2 in Hbeq2. subst...  Qed.
End STLCTypes.

(* ################################################################# *)
(** * The Typechecker *)

(** The typechecker works by walking over the structure of the given
    term, returning either [Some T] or [None].  Each time we make a
    recursive call to find out the types of the subterms, we need to
    pattern-match on the results to make sure that they are not
    [None].  Also, in the [app] case, we use pattern matching to
    extract the left- and right-hand sides of the function's arrow
    type (and fail if the type of the function is not [T11->T12]
    for some [T11] and [T12]). *)

Module FirstTry.
Import STLCTypes.

Fixpoint type_check (Gamma : context) (t : tm) : option ty :=
  match t with
  | tm_var x =>
      Gamma x
  | <{\x:T2, t1}> =>
      match type_check (x |-> T2 ; Gamma) t1 with
      | Some T1 => Some <{T2->T1}>
      | _ => None
      end
  | <{t1 t2}> =>
      match type_check Gamma t1, type_check Gamma t2 with
      | Some <{T11->T12}>, Some T2 =>
          if eqb_ty T11 T2 then Some T12 else None
      | _,_ => None
      end
  | <{true}> =>
      Some <{Bool}>
  | <{false}> =>
      Some <{Bool}>
  | <{if guard then t else f}> =>
      match type_check Gamma guard with
      | Some <{Bool}> =>
          match type_check Gamma t, type_check Gamma f with
          | Some T1, Some T2 =>
              if eqb_ty T1 T2 then Some T1 else None
          | _,_ => None
          end
      | _ => None
      end
  end.

End FirstTry.

(* ################################################################# *)
(** * Digression: Improving the Notation *)

(** Before we consider the properties of this algorithm, let's write
    it out again in a cleaner way, using "monadic" notations in the
    style of Haskell to streamline the plumbing of options.  First, we
    define a notation for composing two potentially failing (i.e.,
    option-returning) computations: *)

Notation " x <- e1 ;; e2" := (match e1 with
                              | Some x => e2
                              | None => None
                              end)
         (right associativity, at level 60).

(** Second, we define [return] and [fail] as synonyms for [Some] and
    [None]: *)

Notation " 'return' e "
  := (Some e) (at level 60).

Notation " 'fail' "
  := None.

Module STLCChecker.
Import STLCTypes.

(** Now we can write the same type-checking function in a more
    imperative-looking style using these notations. *)

Fixpoint type_check (Gamma : context) (t : tm) : option ty :=
  match t with
  | tm_var x =>
    T <- Gamma x ;; return T
  | <{\x:T2, t1}> =>
      T1 <- type_check (x |-> T2 ; Gamma) t1 ;;
      return <{T2->T1}>
  | <{t1 t2}> =>
      T1 <- type_check Gamma t1 ;;
      T2 <- type_check Gamma t2 ;;
      match T1 with
      | <{T11->T12}> =>
          if eqb_ty T11 T2 then return T12 else fail
      | _ => fail
      end
  | <{true}> =>
      return <{ Bool }>
  | <{false}> =>
      return <{ Bool }>
  | <{if guard then t1 else t2}> =>
      Tguard <- type_check Gamma guard ;;
      T1 <- type_check Gamma t1 ;;
      T2 <- type_check Gamma t2 ;;
      match Tguard with
      | <{ Bool }> =>
          if eqb_ty T1 T2 then return T1 else fail
      | _ => fail
      end
  end.

(* ################################################################# *)
(** * Properties *)

(** To verify that the typechecking algorithm is correct, we show that
    it is _sound_ and _complete_ for the original [has_type]
    relation -- that is, [type_check] and [has_type] define the same
    partial function. *)

Theorem type_checking_sound : forall Gamma t T,
  type_check Gamma t = Some T -> has_type Gamma t T.
Proof with eauto.
  intros Gamma t. generalize dependent Gamma.
  induction t; intros Gamma T Htc; inversion Htc.
  - (* var *) rename s into x. destruct (Gamma x) eqn:H.
    rename t into T'. inversion H0. subst. eauto. solve_by_invert.
  - (* app *)
    remember (type_check Gamma t1) as TO1.
    destruct TO1 as [T1|]; try solve_by_invert;
    destruct T1 as [|T11 T12]; try solve_by_invert;
    remember (type_check Gamma t2) as TO2;
    destruct TO2 as [T2|]; try solve_by_invert.
    destruct (eqb_ty T11 T2) eqn: Heqb.
    apply eqb_ty__eq in Heqb.
    inversion H0; subst...
    inversion H0.
  - (* abs *)
    rename s into x, t into T1.
    remember (x |-> T1 ; Gamma) as G'.
    remember (type_check G' t0) as TO2.
    destruct TO2; try solve_by_invert.
    inversion H0; subst...
  - (* tru *) eauto.
  - (* fls *) eauto.
  - (* test *)
    remember (type_check Gamma t1) as TOc.
    remember (type_check Gamma t2) as TO1.
    remember (type_check Gamma t3) as TO2.
    destruct TOc as [Tc|]; try solve_by_invert.
    destruct Tc; try solve_by_invert;
    destruct TO1 as [T1|]; try solve_by_invert;
    destruct TO2 as [T2|]; try solve_by_invert.
    destruct (eqb_ty T1 T2) eqn:Heqb;
    try solve_by_invert.
    apply eqb_ty__eq in Heqb.
    inversion H0. subst. subst...
Qed.

Theorem type_checking_complete : forall Gamma t T,
  has_type Gamma t T -> type_check Gamma t = Some T.
Proof with auto.
  intros Gamma t T Hty.
  induction Hty; simpl.
  - (* T_Var *) destruct (Gamma x0) eqn:H0; assumption.
  - (* T_Abs *) rewrite IHHty...
  - (* T_App *)
    rewrite IHHty1. rewrite IHHty2.
    rewrite (eqb_ty_refl T2)...
  - (* T_True *) eauto.
  - (* T_False *) eauto.
  - (* T_If *) rewrite IHHty1. rewrite IHHty2.
    rewrite IHHty3. rewrite (eqb_ty_refl T1)...
Qed.

End STLCChecker.

(* ################################################################# *)
(** * Exercises *)

(** **** Exercise: 5 stars, standard (typechecker_extensions)

    In this exercise we'll extend the typechecker to deal with the
    extended features discussed in chapter [MoreStlc].  Your job
    is to fill in the omitted cases in the following. *)

Module TypecheckerExtensions.
(* Do not modify the following line: *)
Definition manual_grade_for_type_checking_sound : option (nat*string) := None.
(* Do not modify the following line: *)
Definition manual_grade_for_type_checking_complete : option (nat*string) := None.
Import MoreStlc.
Import STLCExtended.

Fixpoint eqb_ty (T1 T2 : ty) : bool :=
  match T1,T2 with
  | <{{Nat}}>, <{{Nat}}> =>
      true
  | <{{Unit}}>, <{{Unit}}> =>
      true
  | <{{T11 -> T12}}>, <{{T21 -> T22}}> =>
      andb (eqb_ty T11 T21) (eqb_ty T12 T22)
  | <{{T11 * T12}}>, <{{T21 * T22}}> =>
      andb (eqb_ty T11 T21) (eqb_ty T12 T22)
  | <{{T11 + T12}}>, <{{T21 + T22}}> =>
      andb (eqb_ty T11 T21) (eqb_ty T12 T22)
  | <{{List T11}}>, <{{List T21}}> =>
      eqb_ty T11 T21
  | _,_ =>
      false
  end.

Lemma eqb_ty_refl : forall T,
  eqb_ty T T = true.
Proof.
  intros T.
  induction T; simpl; auto;
    rewrite IHT1; rewrite IHT2; reflexivity.  Qed.

Lemma eqb_ty__eq : forall T1 T2,
  eqb_ty T1 T2 = true -> T1 = T2.
Proof.
  intros T1.
  induction T1; intros T2 Hbeq; destruct T2; inversion Hbeq;
    try reflexivity;
    try (rewrite andb_true_iff in H0; inversion H0 as [Hbeq1 Hbeq2];
         apply IHT1_1 in Hbeq1; apply IHT1_2 in Hbeq2; subst; auto);
    try (apply IHT1 in Hbeq; subst; auto).
 Qed.

Fixpoint type_check (Gamma : context) (t : tm) : option ty :=
  match t with
  | tm_var x =>
      match Gamma x with
      | Some T => return T
      | None   => fail
      end
  | <{ \ x1 : T1, t2 }> =>
      T2 <- type_check (x1 |-> T1 ; Gamma) t2 ;;
      return <{{T1 -> T2}}>
  | <{ t1 t2 }> =>
      T1 <- type_check Gamma t1 ;;
      T2 <- type_check Gamma t2 ;;
      match T1 with
      | <{{T11 -> T12}}> =>
          if eqb_ty T11 T2 then return T12 else fail
      | _ => fail
      end
  | tm_const _ =>
      return <{{Nat}}>
  | <{ succ t1 }> =>
      T1 <- type_check Gamma t1 ;;
      match T1 with
      | <{{Nat}}> => return <{{Nat}}>
      | _ => fail
      end
  | <{ pred t1 }> =>
      T1 <- type_check Gamma t1 ;;
      match T1 with
      | <{{Nat}}> => return <{{Nat}}>
      | _ => fail
      end
  | <{ t1 * t2 }> =>
      T1 <- type_check Gamma t1 ;;
      T2 <- type_check Gamma t2 ;;
      match T1, T2 with
      | <{{Nat}}>, <{{Nat}}> => return <{{Nat}}>
      | _,_        => fail
      end
  | <{ t.fst }> =>
    match type_check Gamma t with
    | Some <{{T1 * T2}}> => return T1
    | _ => fail
    end
  | <{ t.snd }> =>
    match type_check Gamma t with
    | Some <{{T1 * T2}}> => return T2
    | _ => fail
    end
  | <{ if0 guard then t else f }> =>
      Tguard <- type_check Gamma guard ;;
      T1 <- type_check Gamma t ;;
      T2 <- type_check Gamma f ;;
      match Tguard with
      | <{{Nat}}> => if eqb_ty T1 T2 then return T1 else fail
      | _ => fail
      end

  (* Complete the following cases. *)

  (* sums *)
  | <{ inl T2 t }> =>
    T1 <- type_check Gamma t ;;
    return <{{T1 + T2}}>
  | <{ inr T1 t }> =>
    T2 <- type_check Gamma t ;;
    return <{{T1 + T2}}>
  | <{ case t of | inl x1 => t1 | inr x2 => t2 }> =>
    match type_check Gamma t with
    | Some <{{T1 + T2}}> =>
      T3  <- type_check (x1 |-> T1; Gamma) t1 ;;
      T3' <- type_check (x2 |-> T2; Gamma) t2 ;;
      if eqb_ty T3 T3' then return T3 else fail
    | _ => fail
    end

  (* lists (the [tlcase] is given for free) *)
  | <{ nil T }> => return <{{List T}}>
  | <{ h :: t }> =>
    Th <- type_check Gamma h ;;
    match type_check Gamma t with
    | Some <{{List T}}> => if eqb_ty T Th then return <{{List T}}> else fail
    | _ => fail
    end
  | <{ case t0 of | nil => t1 | x21 :: x22 => t2 }> =>
      match type_check Gamma t0 with
      | Some <{{List T}}> =>
          match type_check Gamma t1,
                type_check (x21 |-> T ; x22 |-> <{{List T}}> ; Gamma) t2 with
          | Some T1', Some T2' =>
              if eqb_ty T1' T2' then Some T1' else None
          | _,_ => None
          end
      | _ => None
      end
  (* unit *)
  | <{ unit }> => return <{{Unit}}>
  (* pairs *)
  | <{(t1,t2)}> =>
    T1 <- type_check Gamma t1 ;;
    T2 <- type_check Gamma t2 ;;
    return <{{T1 * T2}}>
  (* let *)
  | <{let x=v in t}> =>
    Tv <- type_check Gamma v ;;
    type_check (x |-> Tv ; Gamma) t
  (* fix *)
  | <{ fix t }> =>
    match type_check Gamma t with
    | Some <{{T -> T'}}> => if eqb_ty T T' then return T else fail
    | _ => fail
    end
  end.

(** Just for fun, we'll do the soundness proof with just a bit more
    automation than above, using these "mega-tactics": *)

Ltac invert_typecheck Gamma t T :=
  remember (type_check Gamma t) as TO;
  destruct TO as [T|];
  try solve_by_invert; try (inversion H0; eauto); try (subst; eauto).

Ltac analyze T T1 T2 :=
  destruct T as [T1 T2| |T1 T2|T1| |T1 T2]; try solve_by_invert.

Ltac fully_invert_typecheck Gamma t T T1 T2 :=
  let TX := fresh T in
  remember (type_check Gamma t) as TO;
  destruct TO as [TX|]; try solve_by_invert;
  destruct TX as [T1 T2| |T1 T2|T1| |T1 T2];
  try solve_by_invert; try (inversion H0; eauto); try (subst; eauto).

Ltac case_equality S T :=
  destruct (eqb_ty S T) eqn: Heqb;
  inversion H0; apply eqb_ty__eq in Heqb; subst; subst; eauto.

Theorem type_checking_sound : forall Gamma t T,
  type_check Gamma t = Some T ->
  has_type Gamma t T.
Proof with eauto.
  intros Gamma t. generalize dependent Gamma.
  induction t; intros Gamma T Htc; inversion Htc.
  - (* var *) rename s into x. destruct (Gamma x) eqn:H.
    rename t into T'. inversion H0. subst. eauto. solve_by_invert.
  - (* app *)
    invert_typecheck Gamma t1 T1.
    invert_typecheck Gamma t2 T2.
    analyze T1 T11 T12.
    case_equality T11 T2.
  - (* abs *)
    rename s into x, t into T1.
    remember (x |-> T1 ; Gamma) as Gamma'.
    invert_typecheck Gamma' t0 T0.
  - (* const *) eauto.
  - (* scc *)
    rename t into t1.
    fully_invert_typecheck Gamma t1 T1 T11 T12.
  - (* prd *)
    rename t into t1.
    fully_invert_typecheck Gamma t1 T1 T11 T12.
  - (* mlt *)
    invert_typecheck Gamma t1 T1.
    invert_typecheck Gamma t2 T2.
    analyze T1 T11 T12; analyze T2 T21 T22.
    inversion H0. subst. eauto.
  - (* test0 *)
    invert_typecheck Gamma t1 T1.
    invert_typecheck Gamma t2 T2.
    invert_typecheck Gamma t3 T3.
    destruct T1; try solve_by_invert.
    case_equality T2 T3.
  - (* inl *)
    rename t into T2.
    invert_typecheck Gamma t0 T1.
  - (* inr *)
    rename t into T1.
    invert_typecheck Gamma t0 T2.
  - (* sum case *)
    invert_typecheck Gamma t1 T12.
    analyze T12 T1 T2.
    rename s into x1. rename s0 into x2.
    rename t1 into t0. rename t2 into t1. rename t3 into t2.
    invert_typecheck (x1 |-> T1; Gamma) t1 HT3.
    invert_typecheck (x2 |-> T2; Gamma) t2 HT3'.
    case_equality HT3 HT3'.
  - (* nil *)
    subst...
  - (* cons *)
    invert_typecheck Gamma t1 T1.
    invert_typecheck Gamma t2 T2.
    analyze T2 T2' T2''.
    case_equality T2' T1.
  - (* tlcase *)
    rename s into x31, s0 into x32.
    fully_invert_typecheck Gamma t1 T1 T11 T12.
    invert_typecheck Gamma t2 T2.
    remember (x31 |-> T11 ; x32 |-> <{{List T11}}> ; Gamma) as Gamma'2.
    invert_typecheck Gamma'2 t3 T3.
    case_equality T2 T3.
  - (* unit *)
    eauto.
  - (* pair *)
    invert_typecheck Gamma t1 T1.
    invert_typecheck Gamma t2 T2.
  - (* fst *)
    invert_typecheck Gamma t T'.
    analyze T' T'1 T'2.
    inversion H0. subst. eauto.
  - (* snd *)
    invert_typecheck Gamma t T'.
    analyze T' T'1 T'2.
    inversion H0. subst. eauto.
  - (* let *)
    invert_typecheck Gamma t1 T'.
  - (* fix *)
    invert_typecheck Gamma t T'.
    analyze T' T'1 T'2.
    case_equality T'1 T'2.
Qed.

Theorem type_checking_complete : forall Gamma t T,
  has_type Gamma t T ->
  type_check Gamma t = Some T.
Proof.
  intros Gamma t T Hty.
  induction Hty; simpl;
    try (rewrite IHHty);
    try (rewrite IHHty1);
    try (rewrite IHHty2);
    try (rewrite IHHty3);
    try (rewrite (eqb_ty_refl T0));
    try (rewrite (eqb_ty_refl T1));
    try (rewrite (eqb_ty_refl T2));
    try (rewrite (eqb_ty_refl T3));
    eauto.
  - destruct (Gamma x0); [assumption| solve_by_invert].
  - rewrite eqb_ty_refl.  reflexivity.
Qed.

End TypecheckerExtensions.
(** [] *)

(** **** Exercise: 5 stars, standard, optional (stlc_step_function)

    Above, we showed how to write a typechecking function and prove it
    sound and complete for the typing relation.  Do the same for the
    operational semantics -- i.e., write a _function_ [stepf] of type
    [tm -> option tm] and prove that it is sound and complete with
    respect to [step] from chapter [MoreStlc]. *)

Module StepFunction.
Import MoreStlc.
Import STLCExtended.

Fixpoint valueb (t : tm) : bool :=
  match t with
  | <{\x:T2, t1}> => true
  | tm_const _ => true
  | <{inl _ v}> => valueb v
  | <{inr _ v}> => valueb v
  | <{nil _}> => true
  | <{x :: xs}> => valueb x && valueb xs
  | <{unit}> => true
  | <{(a,b)}> => valueb a && valueb b
  | _ => false
  end.

(* Operational semantics as a Coq function. *)
Fixpoint stepf (t : tm) : option tm :=
  match t with
  | tm_var x => fail
  | <{ \ x1 : T1, t2 }> => fail
  | <{ t1 t2 }> =>
    match valueb t1, valueb t2 with
    | true, true => match t1 with
                   | <{\x1 : T1, t1'}> => Some <{ [x1:=t2]t1' }>
                   | _ => fail
                   end
    | true, false =>
      t2' <- stepf t2 ;; Some <{ t1 t2' }>
    | false, _ =>
      t1' <- stepf t1 ;; Some <{ t1' t2 }>
    end
  | tm_const _ => fail
  | <{ succ t1 }> =>
    match t1 with
    | tm_const n => Some (tm_const (S n))
    | _ => t1' <- stepf t1 ;; Some <{ succ t1' }>
    end
  | <{ pred t1 }> =>
    match t1 with
    | tm_const n => Some (tm_const (n - 1))
    | _ => t1' <- stepf t1 ;; Some <{ pred t1' }>
    end
  | <{ t1 * t2 }> =>
    match valueb t1, valueb t2 with
    | true, true => match t1, t2 with
                   | tm_const n1, tm_const n2 => Some (tm_const (n1 * n2))
                   | _, _ => fail
                   end
    | true, false =>
      t2' <- stepf t2 ;; Some <{ t1 * t2' }>
    | false, _ =>
      t1' <- stepf t1 ;; Some <{ t1' * t2 }>
    end
  | <{(t1,t2)}> =>
    match valueb t1, valueb t2 with
    | false, _ => t1' <- stepf t1 ;; Some <{(t1', t2)}>
    | true, false => t2' <- stepf t2 ;; Some <{(t1, t2')}>
    | _, _ => fail
    end
  | <{ t.fst }> =>
    match stepf t with
    | None => match t with
             | <{(t1,t2)}> => if valueb t1 then
                             if valueb t2 then
                             Some t1
                             else fail
                             else fail
             | _ => fail
             end
    | Some t' => Some <{ t'.fst }>
    end
  | <{ t.snd }> =>
    match stepf t with
    | None => match t with
             | <{(t1,t2)}> => if valueb t1 then
                             if valueb t2 then
                             Some t2
                             else fail
                             else fail
             | _ => fail
             end
    | Some t' => Some <{ t'.snd }>
    end
  | <{ if0 guard then t else f }> =>
    match guard with
    | tm_const 0 => Some t
    | tm_const _ => Some f
    | g => g' <- stepf g ;; Some <{ if0 g' then t else f }>
    end
  | <{ inl T2 t }> => t' <- stepf t ;; Some <{ inl T2 t' }>
  | <{ inr T1 t }> => t' <- stepf t ;; Some <{ inr T1 t' }>
  | <{ case t of | inl x1 => t1 | inr x2 => t2 }> =>
    match stepf t with
    | None => match t with
             | <{ inl T2 t' }> => if valueb t'
                                 then Some <{[x1:=t']t1}>
                                 else fail
             | <{ inr T1 t' }> => if valueb t'
                                 then Some <{[x2:=t']t2}>
                                 else fail
             | _ => fail
             end
    | Some t' => Some <{ case t' of | inl x1 => t1 | inr x2 => t2  }>
    end
  | <{ nil T }> => fail
  | <{ h :: t }> =>
    match valueb h, valueb t with
    | false, _ => h' <- stepf h ;; Some <{ h' :: t }>
    | true, false => t' <- stepf t ;; Some <{ h :: t' }>
    | _, _ => fail
    end
  | <{ case t0 of | nil => t1 | x21 :: x22 => t2 }> =>
    match stepf t0 with
    | None => match t0 with
             | <{ nil _ }> => Some t1
             | <{ h :: t }> => if valueb h then
                             if valueb t then
                             Some <{[x22 := t]([x21 := h]t2)}>
                             else fail
                             else fail
             | _ => fail
             end
    | Some t0' => Some <{ case t0' of | nil => t1 | x21 :: x22 => t2 }>
    end
  | <{ unit }> => fail
  | <{let x=v in t}> =>
    match valueb v with
    | true => Some <{[x:=v]t}>
    | false => v' <- stepf v ;; Some <{let x=v' in t}>
    end
  | <{ fix t }> =>
    match t with
    | <{\xf:T, t}> => Some <{ [xf:=fix (\xf:T, t)]t }>
    | _ => t' <- stepf t ;; Some <{ fix t' }>
    end
  end.

Hint Unfold stepf : core.

Theorem sound_valueb : forall t,
  valueb t = true -> value t.
Proof.
  intro. induction t; intro; eauto; try solve_by_invert.
  - (* cons *)
    inversion H. rewrite andb_true_iff in H1. destruct H1. eauto.
  - (* pair *)
    inversion H. rewrite andb_true_iff in H1. destruct H1. eauto.
Qed.

Theorem complete_valueb : forall t,
  value t -> valueb t = true.
Proof.
  intros. induction H; auto; simpl; rewrite IHvalue1; rewrite IHvalue2; easy.
Qed.

Hint Resolve complete_valueb : core.
Hint Resolve sound_valueb : core.

Theorem no_step_for_values : forall t,
  value t -> stepf t = None.
Proof.
  intros. induction H; auto; simpl;
          try (rewrite IHvalue; auto);
          try (rewrite IHvalue1; rewrite IHvalue2; auto);
          try (apply complete_valueb in H; apply complete_valueb in H0;
               rewrite H; rewrite H0;
               auto).
Qed.

From PLF Require Import MyTactics.

Ltac solve_stepf :=
  repeat (MyTactics.solve_stepf_step tm stepf valueb).

(* Soundness of [stepf]. *)
Theorem sound_stepf : forall t t',
    stepf t = Some t'  ->  t --> t'.
Proof.
  intro.
  induction t; intros t' H; inversion H; clear H; auto.
  (* I intentionally split the cases manually because the proof
   runs for pretty long time and I want to keep track of the progress
   *)
  - (* App *)
    solve_stepf.
  - (* Succ *)
    solve_stepf.
  - (* Pred *)
    solve_stepf.
  - (* Mult *)
    solve_stepf.
  - (* If0 *)
    solve_stepf.
  - (* Inl *)
    solve_stepf.
  - (* Inr *)
    solve_stepf.
  - (* Case Either *)
    solve_stepf.
  - (* Cons *)
    solve_stepf.
  - (* Case List *)
    solve_stepf.
  - (* Pair *)
    solve_stepf.
  - (* Fst *)
    solve_stepf.
  - (* Snd *)
    solve_stepf.
  - (* Let *)
    solve_stepf.
  - (* Fix *)
    solve_stepf.
Qed.

Ltac solve_stepf_converse_step := MyTactics.solve_stepf_converse_step
                                  tm step value stepf valueb
                                  complete_valueb
                                  sound_valueb
                                  no_step_for_values.
Ltac solve_stepf_converse :=
  repeat solve_stepf_converse_step.

(* Completeness of [stepf]. *)
Theorem complete_stepf : forall t t',
    t --> t'  ->  stepf t = Some t'.
Proof.
  intros.
  induction H.
  - (* ST_Abs *)
    solve_stepf_converse.
  - (* ST_App1 *)
    solve_stepf_converse.
  - (* ST_App2 *)
    solve_stepf_converse.
  - (* ST_Succ *)
    solve_stepf_converse.
  - (* ST_SuccNat *)
    solve_stepf_converse.
  - (* ST_Pred *)
    solve_stepf_converse.
  - (* ST_PredNat *)
    solve_stepf_converse.
  - (* ST_Mulconsts *)
    solve_stepf_converse.
  - (* ST_Mult1 *)
    solve_stepf_converse.
  - (* ST_Mult2 *)
    solve_stepf_converse.
  - (* ST_If0 *)
    solve_stepf_converse.
  - (* ST_If0_Zero *)
    solve_stepf_converse.
  - (* ST_If0_Nonzero *)
    solve_stepf_converse.
  - (* ST_Inl *)
    solve_stepf_converse.
  - (* ST_Inr *)
    solve_stepf_converse.
  - (* ST_Case *)
    solve_stepf_converse.
  - (* ST_CaseInl *)
    solve_stepf_converse.
  - (* ST_CaseInr *)
    solve_stepf_converse.
  - (* ST_Cons1 *)
    solve_stepf_converse.
  - (* ST_Cons2 *)
    solve_stepf_converse.
  - (* ST_Lcase1 *)
    solve_stepf_converse.
  - (* ST_LcaseNil *)
    solve_stepf_converse.
  - (* ST_LcaseCons *)
    solve_stepf_converse.
  - (* ST_Pair1 *)
    solve_stepf_converse.
  - (* ST_Pair2 *)
    solve_stepf_converse.
  - (* ST_Fst1 *)
    solve_stepf_converse.
  - (* ST_FstPair *)
    inversion H; subst; solve_stepf_converse.
  - (* ST_Snd1 *)
    solve_stepf_converse.
  - (* ST_SndPair *)
    inversion H; subst; solve_stepf_converse.
  - (* ST_Let1 *)
    solve_stepf_converse.
  - (* ST_LetValue *)
    solve_stepf_converse.
  - (* ST_Fix1 *)
    solve_stepf_converse.
  - (* ST_FixAbs *)
    solve_stepf_converse.
Qed.

End StepFunction.
(** [] *)

(** **** Exercise: 5 stars, standard, optional (stlc_impl)

    Using the Imp parser described in the [ImpParser] chapter
    of _Logical Foundations_ as a guide, build a parser for extended
    STLC programs.  Combine it with the typechecking and stepping
    functions from the above exercises to yield a complete typechecker
    and interpreter for this language. *)

Module StlcImpl.
  Import StepFunction.

  (* I decided not to solve this problem *)

End StlcImpl.
(** [] *)

(* 2020-09-09 21:08 *)


Set Warnings "-extraction-opaque-accessed,-extraction".
Set Warnings "-notation-overridden,-parsing".

Require Import Coq.Classes.Morphisms.
Require Import mathcomp.ssreflect.ssreflect.
From mathcomp Require Import ssrbool ssrnat.
Require Import Sets Tactics.
Require Import Recdef.
Require Import List.

Require Import ZArith ZArith.Znat Arith.

Require Import Producer Generators Enumerators.

Set Bullet Behavior "Strict Subproofs".

Local Open Scope set_scope.

(** Apply a function n times *)
Fixpoint appn {A} (f : A -> A) (n : nat) : A ->  A :=
  fun x =>
    match n with
      | 0%nat => x
      | S n' => f (appn f n' x)
    end.

Infix "^" := appn (at level 30, right associativity) : fun_scope.

(** Instance Hierarchy

   GenSized
      |
      |
     Gen   Shrink
       \    /
        \  /
      Arbitrary
*)

(** * Generator-related classes *)

(* begin gen_sized_class *)
Class GenSized (A : Type) :=
  { arbitrarySized : nat -> G A }.
(* end gen_sized_class *)

(* begin gen_class *)
Class Gen (A : Type) := { arbitrary : G A }.
(* end gen_class *)

(** Shrink class *)
Class Shrink (A : Type) :=
  { shrink : A -> list A }.

(** Arbitrary Class *)
Class Arbitrary (A : Type) `{Gen A} `{Shrink A}.

Class Fuzzy (A : Type) := { fuzz : A -> G A }.

(* TODO: should be able to fail? i.e. outputs G (option A) instead *)
Class Mutate (A : Type) := { mutate : A -> G A }.

Class EnumSized (A : Type) := { enumSized : nat -> E A }.

Class Enum (A : Type) := { enum : E A }.

Class Sized (A : Type) :=
  { size : A -> nat }.

(* ZP: This is not longer usefull *)
(*
(** * Sizes of types *)

Class Sized (A : Type) :=
  { size : A -> nat }.

Class CanonicalSized (A : Type) `{Sized A} :=
  {
    zeroSized : set A;
    succSized : set A -> set A;

    zeroSized_spec : zeroSized <--> [ set x : A | size x = 0 ];
    succSized_spec :
      forall n, succSized [ set x : A | size x <= n ] <--> [ set x : A | size x <= S n ]

  }.

Lemma size_ind (A : Type) `{Hyp : Sized A} :
  forall (P : A -> Prop), (forall y, (forall x, size x < size y -> P x) -> P y) -> (forall x, P x).
Proof.
  intros P H1.
  intros x.
  assert (Hin : [ set y :  A | size y <= size x] x); eauto.
  revert Hin.
  generalize (size x). intros n.
  revert x. induction n.
  - intros x Hl. apply H1. intros x1 Hlt. ssromega.
  - intros x Hleq. eapply H1. intros x1 Hlt.
    eapply IHn. ssromega.
Qed.

Lemma size_lfp (A : Type) `{Hyp : CanonicalSized A} :
  [set x : A | True ] <--> \bigcup_(s : nat) [set x : A | size x <= s ].
Proof.
  intros a; split; eauto. intros _.
  exists (size a). split; eauto. constructor.
Qed.

Lemma succ_lfp (A : Type) `{Hyp : CanonicalSized A}
      `{Proper _ (respectful set_eq set_eq) succSized} s :
  [set x : A | size x <= s ] <-->  (succSized ^ s) zeroSized.
Proof.
  induction s.
  simpl.
  - rewrite zeroSized_spec.
    split; intros; ssromega.
  - simpl. rewrite <- succSized_spec.
    rewrite IHs. reflexivity.
Qed.
*)
(* Lemma succ_lfp' (A : Type) `{Hyp : CanonicalSized A} : *)
(*   \bigcup_(s : nat)  (succSized ^ s) zeroSized <--> [ set x : A | True ]. *)
(* Proof. *)
(*   intros. split; eauto. *)
(*   intros _. *)
(*   eapply set_eq_trans. *)
(*   Focus 2. symmetry. *)
(*   eapply succ_lfp. *)
(*   simpl.  *)
(*   rewrite succ_lfp at 2. *)
(* split. *)
(*     split. rewrite IHs. firstorder. *)
(*     IHs. *)
(*   firstorder. reflexivity. split; intros; eauto. *)
(*   exists (size a). *)
(*   remember (size a) as s. *)
(*   revert a Heqs. induction s; intros. *)
(*   - split. constructor. *)
(*     simpl. eapply zeroSized_spec. now eauto. *)
(*   - split. constructor. *)
(*     simpl. *)
(*     eapply (succSized_spec. *)
(*     eassumption. *)
(*   eapply size_ind. *)

(*   [set x : A | True ] <-->  . *)

(** * Correctness classes *)

(** Correctness of sized generators *)

Class CorrectSized {A : Type} {G} `{Producer G} (g : nat -> G A)  :=
  { prodCorrectSized :
      [ set n | exists s, semProd (g s) n ] <--> [set : A]
  }.

(** Correctness of generators *)
Class Correct (A : Type) {G} `{Producer G} (g : G A)  :=
  {
    prodCorrect : semProd g <--> [set : A]
  }.

(** * Monotonic generators *)

(** Monotonicity of size parametric generators *)
Class GenSizedMonotonic (A : Type) `{GenSized A}
      `{forall s, SizeMonotonic (arbitrarySized s)}.

(** Monotonicity of size parametric generators v2 *)
Class GenSizedSizeMonotonic (A : Type) `{GenSized A}
        `{SizedMonotonic A G arbitrarySized}.

Class GenMonotonic (A : Type) `{Gen A} `{@SizeMonotonic A G ProducerGen arbitrary}.

(** Monotonicity of size parametric generators *)
Class EnumSizedMonotonic (A : Type) `{EnumSized A}
      `{forall s, @SizeMonotonic A E ProducerEnum (enumSized s)}.

(** Monotonicity of size parametric generators v2 *)
Class EnumSizedSizeMonotonic (A : Type) `{EnumSized A}
        `{@SizedMonotonic A E ProducerEnum enumSized}.

Class EnumMonotonic (A : Type) `{Enum A}
        `{@SizeMonotonic A E ProducerEnum enum}.

(** * Correct generators *)

Class GenSizedCorrect (A : Type) `{GenSized A}
      `{@CorrectSized A G ProducerGen arbitrarySized}.

Class GenCorrect (A : Type) `{Gen A}
        `{@Correct A G ProducerGen arbitrary}.

(* Monotonic and Correct generators *)
Class GenMonotonicCorrect (A : Type)
      `{Gen A}
      `{@SizeMonotonic A G ProducerGen arbitrary}
      `{@Correct  A G ProducerGen arbitrary}.

Class EnumSizedCorrect (A : Type)
        `{EnumSized A} `{@CorrectSized A E ProducerEnum enumSized}.

Class EnumCorrect (A : Type) `{Enum A}
        `{@Correct A E ProducerEnum enum}.

(* Monotonic and Correct generators *)
Class EnumMonotonicCorrect (A : Type)
      `{Enum A}
      `{@SizeMonotonic A E ProducerEnum enum}
      `{@Correct A E ProducerEnum enum}.

(** Coercions *)

#[global] Instance GenSizedMonotonicOfSizeMonotonic
         (A : Type) (Hgen : GenSized A) (Hmon : forall s, @SizeMonotonic A G ProducerGen (arbitrarySized s))
: @GenSizedMonotonic A Hgen Hmon := {}.

#[global] Instance GenMonotonicOfSizeMonotonic
         (A : Type) (Hgen : Gen A) (Hmon : @SizeMonotonic A G ProducerGen arbitrary)
: @GenMonotonic A Hgen Hmon := {}.

#[global] Instance GenSizedCorrectOfSizedCorrect
         (A : Type) (Hgen : GenSized A)
         `{Hcor : @CorrectSized A G ProducerGen arbitrarySized}
: @GenSizedCorrect A Hgen Hcor := {}.

#[global] Instance GenMonotonicCorrectOfMonotnicCorrect
         (A : Type) (Hgen : Gen A)
         (Hmon : @SizeMonotonic A G ProducerGen arbitrary)
         `{Hcor : @Correct A G ProducerGen arbitrary}
  : @GenMonotonicCorrect A Hgen Hmon Hcor := {}.

#[global] Instance GenSizedSizeMonotonicOfSizedMonotonic
         (A : Type) (Hgen : GenSized A) (Hmon : @SizedMonotonic A G ProducerGen arbitrarySized)
: @GenSizedSizeMonotonic A Hgen ProducerGen Hmon := {}.

#[global] Instance GenOfGenSized {A} `{GenSized A} : Gen A :=
  {| arbitrary := sized arbitrarySized |}.

#[global] Instance ArbitraryOfGenShrink {A} `{Gen A} `{Shrink A} : Arbitrary A := {}.

Generalizable Variables PSized PMon PSMon PCorr.

#[global] Instance GenMonotonicOfSized (A : Type)
         `{H : GenSized A}
         `{@GenSizedMonotonic A H PMon}
         `{@GenSizedSizeMonotonic A H ProducerGen PSMon}
  : @GenMonotonic A
                  (@GenOfGenSized A H)
                  (@sizedSizeMonotonic G ProducerGen _ A
                                       (@arbitrarySized A H)
                                       PMon PSMon) := {}.

#[global] Instance GenCorrectOfSized (A : Type)
         {H : GenSized A}
         `{@GenSizedMonotonic A H PMon}
         `{@GenSizedSizeMonotonic A H ProducerGen PSMon}
         `{@GenSizedCorrect A H PCorr} : Correct A arbitrary.
Proof.
  constructor. unfold arbitrary, GenOfGenSized.
  eapply set_eq_trans.
  - eapply semSized_alt; now eauto with typeclass_instances.
  - destruct PCorr.
    rewrite <- prodCorrectSized0.

    split. intros [n [? H3]]. eexists. eassumption.

    intros [s H4]. eexists; split; eauto. reflexivity.
Qed.

(*
Lemma nat_set_ind (A : Type) `{GenSized A} `{Hyp : CanonicalSized A} :
  (semProd (arbitrarySized 0) <--> zeroSized) ->
  (forall (s : nat) (elems : set A),
     semProd (arbitrarySized s) <--> elems ->
     semProd (arbitrarySized (s.+1)) <--> succSized elems) ->
  (forall s : nat, semProd (arbitrarySized s) <--> (fun x : A => size x <= s)).
Proof.
  intros HO IH. intros n; induction n.
  - eapply set_eq_trans with (B := (fun x : A => size x = 0)).
    rewrite -zeroSized_spec //=.
    intros s. destruct (size s). now firstorder.
    split; intros; ssromega.
  - rewrite -succSized_spec. eauto.
Qed.
 *)

#[global] Instance EnumSizedMonotonicOfSizeMonotonic
         (A : Type) (Hgen : EnumSized A) (Hmon : forall s, @SizeMonotonic A E ProducerEnum (enumSized s))
: @EnumSizedMonotonic A Hgen Hmon := {}.

#[global] Instance EnumMonotonicOfSizeMonotonic
         (A : Type) (Hgen : Enum A) (Hmon : @SizeMonotonic A E ProducerEnum enum)
: @EnumMonotonic A Hgen Hmon := {}.

#[global] Instance EnumSizedCorrectOfSizedCorrect
         (A : Type) (Hgen : EnumSized A)
         `{Hcor : @CorrectSized A E ProducerEnum enumSized}
: @EnumSizedCorrect A Hgen Hcor := {}.

#[global] Instance EnumCorrectOfCorrect
         (A : Type) (Hgen : Enum A)
         `{Hcor : @Correct A E ProducerEnum enum}
: @EnumCorrect A Hgen Hcor := {}.

#[global] Instance EnumMonotonicCorrectOfMonotnicCorrect
         (A : Type) (Hgen : Enum A)
         (Hmon : @SizeMonotonic A E ProducerEnum enum)
         `{Hcor : @Correct A E ProducerEnum enum}
  : @EnumMonotonicCorrect A Hgen Hmon Hcor := {}.

#[global] Instance EnumSizedSizeMonotonicOfSizedMonotonic
         (A : Type) (Hgen : EnumSized A) (Hmon : @SizedMonotonic A E ProducerEnum enumSized)
: @EnumSizedSizeMonotonic A Hgen Hmon := {}.

#[global] Instance EnumOfEnumSized {A} `{EnumSized A} : Enum A :=
  {| enum := sized enumSized |}.

(*
#[global] Instance EnumOfGenShrink {A} `{Gen A} `{Shrink A} : Arbitrary A := {}.
 *)

#[global] Instance EnumMonotonicOfSized (A : Type)
         `{H : EnumSized A}
         `{@EnumSizedMonotonic A H PMon}
         `{@EnumSizedSizeMonotonic A H PSMon}
  : @EnumMonotonic A
                  (@EnumOfEnumSized A H)
                  (@sizedSizeMonotonic E ProducerEnum _ A
                                       (@enumSized A H)
                                       PMon PSMon) := {}.

#[global] Instance EnumCorrectOfSized' (A : Type)
         {H : EnumSized A}
         `{@EnumSizedMonotonic A H PMon}
         `{@EnumSizedSizeMonotonic A H PSMon}
         `{@EnumSizedCorrect A H PCorr} : Correct A enum.
Proof.
  constructor. unfold arbitrary, EnumOfEnumSized.
  eapply set_eq_trans.
  - eapply semSized_alt; eauto with typeclass_instances.
  - destruct PCorr.
    rewrite <- prodCorrectSized0.

    split. intros [n [? H3]]. eexists. eassumption.

    intros [s H4]. eexists; split; eauto. reflexivity.
Qed.

#[global] Instance EnumCorrectOfSized (A : Type)
         (H1 : EnumSized A)
         (H3 : forall s : nat, SizeMonotonic (enumSized s))
         (H4 : SizedMonotonic enumSized)
         (H5 : CorrectSized enumSized)
  : Correct A enum.
Proof.
  eapply EnumCorrectOfSized'; eauto.
  constructor; eauto.
  constructor; eauto.
  constructor; eauto.
Qed.

Lemma enumerating_complete A (e : E A) {Hc : Correct A e} ch :
  (exists x, ch x = Some true) ->
  exists s, enumerating e ch s = Some true.
Proof.
  intros [x Hch].
  unfold enumerating.
  assert (Hin : semProd e x).
  { eapply Hc. reflexivity. }
  destruct Hin as [s [_ Hsem]]. simpl in *.
  unfold semEnumSize in *.
  exists s. revert Hsem.
  generalize (Enumerators.run e s), false.
  induction l; intros b Hin; inv Hin; simpl.
  - rewrite Hch. reflexivity.
  - destruct (ch a); eauto.
    destruct b0; eauto.
Qed.

Lemma lazylist_backtrack_opt_is_true A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch true = Some b -> b = true.
Proof.
  induction l.
  - simpl. congruence.
  - simpl in *.
    destruct a.
    + destruct (ch a) eqn:Heq. destruct b0. congruence.
      now eauto. now eauto.
    + eauto.
Qed.

Lemma lazylist_backtrack_opt_is_false A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch b = Some false -> b = false.
Proof.
  revert b; induction l; intros b.
  - simpl. destruct b; try congruence.
  - simpl in *.
    destruct a.
    + destruct (ch a) eqn:Heq. destruct b0. congruence.
      now eauto. intros Hin. eapply H in Hin. congruence.
    + intros Hn. eapply lazylist_backtrack_opt_is_true in Hn.
      congruence.
Qed.

Lemma lazylist_backtrack_opt_true A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch b = Some true <->
  exists x, LazyList.In_ll (Some x) l /\ ch x = Some true.
Proof.
  revert b; induction l; intros b.
  - simpl. split; intros Hc; try congruence.
    destruct b; congruence.
    destruct Hc. inv H; exfalso; eauto.
  - split.
    + intros Hl.
      simpl in *. destruct a.
      destruct (ch a) eqn:Heq.
      * destruct b0.
        -- eexists. split; eauto.
        -- eapply H in Hl. destruct Hl. inv H0.
           eexists; split; eauto.
      * eapply H in Hl. destruct Hl. inv H0.
        eexists; split; eauto.
      * eapply H in Hl. destruct Hl. inv H0.
        eexists; split; eauto.
    + intros Hc; inv Hc. inv H0. simpl in *. destruct a.
      * inv H1. inv H3. rewrite H2. reflexivity.
        destruct (ch a) eqn:Heq. destruct b0. reflexivity.
        eapply H. eexists. now split; eauto.
        eapply H. eexists. now split; eauto.
      * inv H1. congruence. eapply H. eexists. now split; eauto.
Qed.

Lemma lazylist_backtrack_opt_false A (l: LazyList.LazyList (option A)) ch :
  lazylist_backtrack_opt l ch false = Some false <->
  (forall x, LazyList.In_ll (Some x) l -> ch x = Some false) /\
  ~ LazyList.In_ll None l.
Proof.
  induction l.
  - simpl. split; intros Hc.
    split; eauto. intros. now inv H.
    reflexivity.
  - split.
    + intros Hl.
      simpl in *.
      destruct a.
      * destruct (ch a) eqn:Heq. destruct b.
        congruence.

        -- eapply H in Hl. inv Hl.
           split.
           intros. inv H2; eauto. inv H3. eassumption.

           intros Hc; inv Hc. congruence. eauto.

        -- eapply lazylist_backtrack_opt_is_false in Hl. congruence.

      * eapply lazylist_backtrack_opt_is_false in Hl. congruence.

    + intros [H1 H2].
      simpl.
      destruct a.
      * destruct (ch a) eqn:Heq. destruct b.
        rewrite H1 in Heq. congruence.
        now left.
        eapply H. split. intros. eapply H1.
        right. eassumption.
        intros Hc. eapply H2. right. eassumption.
        rewrite H1 in Heq. congruence. now left.

      * exfalso. eapply H2. now left.
Qed.

Lemma enumeratingOpt_monotonic A (e : E (option A))
      {Hc : SizeMonotonicOpt e} {Hfp : SizeFP e}
      ch1 ch2 s1 s2 b :
  (s1 <= s2)%coq_nat ->
  (forall x, ch1 x = Some b -> ch2 x = Some b) ->
  enumeratingOpt e ch1 s1 = Some b ->
  enumeratingOpt e ch2 s2 = Some b.
Proof.
  intros Hleq Hall Hen.

  specialize (Hc _ _ Hleq).
  specialize (Hfp _ _ Hleq).

  unfold enumeratingOpt in *.
  unfold semProdSizeOpt in *. simpl in *.
  unfold semEnumSize in *.
  revert Hc Hfp Hen.

  generalize (run e s1), (run e s2).

  intros l1 l2 Hs1 Hs2 Hl.
  destruct b.

  - eapply lazylist_backtrack_opt_true in Hl.
    destruct Hl. destruct H.

    eapply lazylist_backtrack_opt_true.
    eexists. split.
    eapply Hs1. eassumption.
    eapply Hall. eassumption.

  - eapply lazylist_backtrack_opt_false in Hl.
    destruct Hl.

    eapply lazylist_backtrack_opt_false.

    split. intros x Hin. eapply Hall. eapply H.
    eapply Hs2. eassumption. eassumption.

    intros Hc. eapply H0. eapply Hs2. eassumption. eassumption.
Qed.


Load PredMonad.

(* The satisfiability order on monadic predicates, where leqM P1 P2 means that
anything satisfying P1 satisfies P2; we will require this to be a complete
lattice (meets and joins over arbitrary sets) below *)
Class PredMonadOrder (PM: Type -> Type) : Type :=
  leqP : forall {A}, relation (PM A).

(***
 *** Weak Predicate Monad Axioms (as a Typeclass)
 ***)

(* A weak predicate monad, where impliesP, forallP, and existsP are not quite
equivalent to their corresponding logical connectives; I think this becomes
equivalent to the standard PredMonad class if we require the topology to be
preregular (R1), meaning that every pair of non-equal monadic computations
satisfy, respectively, a pair of monadic predicates that are disjoint. *)

Class WeakPredMonad
      {M: Type -> Type} {PM : Type -> Type}
      {PredMonadRet:MonadRet PM} {PredMonadBind:MonadBind PM}
      {PredMonadEquiv:MonadEquiv PM}
      {MonadRet:MonadRet M} {MonadBind:MonadBind M} {MonadEquiv:MonadEquiv M}
      {PredMonadOrder:@PredMonadOrder PM}
      {PredMonadSatisfaction:@PredMonadSatisfaction M PM}
      {PredMonadForall:@PredMonadForall PM}
      {PredMonadExists:@PredMonadExists PM}
      {PredMonadImplication:@PredMonadImplication PM}
  : Prop :=
  {
    (* Both M and PM must be monads *)
    weakpredmonad_monad_M :> Monad M;
    weakpredmonad_monad_PM :> Monad PM;

    (* The satisfiability order is a partial order w.r.t. eqM *)
    weakpredmonad_leqP_preorder :> forall {A}, PreOrder (leqP (A:=A));
    weakpredmonad_leqP_antisymmetric :>
      forall {A}, Antisymmetric (eqM (A:=A)) leqP;
    weakpredmonad_leqP_proper :>
      forall {A}, Proper (eqM (A:=A) ==> eqM (A:=A) ==> impl) leqP;

    (* Satisfaction is monotone with respect to the satisfiability order *)
    weakpredmonad_satisfiesP_monotone :>
      forall {A}, Proper (eqM (A:=A) ==> leqP ==> impl) satisfiesP;

    (* Forall and exists are complete meet and join operators, respectively *)
    weakpredmonad_forallP_lower_bound :
      forall {A B} (f: A -> PM B) a, leqP (forallP f) (f a);
    weakpredmonad_forallP_greatest_lower_bound :
      forall {A B} (f: A -> PM B) P,
        (forall a, leqP P (f a)) -> leqP P (forallP f);
    weakpredmonad_existsP_upper_bound :
      forall {A B} (f: A -> PM B) a, leqP (f a) (existsP f);
    weakpredmonad_existsP_least_upper_bound :
      forall {A B} (f: A -> PM B) P,
        (forall a, leqP (f a) P) ->  leqP (existsP f) P;

    (* The fact that true always holds and false never does is equivalent to the
    existence, for each m, of predicates that m does and does not satisfy,
    respectively. We state the latter property as an axiom, and the former
    properties follow from them. *)
    weakpredmonad_separation :
      forall {A} (m:M A), exists P1 P2, m |= P1 /\ ~(m |= P2);

    (* Implication makes a predicate monad into a complete Heyting algebra,
    which means there is a Galois connection between implication and the andP
    operation.  Note that this is weaker than making PM a Boolean algebra,
    which would be inherently non-constructive. *)
    weakpredmonad_heyting_algegbra :
      forall {A} (P1 P2 P3: PM A),
        leqP (andP P1 P2) P3 <-> leqP P1 (impliesP P2 P3);

    (* Return in the predicate monad precisely characterizes return in M *)
    weakpredmonad_return_return :
      forall {A} (x:A) m, m |= returnM x <-> m == returnM x;

    weakpredmonad_bind_bind :
      forall {A B} (m:M B) (P: PM A) (Q: A -> PM B),
        m |= bindM P Q <->
        (exists (phi:A -> Prop) (m': M {x:A | phi x}) (f : A -> M B),
           (bindM m' (fun x => returnM (proj1_sig x))) |= P /\
           (forall x, phi x -> f x |= Q x) /\
           eqM m (bindM m' (fun x => f (proj1_sig x))))

    (* FIXME: need to commute return and bind with logical operators! *)
  }.

(***
 *** Theorems about predicate monads
 ***)

Section WeakPredMonad_thms.

Context `{WeakPredMonad}.

Theorem existsP_exists_weak {A B} m (P: A -> PM B) :
  (exists x, m |= P x) -> m |= existsP P.
  intro H0; destruct H0.
  apply (weakpredmonad_satisfiesP_monotone m m (reflexivity m) (P x));
    [ apply weakpredmonad_existsP_upper_bound | assumption ].
Qed.

Theorem forallP_forall_weak {A B} m (P: A -> PM B) :
  m |= forallP P -> forall x, m |= P x.
  intros.
  apply (weakpredmonad_satisfiesP_monotone m m (reflexivity m) (forallP P));
    [ apply weakpredmonad_forallP_lower_bound | assumption ].
Qed.

Theorem andP_and_weak {A} m (P1 P2: PM A) :
  m |= andP P1 P2 -> m |= P1 /\ m |= P2.
  intros; split.
  apply (weakpredmonad_satisfiesP_monotone m m (reflexivity m) (andP P1 P2));
    [ | assumption ].
  apply (weakpredmonad_forallP_lower_bound
           (fun (b:bool) => if b then P1 else P2) true); assumption.
  apply (weakpredmonad_satisfiesP_monotone m m (reflexivity m) (andP P1 P2));
    [ | assumption ].
  apply (weakpredmonad_forallP_lower_bound
           (fun (b:bool) => if b then P1 else P2) false); assumption.
Qed.

(* FIXME HERE *)
(*
Theorem impliesP_implies {A} m (P1 P2: PM A) :
  m |= impliesP P1 P2 -> m |= P1 -> m |= P2.
  intros.
*)

End WeakPredMonad_thms.


Inductive RoseTree (A : Type) : Type :=
| L : A -> RoseTree A
| N : list (RoseTree A) -> RoseTree A.

Scheme RoseTree_ind' := Induction for RoseTree Sort Prop.

Check RoseTree_ind'.


Require Import List.
Import ListNotations.
Require Import Nat.
Require Import Coq.Program.Wf.
Require Import Coq.omega.Omega.
Require Import FSets FSets.FMapAVL FSets.FMapFacts.
Require Import Ascii.

From Verbatim Require Import ltac.
From Verbatim Require Import Orders.

Module Type SIGMA.
  Parameter Sigma : Type.
  Parameter SigmaEnum : list Sigma.
  Parameter Sigma_finite : forall a, In a SigmaEnum.
  Parameter Sigma_dec : forall(a a' : Sigma), {a = a'} + {a <> a'}.

  Parameter compareT : Sigma -> Sigma -> comparison.
  Parameter compareT_eq : forall x y : Sigma,
      compareT x y = Eq <-> x = y.

  Parameter compareT_trans : forall c x y z,
      compareT x y = c -> compareT y z = c -> compareT x z = c.

  (* This is a bit superfluous and probably needs some validation unless its identity *)
  Parameter ascii2Sigma : ascii -> Sigma.

End SIGMA.

Module DefsFn (Ty : SIGMA).

  Import Ty.

  Module T_as_UCT <: UsualComparableType.
    Definition t             := Sigma.
    Definition compare       := compareT.
    Definition compare_eq    := compareT_eq.
    Definition compare_trans := compareT_trans.
  End T_as_UCT.

  Module T_as_UOT <: UsualOrderedType := UOT_from_UCT T_as_UCT.

  Module SigFS := FSetAVL.Make T_as_UOT.
  Module SigFSF := FSetFacts.Facts SigFS.

  Module SigFM := FMapAVL.Make T_as_UOT.
  Module SigFMF := FMapFacts.Facts SigFM.

  Module Export Strings.

    Set Warnings "-implicit-core-hint-db,-deprecated".
    Hint Resolve Sigma_dec.

    Definition String : Type := list Sigma.

    Lemma String_dec : forall s s' : String, {s = s'} + {s <> s'}.
    Proof. decide equality. Qed.

    (** * *Should deprecate/move String lemmas *)
    Lemma nil_right : forall(s : String), s ++ [] = s.
    Proof.
      induction s.
      - simpl. reflexivity.
      - simpl. rewrite IHs. reflexivity.
    Qed.

    Lemma nil_left : forall(s : String), [] ++ s = s.
    Proof.
      intros s. simpl. reflexivity.
    Qed.

    Lemma empty_app : forall(s1 s2 : String), s1 ++ s2 = [] -> s1 = [] /\ s2 = [].
    Proof.
      intros s1 s2 H. destruct s1; destruct s2; try(simpl in H; discriminate).
      - split; reflexivity.
    Qed.

    Definition rm_empty (yss : list String) :=
      filter (fun l => match l with | [] => false | _ => true end) yss.

    Lemma rm_empty_mute : forall(yss : list String),
        concat (rm_empty yss) = concat yss.
    Proof.
      intros yss. induction yss.
      - simpl. reflexivity.
      - simpl. destruct a.
        + simpl. apply IHyss.
        + simpl. rewrite IHyss. reflexivity.
    Qed.

    Lemma rm_empty_no_empty : forall(ys : String) (yss : list String),
        In ys (rm_empty yss) -> ys <> [].
    Proof.
      intros ys yss H. unfold not. intros C. rewrite C in H.
      unfold rm_empty in H. induction yss.
      - simpl in H. contradiction.
      - simpl in H. destruct a.
        + apply IHyss in H. contradiction.
        + simpl in H. destruct H.
          * discriminate.
          * apply IHyss in H. contradiction.
    Qed.

    Lemma filtered_subset : forall(xs : String) (yss : list String),
        In xs (rm_empty yss) -> In xs yss.
    Proof.
      intros xs yss. generalize dependent xs. induction yss; intros xs H.
      - simpl in H. contradiction.
      - simpl. destruct a.
        + simpl in H. right. apply IHyss. apply H.
        + simpl in H. destruct H.
          * left. apply H.
          * right. apply IHyss. apply H.
    Qed.

  End Strings.

  Module Export Regexes.

    Inductive regex : Type :=
    | EmptySet
    | EmptyStr
    | Char (t : Sigma)
    | App (r1 r2 : regex)
    | Union (r1 r2 : regex)
    | Star (r : regex).

    Lemma regex_dec : forall r r' : regex, {r = r'} + {r <> r'}.
    Proof. decide equality. Qed.

    (** * *Should rename to re_equal *)
    Fixpoint regex_eq (r1 r2 : regex) : bool :=
      match r1, r2 with
      | EmptyStr, EmptyStr => true
      | EmptySet, EmptySet => true
      | Char a, Char b => if (Ty.Sigma_dec a b) then true else false
      | App x1 y1, App x2 y2 => andb (regex_eq x1 x2) (regex_eq y1 y2)
      | Union x1 y1, Union x2 y2 => andb (regex_eq x1 x2) (regex_eq y1 y2)
      | Star a, Star b => regex_eq a b
      | _, _ => false
      end.

    Lemma regex_eq_correct : forall r1 r2,
        r1 = r2 <-> regex_eq r1 r2 = true.
    Proof.
      induction r1; intros r2; split; intros H; subst; try(auto);
        try(unfold regex_eq in H; repeat dmh; subst; auto; discriminate).
      - simpl. dmg; [reflexivity | contradiction].
      - assert(A1 : regex_eq r1_1 r1_1 = true). apply IHr1_1; reflexivity.
        assert(A2 : regex_eq r1_2 r1_2 = true). apply IHr1_2; reflexivity.
        simpl. rewrite A1. rewrite A2. auto.
      - destruct r2; simpl in H; try(discriminate).
        destruct (regex_eq r1_1 r2_1) eqn:E1; destruct (regex_eq r1_2 r2_2) eqn:E2; try(discriminate).
        apply IHr1_1 in E1. apply IHr1_2 in E2. subst. reflexivity.
      - assert(A1 : regex_eq r1_1 r1_1 = true). apply IHr1_1; reflexivity.
        assert(A2 : regex_eq r1_2 r1_2 = true). apply IHr1_2; reflexivity.
        simpl. rewrite A1. rewrite A2. auto.
      - destruct r2; simpl in H; try(discriminate).
        destruct (regex_eq r1_1 r2_1) eqn:E1; destruct (regex_eq r1_2 r2_2) eqn:E2; try(discriminate).
        apply IHr1_1 in E1. apply IHr1_2 in E2. subst. reflexivity.
      - assert(A : regex_eq r1 r1 = true). apply IHr1; reflexivity.
        simpl. apply A.
      - destruct r2; simpl in H; try(discriminate). apply IHr1 in H. subst; reflexivity.
    Qed.

    Lemma regex_eq_refl : forall e, regex_eq e e = true.
      intros. induction e; auto.
      - simpl. dm.
      - simpl. apply Bool.andb_true_iff. split; auto.
      - simpl. apply Bool.andb_true_iff. split; auto.
    Qed.

    Fixpoint nullable' (r : regex) : bool:=
      match r with
      | EmptySet => false
      | EmptyStr => true
      | Char _ => false
      | App r1 r2 => andb (nullable' r1) (nullable' r2)
      | Union r1 r2 => orb (nullable' r1) (nullable' r2)
      | Star _ => true
      end.

    Fixpoint nullable (r : regex) : bool:=
      match r with
      | EmptySet => false
      | EmptyStr => true
      | Char _ => false
      | App r1 r2 => if negb (nullable r2) (* short circuit and *)
                    then false
                    else (nullable r1)
      | Union r1 r2 => if (nullable r2)    (* short circuit or *)
                      then true
                      else (nullable r1)
      | Star _ => true
      end.

    Fixpoint derivative (a : Sigma) (r : regex) :=
      match r with
      | EmptySet => EmptySet
      | EmptyStr => EmptySet
      | Char x => if Ty.Sigma_dec x a then EmptyStr else EmptySet
      | App r1 r2 => if (nullable r1) then Union (App (derivative a r1) r2) (derivative a r2)
                    else (App (derivative a r1) r2)
      | Union r1 r2 => Union (derivative a r1) (derivative a r2)
      | Star r => App (derivative a r) (Star r)
      end.

    Fixpoint derivative_list (bs : list Sigma) (e : regex) :=
      match bs with
      | [] => e
      | c :: cs => derivative_list cs (derivative c e)
      end.

    Fixpoint re_compare (e1 e2 : regex) : comparison :=
      match e1, e2 with
      | EmptyStr, EmptyStr => Eq
      | EmptyStr, _ => Lt
      | _, EmptyStr => Gt
      | EmptySet, EmptySet => Eq
      | EmptySet, _ => Lt
      | _, EmptySet => Gt
      | Char a, Char b => compareT a b
      | Char _, _ => Lt
      | _, Char _ => Gt
      | App e1 e2, App e3 e4 =>
        match re_compare e1 e3 with
        | Eq => re_compare e2 e4
        | comp => comp
        end
      | App _ _, _ => Lt
      | _, App _ _ => Gt
      | Star e1, Star e2 => re_compare e1 e2
      | Star _, _ => Lt
      | _, Star _ => Gt
      | Union e1 e2, Union e3 e4 =>
        match re_compare e1 e3 with
        | Eq => re_compare e2 e4
        | comp => comp
        end
      end.

    Lemma re_compare_eq : forall x y : regex, re_compare x y = Eq <-> x = y.
    Proof.
      induction x; destruct y; split;
        try(reflexivity);
        try(simpl; intros; discriminate).
      - simpl. intros. apply compareT_eq in H. subst. auto.
      - simpl. intros. injection H; intros; subst. apply compareT_eq. auto.
      - intros. specialize (IHx1 y1). specialize (IHx2 y2). simpl in H.
        destruct (re_compare x1 y1) eqn:E; try(discriminate).
        destruct IHx1. destruct H0. auto. destruct (re_compare x2 y2); try(discriminate).
        destruct IHx2. destruct H0; auto.
      - intros. injection H. intros; subst. simpl.
        destruct (regex_eq y1 y1) eqn:E.
        + apply regex_eq_correct in E. apply IHx1 in E. rewrite E.
          destruct (regex_eq y2 y2) eqn:E1.
          * apply regex_eq_correct in E1. apply IHx2 in E1. auto.
          * apply false_not_true in E1. destruct E1. apply regex_eq_correct. auto.
        + apply false_not_true in E. destruct E. apply regex_eq_correct. auto.
      - intros. specialize (IHx1 y1). specialize (IHx2 y2). simpl in H.
        destruct (re_compare x1 y1) eqn:E; try(discriminate).
        destruct IHx1. destruct H0. auto. destruct (re_compare x2 y2); try(discriminate).
        destruct IHx2. destruct H0; auto.
      - intros. injection H. intros; subst. simpl.
        destruct (regex_eq y1 y1) eqn:E.
        + apply regex_eq_correct in E. apply IHx1 in E. rewrite E.
          destruct (regex_eq y2 y2) eqn:E1.
          * apply regex_eq_correct in E1. apply IHx2 in E1. auto.
          * apply false_not_true in E1. destruct E1. apply regex_eq_correct. auto.
        + apply false_not_true in E. destruct E. apply regex_eq_correct. auto.
      - intros. specialize (IHx y). simpl in H. apply IHx in H. subst. auto.
      - intros. injection H. intros. subst. simpl. apply IHx. auto.
    Qed.

    Lemma re_compare_eq' : forall x,
        re_compare x x = Eq.
    Proof.
      intros. apply re_compare_eq. auto.
    Qed.

    Lemma re_compare_trans : forall c x y z,
        re_compare x y = c -> re_compare y z = c -> re_compare x z = c.
    Proof.
      induction x; destruct y; destruct z; intros; auto;
        try(simpl in *; subst; discriminate).
      - simpl in *. eapply compareT_trans; eauto.
      - simpl in *.
        destruct(re_compare x1 y1) eqn:E; destruct(re_compare x2 y2) eqn:E0;
          destruct(re_compare y1 z1) eqn:E1; destruct(re_compare y2 z2) eqn:E2;
            try(rewrite re_compare_eq in *; subst; repeat rewrite re_compare_eq' in *; auto);
            try(discriminate);
            try(specialize (IHx2 y2 z2); apply IHx2; auto);
            try(rewrite E0; rewrite E1; auto);
            try(rewrite E1; auto);
            try(rewrite E; auto);
            try(specialize (IHx1 y1 z1); rewrite IHx1; [reflexivity|auto|auto]);
            try(subst; specialize (IHx1 y1 z1); specialize (IHx2 y2 z2);
                rewrite IHx1; [reflexivity|auto|auto]); try(discriminate).
      - simpl in *.
        destruct(re_compare x1 y1) eqn:E; destruct(re_compare x2 y2) eqn:E0;
          destruct(re_compare y1 z1) eqn:E1; destruct(re_compare y2 z2) eqn:E2;
            try(rewrite re_compare_eq in *; subst; repeat rewrite re_compare_eq' in *; auto);
            try(discriminate);
            try(specialize (IHx2 y2 z2); apply IHx2; auto);
            try(rewrite E0; rewrite E1; auto);
            try(rewrite E1; auto);
            try(rewrite E; auto);
            try(specialize (IHx1 y1 z1); rewrite IHx1; [reflexivity|auto|auto]);
            try(subst; specialize (IHx1 y1 z1); specialize (IHx2 y2 z2);
                rewrite IHx1; [reflexivity|auto|auto]); try(discriminate).
      - simpl in *. eapply IHx; eauto.
    Qed.

  End Regexes.

  Module regex_as_UCT <: UsualComparableType.
    Definition t := regex.
    Definition compare := re_compare.
    Definition compare_eq := re_compare_eq.
    Definition compare_trans := re_compare_trans.
  End regex_as_UCT.

  Module regex_as_UOT <: UsualOrderedType := UOT_from_UCT regex_as_UCT.

  Module reFS := FSetAVL.Make regex_as_UOT.
  Module reFSF := FSetFacts.Facts reFS.

  Module reFM := FMapAVL.Make regex_as_UOT.
  Module reFMF := FMapFacts.Facts reFM.

  Module Export MatchSpec.

    Inductive exp_match : String -> regex -> Prop :=
    | MEmpty : exp_match [] EmptyStr
    | MChar x : exp_match [x] (Char x)
    | MApp s1 re1 s2 re2
           (H1 : exp_match s1 re1)
           (H2 : exp_match s2 re2) :
        exp_match (s1 ++ s2) (App re1 re2)
    | MUnionL s1 re1 re2
              (H1 : exp_match s1 re1) :
        exp_match s1 (Union re1 re2)
    | MUnionR re1 s2 re2
              (H2 : exp_match s2 re2) :
        exp_match s2 (Union re1 re2)
    | MStar0 re : exp_match [] (Star re)
    | MStarApp s1 s2 re
               (H1 : exp_match s1 re)
               (H2 : exp_match s2 (Star re)) :
        exp_match (s1 ++ s2) (Star re).

    Definition re_equiv (e1 e2 : regex) : Prop :=
      forall z, exp_match z e1 <-> exp_match z e2.

  End MatchSpec.

  Module Export MatchSpecLemmas.

    Theorem nullable_bridge : forall(r : regex),
      true = nullable r <-> exp_match [] r.
    Proof.
      intros r. split; intros H.
      - induction r; try(simpl in H; discriminate).
        + apply MEmpty.
        + simpl in H. destruct (nullable r1); destruct (nullable r2); try(simpl in H; discriminate).
          rewrite <- nil_left with (s := []). apply MApp.
          * apply IHr1. reflexivity.
          * apply IHr2. reflexivity.
        + simpl in H.  destruct (nullable r1); destruct (nullable r2).
          * apply MUnionL. apply IHr1. reflexivity.
          * apply MUnionL. apply IHr1. reflexivity.
          * apply MUnionR. apply IHr2. reflexivity.
          * simpl in H. discriminate.
        + apply MStar0.
      - induction r.
        + inversion H.
        + simpl. reflexivity.
        + inversion H.
        + simpl. inversion H. apply empty_app in H1. destruct H1.
          rewrite H1 in H3. rewrite H5 in H4.
          apply IHr1 in H3. apply IHr2 in H4.
          rewrite <- H3. rewrite <- H4. simpl. reflexivity.
        + simpl. inversion H.
          * apply IHr1 in H2. rewrite <- H2. destruct (nullable r2); simpl; reflexivity.
          * apply IHr2 in H1. rewrite <- H1. destruct (nullable r1); simpl; reflexivity.
        + simpl. reflexivity.
    Qed.

    Theorem nullable_bridge' : forall(r : regex),
        nullable r = true <-> exp_match [] r.
    Proof.
      split; intros.
      - symmetry in H. apply nullable_bridge; auto.
      - symmetry. apply nullable_bridge; auto.
    Qed.

    (* star_concat upto concat_star necessary for hard star case of der_match *)
    Lemma star_concat :
      forall s r',
        exp_match s (Star r')
        -> (exists xss : list (list Sigma),
              s = concat xss
              /\ (forall xs,
                    In xs xss
                    -> exp_match xs r')).
    Proof.
      intros s r' hm.
      remember (Star r') as r. generalize dependent r'.
      induction hm; intros r' heq; inv heq.
      - exists []; split; auto.
        intros xs hi; inv hi.
      - destruct (IHhm2 r') as (xss' & heq & hall); subst; auto.
        exists (s1 :: xss'); split; auto.
        intros xs hi.
        destruct hi as [hh| ht]; subst; auto.
    Qed.

    Lemma star_concat_no_empt : forall(s : String) (r' : regex),
        s <> []
        -> exp_match s (Star r')
        -> (exists xss : list (list Sigma),
              s = concat xss
              /\ (forall xs,
                    In xs xss
                    -> exp_match xs r' /\ xs <> [])).
    Proof.
      intros s r' Hempty Hstar. apply star_concat in Hstar.
      destruct Hstar as (yss & heq & hall).
      exists(rm_empty yss). split.
      - rewrite rm_empty_mute. apply heq.
      - intros xs H. split.
        + apply hall.  apply filtered_subset in H. apply H.
        + apply rm_empty_no_empty in H. apply H.
    Qed.

    Lemma concat_star : forall(xss : list String) (r : regex),
        (forall xs : list Sigma, In xs xss -> exp_match xs r) -> exp_match (concat xss) (Star r).
    Proof.
      intros xss r H. induction xss.
      - simpl. apply MStar0.
      - replace (concat (a :: xss)) with (a ++ (concat xss)).
        + apply MStarApp.
          * apply H. simpl. left. reflexivity.
          * apply IHxss. intros xs H1. apply H. simpl. right. apply H1.
        + simpl. reflexivity.
    Qed.

    Theorem der_match : forall(a : Sigma) (s : String) (r : regex),
        exp_match (a::s) r <-> exp_match s (derivative a r).
    Proof.
      intros a s r.
      split.
      {
        generalize dependent s. induction r; intros s H.
        - inv H.
        - inv H.
        - destruct s.
          + simpl. destruct (Sigma_dec t a).
            * apply MEmpty.
            * inv H. contradiction.
          + inv H.
        - simpl. destruct(nullable r1) eqn:E.
          + inv H. destruct s1.
            * apply MUnionR. apply IHr2. rewrite <- H1. simpl. apply H4.
            * apply MUnionL. simpl in H1. injection H1. intros Happ Hchar.
              rewrite <- Happ. rewrite Hchar in H3. apply IHr1 in H3.
              apply MApp.
              -- apply H3.
              -- apply H4.
          + inv H. destruct s1.
            * apply nullable_bridge in H3. rewrite E in H3. discriminate.
            * simpl in H1. injection H1. intros Happ Hchar.
              rewrite <- Happ. rewrite Hchar in H3. apply IHr1 in H3.
              apply MApp.
              -- apply H3.
              -- apply H4.
        - simpl. inv H.
          + apply IHr1 in H2. apply MUnionL. apply H2.
          + apply IHr2 in H1. apply MUnionR. apply H1.
        (* hard_star: This case was the hard one *)
        - apply star_concat_no_empt in H. destruct H as (xss & heq & hall).
          + assert (H : exists(s1 : String) (yss : list String),
                       ((a :: s1) :: yss) = xss).
            {
              destruct xss.
              - simpl in heq. discriminate.
              - simpl in heq. destruct l eqn:E.
                + apply hall in E.
                  * contradiction.
                  * rewrite E. simpl. left. reflexivity.
                + exists(l0). exists(xss). simpl in heq.
                  injection heq. intros I1 I2. rewrite I2. reflexivity.
            }
            destruct H as [s1]. destruct H as [yss].
            rewrite <- H in hall.
            assert (A : In (a :: s1) ((a :: s1) :: yss)).
            { simpl. left. reflexivity. }
            simpl. replace s with (s1 ++ (concat yss)).
            * apply MApp.
              -- apply IHr. apply hall in A. destruct A. apply H0.
              -- rewrite H in hall.
                 assert (A1 : forall xs : list Sigma, In xs xss -> exp_match xs r).
                 { intros xs. intros HA1. apply hall in HA1. destruct HA1. apply H0. }
                 apply concat_star. intros xs H1. apply A1.
                 rewrite <- H. simpl. right. apply H1.
            * assert (A1 : concat ((a :: s1) :: yss) = concat xss).
              { rewrite H. reflexivity. }
              assert (A2 : concat ((a :: s1) :: yss) = a :: (s1 ++ (concat yss))).
              { simpl. reflexivity. }
              rewrite <- A1 in heq. rewrite A2 in heq. injection heq.
              intros I. symmetry. apply I.
          + unfold not. intros C. discriminate.
      }
      {
        generalize dependent s. induction r; intros s H.
        - inv H.
        - inv H.
        - simpl in H. destruct (Sigma_dec t a); inv H. apply MChar.
        - simpl in H. destruct (nullable r1) eqn:E.
          + inv H.
            * inv H2. replace (a :: s1 ++ s2) with ((a :: s1) ++ s2).
              apply MApp.
              -- apply IHr1. apply H3.
              -- apply H4.
              -- reflexivity.
            * symmetry in E. apply nullable_bridge in E. rewrite <- nil_left with (s := (a :: s)).
              apply MApp.
              -- apply E.
              -- apply IHr2. apply H1.
          + inv H.
            * replace (a :: s1 ++ s2) with ((a :: s1) ++ s2).
              apply MApp.
              -- apply IHr1. apply H3.
              -- apply H4.
              -- reflexivity.
        - inv H.
          + apply MUnionL. apply IHr1. apply H2.
          + apply MUnionR. apply IHr2. apply H1.
        - simpl in H. inv H. replace (a :: s1 ++ s2) with ((a :: s1) ++ s2).
          + apply MStarApp.
            * apply IHr. apply H3.
            * apply H4.
          + reflexivity.
      }
    Qed.

    Lemma derivative_list_cons : forall bs b e,
        derivative_list (b :: bs) e = derivative_list bs (derivative b e).
    Proof.
      auto.
    Qed.

    Lemma derivative_list_str : forall bs e,
        exp_match [] (derivative_list bs e) <-> exp_match bs e.
    Proof.
      induction bs; intros; auto.
      - simpl. split; auto.
      - simpl. rewrite IHbs. split; apply der_match.
    Qed.

  End MatchSpecLemmas.

  Module Notations.

    Notation "e1 = e2 " := (regex_eq e1 e2) (at level 70).
    Notation "a ∂ r" := (derivative a r) (at level 75, right associativity).
    Notation ε := EmptyStr.
    Notation "e1 @ e2" := (App e1 e2) (at level 75).
    Notation "e1 # e2" := (Union e1 e2) (at level 75).

  End Notations.

  Module Helpers.

    Definition Plus (r : regex) : regex := App r (Star r).

    Fixpoint IterUnion (rs : list regex) :=
      match rs with
      | [] => EmptySet
      | [e] => e
      | h::t => Union h (IterUnion t)
      end.

    Fixpoint IterApp (rs : list regex) :=
      match rs with
      | [] => EmptyStr
      | [e] => e
      | h::t => App h (IterApp t)
      end.

    (* r? *)
    Definition Optional (r : regex) := Union EmptyStr r.

    Definition REString (z : String) := IterApp (map Char z).

  End Helpers.

End DefsFn.

Module Type DefsT (Ty : SIGMA).
  Include DefsFn Ty.
End DefsT.

Module Type T.
  Declare Module Ty : SIGMA.
  Declare Module Defs  : DefsT Ty.
  Export Defs.
End T.

Module regexTFn (Ty' : SIGMA) <: T.

  Module Export Ty := Ty'.
  Module Export Defs := DefsFn Ty.

End regexTFn.


Require Import compcert.lib.Coqlib.
Require Import compcert.lib.Integers.
Require Import List. Import ListNotations.
Require Import sha.general_lemmas.
Require Import sha.hmac_pure_lemmas.
Require Import sha.SHA256.
Require Import sha.pure_lemmas.
Require Import sha.functional_prog.
Require Import Lia.

(* Lemma 1: M = Prefix(Pad(M)) *)

(* TODO: replace InWords with InBlocks 4? *)
Inductive InWords : list byte -> Prop :=
  | words_nil : InWords []
  | words_word : forall (a b c d : byte) (msg : list byte),
                   InWords msg -> InWords (a :: b :: c :: d :: msg).

(* *** New definition for this lemma. *)
Definition pad (msg : list byte) : list byte :=
  let n := Zlength msg in
  msg ++ [Byte.repr 128%Z]
      ++ repeat Byte.zero (Z.to_nat (-(n + 9) mod 64))
      ++ intlist_to_bytelist (([Int.repr (n * 8 / Int.modulus); Int.repr (n * 8)])%list).

Definition generate_and_pad' (msg : list byte) : list int :=
  bytelist_to_intlist (pad msg).

(* ----------------- ^ Definitions *)

Lemma fstpad_len :
  forall (msg : list byte),
    Datatypes.length (msg ++ [Byte.repr 128]
                 ++ repeat Byte.zero (Z.to_nat (- (Zlength msg + 9) mod 64)))
= (Datatypes.length msg + (S (Z.to_nat (- (Zlength msg + 9) mod 64))))%nat.
Proof.
  intros msg.
  simpl.
  rewrite -> app_length.
  simpl.
  rewrite -> repeat_length.
  reflexivity.
Qed.

Lemma InWords_len4 : forall (l : list byte),
                       Nat.divide (Z.to_nat WORD) (length l) -> InWords l.
Proof.
  intros l [x H].
  revert l H.
  induction x.
  intros l H. simpl in H.
  destruct l.
    apply words_nil.
    simpl in H. inversion H.
  intros l H.
  destruct l as [ | a [ | b [ | c [ | d ? ]]]].
    inversion H.
    inversion H.
    inversion H.
    inversion H.
    specialize (IHx l).
      apply words_word.
      apply IHx.
      simpl in H. inversion H.
      simpl. apply H1.
Qed.

(* TODO: clear out the SearchAbouts / clean up proof *)
Lemma pad_len_64_mod : forall (msg : list byte),
                           (Zlength (pad msg)) mod 64 = 0.
Proof.
  intros msg.
  unfold pad.
  rewrite -> Zlength_correct.
  repeat rewrite -> app_length.
  simpl.
  assert (succ: forall (n : nat), S n = (n + 1)%nat).
    intros. induction n. reflexivity. lia.
  rewrite -> succ.
  assert ((length msg +
      (length (repeat Byte.zero (Z.to_nat (- (Zlength msg + 9) mod 64))) + 8 +
       1))%nat = (length msg +
      (length (repeat Byte.zero (Z.to_nat (- (Zlength msg + 9) mod 64))) + 9))%nat) by lia.
  rewrite -> H. clear H.

  rewrite -> Zlength_correct.
  rewrite -> repeat_length.

  repeat rewrite -> Nat2Z.inj_add.
  rewrite -> Z2Nat.id.

  assert (move : forall (a b c : Z), a + (b + c) = (a + c) + b).
  intros. lia.

  rewrite -> move.
  rewrite -> Zplus_mod_idemp_r.

  assert (Z_9 : 9 = Z.of_nat (9%nat)). reflexivity.
  rewrite -> Z_9.

  repeat rewrite <- Nat2Z.inj_add.

  assert (forall (x : Z), x + (-x) = 0). intros. lia.

  rewrite -> H.
  reflexivity.

  * apply Z.mod_pos_bound.
    lia.
Qed.

(* more usable versions *)
Lemma pad_len_64 : forall (msg : list byte), exists (n : Z),
                           Zlength (pad msg) = 64 * n /\ n >= 0.
Proof.
  intros msg.
  pose proof pad_len_64_mod msg as pad_len_mod.
  rewrite -> Zmod_divides in *. 2: lia.

  destruct pad_len_mod.
  exists x.
  split.
  apply H.
  specialize (Zlength_nonneg (pad msg)); intros. lia.
Qed.

Lemma pad_len_64_nat : forall (msg : list byte), exists (n : nat),
                           (length (pad msg))%nat = (64 * n)%nat.
Proof.
  intros msg.
  pose proof pad_len_64 msg as pad_len_64.

  rewrite -> Zlength_correct in *.
  destruct pad_len_64.
  exists (Z.to_nat x).
  destruct H.

  assert (app_each : Z.to_nat (Z.of_nat (length (pad msg))) = Z.to_nat (64 * x)).
    rewrite -> H. reflexivity.

  rewrite -> Nat2Z.id in app_each.

  rewrite -> app_each.
(*  SearchAbout (Z.to_nat (_ * _)).*)
  rewrite -> Z2Nat.inj_mul.
  assert (n_64 : Z.to_nat 64 = 64%nat). reflexivity.

  rewrite -> n_64.
  reflexivity.

  * lia.
  * lia.
Qed.

Lemma total_pad_len_Zlist : forall (msg : list byte), exists (n : nat),
     length
       (msg ++ [Byte.repr 128] ++ repeat Byte.zero (Z.to_nat (- (Zlength msg + 9) mod 64)))
     =  (n * Z.to_nat WORD (* 4 *))%nat.
Proof.
  intros msg.
  pose proof pad_len_64_nat msg as pad_len_64_nat.

  unfold pad in *.
  repeat rewrite -> app_length in *.
  destruct pad_len_64_nat.
  assert (sym: (64 * x)%nat = (x * 64)%nat) by lia.
  rewrite -> sym in *. clear sym.

  simpl in *.
  assert (Pos.to_nat 4 = 4%nat) by reflexivity.
  (*rewrite -> H0. clear H0.*)

  rewrite -> repeat_length in *.

  assert (add_both: (length msg + S (Z.to_nat (- (Zlength msg + 9) mod 64) ))%nat =
      (x * 64 - 8)%nat) by lia. clear H.

  rewrite -> add_both.
  assert ((x * 64 - 8)%nat = (4 * (16 * x - 2))%nat) by lia.

  rewrite -> H.
  exists (16 * x - 2)%nat.
  lia.
Qed.

Lemma pad_inwords :
  forall (msg : list byte),
    InWords (msg ++ [Byte.repr 128]
                 ++ repeat Byte.zero (Z.to_nat (- (Zlength msg + 9) mod 64))).
Proof.
  intros msg.
  apply InWords_len4.
  pose proof total_pad_len_Zlist.
  specialize (H msg).
  unfold Nat.divide.
  apply H.
Qed.

Definition fulllen (len : Z) :=
  len + 1%Z + (- (len + 9) mod 64).

Lemma app_left : forall (a b c d : list byte),
   a ++ b ++ c ++ d = (a ++ b ++ c) ++ d.
(* a ++ (b ++ (c ++ d)) = (a ++ (b ++ c)) ++ d *)
Proof.
   intros a b c d.
   assert (b ++ (c ++ d) = (b ++ c) ++ d) as assert1.
     rewrite -> app_assoc. reflexivity.
   rewrite -> assert1.
   rewrite -> app_assoc.
   reflexivity.
Qed.

(* can use extensionality *)
Theorem pad_compose_equal : forall (msg : list byte),
                              generate_and_pad' msg = generate_and_pad msg.
Proof.
  intros msg.
  unfold generate_and_pad'.
  unfold pad.
  unfold generate_and_pad.
  (* need il => ZIL (IZL il), and
     ZIL a ++ Zil b = ZIL (a ++ b) (with length a being a multiple of 4)
   *)
  pose proof pad_inwords as pad_inwords.
  specialize (pad_inwords msg).
  rewrite -> app_left.
  induction pad_inwords.
  (* case none *)
    assert (forall l : list byte, [] ++ l = l) as Happend. reflexivity.
    specialize (Happend (intlist_to_bytelist
        [Int.repr (Zlength msg * 8 / Int.modulus);
        Int.repr (Zlength msg * 8)])).
    rewrite -> Happend.
    rewrite -> intlist_to_bytelist_to_intlist.
    reflexivity.
  (* case a :: b :: c :: d :: msg0 *)
    Opaque intlist_to_bytelist.
    simpl.
    apply f_equal.
    apply IHpad_inwords.
Qed.

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

(* Lemma 2: |M1| = |M2| -> |Pad(M1)| = |Pad(M2)| *)

Theorem length_equal_pad_length : forall (msg1 : list byte) (msg2 : list byte),
     Zlength msg1  = Zlength msg2 ->
     Zlength (generate_and_pad msg1) = Zlength (generate_and_pad msg2).
Proof.
  intros m1 m2 H.
  repeat rewrite -> functional_prog.length_generate_and_pad.
  rewrite -> H.
  reflexivity.
Qed.

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

(* Lemma 3: |M1| =/= |M2| ->
last block of Pad(M1) =/= last block of Pad(M2)

or, if one-to-one property is desired (for HMAC), only need to prove that
the padded messages differ
*)


Require Export ComponentSM2.
Require Export ComponentSM4.
Require Export DTimeQ.

Inductive exmsg :=
| ADD1 (n : nat)
| ADD2 (n : nat)
| TOT  (n : nat).

Global Instance EX1Node : Node := MkNode nat deq_nat.
Global Instance EX1Key : Key := MkKey nat nat.
Global Instance EX1Msg : Msg := MkMsg exmsg.

Definition bstatefun (cn : CompName) : Type :=
  if String.string_dec (comp_name_kind cn) "STATE" then nat else unit.

Global Instance EX1baseFunIO : baseFunIO := MkBaseFunIO (fun _ => CIOnat).
Global Instance EX1baseStateFun : baseStateFun := MkBaseStateFun bstatefun.
Global Instance EX1IOTrustedFun : IOTrustedFun := MkIOTrustedFun (fun _ => MkIOTrusted nat nat 0).
Global Instance EX1trustedStateFun : trustedStateFun := MkTrustedStateFun (fun _ => nat).

(* ====== STATE ====== *)

Definition STATEname : CompName := MkCN "STATE" 0 true.

Definition STATEupdate : M_Update 0 STATEname _ :=
  fun (s : nat) i =>
    interp_s_proc
      ([R] (s + i, s + i)).

Definition STATE : n_proc _ _ :=
  build_m_sm STATEupdate 0.

(* ====== OP1 ====== *)

Definition OP1name : CompName := MkCN "OP" 1 false.

Definition OP1update : M_Update 1 OP1name _ :=
  fun s i =>
    interp_s_proc
      ((STATEname [C] i)
       [>>=] fun out =>
       [R] (tt, out)).

Definition OP1 : n_proc _ _ :=
  build_m_sm OP1update tt.

(* ====== OP2 ====== *)

Definition OP2name : CompName := MkCN "OP" 2 false.

Definition OP2update : M_Update 1 OP2name _ :=
  fun s i =>
    interp_s_proc
      ((STATEname [C] i)
       [>>=] fun _ =>
       (STATEname [C] i)
       [>>=] fun out =>
       [R] (tt, out)).

Definition OP2 : n_proc _ _ :=
  build_m_sm OP2update tt.

(* ====== Main ====== *)

Definition MAINname := msg_comp_name 0.

Definition MAINupdate : M_Update 2 MAINname _ :=
  fun s i =>
    interp_s_proc
      ((match i with
        | ADD1 n => (OP1name [C] n)
        | ADD2 n => (OP2name [C] n)
        | TOT  n => [R] n
        end) [>>=] (fun out => [R] (tt, [MkDMsg (TOT out) [] ('0)]))).

Definition MAIN : n_proc _ _ :=
  build_m_sm MAINupdate tt.

(* ====== Local System ====== *)

Definition progs : n_procs _ :=
  [
    MkPProc _ MAIN,
    MkPProc _ (lift_n_proc 1 OP1),
    MkPProc _ (lift_n_proc 1 OP2),
    MkPProc _ (lift_n_proc 2 STATE)
  ].

Definition ls : LocalSystem 3 0 := progs.

Definition test1 := call_procs_out ls MAINname [ADD1 2, ADD2 3] (ADD1 1).
Eval compute in test1.