import TsmLean.Core.TypeSoundness

namespace TsmLean.Core

/-! # Preservation and progress for TSM.

Local (per-instruction) preservation: if the stack matches an
instruction's input type and that instruction succeeds, the post-stack
matches its output type.

Global type soundness — that *every* reachable PC has a consistent
stackmap — requires program-wide code typing (JVM-style stackmaps).
That's a layer above; this file proves the per-instruction theorem
on which the global one is built.

Progress: well-typed non-halt instructions always make a step. -/

theorem stack_preservation
    {s s' : State} {in_ty out_ty : StackTy}
    (h_pc    : s.pc < s.code.size)
    (h_typed : HasTypeInstr (s.code[s.pc]'h_pc) in_ty out_ty)
    (h_stack : HasTypeStack s.stack in_ty)
    (h_step  : step s = some s') :
    HasTypeStack s'.stack out_ty := by
  unfold step at h_step
  rw [dif_pos h_pc] at h_step
  generalize h_at : s.code[s.pc]'h_pc = instr at h_typed h_step
  generalize h_stk : s.stack = stk at h_stack h_step
  cases h_typed with
  | push n =>
    simp at h_step
    obtain ⟨_, rfl⟩ := h_step
    exact .cons (.vInt n) h_stack
  | pushB b =>
    simp at h_step
    obtain ⟨_, rfl⟩ := h_step
    exact .cons (.vBool b) h_stack
  | pop =>
    cases h_stack with
    | cons _ hs =>
      simp at h_step
      obtain ⟨_, rfl⟩ := h_step
      exact hs
  | dup =>
    cases h_stack with
    | cons hv hs =>
      simp at h_step
      obtain ⟨_, rfl⟩ := h_step
      exact .cons hv (.cons hv hs)
  | swap =>
    cases h_stack with
    | cons hv1 h_rest =>
      cases h_rest with
      | cons hv2 hs =>
        simp at h_step
        obtain ⟨_, rfl⟩ := h_step
        exact .cons hv2 (.cons hv1 hs)
  | add =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt a =>
        cases h1 with
        | cons hv2 hs =>
          cases hv2 with
          | vInt b =>
            simp at h_step
            obtain ⟨_, rfl⟩ := h_step
            exact .cons (.vInt _) hs
  | sub =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt a =>
        cases h1 with
        | cons hv2 hs =>
          cases hv2 with
          | vInt b =>
            simp at h_step
            obtain ⟨_, rfl⟩ := h_step
            exact .cons (.vInt _) hs
  | mul =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt a =>
        cases h1 with
        | cons hv2 hs =>
          cases hv2 with
          | vInt b =>
            simp at h_step
            obtain ⟨_, rfl⟩ := h_step
            exact .cons (.vInt _) hs
  | eq_int =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt a =>
        cases h1 with
        | cons hv2 hs =>
          cases hv2 with
          | vInt b =>
            simp at h_step
            obtain ⟨_, rfl⟩ := h_step
            exact .cons (.vBool _) hs
  | lt_int =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt a =>
        cases h1 with
        | cons hv2 hs =>
          cases hv2 with
          | vInt b =>
            simp at h_step
            obtain ⟨_, rfl⟩ := h_step
            exact .cons (.vBool _) hs
  | jmp =>
    simp at h_step
    obtain ⟨_, rfl⟩ := h_step
    exact h_stack
  | jmpFalse =>
    cases h_stack with
    | cons hv hs =>
      cases hv with
      | vBool b =>
        cases b with
        | false =>
          simp at h_step
          obtain ⟨_, rfl⟩ := h_step
          exact hs
        | true =>
          simp at h_step
          obtain ⟨_, rfl⟩ := h_step
          exact hs
  | halt =>
    simp at h_step

theorem progress
    {s : State} {in_ty out_ty : StackTy}
    (h_pc    : s.pc < s.code.size)
    (h_typed : HasTypeInstr (s.code[s.pc]'h_pc) in_ty out_ty)
    (h_stack : HasTypeStack s.stack in_ty)
    (h_not_halt : s.code[s.pc]'h_pc ≠ .halt) :
    ∃ s', step s = some s' := by
  unfold step
  rw [dif_pos h_pc]
  generalize h_at : s.code[s.pc]'h_pc = instr at h_typed h_not_halt
  generalize h_stk : s.stack = stk at h_stack
  cases h_typed with
  | push n  => exact ⟨_, rfl⟩
  | pushB b => exact ⟨_, rfl⟩
  | pop =>
    cases h_stack with
    | cons _ _ => exact ⟨_, rfl⟩
  | dup =>
    cases h_stack with
    | cons _ _ => exact ⟨_, rfl⟩
  | swap =>
    cases h_stack with
    | cons _ h1 => cases h1 with | cons _ _ => exact ⟨_, rfl⟩
  | add =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt _ =>
        cases h1 with
        | cons hv2 _ => cases hv2 with | vInt _ => exact ⟨_, rfl⟩
  | sub =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt _ =>
        cases h1 with
        | cons hv2 _ => cases hv2 with | vInt _ => exact ⟨_, rfl⟩
  | mul =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt _ =>
        cases h1 with
        | cons hv2 _ => cases hv2 with | vInt _ => exact ⟨_, rfl⟩
  | eq_int =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt _ =>
        cases h1 with
        | cons hv2 _ => cases hv2 with | vInt _ => exact ⟨_, rfl⟩
  | lt_int =>
    cases h_stack with
    | cons hv1 h1 =>
      cases hv1 with
      | vInt _ =>
        cases h1 with
        | cons hv2 _ => cases hv2 with | vInt _ => exact ⟨_, rfl⟩
  | jmp => exact ⟨_, rfl⟩
  | jmpFalse =>
    cases h_stack with
    | cons hv _ => cases hv with | vBool b => cases b <;> exact ⟨_, rfl⟩
  | halt => exact absurd rfl h_not_halt

end TsmLean.Core