-- LeroyCompilerVerificationCourse/Deadcode.lean

import LeroyCompilerVerificationCourse.Imp

import Std.Data.HashSet

abbrev IdentSet := Std.HashSet ident

@[grind] def fv_aexp (a : aexp) : IdentSet :=
  match a with
  | .CONST _ => ∅
  | .VAR v => Std.HashSet.instSingleton.singleton v
  | .PLUS a1 a2 =>  (fv_aexp a1) ∪ (fv_aexp a2)
  | .MINUS a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)

@[grind] def fv_bexp (b : bexp) : IdentSet :=
  match b with
  | .TRUE => ∅
  | .FALSE => ∅
  | .EQUAL a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)
  | .LESSEQUAL a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)
  | .NOT b1 => fv_bexp b1
  | .AND b1 b2 => (fv_bexp b1) ∪ (fv_bexp b2)

@[grind] def fv_com (c : com) : IdentSet :=
  match c with
  | .SKIP => ∅
  | .ASSIGN _ a => fv_aexp a
  | .SEQ c1 c2 => (fv_com c1) ∪ (fv_com c2)
  | .IFTHENELSE b c1 c2 => (fv_bexp b) ∪ ((fv_com c1) ∪ (fv_com c2))
  | .WHILE b c => (fv_bexp b) ∪ (fv_com c)

@[grind] noncomputable def deadcode_fixpoint_rec (F : IdentSet → IdentSet) (default : IdentSet) (fuel : Nat) (x : IdentSet) : IdentSet :=
  match fuel with
  | 0 => default
  | fuel + 1 =>
      let x' := F x
      if x' ⊆ x then x else deadcode_fixpoint_rec F default fuel x'

@[grind] noncomputable def deadcode_fixpoint (F : IdentSet → IdentSet) (default : IdentSet) : IdentSet :=
  deadcode_fixpoint_rec F default 20 ∅

@[grind] noncomputable def live (c : com) (L : IdentSet) : IdentSet :=
  match c with
  | .SKIP => L
  | .ASSIGN x a =>
      if x ∈ L
      then (L.erase x) ∪ (fv_aexp a)
      else L
  | .SEQ c1 c2 =>
      live c1 (live c2 L)
  | .IFTHENELSE b c1 c2 =>
       (fv_bexp b) ∪ ((live c1 L) ∪ (live c2 L))
  | .WHILE b c =>
      let L' := (fv_bexp b) ∪  L
      let default := (fv_com (.WHILE b c)) ∪ L
      deadcode_fixpoint (fun x => L' ∪ (live c x)) default

@[grind] noncomputable def dce (c : com) (L : IdentSet) : com :=
  match c with
  | .SKIP => .SKIP
  | .ASSIGN x a => if x ∈ L then .ASSIGN x a else .SKIP
  | .SEQ c1 c2 => .SEQ (dce c1 (live c2 L)) (dce c2 L)
  | .IFTHENELSE b c1 c2 => .IFTHENELSE b (dce c1 L) (dce c2 L)
  | .WHILE b c => .WHILE b (dce c (live (.WHILE b c) L))

@[grind] def agree (L : IdentSet) (s1 s2 : store) : Prop :=
  ∀ x, x  ∈ L -> s1 x = s2 x

-- 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: abbrev IdentSet := Std.HashSet ident

content: @[grind] def fv_aexp (a : aexp) : IdentSet :=
  match a with
  | .CONST _ => ∅
  | .VAR v => Std.HashSet.instSingleton.singleton v
  | .PLUS a1 a2 =>  (fv_aexp a1) ∪ (fv_aexp a2)
  | .MINUS a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)

content: @[grind] def fv_bexp (b : bexp) : IdentSet :=
  match b with
  | .TRUE => ∅
  | .FALSE => ∅
  | .EQUAL a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)
  | .LESSEQUAL a1 a2 => (fv_aexp a1) ∪ (fv_aexp a2)
  | .NOT b1 => fv_bexp b1
  | .AND b1 b2 => (fv_bexp b1) ∪ (fv_bexp b2)

content: @[grind] def fv_com (c : com) : IdentSet :=
  match c with
  | .SKIP => ∅
  | .ASSIGN _ a => fv_aexp a
  | .SEQ c1 c2 => (fv_com c1) ∪ (fv_com c2)
  | .IFTHENELSE b c1 c2 => (fv_bexp b) ∪ ((fv_com c1) ∪ (fv_com c2))
  | .WHILE b c => (fv_bexp b) ∪ (fv_com c)

