-- LeroyCompilerVerificationCourse/Constprop.lean

import LeroyCompilerVerificationCourse.Imp

import Std.Data.HashMap

open Classical in

@[grind] def mk_PLUS_CONST (a : aexp) (n : Int) : aexp :=
  if n = 0 then a else
  match a with
  | .CONST m => .CONST (m + n)
  | .PLUS a (.CONST m) => .PLUS a (.CONST (m + n))
  | _ => .PLUS a (.CONST n)

@[grind] def mk_PLUS (a1 a2 : aexp) : aexp :=
  match a1, a2 with
  | .CONST m, _ => mk_PLUS_CONST a2 m
  | _, .CONST m => mk_PLUS_CONST a1 m
  | .PLUS a1 (.CONST m1), .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.PLUS a1 a2) (m1 + m2)
  | .PLUS a1 (.CONST m1), _ => mk_PLUS_CONST (.PLUS a1 a2) m1
  | _, .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.PLUS a1 a2) m2
  | _, _ => .PLUS a1 a2

@[grind] def mk_MINUS (a1 a2 : aexp) : aexp :=
  match a1, a2 with
  | _, .CONST m => mk_PLUS_CONST a1 (-m)
  | .PLUS a1 (.CONST m1), .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.MINUS a1 a2) (m1 - m2)
  | .PLUS a1 (.CONST m1), _ => mk_PLUS_CONST (.MINUS a1 a2) m1
  | _, .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.MINUS a1 a2) (-m2)
  | _, _ => .MINUS a1 a2

@[grind] def mk_EQUAL (a1 a2 : aexp) : bexp :=
  match a1, a2 with
  | .CONST n1, .CONST n2 => if n1 = n2 then .TRUE else .FALSE
  | .PLUS a1 (.CONST n1), .CONST n2 => .EQUAL a1 (.CONST (n2 - n1))
  | _, _ => .EQUAL a1 a2

@[grind] def mk_LESSEQUAL (a1 a2 : aexp) : bexp :=
  match a1, a2 with
  | .CONST n1, .CONST n2 => if n1 <= n2 then .TRUE else .FALSE
  | .PLUS a1 (.CONST n1), .CONST n2 => .LESSEQUAL a1 (.CONST (n2 - n1))
  | _, _ => .LESSEQUAL a1 a2

@[grind] def mk_NOT (b : bexp) : bexp :=
  match b with
  | .TRUE => .FALSE
  | .FALSE => .TRUE
  | .NOT b => b
  | _ => .NOT b

@[grind] def mk_AND (b1 b2 : bexp) : bexp :=
  match b1, b2 with
  | .TRUE, _ => b2
  | _, .TRUE => b1
  | .FALSE, _ => .FALSE
  | _, .FALSE => .FALSE
  | _, _ => .AND b1 b2

@[grind] def mk_IFTHENELSE (b : bexp) (c1 c2 : com) : com :=
  match b with
  | .TRUE => c1
  | .FALSE => c2
  | _ => .IFTHENELSE b c1 c2

@[grind] def mk_WHILE (b : bexp) (c : com) : com :=
  match b with
  | .FALSE => .SKIP
  | _ => .WHILE b c

def Store := Std.HashMap ident Int

@[grind] def matches' (s : store) (S : Store) : Prop :=
  ∀ x n, S.get? x = .some n -> s x = n

