From Stdlib Require Import List Arith Lia String.
Notation var := string.
Import ListNotations.
Open Scope string_scope.

Inductive arith : Set :=
| Const (n : nat)
| Var (x : var)
| Plus (e1 e2 : arith)
| Minus (e1 e2 : arith)
| Times (e1 e2 : arith).

Example ex1 := Const 42.
Example ex2 := Plus (Var "y") (Times (Var "x") (Const 3)).

(* The above definition only explains what programs *look like*.
 * We also care about what they *mean*.
 * The natural meaning of an expression is the number it evaluates to.
 * Actually, it's not quite that simple.
 * We need to consider the meaning to be a function over a valuation
 * to the variables. aluation is a function from variable
 * names to numbers. *)
 
Definition valuation := var -> option nat.

Definition empty : valuation := fun _ => None.

Definition lookup (x:var) (v: valuation) : option nat := v x.

Definition update (x:var) (n:nat) (v: valuation) : valuation :=
  fun y =>
    match string_dec x y with
        | left H => Some n
        | right H' => v y
    end.


(* That is, the domain is [var] (a synonym for [string]) and the codomain/range
 * is [nat]. *)

(* The interpreter is a fairly innocuous-looking recursive function. *)
Fixpoint interp (e : arith) (v : valuation) : nat :=
  match e with
  | Const n => n
  | Var x =>
    match v x with
    | None => 0 (* goofy default value! *)
    | Some n => n
    end
  | Plus e1 e2 => interp e1 v + interp e2 v
  | Minus e1 e2 => interp e1 v - interp e2 v
  | Times e1 e2 => interp e1 v * interp e2 v
  end.

