import TsmLean.Compile.Compile
import TsmLean.Core.Semantics

namespace TsmLean.Compile

open TsmLean.Core

/-! # Compiler-correctness theorem.

  `Source.Eval e v  ⟹  TSM.MultiStep (start of compile e) (end of compile e, with v on stack)`

The CompCert-flavored bridge: source-level evaluation and target-level
execution agree on the value produced. -/

/-! ## Multi-step utilities. -/

theorem MultiStep.trans
    {s₁ s₂ s₃ : State}
    (h₁ : MultiStep s₁ s₂) (h₂ : MultiStep s₂ s₃) :
    MultiStep s₁ s₃ := by
  induction h₁ with
  | refl => exact h₂
  | cons hs _ ih => exact .cons hs (ih h₂)

theorem MultiStep.single
    {s s' : State} (h : step s = some s') :
    MultiStep s s' := .cons h (.refl _)

/-! ## Single-step reduction lemmas. -/

theorem step_push
    {code : Code} {pc : Nat} {stack : List Value} {n : Int}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .push n) :
    step { code := code, pc := pc, stack := stack }
    = some { code := code, pc := pc + 1, stack := .vInt n :: stack } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_pushB
    {code : Code} {pc : Nat} {stack : List Value} {b : Bool}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .pushB b) :
    step { code := code, pc := pc, stack := stack }
    = some { code := code, pc := pc + 1, stack := .vBool b :: stack } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_add
    {code : Code} {pc : Nat} {a b : Int} {rest : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .add) :
    step { code := code, pc := pc, stack := .vInt b :: .vInt a :: rest }
    = some { code := code, pc := pc + 1, stack := .vInt (a + b) :: rest } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_sub
    {code : Code} {pc : Nat} {a b : Int} {rest : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .sub) :
    step { code := code, pc := pc, stack := .vInt b :: .vInt a :: rest }
    = some { code := code, pc := pc + 1, stack := .vInt (a - b) :: rest } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_mul
    {code : Code} {pc : Nat} {a b : Int} {rest : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .mul) :
    step { code := code, pc := pc, stack := .vInt b :: .vInt a :: rest }
    = some { code := code, pc := pc + 1, stack := .vInt (a * b) :: rest } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_jmp
    {code : Code} {pc k : Nat} {stack : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .jmp k) :
    step { code := code, pc := pc, stack := stack }
    = some { code := code, pc := k, stack := stack } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_jmpFalse_true
    {code : Code} {pc k : Nat} {rest : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .jmpFalse k) :
    step { code := code, pc := pc, stack := .vBool true :: rest }
    = some { code := code, pc := pc + 1, stack := rest } := by
  unfold step; rw [dif_pos h_pc, h_get]

