-- Frap/Trans.lean

import Frap.Equiv

import Frap.Exercises.Equiv

namespace Imp

open AExp

open BExp

open Com

open CEval

def atrans_sound (atrans : AExp → AExp) :=
  ∀ (a : AExp), aequiv a (atrans a)

def btrans_sound (btrans : BExp → BExp) :=
  ∀ (b : BExp), bequiv b (btrans b)

def ctrans_sound (ctrans : Com → Com) :=
  ∀ (c : Com), cequiv c (ctrans c)

def fold_constants_aexp (a : AExp) : AExp :=
  match a with
  | a_num _
  | a_id _
    => a
  | a_plus a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 + n2)
    | (a1', a2') => a_plus a1' a2'
  | a_minus a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 - n2)
    | (a1', a2') => a_minus a1' a2'
  | a_mult a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 * n2)
    | (a1', a2') => a_mult a1' a2'

example : fold_constants_aexp <{(1 + 2) * x}> = <{3 * x}> := by admit /- proof elided -/

def fold_constants_bexp (b : BExp) : BExp :=
  match b with
  | b_true
  | b_false
    => b
  | b_eq a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 == n2 then b_true else b_false
    | (a1', a2') => b_eq a1' a2'
  | b_neq a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 != n2 then b_true else b_false
    | (a1', a2') => b_neq a1' a2'
  | b_le a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 <= n2 then b_true else b_false
    | (a1', a2') => b_le a1' a2'
  | b_not b1 =>
    match fold_constants_bexp b1 with
    | b_true => b_false
    | b_false => b_true
    | b1' => b_not b1'
  | b_and b1 b2 =>
    match (fold_constants_bexp b1, fold_constants_bexp b2) with
    | (b_true, b_true) => b_true
    | (b_true, b_false) => b_false
    | (b_false, b_true) => b_false
    | (b_false, b_false) => b_false
    | (b1', b2') => b_and b1' b2'
  | b_or b1 b2 =>
    match (fold_constants_bexp b1, fold_constants_bexp b2) with
    | (b_true, b_true) => b_true
    | (b_true, b_false) => b_true
    | (b_false, b_true) => b_true
    | (b_false, b_false) => b_false
    | (b1', b2') => b_or b1' b2'

example : fold_constants_bexp <{true && !(false && true)}> = <{true}> := by admit /- proof elided -/

def fold_constants_com (c : Com) : Com :=
  match c with
  | c_skip => c
  | c_asgn x a => c_asgn x (fold_constants_aexp a)
  | c_seq c1 c2 => c_seq (fold_constants_com c1) (fold_constants_com c2)
  | c_if b c1 c2 =>
    match fold_constants_bexp b with
    | b_true => fold_constants_com c1
    | b_false => fold_constants_com c2
    | b' => c_if b' (fold_constants_com c1) (fold_constants_com c2)
  | c_while b c1 =>
    match fold_constants_bexp b with
    | b_true => loop
    | b_false => c_skip
    | b' => c_while b' (fold_constants_com c1)

example : fold_constants_com
    
    <{ x := 4 + 5;
       y := x - 3;
       if (x - y) = (2 + 4) then skip
       else y := 0 end;
       if 0 <= (4 - (2 + 1)) then y := 0
       else skip end;
       while y = 0 do
         x := x + 1
       end }>
    = 
    <{ x := 9;
       y := x - 3;
       if (x - y) = 6 then skip
       else y := 0 end;
       y := 0;
       while y = 0 do
         x := x + 1
       end }> := by admit /- proof elided -/

-- Repository definitions used:
content: syntax "<[" term "]>" : imp

syntax "<{" imp "}>" : term

syntax "false" : imp

syntax "false" : term

syntax:21 "if" imp:20 "then" imp:20 "else" imp:20 "end" : imp

syntax "true" : imp

syntax "true" : term

syntax "skip" : imp

syntax:21 "while" imp:20 "do" imp:20 "end" : imp