(* Here's an example aluation, using an infix operator for map extension. *)
Definition valuation0 : valuation := 
  update "y" 3 (update "x" 17 empty).

Theorem interp_ex1 : interp ex1 valuation0 = 42.
Proof.
  Print ex1.
  simpl.
  trivial.
Qed.

Theorem interp_ex2 : interp ex2 valuation0 = 54.
Proof.
  unfold valuation0.
  simpl. trivial.
Qed.

(* Here's the silly transformation we defined last time. *)
Fixpoint commuter (e : arith) : arith :=
  match e with
  | Const _ => e
  | Var _ => e
  | Plus e1 e2 => Plus (commuter e2) (commuter e1)
  | Minus e1 e2 => Minus (commuter e1) (commuter e2)
                   (* ^-- NB: didn't change the operand order here! *)
  | Times e1 e2 => Times (commuter e2) (commuter e1)
  end.

(* Instead of proving various odds-and-ends properties about it,
 * let's show what we *really* care about: it preserves the
 * *meanings* of expressions! *)
Theorem commuter_ok : forall v e, interp (commuter e) v = interp e v.
Proof.
  induction e; simpl. 

  - trivial.
  (*- destruct (v x).
   + reflexivity.
   + reflexivity.      
   *)
  - lia. 
  - rewrite IHe1, IHe2. lia. 
  - rewrite IHe1, IHe2. lia.
  - rewrite IHe1, IHe2. lia. 
Qed.

(* Let's also revisit substitution. *)
Fixpoint substitute (inThis : arith) (replaceThis : var) (withThis : arith) : arith :=
  match inThis with
  | Const _ => inThis
  | Var x => if (string_dec x replaceThis) then withThis else inThis
  | Plus e1 e2 => Plus (substitute e1 replaceThis withThis) (substitute e2 replaceThis withThis)
  | Minus e1 e2 => Minus (substitute e1 replaceThis withThis) (substitute e2 replaceThis withThis)
  | Times e1 e2 => Times (substitute e1 replaceThis withThis) (substitute e2 replaceThis withThis)
  end.

Theorem substitute_ok : forall v replaceThis withThis inThis,
  interp (substitute inThis replaceThis withThis) v
  = interp inThis (update replaceThis (interp withThis v) v).
Proof.
  induction inThis; simpl. 
  - trivial.
  - unfold update, lookup.
    destruct (string_dec replaceThis x).
    + rewrite e. destruct (string_dec x x).
      * congruence.
      * contradiction. 
    + destruct (string_dec x replaceThis).
      * rewrite <-e in n. contradiction.
      * simpl. congruence.
  - rewrite IHinThis1. rewrite IHinThis2. reflexivity.
  - rewrite IHinThis1. rewrite IHinThis2. reflexivity.
  - congruence.

Qed.

Theorem substitute_ok' : forall v replaceThis withThis inThis,
  interp (substitute inThis replaceThis withThis) v
  = interp inThis (update replaceThis (interp withThis v) v).
Proof.
  induction inThis; simpl; 
   try (unfold update, lookup;
    destruct (string_dec replaceThis x); 
     destruct (string_dec x replaceThis); simpl); congruence.
Qed.

(* Let's also define a pared-down version of the expression-simplification
 * functions from last chapter. *)
Fixpoint doSomeArithmetic (e : arith) : arith :=
  match e with
  | Const _ => e
  | Var _ => e
  | Plus (Const n1) (Const n2) => Const (n1 + n2)
  | Plus e1 e2 => Plus (doSomeArithmetic e1) (doSomeArithmetic e2)
  | Minus e1 e2 => Minus (doSomeArithmetic e1) (doSomeArithmetic e2)
  | Times (Const n1) (Const n2) => Const (n1 * n2)
  | Times e1 e2 => Times (doSomeArithmetic e1) (doSomeArithmetic e2)
  end.

Compute (doSomeArithmetic (Plus (Const 10) (Const 20))).

Theorem doSomeArithmetic_ok : forall e v, 
               interp (doSomeArithmetic e) v = interp e v.
Proof.
  induction e.
  - simpl. intros. reflexivity. 
  - simpl. reflexivity. 
  - destruct e1, e2; simpl in *; try congruence.
  - destruct e1, e2; simpl in *; try congruence.
  - destruct e1, e2; simpl in *; try congruence.
Qed.

(* Of course, we're going to get bored if we confine ourselves to arithmetic
 * expressions for the rest of our journey.  Let's get a bit fancier and define
 * a *stack machine*, related to postfix calculators that some of you may have
 * experienced. *)


Inductive instruction :=
| PushConst (n : nat)
| PushVar (x : var)
| Add
| Subtract
| Multiply.

(* What does it all mean?  An interpreter tells us unambiguously! *)
Definition run1 (i : instruction) (v : valuation) (stack : list nat) :
       list nat :=
  match i with
  | PushConst n => n :: stack
  | PushVar x => (match v x with
                  | None => 0
                  | Some n => n
                  end) :: stack
  | Add =>
    match stack with
    | arg2 :: arg1 :: stack' => arg1 + arg2 :: stack'
    | _ => stack (* arbitrary behavior in erroneous case (stack underflow) *)
    end
  | Subtract =>
    match stack with
    | arg2 :: arg1 :: stack' => arg1 - arg2 :: stack'
    | _ => stack (* arbitrary behavior in erroneous case *)
    end
  | Multiply =>
    match stack with
    | arg2 :: arg1 :: stack' => arg1 * arg2 :: stack'
    | _ => stack (* arbitrary behavior in erroneous case *)
    end
  end.

Compute (run1 (PushConst 1) empty []).
Compute (run1 (PushVar "x") valuation0 []).
Compute (run1 Add empty [1;2;10;20]).

Compute (run1 Multiply empty [5;2]).


(* That function explained how to run one instruction.
 * Here's how to run several of them. *)
Fixpoint run (is : list instruction) (v : valuation) (stack : list nat) : list nat :=
  match is with
  | nil => stack
  | i :: is' => run is' v (run1 i v stack)
  end.

  (* (1 + 2)  * 5 *)
Compute (run [(PushConst 1); (PushConst 2); Add;
              (PushConst 5); Multiply]
           empty []).
(*
x = 5
y = 2
"x * 3 - (y + 4)" ==> Parser ==>
Minus(Mult(x,3), Plus(y,4))

*)
Definition valuation1 : valuation := 
  update "y" 2 (update "x" 5 empty).

Compute (run 
[PushVar "x"; 
 PushConst 3; 
 Multiply; 
 PushVar "y"; 
 PushConst 4;
 Add;
 Subtract] valuation1
 []).


(* Instead of writing fiddly stack programs ourselves, let's *compile*
 * arithmetic expressions into equivalent stack programs. *)
Fixpoint compile (e : arith) : list instruction :=
  match e with
  | Const n => PushConst n :: nil
  | Var x => PushVar x :: nil
  | Plus e1 e2 => compile e1 ++ compile e2 ++ Add :: nil
  | Minus e1 e2 => compile e1 ++ compile e2 ++ Subtract :: nil
  | Times e1 e2 => compile e1 ++ compile e2 ++ Multiply :: nil
  end.

Compute (compile (Const 20)).

Compute (interp (Plus (Const 2) (Const 5))). (* python*)


Compute (run (compile (Plus (Const 2) (Const 5)))). (* C *)

Example ex101: compile (Minus
                         (Times (Var "x") (Const 3)) 
                         (Plus (Var "y")  (Const 4))) 
                         = 
[PushVar "x"; 
 PushConst 3; 
 Multiply; 
 PushVar "y"; 
 PushConst 4;
 Add;
 Subtract].
Proof. simpl. reflexivity. Qed.   


(* gcc: 
1. C definitional interpreter: Operatioinal semantics
    Oracle. (eval (C program))= result: meaning of my c program

    slow, 

  compiler: optimize, fast, c==> x86 assembly. run the generated 
  assembly code. How do you know it is correct?

  1) Testing: run (compile x) == eval x.
 
  prove: run (compile x) == eval x.

  Auth: 1) eval (functional code)
        2) fast version
*)


(* Now, of course, we should prove our compiler correct.
 * Skip down to the next theorem to see the overall correctness statement.
 * It turns out that we need to strengthen the induction hypothesis with a
 * lemma, to push the proof through. *)
Lemma compile_ok' : forall e v is stack,
    run (compile e ++ is) v stack = run is v (interp e v :: stack).
Proof.
  induction e; intros; simpl.
  - reflexivity.
  - reflexivity.
  (* Here we want to use associativity of [++], to get the conclusion to match
   * an induction hypothesis.  Let's ask Rocq to search its library for lemmas
   * that would justify such a rewrite, giving a pattern with wildcards, to
   * specify the essential structure that the rewrite should match. *)
  Search ((_ ++ _) ++ _).
  - rewrite <- app_assoc.
    rewrite IHe1.
    rewrite <- app_assoc.
    rewrite IHe2.
    simpl.
    reflexivity.

- rewrite <- app_assoc.
  rewrite IHe1.
  rewrite <- app_assoc.
  rewrite IHe2.
  simpl.
  trivial.

- rewrite <- app_assoc.
  rewrite IHe1.
  rewrite <- app_assoc.
  rewrite IHe2.
  simpl.
  trivial.
Qed.

Lemma compile_ok'' : forall e v is stack,
    run (compile e ++ is) v stack = run is v (interp e v :: stack).
Proof.
  induction e; intros; simpl;
  try (rewrite <- app_assoc; rewrite IHe1; 
      rewrite <- app_assoc; rewrite IHe2);
  reflexivity.
Qed. 


(* The overall theorem follows as a simple corollary. *)
Theorem compile_ok : forall e v, 

run (compile e) v nil =  interp e v :: nil.

Proof.
  intros. simpl.

  (* To match the form of our lemma, we need to replace [compile e] with
   * [compile e ++ nil], adding a "pointless" concatenation of the empty list.
   * [Search] again helps us find a library lemma. *)
  
   Search (_ ++ nil).
  rewrite <- (app_nil_r (compile e)).
  (* Note that we can use [rewrite] with explicit values of the first few
   * quantified variables of a lemma.  Otherwise, [rewrite] picks an
   * unhelpful place to rewrite.  (Try it and see!) *)
  Check compile_ok'.
  apply compile_ok'.
  (* Direct appeal to a previously proved lemma *)
Qed.

(* Let's get a bit fancier, moving toward the level of general-purpose
 * imperative languages.  Here's a language of commands, building on the
 * language of expressions we have defined. *)
Inductive cmd :=
| Skip
| Assign (x : var) (e : arith)
| Sequence (c1 c2 : cmd)
| Repeat (e : arith) (body : cmd).

(* That last constructor is for repeating a body command some number of
 * times.  Note that we sneakily avoid constructs that could introduce
 * nontermination, since Rocq only accepts terminating programs, and we want to
 * write an interpreter for commands.
 * In contrast to our last one, this interpreter *transforms valuations*.
 * We use a helper function for self-composing a function some number of
 * times. *)

Fixpoint selfCompose {A} (f : A -> A) (n : nat) : A -> A :=
  match n with
  | O => fun x => x
  | S n' => fun x => selfCompose f n' (f x)
  end.

(* 
selfCompose f 0 = identity function
selfCompose f 1 = f
selfCompose f 2 = f ∘ f
selfCompose f 3 = f ∘ f ∘ f
*)
(*
(selfCompose f 1).
fun x => selfCompose f 0 (f x)
fun x => (fun x => x) (f x)
fun x =>  (f x)
*)
(*
selfCompose f 1
= fun x => selfCompose f 0 (f x)
= fun x => (fun y => y) (f x)
= fun x => f x
*)
(*
selfCompose f 2
= fun x => selfCompose f 1 (f x)
= fun x => (fun y => selfCompose f 0 (f y)) (f x)
 subst y with (f x)
= fun x => selfCompose f 0 (f (f x))
= fun x => (fun z => z) (f (f x))
= fun x => f (f x)
*)

Definition t1 := (selfCompose (fun x => x + 1) 5 ).
Print t1.
Compute ((selfCompose (fun x => x + 1) 5 ) 10). (* 10 + 1 + 1 + 1 + 1= 15 *)


Fixpoint exec (c : cmd) (v : valuation) : valuation :=
  match c with
  | Skip => v
  | Assign x e => update x (interp e v) v
  | Sequence c1 c2 => exec c2 (exec c1 v)
  | Repeat e body => selfCompose (exec body) (interp e v) v
  end.

(* Let's define some programs and prove that they operate in certain ways. *)

Example factorial_ugly :=
  Sequence
    (Assign "output" (Const 1))
    (Repeat (Var "input")
            (Sequence
               (Assign "output" (Times (Var "output") (Var "input")))
               (Assign "input" (Minus (Var "input") (Const 1))))).

(* Ouch; that code is hard to read.  Let's introduce some notations to make the
 * concrete syntax more palatable.  We won't explain the general mechanisms on
 * display here, but see the Rocq manual for details, or try to reverse-engineer
 * them from our examples. *)
Coercion Const : nat >-> arith.
Coercion Var : var >-> arith.
(*Declare Scope arith_scope.*)
Infix "+" := Plus : arith_scope.
Infix "-" := Minus : arith_scope.
Infix "*" := Times : arith_scope.
Delimit Scope arith_scope with arith.
Notation "x <- e" := (Assign x e%arith) (at level 75).
Infix ";" := Sequence (at level 76).
Notation "'repeat' e 'doing' body 'done'" := (Repeat e%arith body) (at level 75).

(* OK, let's try that program again. *)
Example factorial :=
  "output" <- 1;
  repeat "input" doing
    "output" <- "output" * "input";
    "input" <- "input" - 1
  done.

Compute (exec factorial (update "input" 5 empty) "output").

Example fact5: (exec factorial (update "input" 5 empty) "output") = Some 120.
Proof. simpl. reflexivity.  Qed.

Example sum :=
  "output" <- 0;
  repeat "input" doing
    "output" <- "output" + "input";
    "input" <- "input" - 1
  done.

Compute (exec sum (update "input" 10 empty) "output").

Example sum10: (exec sum (update "input" 10 empty) "output") = Some 55.
Proof. 
  simpl. reflexivity.  
Qed.


(* Now we prove that it really computes factorial.
 * First, a reference implementation as a functional program. *)
Fixpoint fact (n : nat) : nat :=
  match n with
  | O => 1
  | S n' => n * fact n'
  end.

Theorem factorial_ok : forall v input,
  v "input" = Some input
  -> exec factorial v "output" = Some (fact input).
Proof. Admitted.

(* One last example: let's try to do loop unrolling, for constant iteration
 * counts.  That is, we can duplicate the loop body instead of using an explicit
 * loop. *)

Fixpoint seqself (c : cmd) (n : nat) : cmd :=
  match n with
  | O => Skip
  | S n' => Sequence c (seqself c n')
  end.

Compute (seqself ( "x" <- "x" - 1) 5).

Compute ((exec (seqself ( "x" <- "x" - 1) 5) (update "x" 20 empty)) "x").

Fixpoint unroll (c : cmd) : cmd :=
  match c with
  | Skip => c
  | Assign _ _ => c
  | Sequence c1 c2 => Sequence (unroll c1) (unroll c2)
  | Repeat (Const n) c1 => seqself (unroll c1) n
  (* ^-- the crucial case! *)
  | Repeat e c1 => Repeat e (unroll c1)
  end.

Example ex5 := repeat 5 doing Skip done.
Compute (unroll ex5).

Example ex6 := repeat 3 doing ( "x" <- "x" + 1)  done.
Compute (unroll ex6).
(*
"x" <- "x" + 1;
("x" <- "x" + 1; 
("x" <- "x" + 1; Skip))
*)
Compute (exec (unroll ex6) (update "x" 10 empty)) "x".

(* This obvious-sounding fact will come in handy: self-composition gives the
 * same result, when passed two functions that map equal inputs to equal
 * outputs. *)
Lemma selfCompose_extensional : forall {A} (f g : A -> A) n x,
  (forall y, f y = g y)
  -> selfCompose f n x = selfCompose g n x.
Proof. 
  induction n.
  - simpl. reflexivity.
  - intros. simpl. rewrite H.  apply IHn. assumption.       
Qed.

(* Crucial lemma: [seqself] is acting just like [selfCompose], in a suitable
 * sense. *)
Lemma seqself_ok : forall c n v,
  exec (seqself c n) v = selfCompose (exec c) n v.
Proof.

  induction n; simpl; intros.
  - reflexivity.
  - rewrite IHn. reflexivity.      
 (* induction n; simpl; congruence. *)
Qed.

(* The two lemmas we just proved are the main ingredients to prove the natural
 * correctness condition for [unroll]. *)
Theorem unroll_ok : forall c v, exec (unroll c) v = exec c v.
Proof.
  induction c; simpl; try reflexivity.
  - intros. rewrite IHc1. rewrite IHc2. reflexivity.
  - intros. destruct e; simpl.
    + rewrite seqself_ok. apply selfCompose_extensional. assumption.
    + apply selfCompose_extensional. assumption.
    + apply selfCompose_extensional. assumption.
    + apply selfCompose_extensional. assumption.  
    + apply selfCompose_extensional. assumption. 

   (*
   congruence.
  - intros. destruct e; simpl; try congruence.
   + rewrite seqself_ok. 
  apply selfCompose_extensional.
  trivial.

  + apply selfCompose_extensional.
  trivial.

  + apply selfCompose_extensional.
  trivial.

  + apply selfCompose_extensional.
  trivial.

  + apply selfCompose_extensional.
  trivial.
  *)
Qed.

Theorem unroll_ok' : forall c v, exec (unroll c) v = exec c v.
Proof.
  induction c; simpl; 
  try destruct e; intros;
  try (rewrite IHc1; rewrite IHc2);
  try rewrite seqself_ok; 
  try  apply selfCompose_extensional; 
  trivial.
Qed.