-- LeroyCompilerVerificationCourse/Constprop.lean

import LeroyCompilerVerificationCourse.Imp

import Std.Data.HashMap

open Classical in

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

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

content: def store : Type := ident → Int

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: @[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: 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

-- Local lemmas used:
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

-- Target theorem and proof:
@[grind] theorem Cexec_sound :
  ∀ c s1 s2 S1,
  cexec s1 c s2 -> matches' s1 S1 -> matches' s2 (Cexec S1 c)
:= by
    intro c
    induction c
    next =>
      intro s1 s2 S1 EXEC
      cases EXEC
      grind
    next x a =>
      intro s1 s2 S1 EXEC
      cases EXEC
      grind
    next c1 c2 c1_ih c2_ih =>
      grind
    next b c1 c2 c1_ih c2_ih =>
      intro s1 s2 S1 EXEC
      cases EXEC
      next EXEC =>
        by_cases beval s1 b
        case pos h =>
          unfold Cexec
          intro h2
          have := Beval_sound s1 S1 h2 b
          split <;> grind
        case neg h =>
          simp [h] at EXEC
          intro h2
          unfold Cexec
          have := Beval_sound s1 S1 h2 b
          split <;> grind
    case WHILE b c c1_ih =>
      intro s1 s2 S1 EXEC MATCHES
      generalize eq1 : (fun x => Join S1 (Cexec x c)) = F
      generalize eq2 : fixpoint F S1 = X
      have INNER : ∀ s1 c1 s2,
                 cexec s1 c1 s2 ->
                 c1 = .WHILE b c ->
                 matches' s1 X ->
                 matches' s2 X := by
                  intro s3 c1 s4 EXEC2 EQ AG
                  induction EXEC2
                  any_goals grind
                  case cexec_while_loop s' b' c' s5 s6 EXEC2 EXEC3 EXEC4 _ a_ih2 =>
                    apply a_ih2
                    · grind
                    · apply matches_Le
                      rotate_right
                      · exact F X
                      · exact @fixpoint_sound X F S1 (by grind)
                      · rw [←eq2, ←eq1]
                        simp
                        apply matches_Le
                        apply Le_Join_r
                        rw [eq1, eq2]
                        apply c1_ih
                        injections EQ
                        rename_i eq1 eq2
                        rw [eq2] at EXEC3
                        exact EXEC3
                        exact AG
      unfold Cexec
      rw [eq1, eq2]
      apply INNER
      · apply EXEC
      · rfl
      · apply matches_Le
        have := @fixpoint_sound X F
        apply this
        rotate_left
        · exact S1
        rotate_left
        · grind
        · rw [←eq1]
          simp
          apply matches_Le
          · apply Le_Join_l
          · exact MATCHES