theorem step_jmpFalse_false
    {code : Code} {pc k : Nat} {rest : List Value}
    (h_pc : pc < code.size) (h_get : code[pc]'h_pc = .jmpFalse k) :
    step { code := code, pc := pc, stack := .vBool false :: rest }
    = some { code := code, pc := k, stack := rest } := by
  unfold step; rw [dif_pos h_pc, h_get]

/-! ## Generic array-lookup helpers. -/

/-- The instruction at the boundary of a `c1 ++ c2 ++ #[op]` arrangement. -/
theorem getElem_at_op_boundary
    (c1 c2 : Code) (op : Instr)
    (h : c1.size + c2.size < (c1 ++ c2 ++ #[op]).size) :
    (c1 ++ c2 ++ #[op])[c1.size + c2.size]'h = op := by
  show ((c1 ++ c2) ++ #[op])[c1.size + c2.size]'h = op
  have hle : (c1 ++ c2).size ≤ c1.size + c2.size := Nat.le_of_eq Array.size_append
  rw [Array.getElem_append_right hle]
  simp [Array.size_append, Nat.sub_self]

/-- Lookup at offset `pre.size + i` within a `pre ++ X ++ suf` array reduces
    to lookup at `i` within `X` itself, when `i < X.size`. -/
theorem getElem_at_offset
    (pre X suf : Code) (i : Nat)
    (h_lt : i < X.size)
    (h_pc : pre.size + i < (pre ++ X ++ suf).size) :
    (pre ++ X ++ suf)[pre.size + i]'h_pc = X[i]'h_lt := by
  have h_pre_X : pre.size + i < (pre ++ X).size := by
    rw [Array.size_append]; omega
  rw [Array.getElem_append_left h_pre_X]
  rw [Array.getElem_append_right (Nat.le_add_right _ _)]
  congr 1; omega

/-! ## Per-construct compile-output lookup lemmas. -/

theorem compile_add_get_op (offset : Nat) (e1 e2 : Source.Expr)
    (h : (compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size
       < (compile offset (Source.Expr.add e1 e2)).size) :
    (compile offset (Source.Expr.add e1 e2))[(compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size]'h = .add :=
  getElem_at_op_boundary (compile offset e1) (compile (offset + (compile offset e1).size) e2) _ h

theorem compile_sub_get_op (offset : Nat) (e1 e2 : Source.Expr)
    (h : (compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size
       < (compile offset (Source.Expr.sub e1 e2)).size) :
    (compile offset (Source.Expr.sub e1 e2))[(compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size]'h = .sub :=
  getElem_at_op_boundary (compile offset e1) (compile (offset + (compile offset e1).size) e2) _ h

theorem compile_mul_get_op (offset : Nat) (e1 e2 : Source.Expr)
    (h : (compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size
       < (compile offset (Source.Expr.mul e1 e2)).size) :
    (compile offset (Source.Expr.mul e1 e2))[(compile offset e1).size + (compile (offset + (compile offset e1).size) e2).size]'h = .mul :=
  getElem_at_op_boundary (compile offset e1) (compile (offset + (compile offset e1).size) e2) _ h


/-! ## Main theorem.

The compile is offset-aware: `compile pre.size e` produces bytecode
correctly placed at position `pre.size` in `pre ++ ... ++ suf`. -/

theorem compile_correct
    {e : Source.Expr} {v : Source.Value}
    (h_eval : Source.Eval e v) :
    ∀ (pre suf : Code) (rest : List Value),
      MultiStep
        { code := pre ++ compile pre.size e ++ suf, pc := pre.size, stack := rest }
        { code := pre ++ compile pre.size e ++ suf,
          pc := pre.size + (compile pre.size e).size, stack := v :: rest } := by
  induction h_eval with
  | intLit n =>
    intros pre suf rest
    apply MultiStep.single
    have h_pc : pre.size < (pre ++ compile pre.size (Source.Expr.intLit n) ++ suf).size := by
      simp only [compile, Array.size_append, Array.size_singleton]; omega
    have h_get : (pre ++ compile pre.size (Source.Expr.intLit n) ++ suf)[pre.size]'h_pc = .push n := by
      have h_pre_ce : pre.size < (pre ++ compile pre.size (Source.Expr.intLit n)).size := by
        simp only [compile, Array.size_append, Array.size_singleton]; omega
      rw [Array.getElem_append_left h_pre_ce]
      rw [Array.getElem_append_right (Nat.le_refl _)]
      simp [compile, Nat.sub_self]
    have step_thm := step_push h_pc h_get (stack := rest)
    have h_size : (compile pre.size (Source.Expr.intLit n)).size = 1 := by simp [compile]
    rw [h_size]
    exact step_thm
  | boolLit b =>
    intros pre suf rest
    apply MultiStep.single
    have h_pc : pre.size < (pre ++ compile pre.size (Source.Expr.boolLit b) ++ suf).size := by
      simp only [compile, Array.size_append, Array.size_singleton]; omega
    have h_get : (pre ++ compile pre.size (Source.Expr.boolLit b) ++ suf)[pre.size]'h_pc = .pushB b := by
      have h_pre_ce : pre.size < (pre ++ compile pre.size (Source.Expr.boolLit b)).size := by
        simp only [compile, Array.size_append, Array.size_singleton]; omega
      rw [Array.getElem_append_left h_pre_ce]
      rw [Array.getElem_append_right (Nat.le_refl _)]
      simp [compile, Nat.sub_self]
    have step_thm := step_pushB h_pc h_get (stack := rest)
    have h_size : (compile pre.size (Source.Expr.boolLit b)).size = 1 := by simp [compile]
    rw [h_size]
    exact step_thm
  | @add e1 e2 a b _ _ ih1 ih2 =>
    intros pre suf rest
    have stepA := ih1 pre (compile (pre.size + (compile pre.size e1).size) e2 ++ #[.add] ++ suf) rest
    have stepB := ih2 (pre ++ compile pre.size e1) (#[.add] ++ suf) (.vInt a :: rest)
    have h_pre_e1_size : (pre ++ compile pre.size e1).size = pre.size + (compile pre.size e1).size := by
      simp [Array.size_append]
    have h_code_A : pre ++ compile pre.size e1 ++ (compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.add] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.add e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    have h_code_B : pre ++ compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ (#[Instr.add] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.add e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    rw [h_code_A] at stepA
    rw [h_pre_e1_size] at stepB
    rw [h_code_B] at stepB
    apply MultiStep.trans stepA
    apply MultiStep.trans stepB
    apply MultiStep.single
    have h_total_size : (compile pre.size (Source.Expr.add e1 e2)).size
                      = (compile pre.size e1).size
                      + (compile (pre.size + (compile pre.size e1).size) e2).size + 1 := by
      show (compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.add]).size
         = (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size + 1
      simp [Array.size_append]; omega
    have h_in_comp : (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                     < (compile pre.size (Source.Expr.add e1 e2)).size := by
      simp [compile, Array.size_append]
    have h_full_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)
                     < (pre ++ compile pre.size (Source.Expr.add e1 e2) ++ suf).size := by
      simp only [Array.size_append, h_total_size]; omega
    have h_op_get : (pre ++ compile pre.size (Source.Expr.add e1 e2) ++ suf)[pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)]'h_full_pc = .add := by
      rw [getElem_at_offset pre (compile pre.size (Source.Expr.add e1 e2)) suf
            ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) h_in_comp h_full_pc]
      exact compile_add_get_op pre.size e1 e2 h_in_comp
    have h_step := step_add (a := a) (b := b) (rest := rest) h_full_pc h_op_get
    have h_pre_pc : pre.size + (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                  = pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) := by omega
    have h_post_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) + 1
                   = pre.size + (compile pre.size (Source.Expr.add e1 e2)).size := by
      rw [h_total_size]; omega
    rw [h_pre_pc, ← h_post_pc]
    exact h_step
  | @sub e1 e2 a b _ _ ih1 ih2 =>
    intros pre suf rest
    have stepA := ih1 pre (compile (pre.size + (compile pre.size e1).size) e2 ++ #[.sub] ++ suf) rest
    have stepB := ih2 (pre ++ compile pre.size e1) (#[.sub] ++ suf) (.vInt a :: rest)
    have h_pre_e1_size : (pre ++ compile pre.size e1).size = pre.size + (compile pre.size e1).size := by
      simp [Array.size_append]
    have h_code_A : pre ++ compile pre.size e1 ++ (compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.sub] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.sub e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    have h_code_B : pre ++ compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ (#[Instr.sub] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.sub e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    rw [h_code_A] at stepA
    rw [h_pre_e1_size] at stepB
    rw [h_code_B] at stepB
    apply MultiStep.trans stepA
    apply MultiStep.trans stepB
    apply MultiStep.single
    have h_total_size : (compile pre.size (Source.Expr.sub e1 e2)).size
                      = (compile pre.size e1).size
                      + (compile (pre.size + (compile pre.size e1).size) e2).size + 1 := by
      show (compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.sub]).size
         = (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size + 1
      simp [Array.size_append]; omega
    have h_in_comp : (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                     < (compile pre.size (Source.Expr.sub e1 e2)).size := by
      simp [compile, Array.size_append]
    have h_full_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)
                     < (pre ++ compile pre.size (Source.Expr.sub e1 e2) ++ suf).size := by
      simp only [Array.size_append, h_total_size]; omega
    have h_op_get : (pre ++ compile pre.size (Source.Expr.sub e1 e2) ++ suf)[pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)]'h_full_pc = .sub := by
      rw [getElem_at_offset pre (compile pre.size (Source.Expr.sub e1 e2)) suf
            ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) h_in_comp h_full_pc]
      exact compile_sub_get_op pre.size e1 e2 h_in_comp
    have h_step := step_sub (a := a) (b := b) (rest := rest) h_full_pc h_op_get
    have h_pre_pc : pre.size + (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                  = pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) := by omega
    have h_post_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) + 1
                   = pre.size + (compile pre.size (Source.Expr.sub e1 e2)).size := by
      rw [h_total_size]; omega
    rw [h_pre_pc, ← h_post_pc]
    exact h_step
  | @mul e1 e2 a b _ _ ih1 ih2 =>
    intros pre suf rest
    have stepA := ih1 pre (compile (pre.size + (compile pre.size e1).size) e2 ++ #[.mul] ++ suf) rest
    have stepB := ih2 (pre ++ compile pre.size e1) (#[.mul] ++ suf) (.vInt a :: rest)
    have h_pre_e1_size : (pre ++ compile pre.size e1).size = pre.size + (compile pre.size e1).size := by
      simp [Array.size_append]
    have h_code_A : pre ++ compile pre.size e1 ++ (compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.mul] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.mul e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    have h_code_B : pre ++ compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ (#[Instr.mul] ++ suf)
                  = pre ++ compile pre.size (Source.Expr.mul e1 e2) ++ suf := by
      simp only [compile, Array.append_assoc]
    rw [h_code_A] at stepA
    rw [h_pre_e1_size] at stepB
    rw [h_code_B] at stepB
    apply MultiStep.trans stepA
    apply MultiStep.trans stepB
    apply MultiStep.single
    have h_total_size : (compile pre.size (Source.Expr.mul e1 e2)).size
                      = (compile pre.size e1).size
                      + (compile (pre.size + (compile pre.size e1).size) e2).size + 1 := by
      show (compile pre.size e1 ++ compile (pre.size + (compile pre.size e1).size) e2 ++ #[Instr.mul]).size
         = (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size + 1
      simp [Array.size_append]; omega
    have h_in_comp : (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                     < (compile pre.size (Source.Expr.mul e1 e2)).size := by
      simp [compile, Array.size_append]
    have h_full_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)
                     < (pre ++ compile pre.size (Source.Expr.mul e1 e2) ++ suf).size := by
      simp only [Array.size_append, h_total_size]; omega
    have h_op_get : (pre ++ compile pre.size (Source.Expr.mul e1 e2) ++ suf)[pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size)]'h_full_pc = .mul := by
      rw [getElem_at_offset pre (compile pre.size (Source.Expr.mul e1 e2)) suf
            ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) h_in_comp h_full_pc]
      exact compile_mul_get_op pre.size e1 e2 h_in_comp
    have h_step := step_mul (a := a) (b := b) (rest := rest) h_full_pc h_op_get
    have h_pre_pc : pre.size + (compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size
                  = pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) := by omega
    have h_post_pc : pre.size + ((compile pre.size e1).size + (compile (pre.size + (compile pre.size e1).size) e2).size) + 1
                   = pre.size + (compile pre.size (Source.Expr.mul e1 e2)).size := by
      rw [h_total_size]; omega
    rw [h_pre_pc, ← h_post_pc]
    exact h_step

end TsmLean.Compile