content: @[grind] noncomputable def deadcode_fixpoint_rec (F : IdentSet → IdentSet) (default : IdentSet) (fuel : Nat) (x : IdentSet) : IdentSet :=
  match fuel with
  | 0 => default
  | fuel + 1 =>
      let x' := F x
      if x' ⊆ x then x else deadcode_fixpoint_rec F default fuel x'

content: @[grind] noncomputable def deadcode_fixpoint (F : IdentSet → IdentSet) (default : IdentSet) : IdentSet :=
  deadcode_fixpoint_rec F default 20 ∅

content: @[grind] noncomputable def live (c : com) (L : IdentSet) : IdentSet :=
  match c with
  | .SKIP => L
  | .ASSIGN x a =>
      if x ∈ L
      then (L.erase x) ∪ (fv_aexp a)
      else L
  | .SEQ c1 c2 =>
      live c1 (live c2 L)
  | .IFTHENELSE b c1 c2 =>
       (fv_bexp b) ∪ ((live c1 L) ∪ (live c2 L))
  | .WHILE b c =>
      let L' := (fv_bexp b) ∪  L
      let default := (fv_com (.WHILE b c)) ∪ L
      deadcode_fixpoint (fun x => L' ∪ (live c x)) default

content: @[grind] noncomputable def dce (c : com) (L : IdentSet) : com :=
  match c with
  | .SKIP => .SKIP
  | .ASSIGN x a => if x ∈ L then .ASSIGN x a else .SKIP
  | .SEQ c1 c2 => .SEQ (dce c1 (live c2 L)) (dce c2 L)
  | .IFTHENELSE b c1 c2 => .IFTHENELSE b (dce c1 L) (dce c2 L)
  | .WHILE b c => .WHILE b (dce c (live (.WHILE b c) L))

content: @[grind] def agree (L : IdentSet) (s1 s2 : store) : Prop :=
  ∀ x, x  ∈ L -> s1 x = s2 x

-- Local lemmas used:
content: @[grind =] theorem subset_def (a b : IdentSet) : a ⊆ b ↔ ∀ x ∈ a, x ∈ b

content: theorem live_while_charact (b : bexp) (c : com) (L L' : IdentSet)
  (eq : L' = live (.WHILE b c) L) :
  (fv_bexp b) ⊆ L' ∧ L ⊆ L' ∧ (live c L') ⊆ L'

content: theorem beval_agree :
  ∀ L s1 s2, agree L s1 s2 ->
  ∀ b, (fv_bexp b) ⊆ L -> beval s1 b = beval s2 b

-- Target theorem and proof:
theorem dce_correct_terminating :
  ∀ s c s', cexec s c s' ->
  ∀ L s1, agree (live c L) s s1 ->
  ∃ s1', cexec s1 (dce c L) s1' /\ agree L s' s1'
:= by
    intro s c s' EXEC
    induction EXEC
    any_goals grind
    case cexec_while_loop s1 b c1 s2 s3 isTrue EX1 EX2 a_ih a_ih2 =>
      intro L s4 hyp
      have ⟨t1, ht1, ht2⟩ :=  a_ih (live (.WHILE b c1) L) s4 (by grind)
      have ⟨u1, hu1, hu2⟩ :=  a_ih2 L t1 ht2
      exists u1
      constructor
      rotate_right
      · exact hu2
      · apply cexec.cexec_while_loop
        · have := beval_agree (live (com.WHILE b c1) L) s1 s4
          grind
        · exact ht1
        · grind
    case cexec_assign s2 x a=>
      intro L s3 AG
      simp [live] at AG
      by_cases x ∈ L
      case neg notIn =>
        exists s3
        grind
      case pos isIn =>
        exists (update x (aeval s3 a) s3)
        grind
    case cexec_ifthenelse s2 b c1 c2 s3 EXEC ih =>
      intro L s4 AG
      simp [dce]
      have EQ : beval s2 b = beval s4 b := by
        apply beval_agree
        · apply AG
        · grind
      by_cases beval s2 b = true
      case pos isTrue =>
        specialize ih L s4
        grind
      case neg isFalse =>
        specialize ih L s4
        grind
    case cexec_while_done s2 b c isFalse =>
      intro L s1 AG
      have ⟨h1, h2, h3⟩ := live_while_charact b c L (live (com.WHILE b c) L) (by grind)
      have EQ : beval s2 b = beval s1 b := by
        apply beval_agree
        · apply AG
        · grind
      exists s1
      grind
    case cexec_seq => grind [-subset_def] -- TODO: what does wrong?