From Coq Require Import Bool.Bool.
From Coq Require Import Init.Nat.
From Coq Require Import Arith.Arith.
From Coq Require Import Arith.EqNat. Import Nat.
From Coq Require Import Lia.
From Coq Require Import Lists.List. Import ListNotations.
From Coq Require Import Strings.String.
From LF Require Import Maps.
Set Default Goal Selector "!".

Module AExp.

  Inductive aexp : Type :=
  | ANum (n : nat)
  | AId (x : string)
  | APlus (a1 a2 : aexp)
  | AMinus (a1 a2 : aexp)
  | AMult (a1 a2 : aexp).

  Definition W : string := "W".
  Definition X : string := "X".
  Definition Y : string := "Y".
  Definition Z : string := "Z".

  Inductive bexp : Type :=
  | BTrue
  | BFalse
  | BEq (a1 a2 : aexp)
  | BNeq (a1 a2 : aexp)
  | BLe (a1 a2 : aexp)
  | BGt (a1 a2 : aexp)
  | BNot (b : bexp)
  | BAnd (b1 b2 : bexp).

  Coercion AId : string >-> aexp.
  Coercion ANum : nat >-> aexp.
  Declare Custom Entry com.
  Declare Scope com_scope.
  Declare Custom Entry com_aux.
  Notation "<{ e }>" := e (e custom com_aux) : com_scope.
  Notation "e" := e (in custom com_aux at level 0, e custom com) : com_scope.
  Notation "( x )" := x (in custom com, x at level 99) : com_scope.
  Notation "x" := x (in custom com at level 0, x constr at level 0) : com_scope.
  Notation "f x .. y" := (.. (f x) .. y)
                           (in custom com at level 0, only parsing,
                               f constr at level 0, x constr at level 9,
                               y constr at level 9) : com_scope.
  Notation "x + y" := (APlus x y) (in custom com at level 50, left associativity).
  Notation "x - y" := (AMinus x y) (in custom com at level 50, left associativity).
  Notation "x * y" := (AMult x y) (in custom com at level 40, left associativity).
  Notation "'true'" := true (at level 1).
  Notation "'true'" := BTrue (in custom com at level 0).
  Notation "'false'" := false (at level 1).
  Notation "'false'" := BFalse (in custom com at level 0).
  Notation "x <= y" := (BLe x y) (in custom com at level 70, no associativity).
  Notation "x > y" := (BGt x y) (in custom com at level 70, no associativity).
  Notation "x = y" := (BEq x y) (in custom com at level 70, no associativity).
  Notation "x <> y" := (BNeq x y) (in custom com at level 70, no associativity).
  Notation "x && y" := (BAnd x y) (in custom com at level 80, left associativity).
  Notation "'~' b" := (BNot b) (in custom com at level 75, right associativity).
  Open Scope com_scope.

  Definition example_aexp : aexp := <{ 3 + (X * 2) }>.
  Definition example_bexp := <{ true && ~(X <= 4) }>.

  Definition state := total_map nat.

  Fixpoint aeval (st : state) (a : aexp) : nat :=
    match a with
    | ANum n => n
    | AId x => st x
    | <{a1 + a2}> => (aeval st a1) + (aeval st a2)
    | <{a1 - a2}> => (aeval st a1) - (aeval st a2)
    | <{a1 * a2}> => (aeval st a1) * (aeval st a2)
    end.

  Definition empty_st := (_ !-> 0).

  Example test_aeval1 :
    aeval empty_st <{2 + 2}> = 4.
  Proof. reflexivity. Qed.

  Fixpoint beval (st : state) (b : bexp) : bool :=
    match b with
    | <{true}> => true
    | <{false}> => false
    | <{a1 = a2}> => (aeval st a1) =? (aeval st a2)
    | <{a1 <> a2}> => negb ((aeval st a1) =? (aeval st a2))
    | <{a1 <= a2}> => (aeval st a1) <=? (aeval st a2)
    | <{a1 > a2}> => negb ((aeval st a1) <=? (aeval st a2))
    | <{~ b1}> => negb (beval st b1)
    | <{b1 && b2}> => andb (beval st b1) (beval st b2)
    end.

  Notation "x '!->' v" := (x !-> v ; empty_st) (at level 100).

  Example aexp1 :
    aeval (X !-> 5) <{ 3 + (X * 2) }> = 13.
  Proof.
    reflexivity. Qed.

  Inductive com : Type :=
  | CSkip
  | CAsgn (x : string) (a : aexp)
  | CSeq (c1 c2 : com)
  | CIf (b : bexp) (c1 c2 : com)
  | CWhile (b : bexp) (c : com).

  Notation "'skip'"  :=
    CSkip (in custom com at level 0) : com_scope.
  Notation "x := y"  :=
    (CAsgn x y)
      (in custom com at level 0, x constr at level 0,
          y at level 85, no associativity) : com_scope.
  Notation "x ; y" :=
    (CSeq x y)
      (in custom com at level 90,
          right associativity) : com_scope.
  Notation "'if' x 'then' y 'else' z 'end'" :=
    (CIf x y z)
      (in custom com at level 89, x at level 99,
          y at level 99, z at level 99) : com_scope.
  Notation "'while' x 'do' y 'end'" :=
    (CWhile x y)
      (in custom com at level 89, x at level 99,
          y at level 99) : com_scope.

  Reserved Notation "st '=[' c ']=>' st'"
    (at level 40, c custom com at level 99, st' constr at next level).

  Inductive ceval : com -> state -> state -> Prop :=
  | E_Skip : forall st,
      st =[ skip ]=> st
  | E_Asgn : forall st a n x,
      aeval st a = n -> st =[ x := a ]=> (x !-> n ; st)
  | E_Seq : forall c1 c2 st st' st'',
      st =[ c1 ]=> st' ->
      st' =[ c2 ]=> st'' ->
      st =[ c1 ; c2 ]=> st''
  | E_IfTrue : forall st st' b c1 c2,
      beval st b = true ->
      st =[ c1 ]=> st' ->
      st =[ if b then c1 else c2 end]=> st'
  | E_IfFalse : forall st st' b c1 c2,
      beval st b = false ->
      st =[ c2 ]=> st' ->
      st =[ if b then c1 else c2 end]=> st'
  | E_WhileFalse : forall b st c,
      beval st b = false ->
      st =[ while b do c end ]=> st
  | E_WhileTrue : forall st st' st'' b c,
      beval st b = true ->
      st =[ c ]=> st' ->
      st' =[ while b do c end ]=> st'' ->
      st =[ while b do c end ]=> st''

  where "st =[ c ]=> st'" := (ceval c st st').

  Example ceval_example1 :
    empty_st =[
        X := 2;
        if (X <= 1)
            then Y := 3
            else Z := 4
        end
    ]=> (Z !-> 4 ; X !-> 2).
  Proof.
    apply E_Seq with (X !-> 2).
    - apply E_Asgn. reflexivity.
    - apply E_IfFalse.
      + reflexivity.
      + apply E_Asgn.
        reflexivity.
  Qed.

  Example ceval_example2 :
    empty_st =[
        X := 0;
        Y := 1;
        Z := 2
      ]=> (Z !-> 2 ; Y !-> 1 ; X !-> 0).
  Proof.
    apply E_Seq with (X !-> 0).
    { apply E_Asgn. reflexivity. }
    apply E_Seq with (Y !-> 1 ; X !-> 0).
    { apply E_Asgn. reflexivity. }
    apply E_Asgn. reflexivity.
  Qed.

  Definition pup_to_n : com :=
    <{ while X <> 0 do
           Y := Y + X;
           X := X - 1
       end
      }>.

  Theorem pup_to_2_ceval :
    (X !-> 2) =[
        pup_to_n
      ]=> (X !-> 0 ; Y !-> 3 ; X !-> 1 ; Y !-> 2 ; X !-> 2).
  Proof.
    apply E_WhileTrue with (X !-> 1 ; Y !-> 2 ; X !-> 2).
    { reflexivity. }
    { apply E_Seq with (Y !-> 2 ; X !-> 2).
      { apply E_Asgn. reflexivity. }
      apply E_Asgn. reflexivity. }
    apply E_WhileTrue with (X !-> 0 ; Y !-> 3 ; X !-> 1 ; Y !-> 2 ; X !-> 2).
    { reflexivity. }
    { apply E_Seq with (Y !-> 3 ; X !-> 1 ; Y !-> 2 ; X !-> 2).
      { apply E_Asgn. reflexivity. }
      apply E_Asgn. reflexivity. }

    apply E_WhileFalse.
    reflexivity.
  Qed.

  Theorem ceval_deterministic : forall c st st1 st2,
      st =[ c ]=> st1 ->
      st =[ c ]=> st2 ->
      st1 = st2.
  Proof.
    intros c st st1 st2 E1 E2.
    generalize dependent st2.
    induction E1; intros st2 E2; inversion E2; subst.
    - reflexivity.
    - reflexivity.
    - rewrite (IHE1_1 st'0 H1) in *.
      apply IHE1_2.
      assumption.
    - apply IHE1. assumption.
    - rewrite H in H5. discriminate.
    - rewrite H in H5. discriminate.
    - apply IHE1. assumption.
    - reflexivity.
    - rewrite H in H2. discriminate.
    - rewrite H in H4. discriminate.
    - rewrite (IHE1_1 st'0 H3) in *.
      apply IHE1_2. assumption. Qed.

  Definition plus2 : com :=
    <{ X := X + 2 }>.

  Theorem plus2_spec : forall st n st',
      st X = n ->
      st =[ plus2 ]=> st' ->
      st' X = n + 2.
  Proof.
    intros st n st' HX Heval.
    inversion Heval.
    subst. clear Heval. simpl.
    apply t_update_eq.
  Qed.

  Definition XtimesYinZ : com :=
    <{ Z := X * Y }>.

  Theorem XtimesYinZ_spec : forall st n m st',
      st X = n ->
      st Y = m ->
      st =[ XtimesYinZ ]=> st' ->
      st' Z = n * m.
  Proof.
    intros st n m st' HX HY Heval.
    inversion Heval.
    subst.
    simpl.
    apply t_update_eq.
  Qed.

  Definition loop : com :=
    <{ while true do
         skip
       end }>.

  Theorem loop_never_stops : forall st st',
      ~(st =[ loop ]=> st').
  Proof.
    intros st st' contra.
    unfold loop in contra.
    remember <{ while true do skip end }> as loopdef eqn:Heqloopdef.
    induction contra ; try discriminate.
    - inversion Heqloopdef; subst.
      simpl in H.
      discriminate.
    - inversion Heqloopdef; subst.
      apply IHcontra2.
      reflexivity.
  Qed.

  Fixpoint no_whiles (c : com) : bool :=
    match c with
    | <{ skip }> => true
    | <{ _ := _ }> => true
    | <{ c1 ; c2 }> =>
        andb (no_whiles c1) (no_whiles c2)
    | <{ if _ then ct else cf end }> =>
        andb (no_whiles ct) (no_whiles cf)
    | <{ while _ do _ end }> => false
    end.

  Inductive no_whilesR : com -> Prop :=
  | no_whilesSkip : no_whilesR <{ skip }>
  | no_whilesAsgn (x : string) (a : aexp) : no_whilesR <{ x := a }>
  | no_whilesSeq (c1 c2 : com) (H1 : no_whilesR c1) (H2 : no_whilesR c2) :
    no_whilesR <{ c1 ; c2 }>
  | no_whilesIf (b : bexp) (ct cf : com) (H1 : no_whilesR ct) (H2 : no_whilesR cf) :
    no_whilesR <{ if b then ct else cf end }>.

  Theorem no_whiles_eqv:
    forall c, no_whiles c = true <-> no_whilesR c.
  Proof.
    intros c.
    split.
    - intros H.
      induction c;
        try (constructor ; reflexivity) ;
        try (constructor ;
             [ apply IHc1 ; reflexivity
             | apply IHc2 ; reflexivity ]) ;
        try (constructor ; [ apply IHc1 | apply IHc2 ] ;
             inversion H;
             apply andb_true_iff
               with (b1 := no_whiles c1)
                    (b2 := no_whiles c2)
               in H1 as [G1 G2];
             rewrite G1;
             rewrite G2;
             reflexivity).
      + simpl in H. discriminate.
    - intros H.
      induction H;
        try reflexivity;
        try (simpl; rewrite IHno_whilesR1;
             rewrite IHno_whilesR2; reflexivity).
  Qed.

  Theorem no_whiles_terminating : forall (st : state) (c : com),
      no_whilesR c -> exists (st' : state),
        st =[ c ]=> st'.
  Proof.
    intros st c H.
    generalize dependent st.
    induction c.
    - intros.
      exists st.
      constructor.
    - intros.
      exists (x !-> (aeval st a) ; st).
      constructor.
      reflexivity.
    - intros.
      inversion H.
      subst.
      apply IHc1 with (st := st) in H0.
      destruct H0 as [st' H1].
      apply IHc2 with (st := st') in H3.
      destruct H3 as [st'' H2].
      exists st''.
      apply (E_Seq c1 c2 st st' st'' H1 H2).
    - intros.
      inversion H.
      subst.
      destruct (beval st b) eqn:Eb ;
        [ apply IHc1 with (st := st) in H0 as G
        | apply IHc2 with (st := st) in H4 as G] ;
        destruct G as [st' G1] ;
        exists st' ;
        constructor ;
        assumption.
    - intros.
      inversion H.
  Qed.


  Inductive sinstr : Type :=
  | SPush (n : nat)
  | SLoad (x : string)
  | SPlus
  | SMinus
  | SMult.

  Fixpoint s_execute (st : state) (stack : list nat)
    (prog : list sinstr)
    : list nat :=
    match prog with
    | [] => stack
    | (SPush n :: rest) => s_execute st (n :: stack) rest
    | (SLoad x :: rest) => s_execute st (st x :: stack) rest
    | (SPlus :: rest) => match stack with
                         | x :: y :: ss =>
                             s_execute st ((y + x) :: ss) rest
                         | _ => s_execute st stack rest
                         end
    | (SMinus :: rest) => match stack with
                          | x :: y :: ss =>
                              s_execute st ((y - x) :: ss) rest
                          | _ => s_execute st stack rest
                          end
    | (SMult :: rest) => match stack with
                         | x :: y :: ss =>
                             s_execute st ((y * x) :: ss) rest
                         | _ => s_execute st stack rest
                         end
    end.

  Example s_execute1 :
    s_execute empty_st []
      [SPush 5; SPush 3; SPush 1; SMinus]
    = [2; 5].
  Proof. reflexivity. Qed.

  Fixpoint s_compile (e : aexp) : list sinstr :=
    match e with
    | ANum n => [SPush n]
    | AId x => [SLoad x]
    | APlus a1 a2 => s_compile a1 ++ s_compile a2 ++ [SPlus]
    | AMinus a1 a2 => s_compile a1 ++ s_compile a2 ++ [SMinus]
    | AMult a1 a2 => s_compile a1 ++ s_compile a2 ++ [SMult]
    end.

  Example s_compile1 :
    s_compile <{ X - (2 * Y) }>
      = [SLoad X; SPush 2; SLoad Y; SMult; SMinus].
  Proof. reflexivity.
  Qed.

  Theorem execute_app : forall st p1 p2 stack,
      s_execute st stack (p1 ++ p2) = s_execute st (s_execute st stack p1) p2.
  Proof.
    intros st p1 p2.
    induction p1.
    - intros.
      reflexivity.
    - intros.
      destruct a eqn:Ea.
      + simpl. apply (IHp1 (n :: stack)).
      + simpl. apply (IHp1 (st x :: stack)).
      + simpl. destruct stack.
        { apply IHp1. }
        destruct stack.
        { apply IHp1. }