content: macro_rules
  | `(term|true) => `(Bool.true)
  | `(term|false) => `(Bool.false)
  | `(term|<{$x}>) => `(imp|$x)
  | `(imp|$n:num) => `(a_num $n)
  | `(imp|$s:str) => `(a_id $s)
  | `(imp|$x + $y) => `(a_plus <{$x}> <{$y}>)
  | `(imp|$x - $y) => `(a_minus <{$x}> <{$y}>)
  | `(imp|$x * $y) => `(a_mult <{$x}> <{$y}>)
  | `(imp|true) => `(b_true)
  | `(imp|false) => `(b_false)
  | `(imp|$x = $y) => `(b_eq <{$x}> <{$y}>)
  | `(imp|$x != $y) => `(b_neq <{$x}> <{$y}>)
  | `(imp|$x <= $y) => `(b_le <{$x}> <{$y}>)
  | `(imp|!$x) => `(b_not <{$x}>)
  | `(imp|$x && $y) => `(b_and <{$x}> <{$y}>)
  | `(imp|$x || $y) => `(b_or <{$x}> <{$y}>)
  | `(imp|($x)) => `(<{$x}>)
  | `(imp|$x:ident) => `(a_id $(Lean.quote (toString x.getId)))
  | `(imp|<[$t:term]>) => pure t

content: macro_rules
  | `(imp|skip) => `(c_skip)
  | `(imp|$x:str := $y) => `(c_asgn $x <{$y}>)
  | `(imp|$x:ident := $y) => `(c_asgn $x <{$y}>)
  | `(imp|$c1 ; $c2) => `(c_seq <{$c1}> <{$c2}>)
  | `(imp|if $b then $c1 else $c2 end) => `(c_if <{$b}> <{$c1}> <{$c2}>)
  | `(imp|while $b do $c end) => `(c_while <{$b}> <{$c}>)

content: def cequiv (c₁ c₂ : Com) : Prop :=
  ∀ st st' : State, (st =[<[c₁]>]=> st') ↔ (st =[<[c₂]>]=> st')

content: abbrev State := String → Nat

content: inductive Com :=
  | c_skip : Com
  | c_asgn : String → AExp → Com
  | c_seq : Com → Com → Com
  | c_if : BExp → Com → Com → Com
  | c_while : BExp → Com → Com

content: inductive AExp where
  | a_num : Nat → AExp
  | a_id : String → AExp  
  | a_plus : AExp → AExp → AExp
  | a_minus : AExp → AExp → AExp
  | a_mult : AExp → AExp → AExp

content: inductive AExp where
  | a_num : Nat → AExp
  | a_plus : AExp → AExp → AExp
  | a_minus : AExp → AExp → AExp
  | a_mult : AExp → AExp → AExp

content: inductive Com : Type :=
  | c_skip : Com
  | c_asgn : String → AExp → Com
  | c_seq : Com → Com → Com
  | c_if : BExp → Com → Com → Com
  | c_while : BExp → Com → Com
  | c_par : Com → Com → Com   

content: inductive BExp where
  | b_true : BExp
  | b_false : BExp
  | b_eq : AExp → AExp → BExp
  | b_neq : AExp → AExp → BExp
  | b_le : AExp → AExp → BExp
  | b_not : BExp → BExp
  | b_and : BExp → BExp → BExp
  | b_or : BExp → BExp → BExp

content: abbrev x := "x"

content: abbrev y := "y"

content: def loop : Com :=
  <{ while true do
       skip
     end }>

content: def bequiv (b₁ b₂ : BExp) : Prop :=
  ∀ (st : State), beval st b₁ = beval st b₂

content: def beval (st : State) (b : BExp) : Bool :=
  match b with
  | b_true => true
  | b_false => false
  | b_eq a₁ a₂ => (aeval st a₁) == (aeval st a₂)
  | b_neq a₁ a₂ => (aeval st a₁) != (aeval st a₂)
  | b_le a₁ a₂ => (aeval st a₁) <= (aeval st a₂)
  | b_not b₁ => not (beval st b₁)
  | b_and b₁ b₂ => and (beval st b₁) (beval st b₂)
  | b_or b₁ b₂ => or (beval st b₁) (beval st b₂)

example : aeval (update empty x 5)
    <{3 + x * 2}>
    
    = 13 := by admit /- proof elided -/

content: def empty : State := fun _ => 0

content: abbrev z := "z"

content: def aeval (st : State) (a : AExp) : Nat :=
  match a with
  | a_num n => n
  | a_id x => st x
  | a_plus a₁ a₂ => (aeval st a₁) + (aeval st a₂)
  | a_minus a₁ a₂ => (aeval st a₁) - (aeval st a₂)
  | a_mult a₁ a₂ => (aeval st a₁) * (aeval st a₂)

content: def update (st : State) (k : String) (v : Nat) : State :=
  fun x => if x == k then v else st x

content: def aequiv (a₁ a₂ : AExp) : Prop :=
  ∀ (st : State), aeval st a₁ = aeval st a₂

content: inductive CEval : Com → State → State → Prop :=
  | e_skip : ∀ st,
      CEval c_skip st st
  | e_asgn : ∀ a n x st,
      aeval st a = n
      → CEval (c_asgn x a) st (update st x n)
  | e_seq : ∀ c₁ c₂ st st' st'',
      CEval c₁ st st' → CEval c₂ st' st''
      → CEval (c_seq c₁ c₂) st st''
  | e_ifTrue : ∀ b c₁ c₂ st st',
      beval st b = true → CEval c₁ st st'
      → CEval (c_if b c₁ c₂) st st'
  | e_ifFalse : ∀ b c₁ c₂ st st',
      beval st b = false → CEval c₂ st st'
      → CEval (c_if b c₁ c₂) st st'
  | e_whileFalse : ∀ b c st,
      beval st b = false
      → CEval (c_while b c) st st
  | e_whileTrue : ∀ b c st st' st'',
      beval st b = true
      → CEval c st st'
      → CEval (c_while b c) st' st''
      → CEval (c_while b c) st st''

-- Repository lemmas used:
content: theorem if_true b c₁ c₂
    : bequiv b <{true}> → cequiv <{if <[b]> then <[c₁]> else <[c₂]> end}> c₁

content: theorem if_false b c₁ c₂
    : bequiv b <{false}> → cequiv <{if <[b]> then <[c₁]> else <[c₂]> end}> c₂

content: theorem while_false b c : bequiv b <{false}> →
    cequiv <{while <[b]> do <[c]> end}> <{skip}>

-- Local definitions used:
content: def atrans_sound (atrans : AExp → AExp) :=
  ∀ (a : AExp), aequiv a (atrans a)

content: def btrans_sound (btrans : BExp → BExp) :=
  ∀ (b : BExp), bequiv b (btrans b)

content: def ctrans_sound (ctrans : Com → Com) :=
  ∀ (c : Com), cequiv c (ctrans c)

content: def fold_constants_aexp (a : AExp) : AExp :=
  match a with
  | a_num _
  | a_id _
    => a
  | a_plus a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 + n2)
    | (a1', a2') => a_plus a1' a2'
  | a_minus a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 - n2)
    | (a1', a2') => a_minus a1' a2'
  | a_mult a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => a_num (n1 * n2)
    | (a1', a2') => a_mult a1' a2'

example : fold_constants_aexp <{(1 + 2) * x}> = <{3 * x}> := by admit /- proof elided -/

content: def fold_constants_bexp (b : BExp) : BExp :=
  match b with
  | b_true
  | b_false
    => b
  | b_eq a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 == n2 then b_true else b_false
    | (a1', a2') => b_eq a1' a2'
  | b_neq a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 != n2 then b_true else b_false
    | (a1', a2') => b_neq a1' a2'
  | b_le a1 a2 =>
    match (fold_constants_aexp a1, fold_constants_aexp a2) with
    | (a_num n1, a_num n2) => if n1 <= n2 then b_true else b_false
    | (a1', a2') => b_le a1' a2'
  | b_not b1 =>
    match fold_constants_bexp b1 with
    | b_true => b_false
    | b_false => b_true
    | b1' => b_not b1'
  | b_and b1 b2 =>
    match (fold_constants_bexp b1, fold_constants_bexp b2) with
    | (b_true, b_true) => b_true
    | (b_true, b_false) => b_false
    | (b_false, b_true) => b_false
    | (b_false, b_false) => b_false
    | (b1', b2') => b_and b1' b2'
  | b_or b1 b2 =>
    match (fold_constants_bexp b1, fold_constants_bexp b2) with
    | (b_true, b_true) => b_true
    | (b_true, b_false) => b_true
    | (b_false, b_true) => b_true
    | (b_false, b_false) => b_false
    | (b1', b2') => b_or b1' b2'

example : fold_constants_bexp <{true && !(false && true)}> = <{true}> := by admit /- proof elided -/

content: def fold_constants_com (c : Com) : Com :=
  match c with
  | c_skip => c
  | c_asgn x a => c_asgn x (fold_constants_aexp a)
  | c_seq c1 c2 => c_seq (fold_constants_com c1) (fold_constants_com c2)
  | c_if b c1 c2 =>
    match fold_constants_bexp b with
    | b_true => fold_constants_com c1
    | b_false => fold_constants_com c2
    | b' => c_if b' (fold_constants_com c1) (fold_constants_com c2)
  | c_while b c1 =>
    match fold_constants_bexp b with
    | b_true => loop
    | b_false => c_skip
    | b' => c_while b' (fold_constants_com c1)

example : fold_constants_com
    
    <{ x := 4 + 5;
       y := x - 3;
       if (x - y) = (2 + 4) then skip
       else y := 0 end;
       if 0 <= (4 - (2 + 1)) then y := 0
       else skip end;
       while y = 0 do
         x := x + 1
       end }>
    = 
    <{ x := 9;
       y := x - 3;
       if (x - y) = 6 then skip
       else y := 0 end;
       y := 0;
       while y = 0 do
         x := x + 1
       end }> := by admit /- proof elided -/

-- Local lemmas used:
content: theorem while_true b c :
    bequiv b <{true}> → cequiv <{while <[b]> do <[c]> end}> loop

content: theorem sym_bequiv b₁ b₂ : bequiv b₁ b₂ → bequiv b₂ b₁

content: theorem sym_cequiv c₁ c₂ : cequiv c₁ c₂ → cequiv c₂ c₁

content: theorem trans_cequiv c₁ c₂ c₃ : cequiv c₁ c₂ → cequiv c₂ c₃ → cequiv c₁ c₃

content: theorem c_asgn_congruence x a a'
    : aequiv a a' → cequiv <{x := <[a]>}> <{x := <[a']>}>

content: theorem c_while_congruence b b' c c'
    : bequiv b b' → cequiv c c'
      → cequiv <{while <[b]> do <[c]> end}> <{while <[b']> do <[c']> end}>

content: theorem c_seq_congruence c₁ c₁' c₂ c₂'
    : cequiv c₁ c₁' → cequiv c₂ c₂' → cequiv (c_seq c₁ c₂) (c_seq c₁' c₂')

content: theorem c_if_congruence b b' c₁ c₁' c₂ c₂'
    : bequiv b b' → cequiv c₁ c₁' → cequiv c₂ c₂'
      → cequiv (c_if b c₁ c₂) (c_if b' c₁' c₂')

content: theorem fold_constants_aexp_sound : atrans_sound fold_constants_aexp

content: theorem fold_constants_bexp_sound : btrans_sound fold_constants_bexp

-- Target theorem and proof:
theorem fold_constants_com_sound : ctrans_sound fold_constants_com
:= by
  intro c; induction c with simp [fold_constants_com]
  | c_asgn =>
    apply c_asgn_congruence
    apply fold_constants_aexp_sound
  | c_seq =>
    apply c_seq_congruence <;> assumption
  | c_if b c₁ c₂ ih1 ih2 =>
    split
    . apply trans_cequiv _ c₁
      . apply if_true
        rename_i heq; rw [← heq]
        apply fold_constants_bexp_sound
      . assumption
    . apply trans_cequiv _ c₂
      . apply if_false
        rename_i heq; rw [← heq]
        apply fold_constants_bexp_sound
      . assumption
    . apply c_if_congruence
      . apply fold_constants_bexp_sound
      . assumption
      . assumption
  | c_while b c' ih =>
    split
    . apply while_true
      rename_i heq; rw [← heq]
      apply fold_constants_bexp_sound
    . apply while_false
      rename_i heq; rw [← heq]
      apply fold_constants_bexp_sound
    . apply c_while_congruence
      . apply fold_constants_bexp_sound
      . assumption