From Stdlib Require Import Arith Lia Strings.String.

Module ArithWithConstants.

  (* The following definition closely mirrors a standard BNF grammar for expressions.
   * It defines abstract syntax trees of arithmetic expressions. *)
  Inductive arith : Set :=
  | Const (n : nat)
  | Plus (e1 e2 : arith)
  | Times (e1 e2 : arith).

  (* Here are a few examples of specific expressions. *)
  Example ex1 := Const 42.
  Example ex2 := Plus (Const 1) (Times (Const 2) (Const 3)).
  (*
         + 
        /  \
       1    *
           / \
          2   3
  *)

  Fixpoint size (e : arith) : nat :=
    match e with
    | Const _ => 1
    | Plus e1 e2 => 1 + size e1 + size e2
    | Times e1 e2 => 1 + size e1 + size e2
    end.

  (* Here's how to run a program (evaluate a term) in Rocq. *)
  Compute size ex1.
  Compute size ex2.

  (* What's the longest path from the root of a syntax tree to a leaf? *)
  Fixpoint depth (e : arith) : nat :=
    match e with
    | Const _ => 1
    | Plus e1 e2 => 1 + max (depth e1) (depth e2)
    | Times e1 e2 => 1 + max (depth e1) (depth e2)
    end.

  Compute depth ex1.
  Compute depth ex2.

  (* Our first proof!
   * Size is an upper bound on depth. *)

(* Fixpoint max n m {struct n} : nat :=
  match n, m with
    | O, _ => m
    | S n', O => n
    | S n', S m' => S (max n' m')
  end. *)

Definition max2 (n m:nat):nat := 
  (if le_ge_dec n m then m else n).
Compute (max 10 20).
Compute (max 10 1).
Compute (le_ge_dec 1 2).


Lemma max_add : forall n m : nat,
  max n m <= n + m.
Proof.
  intros n m.
  destruct (le_ge_dec n m) as [Hle | Hge].
  - (* n ≤ m *)
   Check max_r. 
    rewrite max_r; try assumption.
    apply (Nat.le_add_l m n).
  - (* n ≥ m *)
    rewrite max_l; try assumption.
    apply Nat.le_add_r.
Qed.

Theorem depth_le_size : forall e, depth e <= size e.
  Proof.
    intros.
    induction e.
    
   -  simpl. 
     trivial.  
    - simpl.   apply le_n_S. 
      apply (Nat.le_trans (Nat.max (depth e1) (depth e2)) 
             (depth e1 + depth e2) (size e1 + size e2)).
     + 
     apply max_add.
     + apply Nat.add_le_mono; assumption. 
    - apply le_n_S. 
      apply (Nat.le_trans (Nat.max (depth e1) (depth e2)) (depth e1 + depth e2) (size e1 + size e2)).
     +  apply max_add.
     +  apply Nat.add_le_mono; assumption. 
  Qed.


  Theorem depth_le_size_snazzy : forall e, depth e <= size e.
  Proof.
      induction e; simpl; intros;
      lia.    
  Qed. 
  (* A silly recursive function: swap the operand orders of all binary operators. *)
  Fixpoint commuter (e : arith) : arith :=
    match e with
    | Const _ => e
    | Plus e1 e2 => Plus (commuter e2) (commuter e1)
    | Times e1 e2 => Times (commuter e2) (commuter e1)
    end.

  Compute commuter ex1.
  Print ex2.
  Compute commuter ex2.
(*
         + 
        /  \
       1    *
           / \
          2   3
  *)

   (*
         + 
        /  \
       *    1
      / \
     3   2
  *)
  (* [commuter] has all the appropriate interactions with other functions (and itself). *)

  Theorem size_commuter : forall e, size (commuter e) = size e.
  Proof.
      induction e; simpl; intros. 
    - reflexivity.
    - rewrite IHe1. rewrite IHe2. f_equal. rewrite Nat.add_comm. reflexivity. 
    - rewrite IHe1. rewrite IHe2. f_equal. rewrite Nat.add_comm. reflexivity.
  Qed.  

  
Theorem depth_commuter_snazzy : forall e, depth (commuter e) = depth e.
 Proof.   
  induction e; simpl; lia.
Qed.  

(* Here is a proof that does not use lia *)
Lemma max_comm2: forall (n m:nat),
  max n m = max m n.
  Proof.
    induction n, m. 
    - simpl. trivial.  
    - simpl. trivial.
    - simpl. trivial.
    - rewrite <- Nat.succ_max_distr.
    rewrite <- Nat.succ_max_distr.
    f_equal. rewrite IHn. trivial.   
Qed. 


Theorem depth_commuter : forall e, depth (commuter e) = depth e.
  Proof.
    induction e; simpl; intros.
    - reflexivity.
    - rewrite IHe1. rewrite IHe2. f_equal. apply max_comm2. 
    - rewrite IHe1. rewrite IHe2. f_equal. apply max_comm2.        
  Qed. 


Theorem commuter_inverse : forall e, commuter (commuter e) = e.
  Proof.
    induction e; simpl; intros. 
    - reflexivity.
    - rewrite IHe1. rewrite IHe2. reflexivity.
    -  rewrite IHe1. rewrite IHe2. reflexivity. 
  Qed. 
    
End ArithWithConstants.

(* Let's shake things up a bit by adding variables to expressions.
 * Note that all of the automated proof scripts from before will keep working
 * with no changes!  That sort of "free" proof evolution is invaluable for
 * theorems about real-world compilers, say. *)

 Module ArithWithVariables.

  Notation var := string.

  Inductive arith : Set :=
  | Const (n : nat)
  | Var (x : var) (* <-- this is the new constructor! *)
  | Plus (e1 e2 : arith)
  | Times (e1 e2 : arith).

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

  Fixpoint size (e : arith) : nat :=
    match e with
    | Const _ => 1
    | Var _ => 1
    | Plus e1 e2 => 1 + size e1 + size e2
    | Times e1 e2 => 1 + size e1 + size e2
    end.

  Compute size ex1.
  Compute size ex2.

  Fixpoint depth (e : arith) : nat :=
    match e with
    | Const _ => 1
    | Var _ => 1
    | Plus e1 e2 => 1 + max (depth e1) (depth e2)
    | Times e1 e2 => 1 + max (depth e1) (depth e2)
    end.

  Compute depth ex1.
  Compute depth ex2.

  Theorem depth_le_size : forall e, depth e <= size e.
  Proof.
    induction e; simpl; lia.
  Qed.

  Fixpoint commuter (e : arith) : arith :=
    match e with
    | Const _ => e
    | Var _ => e
    | Plus e1 e2 => Plus (commuter e2) (commuter e1)
    | Times e1 e2 => Times (commuter e2) (commuter e1)
    end.

  Compute commuter ex1.
  Compute commuter ex2.

  Theorem size_commuter : forall e, size (commuter e) = size e.
  Proof.
     induction e; simpl; lia.
  Qed. 

  Theorem depth_commuter : forall e, depth (commuter e) = depth e.
  Proof.
     induction e; simpl; lia.
  Qed. 

Theorem commuter_inverse : forall e, commuter (commuter e) = e.
Proof.
    (*
    induction e; simpl. 
      - trivial.
      - trivial.
      - rewrite IHe1. rewrite IHe2. trivial.
      - rewrite IHe1. rewrite IHe2. trivial.
    *)
    induction e; simpl; (try rewrite IHe1; try rewrite IHe2); trivial. 
  Qed.


  (* Now that we have variables, we can consider new operations,
   * like substituting an expression for a variable.
   * We use an infix operator [==v] for equality tests on strings.
   * It has a somewhat funny and very expressive type,
   * whose details we will try to gloss over.
   * (We later return to them in SubsetTypes.v.) *)
  Infix "==" := string_dec (at level 70).
  
   Fixpoint substitute (inThis : arith) (replaceThis : var) (withThis : arith) : arith :=
    match inThis with
    | Const _ => inThis
    | Var x => if x == replaceThis then withThis else inThis
    | Plus e1 e2 => Plus (substitute e1 replaceThis withThis) (substitute e2 replaceThis withThis)
    | Times e1 e2 => Times (substitute e1 replaceThis withThis) (substitute e2 replaceThis withThis)
    end.

  (* An intuitive property about how much [substitute] might increase depth. *)
  Theorem substitute_depth : forall replaceThis withThis inThis,
    depth (substitute inThis replaceThis withThis) <= depth inThis + depth withThis.
  Proof.
    induction inThis; simpl.
    - lia.
    - destruct (x == replaceThis). 
      + lia.
      + simpl. lia.
    - lia.           
    - lia. 
  Qed. 
  (* Let's get fancier about automation, using [match goal] to pattern-match the goal
   * and decide what to do next!
   * The [|-] syntax separates hypotheses and conclusion in a goal.
   * The [context] syntax is for matching against *any subterm* of a term.
   * The construct [try] is also useful, for attempting a tactic and rolling back
   * the effect if any error is encountered. *)
  Theorem substitute_depth_snazzy : forall replaceThis withThis inThis,
    depth (substitute inThis replaceThis withThis) <= depth inThis + depth withThis.
  Proof.
    induction inThis; simpl;
   try match goal with
    | [ |- context[if ?a == ?b then _ else _] ] =>
        destruct (a == b); simpl
    end; lia.
  Qed. 
  (* A silly self-substitution has no effect. *)
  Theorem substitute_self : forall replaceThis inThis,
    substitute inThis replaceThis (Var replaceThis) = inThis.
  Proof.
    induction inThis; simpl;
    try match goal with
    | [ |- context[if ?a == ?b then _ else _] ] =>
        destruct (a == b); simpl
    end; congruence.
Qed.
    
  (* We can do substitution and commuting in either order. *)
  Theorem substitute_commuter : forall replaceThis withThis inThis,
    commuter (substitute inThis replaceThis withThis)
    = substitute (commuter inThis) replaceThis (commuter withThis).
  Proof.
    induction inThis; simpl;
    try match goal with
    | [ |- context[if ?a == ?b then _ else _] ] =>
        destruct (a == b); simpl
    end; congruence.
Qed.
    
  (* *Constant folding* is one of the classic compiler optimizations.
   * We repeatedly find opportunities to replace fancier expressions
   * with known constant values. *)
  Fixpoint constantFold (e : arith) : arith :=
    match e with
    | Const _ => e
    | Var _ => e
    | Plus e1 e2 =>
      let e1' := constantFold e1 in
      let e2' := constantFold e2 in
      match e1', e2' with
      | Const n1, Const n2 => Const (n1 + n2)
      | Const 0, _ => e2'
      | _, Const 0 => e1'
      | _, _ => Plus e1' e2'
      end
    | Times e1 e2 =>
      let e1' := constantFold e1 in
      let e2' := constantFold e2 in
      match e1', e2' with
      | Const n1, Const n2 => Const (n1 * n2)
      | Const 1, _ => e2'
      | _, Const 1 => e1'
      | Const 0, _ => Const 0
      | _, Const 0 => Const 0
      | _, _ => Times e1' e2'
      end
    end.

  (* This is supposed to be an *optimization*, so it had better not *increase*
   * the size of an expression!
   * There are enough cases to consider here that we skip straight to
   * the automation.
   * A new scripting construct is [match] patterns with dummy bodies.
   * Such a pattern matches *any* [match] in a goal, over any type! *)
  Theorem size_constantFold : forall e, size (constantFold e) <= size e.
  Proof.
      (* induction e; simpl.
      - lia.
      - lia. 
      - destruct (constantFold e1); simpl in *.
       + destruct n; simpl. 
         * destruct (constantFold e2); simpl in *; lia. 
         * destruct (constantFold e2); simpl in *; lia. 
       + destruct (constantFold e2); simpl in *. 
         **  destruct n; simpl; lia.
         **  lia.
         ** lia.
         ** lia.
      +  destruct (constantFold e2); simpl in *;
           try (destruct n; simpl; lia); lia. 
      +   destruct (constantFold e2); simpl in *;
           try (destruct n; simpl; lia); lia. 
    -  .... *)
    induction e; simpl;
   repeat match goal with
    | [ |- context[match ?E with _ => _ end] ] =>
        destruct E; simpl in *
    end; lia.
Qed.

  (* Business as usual, with another commuting law *)
  Theorem commuter_constantFold : forall e, 
      commuter (constantFold e) = constantFold (commuter e).
  Proof.
  induction e; simpl;
   repeat match goal with
    | [ |- context[match ?E with _ => _ end] ] =>
        destruct E; simpl in *
    | [ H : ?f _ = ?f _ |- _ ] =>
        injection H as H; subst
    | [ |- ?f _ = ?f _ ] =>
        f_equal
    end; try (lia || congruence).
  Qed.

  (* To define a further transformation, we first write a roundabout way of
   * testing whether an expression is a constant.
   * This detour happens to be useful to avoid overhead in concert with
   * pattern matching, since Rocq internally elaborates wildcard [_] patterns
   * into separate cases for all constructors not considered beforehand.
   * That expansion can create serious code blow-ups, leading to serious
   * proof blow-ups! *)
  Definition isConst (e : arith) : option nat :=
    match e with
    | Const n => Some n
    | _ => None
    end.

  (* Our next target is a function that finds multiplications by constants
   * and pushes the multiplications to the leaves of syntax trees,
   * ideally finding constants, which can be replaced by larger constants,
   * not affecting the meanings of expressions.
   * This helper function takes a coefficient [multiplyBy] that should be
   * applied to an expression. *)
  Fixpoint pushMultiplicationInside' (multiplyBy : nat) (e : arith) : arith :=
    match e with
    | Const n => Const (multiplyBy * n)
    | Var _ => Times (Const multiplyBy) e
    | Plus e1 e2 => Plus (pushMultiplicationInside' multiplyBy e1)
                         (pushMultiplicationInside' multiplyBy e2)
    | Times e1 e2 =>
      match isConst e1 with
      | Some k => pushMultiplicationInside' (k * multiplyBy) e2
      | None => Times (pushMultiplicationInside' multiplyBy e1) e2
      end
    end.

  Compute (pushMultiplicationInside'  5 (Const 10)). (* 50 *)
  Compute (pushMultiplicationInside'  5 (Var "x")). (* 5 * x *)
  Compute (pushMultiplicationInside'  5 (Plus(Var "x") (Const 2))). (*  (5 * x) + (2 * 5)  *)
  Compute (pushMultiplicationInside'  5 (Times (Const 2) (Var "x"))). (*  (5 * 2). * x   *)

  (* The overall transformation just fixes the initial coefficient as [1]. *)
  Definition pushMultiplicationInside (e : arith) : arith :=
    pushMultiplicationInside' 1 e.

  (* Let's prove this boring arithmetic property, so that we may use it below. *)
  Lemma n_times_0 : forall n, n * 0 = 0.
Proof. 
    lia. 
Qed.
    

  (* A fun fact about pushing multiplication inside:
   * the coefficient has no effect on depth!
   * Let's start by showing any coefficient is equivalent to coefficient 0. *)
  Lemma depth_pushMultiplicationInside'_irrelevance0 : forall e multiplyBy,
    depth (pushMultiplicationInside' multiplyBy e)
    = depth (pushMultiplicationInside' 0 e).
  Proof.
    induction e; intros; simpl in *.
    - lia.
    - lia.
    - rewrite IHe1. rewrite IHe2. lia.
    - destruct (isConst e1); simpl.
      + rewrite IHe2. rewrite Nat.mul_0_r. lia.
      + rewrite IHe1. lia.     
  Qed. 
  (* It can be remarkably hard to get Rocq's automation to be dumb enough to
   * help us demonstrate all of the primitive tactics. ;-)
   * In particular, we can redo the proof in an automated way, without the
   * explicit rewrites. *)
  Lemma depth_pushMultiplicationInside'_irrelevance0_snazzy : forall e multiplyBy,
    depth (pushMultiplicationInside' multiplyBy e)
    = depth (pushMultiplicationInside' 0 e).
  Proof.
   induction e; simpl;
   try match goal with
      | [ |- context[match ?E with _ => _ end] ] => destruct E; simpl
      end; congruence.
  Qed. 

  (* Now the general corollary about irrelevance of coefficients for depth. *)
  Lemma depth_pushMultiplicationInside'_irrelevance : forall e multiplyBy1 multiplyBy2,
    depth (pushMultiplicationInside' multiplyBy1 e)
    = depth (pushMultiplicationInside' multiplyBy2 e).
  Proof.
    intros.
    transitivity (depth (pushMultiplicationInside' 0 e)).
    apply depth_pushMultiplicationInside'_irrelevance0.
    symmetry.
    apply depth_pushMultiplicationInside'_irrelevance0.
  Qed. 
  (* Let's prove that pushing-inside has only a small effect on depth,
   * considering for now only coefficient 0. *)
  Lemma depth_pushMultiplicationInside' : forall e,
    depth (pushMultiplicationInside' 0 e) <= S (depth e).
  Proof.
    induction e; simpl. lia. lia. lia.
    destruct (isConst e1); simpl.
    rewrite n_times_0.
    lia.
    lia.
  Qed.

  (*Local Hint Rewrite n_times_0.*)
  (* Registering rewrite hints will get [simplify] to apply them for us
   * automatically! *)
#[local] Hint Rewrite Nat.mul_0_r : arith.
  Lemma depth_pushMultiplicationInside'_snazzy : forall e,
    depth (pushMultiplicationInside' 0 e) <= S (depth e).
  Proof.
    induction e; intros; simpl in *;
    try match goal with
        | [ |- context[match ?E with _ => _ end] ] =>
        destruct E; simpl; autorewrite with arith
        end;lia.
  Qed.

  Theorem depth_pushMultiplicationInside : forall e,
    depth (pushMultiplicationInside e) <= S (depth e).
  Proof.
   intros.
    unfold pushMultiplicationInside.
    (* [unfold X]: replace [X] by its definition. *)
    rewrite depth_pushMultiplicationInside'_irrelevance0.
    apply depth_pushMultiplicationInside'.
  Qed.

End ArithWithVariables.