-- LeroyCompilerVerificationCourse/Compil.lean

import LeroyCompilerVerificationCourse.Imp

@[grind] inductive instr : Type where
  | Iconst (n : Int)             
  | Ivar (x : ident)             
  | Isetvar (x : ident)          
  | Iadd                         
  | Iopp                         
  | Ibranch (d : Int)            
  | Ibeq (d1 : Int) (d0 : Int)   
  | Ible (d1 : Int) (d0 : Int)   
  | Ihalt                        
    deriving Repr

@[grind] def codelen (c : List instr) : Int := c.length

def stack : Type := List Int

def config : Type := Int × stack × store

@[grind] def instr_at (C : List instr) (pc : Int) : Option instr :=
  match C with
  | [] => .none
  | i :: C' => if pc = 0 then .some i else instr_at C' (pc - 1)

@[grind] inductive transition (C : List instr) : config → config → Prop where
  | trans_const : ∀ pc stk s n,
      instr_at C pc = .some (.Iconst n) →
      transition C (pc    , stk     , s)
                   (pc + 1, n :: stk, s)
  | trans_var : ∀ pc stk s x,
      instr_at C pc = .some (.Ivar x) ->
      transition C (pc    , stk       , s)
                   (pc + 1, s x :: stk, s)
  | trans_setvar : ∀ pc stk s x n,
      instr_at C pc = .some (.Isetvar x) ->
      transition C (pc    , n :: stk, s)
                   (pc + 1, stk     , update x n s)
  | trans_add : ∀ pc stk s n1 n2,
      instr_at C pc = .some (.Iadd) ->
      transition C (pc    , n2 :: n1 :: stk , s)
                   (pc + 1, (n1 + n2) :: stk, s)
  | trans_opp : ∀ pc stk s n,
      instr_at C pc = .some (.Iopp) ->
      transition C (pc    , n :: stk    , s)
                   (pc + 1, (- n) :: stk, s)
  | trans_branch : ∀ pc stk s d pc',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      transition C (pc , stk, s)
                   (pc', stk, s)
  | trans_beq : ∀ pc stk s d1 d0 n1 n2 pc',
      instr_at C pc = .some (.Ibeq d1 d0) ->
      pc' = pc + 1 + (if n1 = n2 then d1 else d0) ->
      transition C (pc , n2 :: n1 :: stk, s)
                   (pc', stk            , s)
  | trans_ble : ∀ pc stk s d1 d0 n1 n2 pc',
      instr_at C pc = .some (.Ible d1 d0) ->
      pc' = pc + 1 + (if n1 ≤ n2 then d1 else d0) ->
      transition C (pc , n2 :: n1 :: stk, s)
                   (pc', stk            , s)

@[grind] def transitions (C : List instr) : config → config → Prop :=
  star (transition C)

@[grind] def compile_aexp (a : aexp) : List instr :=
  match a with
  | .CONST n => .Iconst n :: []
  | .VAR x => .Ivar x :: []
  | .PLUS a1 a2  => (compile_aexp a1) ++ (compile_aexp a2) ++ (.Iadd :: [])
  | .MINUS a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ (.Iopp :: .Iadd :: [])

@[grind] def compile_bexp (b : bexp) (d1 : Int) (d0 : Int) : List instr :=
  match b with
  | .TRUE => if d1 = 0 then [] else .Ibranch d1 :: []
  | .FALSE => if d0 = 0 then [] else .Ibranch d0 :: []
  | .EQUAL a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ .Ibeq d1 d0 :: []
  | .LESSEQUAL a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ .Ible d1 d0 :: []
  | .NOT b1 => compile_bexp b1 d0 d1
  | .AND b1 b2 =>
      let code2 := compile_bexp b2 d1 d0
      let code1 := compile_bexp b1 0 (codelen code2 + d0)
      code1 ++ code2

@[grind] def compile_com (c : com) : List instr :=
  match c with
  | .SKIP =>
      []
  | .ASSIGN x a =>
      compile_aexp a ++ .Isetvar x :: []
  | .SEQ c1 c2 =>
      compile_com c1 ++ compile_com c2
  | .IFTHENELSE b ifso ifnot =>
      let code_ifso := compile_com ifso
      let code_ifnot := compile_com ifnot
      compile_bexp b 0 (codelen code_ifso + 1)
      ++ code_ifso
      ++ .Ibranch (codelen code_ifnot)
      :: code_ifnot
  | .WHILE b body =>
      let code_body := compile_com body
      let code_test := compile_bexp b 0 (codelen code_body + 1)
      code_test
      ++ code_body
      ++ .Ibranch (- (codelen code_test + codelen code_body + 1)) :: []

@[grind] inductive code_at : List instr → Int → List instr → Prop where
  | code_at_intro : ∀ C1 C2 C3 pc,
      pc = codelen C1 ->
      code_at (C1 ++ C2 ++ C3) pc C2

-- Repository definitions used:
content: def store : Type := ident → Int

content: def ident := String deriving BEq, Repr, Hashable

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

content: @[grind] inductive star (R : α → α → Prop) : α → α → Prop where
  | star_refl : ∀ x : α, star R x x
  | star_step : ∀ {x y z}, R x y → star R y z → star R x z

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: notation:10 l:10 " ;; " r:11 => com.SEQ l r

-- Repository lemmas used:
content: @[grind] theorem star_trans {α} (R : α → α → Prop) (a b : α) (sab : star R a b) : ∀ c : α, star R b c → star R a c

-- Local definitions used:
content: @[grind] inductive instr : Type where
  | Iconst (n : Int)             
  | Ivar (x : ident)             
  | Isetvar (x : ident)          
  | Iadd                         
  | Iopp                         
  | Ibranch (d : Int)            
  | Ibeq (d1 : Int) (d0 : Int)   
  | Ible (d1 : Int) (d0 : Int)   
  | Ihalt                        
    deriving Repr

content: @[grind] def codelen (c : List instr) : Int := c.length

content: def stack : Type := List Int

content: def config : Type := Int × stack × store

content: @[grind] def instr_at (C : List instr) (pc : Int) : Option instr :=
  match C with
  | [] => .none
  | i :: C' => if pc = 0 then .some i else instr_at C' (pc - 1)

content: @[grind] inductive transition (C : List instr) : config → config → Prop where
  | trans_const : ∀ pc stk s n,
      instr_at C pc = .some (.Iconst n) →
      transition C (pc    , stk     , s)
                   (pc + 1, n :: stk, s)
  | trans_var : ∀ pc stk s x,
      instr_at C pc = .some (.Ivar x) ->
      transition C (pc    , stk       , s)
                   (pc + 1, s x :: stk, s)
  | trans_setvar : ∀ pc stk s x n,
      instr_at C pc = .some (.Isetvar x) ->
      transition C (pc    , n :: stk, s)
                   (pc + 1, stk     , update x n s)
  | trans_add : ∀ pc stk s n1 n2,
      instr_at C pc = .some (.Iadd) ->
      transition C (pc    , n2 :: n1 :: stk , s)
                   (pc + 1, (n1 + n2) :: stk, s)
  | trans_opp : ∀ pc stk s n,
      instr_at C pc = .some (.Iopp) ->
      transition C (pc    , n :: stk    , s)
                   (pc + 1, (- n) :: stk, s)
  | trans_branch : ∀ pc stk s d pc',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      transition C (pc , stk, s)
                   (pc', stk, s)
  | trans_beq : ∀ pc stk s d1 d0 n1 n2 pc',
      instr_at C pc = .some (.Ibeq d1 d0) ->
      pc' = pc + 1 + (if n1 = n2 then d1 else d0) ->
      transition C (pc , n2 :: n1 :: stk, s)
                   (pc', stk            , s)
  | trans_ble : ∀ pc stk s d1 d0 n1 n2 pc',
      instr_at C pc = .some (.Ible d1 d0) ->
      pc' = pc + 1 + (if n1 ≤ n2 then d1 else d0) ->
      transition C (pc , n2 :: n1 :: stk, s)
                   (pc', stk            , s)

content: @[grind] def transitions (C : List instr) : config → config → Prop :=
  star (transition C)

content: @[grind] def compile_aexp (a : aexp) : List instr :=
  match a with
  | .CONST n => .Iconst n :: []
  | .VAR x => .Ivar x :: []
  | .PLUS a1 a2  => (compile_aexp a1) ++ (compile_aexp a2) ++ (.Iadd :: [])
  | .MINUS a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ (.Iopp :: .Iadd :: [])

content: @[grind] def compile_bexp (b : bexp) (d1 : Int) (d0 : Int) : List instr :=
  match b with
  | .TRUE => if d1 = 0 then [] else .Ibranch d1 :: []
  | .FALSE => if d0 = 0 then [] else .Ibranch d0 :: []
  | .EQUAL a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ .Ibeq d1 d0 :: []
  | .LESSEQUAL a1 a2 => compile_aexp a1 ++ compile_aexp a2 ++ .Ible d1 d0 :: []
  | .NOT b1 => compile_bexp b1 d0 d1
  | .AND b1 b2 =>
      let code2 := compile_bexp b2 d1 d0
      let code1 := compile_bexp b1 0 (codelen code2 + d0)
      code1 ++ code2

content: @[grind] def compile_com (c : com) : List instr :=
  match c with
  | .SKIP =>
      []
  | .ASSIGN x a =>
      compile_aexp a ++ .Isetvar x :: []
  | .SEQ c1 c2 =>
      compile_com c1 ++ compile_com c2
  | .IFTHENELSE b ifso ifnot =>
      let code_ifso := compile_com ifso
      let code_ifnot := compile_com ifnot
      compile_bexp b 0 (codelen code_ifso + 1)
      ++ code_ifso
      ++ .Ibranch (codelen code_ifnot)
      :: code_ifnot
  | .WHILE b body =>
      let code_body := compile_com body
      let code_test := compile_bexp b 0 (codelen code_body + 1)
      code_test
      ++ code_body
      ++ .Ibranch (- (codelen code_test + codelen code_body + 1)) :: []

content: @[grind] inductive code_at : List instr → Int → List instr → Prop where
  | code_at_intro : ∀ C1 C2 C3 pc,
      pc = codelen C1 ->
      code_at (C1 ++ C2 ++ C3) pc C2

-- Local lemmas used:
content: @[grind =] theorem codelen_cons :
  ∀ i c, codelen (i :: c) = codelen c + 1

content: @[grind =] theorem codelen_singleton : codelen [i] = 1

content: @[grind =] theorem codelen_app :
  ∀ c1 c2, codelen (c1 ++ c2) = codelen c1 + codelen c2

content: @[grind =>] theorem instr_a : ∀ i c2 c1 pc,
  pc = codelen c1 ->
  instr_at (c1 ++ (i :: c2) ) pc = .some i

content: @[grind] theorem code_at_app_right :
  ∀ C pc C1 C2,
  code_at C pc (C1 ++ C2) ->
  code_at C (pc + codelen C1) C2

content: @[grind] theorem code_at_to_instr_at : code_at C pc (c1 ++ i :: c2) → instr_at C (pc + codelen c1) = .some i

content: theorem compile_aexp_correct (C : List instr) (s : store) (a : aexp) (pc : Int) (stk : stack) :
  code_at C pc (compile_aexp a) →
  transitions C (pc, stk, s) (pc + codelen (compile_aexp a), aeval s a :: stk, s)

content: theorem compile_bexp_correct (C : List instr) (s : store) (b : bexp) (d1 d0 : Int) (pc : Int) (stk : stack) (h : code_at C pc (compile_bexp b d1 d0))  :
   transitions C
       (pc, stk, s)
       (pc + codelen (compile_bexp b d1 d0) + (if beval s b then d1 else d0), stk, s)

-- Target theorem and proof:
theorem compile_com_correct_terminating (s s' : store) (c : com) (h₁ : cexec s c s') :
 ∀ C pc stk, code_at C pc (compile_com c) →
  transitions C
      (pc, stk, s)
      (pc + codelen (compile_com c), stk, s')
:= by
  induction h₁
  case cexec_skip =>
    intro C pc stk h
    unfold compile_com
    dsimp [codelen]
    simp only [Int.add_zero]
    apply star.star_refl
  case cexec_assign s' x a =>
    intro C pc stk h
    unfold compile_com
    apply star_trans
    · apply compile_aexp_correct (a := a)
      grind
    · apply star_one
      · have := @transition.trans_setvar C
        grind
  case cexec_seq s'2 c1 s1 c2 s2 cexec1 cexec2 c1_ih c2_ih =>
    intro C pc stk h
    apply star_trans
    · apply c1_ih
      grind
    · specialize c2_ih C (pc + codelen (compile_com c1)) stk
      simp [compile_com, codelen_app]
      simp [Int.add_assoc] at c2_ih
      apply c2_ih
      grind
  case cexec_ifthenelse s b c1 c2 s' cexec_h ih =>
    intro C pc stk
    generalize heq1 : compile_com c1 = code1
    generalize heq2 : compile_com c2 = code2
    generalize heq3 : compile_bexp b 0 (codelen code1 + 1) = code3
    simp [compile_com]
    rw [heq1, heq2, heq3]
    intro h
    apply star_trans
    · have := compile_bexp_correct C s b 0 (codelen code1 + 1) pc stk (by grind)
      apply this
    · by_cases beval s b = true
      case pos isTrue =>
        simp [isTrue]
        apply star_trans
        · apply ih
          grind
        · apply star_one
          · apply transition.trans_branch (d := codelen code2) <;> grind
      case neg isFalse =>
        simp [isFalse]
        rw [heq3]
        specialize ih C (pc + codelen code3 + (codelen code1 + 1)) stk
        simp [isFalse] at ih
        suffices h2 : code_at C (pc + codelen code3 + (codelen code1 + 1)) (compile_com c2) from by
          specialize ih h2
          simp [codelen_app, codelen_cons]
          have : (pc + codelen code3 + (codelen code1 + 1) + codelen (compile_com c2)) = (pc + (codelen code3 + (codelen code1 + (codelen code2 + 1)))) := by grind
          rw [this] at ih
          apply ih
        have := @code_at_app_right C pc (code3 ++ code1 ++ [instr.Ibranch (codelen code2)]) code2 (by simp[h])
        grind
  case cexec_while_done s b c1 isFalse =>
    intro C pc stk h
    generalize heq1 : compile_com c1 = code_body
    generalize heq2 : compile_bexp b 0 (codelen code_body + 1) = code_branch
    generalize heq3 : - (codelen code_branch + codelen code_body + 1) = d
    simp [compile_com]
    rw [heq1, heq2, heq3]
    simp [codelen_app, codelen_singleton]
    apply star_trans
    · apply compile_bexp_correct C s b 0 (codelen code_body + 1) pc stk (by grind)
    · grind
  case cexec_while_loop s b c1 s_intermediate s' isTrue cexec1 cexec2 ih1 ih2 =>
    intro C pc stk
    generalize heq1 : compile_com c1 = code_body
    generalize heq2 : compile_bexp b 0 (codelen code_body + 1) = code_branch
    generalize heq3 : - (codelen code_branch + codelen code_body + 1) = d
    simp [compile_com]
    rw [heq1, heq2, heq3]
    intro h
    apply star_trans
    · apply compile_bexp_correct C s b 0 (codelen code_body + 1) pc stk (by grind)
    · apply star_trans
      · apply ih1
        grind
      · apply star_trans
        · apply star_one
          apply transition.trans_branch (d := d)
          rotate_left
          rotate_left
          · exact (pc + codelen code_branch + codelen code_body + 1 + d)
          · grind
          · grind
        · specialize ih2 C (pc + codelen code_branch + codelen code_body + 1 + d) stk
          suffices h2 : code_at C (pc + codelen code_branch + codelen code_body + 1 + d) (compile_com (com.WHILE b c1)) from by
            specialize ih2 h2
            simp [compile_com] at ih2
            rw [heq1, heq2, heq3] at ih2
            simp [codelen_app]
            simp [codelen_app] at ih2
            have : (pc + codelen code_branch + codelen code_body + 1 + d +
      (codelen code_branch + (codelen code_body + codelen [instr.Ibranch d])) ) = (pc + (codelen code_branch + (codelen code_body + codelen [instr.Ibranch d]))) := by grind
            rw [←this]
            apply ih2
          grind