@[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

@[grind] def Top : Store := Std.HashMap.emptyWithCapacity

@[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

@[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

@[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

@[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

@[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

@[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

@[grind] noncomputable def fixpoint_rec (F : Store -> Store) (fuel : Nat) (S : Store) : Store :=
  match fuel with
  | 0 => Top
  | fuel + 1 =>
      let S' := F S
      if Equal S' S then S else fixpoint_rec F fuel S'

@[grind] def num_iter : Nat := 20

@[grind] noncomputable def fixpoint (F : Store -> Store) (init_S : Store) : Store :=
  fixpoint_rec F num_iter init_S

@[grind] noncomputable def Cexec (S : Store) (c : com) : Store :=
  match c with
  | .SKIP => S
  | .ASSIGN x a => Update x (Aeval S a) S
  | .SEQ c1 c2 => Cexec (Cexec S c1) c2
  | .IFTHENELSE b c1 c2 =>
      match Beval S b with
      | .some true => Cexec S c1
      | .some false => Cexec S c2
      | .none => Join (Cexec S c1) (Cexec S c2)
  | .WHILE _ c1 =>
      fixpoint (fun x => Join S (Cexec x c1)) S

@[grind =] def cp_aexp (S : Store) (a : aexp) : aexp :=
  match a with
  | .CONST n => .CONST n
  | .VAR x => match S.get? x with
    | .some n => .CONST n
    | .none => .VAR x
  | .PLUS a1 a2 => mk_PLUS (cp_aexp S a1) (cp_aexp S a2)
  | .MINUS a1 a2 => mk_MINUS (cp_aexp S a1) (cp_aexp S a2)

@[grind] def cp_bexp (S : Store) (b : bexp) : bexp :=
  match b with
  | .TRUE => .TRUE
  | .FALSE => .FALSE
  | .EQUAL a1 a2 => mk_EQUAL (cp_aexp S a1) (cp_aexp S a2)
  | .LESSEQUAL a1 a2 => mk_LESSEQUAL (cp_aexp S a1) (cp_aexp S a2)
  | .NOT b => mk_NOT (cp_bexp S b)
  | .AND b1 b2 => mk_AND (cp_bexp S b1) (cp_bexp S b2)

@[grind] noncomputable def cp_com (S : Store) (c : com) : com :=
  match c with
  | .SKIP => .SKIP
  | .ASSIGN x a =>
      .ASSIGN x (cp_aexp S a)
  | .SEQ c1 c2 =>
      .SEQ (cp_com S c1) (cp_com (Cexec S c1) c2)
  | .IFTHENELSE b c1 c2 =>
      mk_IFTHENELSE (cp_bexp S b) (cp_com S c1) (cp_com S c2)
  | .WHILE b c =>
      let sfix := Cexec S (.WHILE b c)
      mk_WHILE (cp_bexp sfix b) (cp_com sfix c)

-- Repository definitions used:
content: def ident := String deriving BEq, Repr, Hashable

content: inductive aexp : Type where
  | CONST (n : Int)               
  | VAR (x : ident)               
  | PLUS (a1 : aexp) (a2 : aexp)  
  | MINUS (a1 : aexp) (s2 : aexp) 

content: inductive bexp : Type where
  | TRUE                              
  | FALSE                             
  | EQUAL (a1 : aexp) (a2 : aexp)     
  | LESSEQUAL (a1 : aexp) (a2 : aexp) 
  | NOT (b1 : bexp)                   
  | AND (b1 : bexp) (b2 : bexp)       

content: inductive com : Type where
  | SKIP                                        
  | ASSIGN (x : ident) (a : aexp)               
  | SEQ (c1 : com) (c2 : com)                   
  | IFTHENELSE (b : bexp) (c1 : com) (c2 : com) 
  | WHILE (b : bexp) (c1 : com)                 

content: def store : Type := ident → Int

content: @[grind] inductive cexec : store → com → store → Prop where
  | cexec_skip :
      cexec s .SKIP s
  | cexec_assign :
      cexec s (.ASSIGN x a) (update x (aeval s a) s)
  | cexec_seq :
      cexec s c1 s' -> cexec s' c2 s'' ->
      cexec s (.SEQ c1 c2) s''
  | cexec_ifthenelse :
      cexec s (if beval s b then c1 else c2) s' ->
      cexec s (.IFTHENELSE b c1 c2) s'
  | cexec_while_done :
      beval s b = false ->
      cexec s (.WHILE b c) s
  | cexec_while_loop :
      beval s b = true -> cexec s c s' -> cexec s' (.WHILE b c) s'' ->
      cexec s (.WHILE b c) s''

content: @[grind] def beval (s : store) (b : bexp) : Bool :=
  match b with
  | .TRUE => true
  | .FALSE => false
  | .EQUAL a1 a2 => aeval s a1 = aeval s a2
  | .LESSEQUAL a1 a2 => aeval s a1 <= aeval s a2
  | .NOT b1 =>  !(beval s b1)
  | .AND b1 b2 => beval s b1 && beval s b2

content: @[grind] def aeval (s : store) (a : aexp) : Int :=
  match a with
  | .CONST n => n
  | .VAR x => s x
  | .PLUS a1 a2 => aeval s a1 + aeval s a2
  | .MINUS a1 a2 => aeval s a1 - aeval s a2

content: @[grind] def update (x : ident) (v : Int) (s : store) : store :=
  fun y => if x == y then v else s y

content: notation:10 l:10 " ;; " r:11 => com.SEQ l r

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: @[grind] def mk_PLUS_CONST (a : aexp) (n : Int) : aexp :=
  if n = 0 then a else
  match a with
  | .CONST m => .CONST (m + n)
  | .PLUS a (.CONST m) => .PLUS a (.CONST (m + n))
  | _ => .PLUS a (.CONST n)

content: @[grind] def mk_PLUS (a1 a2 : aexp) : aexp :=
  match a1, a2 with
  | .CONST m, _ => mk_PLUS_CONST a2 m
  | _, .CONST m => mk_PLUS_CONST a1 m
  | .PLUS a1 (.CONST m1), .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.PLUS a1 a2) (m1 + m2)
  | .PLUS a1 (.CONST m1), _ => mk_PLUS_CONST (.PLUS a1 a2) m1
  | _, .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.PLUS a1 a2) m2
  | _, _ => .PLUS a1 a2

content: @[grind] def mk_MINUS (a1 a2 : aexp) : aexp :=
  match a1, a2 with
  | _, .CONST m => mk_PLUS_CONST a1 (-m)
  | .PLUS a1 (.CONST m1), .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.MINUS a1 a2) (m1 - m2)
  | .PLUS a1 (.CONST m1), _ => mk_PLUS_CONST (.MINUS a1 a2) m1
  | _, .PLUS a2 (.CONST m2) => mk_PLUS_CONST (.MINUS a1 a2) (-m2)
  | _, _ => .MINUS a1 a2

content: @[grind] def mk_EQUAL (a1 a2 : aexp) : bexp :=
  match a1, a2 with
  | .CONST n1, .CONST n2 => if n1 = n2 then .TRUE else .FALSE
  | .PLUS a1 (.CONST n1), .CONST n2 => .EQUAL a1 (.CONST (n2 - n1))
  | _, _ => .EQUAL a1 a2

content: @[grind] def mk_LESSEQUAL (a1 a2 : aexp) : bexp :=
  match a1, a2 with
  | .CONST n1, .CONST n2 => if n1 <= n2 then .TRUE else .FALSE
  | .PLUS a1 (.CONST n1), .CONST n2 => .LESSEQUAL a1 (.CONST (n2 - n1))
  | _, _ => .LESSEQUAL a1 a2

content: @[grind] def mk_NOT (b : bexp) : bexp :=
  match b with
  | .TRUE => .FALSE
  | .FALSE => .TRUE
  | .NOT b => b
  | _ => .NOT b

content: @[grind] def mk_AND (b1 b2 : bexp) : bexp :=
  match b1, b2 with
  | .TRUE, _ => b2
  | _, .TRUE => b1
  | .FALSE, _ => .FALSE
  | _, .FALSE => .FALSE
  | _, _ => .AND b1 b2

content: @[grind] def mk_IFTHENELSE (b : bexp) (c1 c2 : com) : com :=
  match b with
  | .TRUE => c1
  | .FALSE => c2
  | _ => .IFTHENELSE b c1 c2

content: @[grind] def mk_WHILE (b : bexp) (c : com) : com :=
  match b with
  | .FALSE => .SKIP
  | _ => .WHILE b c

content: def Store := Std.HashMap ident Int

content: @[grind] def matches' (s : store) (S : Store) : Prop :=
  ∀ x n, S.get? x = .some n -> s x = n

content: @[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

content: @[grind] def Top : Store := Std.HashMap.emptyWithCapacity

content: @[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

content: def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

content: @[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

content: @[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

content: @[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

content: @[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

content: @[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

content: @[grind] noncomputable def fixpoint_rec (F : Store -> Store) (fuel : Nat) (S : Store) : Store :=
  match fuel with
  | 0 => Top
  | fuel + 1 =>
      let S' := F S
      if Equal S' S then S else fixpoint_rec F fuel S'

content: @[grind] def num_iter : Nat := 20

content: @[grind] noncomputable def fixpoint (F : Store -> Store) (init_S : Store) : Store :=
  fixpoint_rec F num_iter init_S

content: @[grind] noncomputable def Cexec (S : Store) (c : com) : Store :=
  match c with
  | .SKIP => S
  | .ASSIGN x a => Update x (Aeval S a) S
  | .SEQ c1 c2 => Cexec (Cexec S c1) c2
  | .IFTHENELSE b c1 c2 =>
      match Beval S b with
      | .some true => Cexec S c1
      | .some false => Cexec S c2
      | .none => Join (Cexec S c1) (Cexec S c2)
  | .WHILE _ c1 =>
      fixpoint (fun x => Join S (Cexec x c1)) S

content: @[grind =] def cp_aexp (S : Store) (a : aexp) : aexp :=
  match a with
  | .CONST n => .CONST n
  | .VAR x => match S.get? x with
    | .some n => .CONST n
    | .none => .VAR x
  | .PLUS a1 a2 => mk_PLUS (cp_aexp S a1) (cp_aexp S a2)
  | .MINUS a1 a2 => mk_MINUS (cp_aexp S a1) (cp_aexp S a2)

content: @[grind] def cp_bexp (S : Store) (b : bexp) : bexp :=
  match b with
  | .TRUE => .TRUE
  | .FALSE => .FALSE
  | .EQUAL a1 a2 => mk_EQUAL (cp_aexp S a1) (cp_aexp S a2)
  | .LESSEQUAL a1 a2 => mk_LESSEQUAL (cp_aexp S a1) (cp_aexp S a2)
  | .NOT b => mk_NOT (cp_bexp S b)
  | .AND b1 b2 => mk_AND (cp_bexp S b1) (cp_bexp S b2)

content: @[grind] noncomputable def cp_com (S : Store) (c : com) : com :=
  match c with
  | .SKIP => .SKIP
  | .ASSIGN x a =>
      .ASSIGN x (cp_aexp S a)
  | .SEQ c1 c2 =>
      .SEQ (cp_com S c1) (cp_com (Cexec S c1) c2)
  | .IFTHENELSE b c1 c2 =>
      mk_IFTHENELSE (cp_bexp S b) (cp_com S c1) (cp_com S c2)
  | .WHILE b c =>
      let sfix := Cexec S (.WHILE b c)
      mk_WHILE (cp_bexp sfix b) (cp_com sfix c)

-- Local lemmas used:
content: theorem mk_EQUAL_sound :
  ∀ s a1 a2, beval s (mk_EQUAL a1 a2) = (aeval s a1 = aeval s a2)

content: theorem mk_LESSEQUAL_sound :
  ∀ s a1 a2, beval s (mk_LESSEQUAL a1 a2) = (aeval s a1 <= aeval s a2)

content: theorem mk_NOT_sound :
  ∀ s b, beval s (mk_NOT b) = ¬ (beval s b)

content: theorem mk_AND_sound :
  ∀ s b1 b2, beval s (mk_AND b1 b2) = (beval s b1 ∧ beval s b2)

content: theorem cexec_mk_IFTHENELSE : ∀ s1 b c1 c2 s2,
  cexec s1 (if beval s1 b then c1 else c2) s2 ->
  cexec s1 (mk_IFTHENELSE b c1 c2) s2

content: theorem cexec_mk_WHILE_done : ∀ s1 b c,
  beval s1 b = false ->
  cexec s1 (mk_WHILE b c) s1

content: theorem cexec_mk_WHILE_loop : ∀ b c s1 s2 s3,
  beval s1 b = true -> cexec s1 c s2 -> cexec s2 (mk_WHILE b c) s3 ->
  cexec s1 (mk_WHILE b c) s3

content: theorem matches_Le : ∀ s S1 S2, Le S1 S2 -> matches' s S1 -> matches' s S2

content: theorem Le_Join_l : ∀ S1 S2, Le S1 (Join S1 S2)

content: theorem  Le_Join_r : ∀ S1 S2, Le S2 (Join S1 S2)

content: theorem Equal_Le : ∀ S1 S2, Equal S1 S2 -> Le S1 S2

content: theorem Beval_sound :
  ∀ s S, matches' s S ->
  ∀ b n, Beval S b = .some n -> beval s b = n

content: theorem fixpoint_sound (F : Store → Store) (init_S : Store) (h : S = fixpoint F init_S) :
  Le (F S) S

content: @[grind] theorem Cexec_sound :
  ∀ c s1 s2 S1,
  cexec s1 c s2 -> matches' s1 S1 -> matches' s2 (Cexec S1 c)

content: theorem cp_bexp_sound :
  ∀ s S, matches' s S ->
  ∀ b, beval s (cp_bexp S b) = beval s b

-- Target theorem and proof:
theorem cp_com_correct_terminating :
  ∀ c s1 s2 S1,
  cexec s1 c s2 -> matches' s1 S1 -> cexec s1 (cp_com S1 c) s2
:= by
    intro c s1 s2 S1 EXEC AG
    induction c generalizing s1 s2 S1
    any_goals grind
    case ASSIGN x a =>
      cases EXEC
      next =>
        have := @cexec.cexec_assign s1 x (cp_aexp S1 a)
        grind
    case IFTHENELSE b c1 c2 c1_ih c2_ih =>
      cases EXEC
      next =>
        apply cexec_mk_IFTHENELSE
        grind [cp_bexp_sound]
    case WHILE b c c_ih =>
      generalize heq1 : com.WHILE b c = loop
      generalize heq2 : Cexec S1 (.WHILE b c) = X
      have INNER : ∀ s1 c1 s2,
                 cexec s1 c1 s2 ->
                 c1 = .WHILE b c ->
                 matches' s1 X ->
                 cexec s1 (mk_WHILE (cp_bexp X b) (cp_com X c)) s2 := by
                  intro s1 c1
                  induction c1 generalizing s1 c
                  any_goals grind
                  case WHILE b1 c1 c1_ih  =>
                    intro s2 EXEC EQ AG1
                    injections heq1' heq2'
                    generalize heq : (com.WHILE b1 c1) = loop
                    rw [heq] at EXEC
                    induction EXEC
                    any_goals grind
                    case cexec_while_done isFalse =>
                      apply cexec_mk_WHILE_done
                      · grind [cp_bexp_sound]
                    case cexec_while_loop s3 b' c' s4 s5 isTrue EXEC2 EXEC3 a_ih a_ih2 =>
                      apply cexec_mk_WHILE_loop
                      · grind [cp_bexp_sound]
                      · apply c_ih
                        · injections heq4 heq5
                          rw [heq2'] at heq5
                          rw [←heq5] at EXEC2
                          exact EXEC2
                        · exact AG1
                      · apply a_ih2
                        rotate_left
                        · grind
                        · apply matches_Le
                          rw [←heq2]
                          simp [Cexec]
                          apply fixpoint_sound
                          rotate_left
                          · exact (fun x => Join S1 (Cexec x c))
                          · exact S1
                          rotate_left
                          · grind
                          · apply matches_Le
                            · apply Le_Join_r
                            · apply Cexec_sound
                              · rw [←heq2']
                                injections heq3 heq4
                                rw [heq4]
                                exact EXEC2
                              · grind
      rw [heq1] at EXEC
      induction EXEC
      any_goals grind
      case cexec_while_loop s3 b' c' s4 s5 isTrue EXEC1 EXEC2 ih1 ih2 =>
        injections heq3 heq4
        simp [cp_com]
        specialize INNER s3 (.WHILE b' c') s5
        rw [←heq3, ←heq4, heq2]
        apply INNER
        any_goals grind
        · rw [←heq2]
          simp [Cexec]
          apply matches_Le
          · apply fixpoint_sound
            rotate_left
            · exact (fun x => Join S1 (Cexec x c))
            · exact S1
            · grind
          · grind