import GolangLean.Core.TypeSoundness

namespace GolangLean.Core

/-! # Preservation theorem for big-step semantics.

`Γ ⊢ e : T  ∧  HasTypeH ht h  ∧  HasTypeEnv ht env Γ  ∧  BigStep h env e v h'
   ⟹  ∃ ht', HeapTy.extends ht ht'  ∧  HasTypeH ht' h'  ∧  HasTypeV ht' v T  ∧  HasTypeEnv ht' env Γ`

By induction on the big-step derivation, with case-splits on the typing
hypothesis. Each rule case threads heap-typings through its sub-derivations
and weakens earlier results across the extension.

Big-step preservation has no progress analogue — the theorem only speaks
about derivations that exist. -/

/-! ## Heap extension by push. -/

theorem HeapTy.extends_push (ht : HeapTy) (T : Ty) : ht.extends (ht.push T) := by
  refine ⟨by simp [Array.size_push], ?_⟩
  intros i T' hL
  rw [Array.getElem?_push]
  by_cases heq : i = ht.size
  · subst heq
    have hnone : ht[ht.size]? = none := Array.getElem?_eq_none (Nat.le_refl _)
    rw [hnone] at hL
    contradiction
  · simp [heq]; exact hL

/-! ## Heap-conformance preservation. -/

theorem HasTypeH.push
    {ht : HeapTy} {h : Heap} {v : Value} {T : Ty}
    (hH : HasTypeH ht h) (hV : HasTypeV ht v T) :
    HasTypeH (ht.push T) (h.push v) := by
  refine ⟨by simp [Array.size_push, hH.1], ?_⟩
  intros i T' hL
  rw [Array.getElem?_push] at hL
  by_cases heq : i = ht.size
  · subst heq
    simp at hL
    subst hL
    refine ⟨v, ?_, hV.weaken (HeapTy.extends_push ht T)⟩
    rw [Array.getElem?_push]
    simp [hH.1]
  · simp [heq] at hL
    have ⟨v', hgv, hVT⟩ := hH.2 i T' hL
    refine ⟨v', ?_, hVT.weaken (HeapTy.extends_push ht T)⟩
    rw [Array.getElem?_push]
    have hisize : i ≠ h.size := hH.1 ▸ heq
    simp [hisize, hgv]

theorem HasTypeH.setIfInBounds
    {ht : HeapTy} {h : Heap} {loc : Nat} {v : Value} {T : Ty}
    (hH : HasTypeH ht h) (hV : HasTypeV ht v T)
    (hLoc : (ht[loc]? : Option Ty) = some T) :
    HasTypeH ht (h.setIfInBounds loc v) := by
  refine ⟨by rw [Array.size_setIfInBounds]; exact hH.1, ?_⟩
  intros i T' hL
  by_cases hi : loc = i
  · subst hi
    rw [hL] at hLoc
    simp at hLoc
    subst hLoc
    have hloc_lt : loc < h.size := by
      rw [← hH.1]
      rcases Nat.lt_or_ge loc ht.size with hlt | hge
      · exact hlt
      · rw [Array.getElem?_eq_none hge] at hL
        contradiction
    have heq : (h.setIfInBounds loc v)[loc]? = some v := by
      rw [Array.getElem?_setIfInBounds]; simp [hloc_lt]
    exact ⟨v, heq, hV⟩
  · have ⟨v', hgv, hVT⟩ := hH.2 i T' hL
    have heq : (h.setIfInBounds loc v)[i]? = some v' := by
      rw [Array.getElem?_setIfInBounds, if_neg hi]; exact hgv
    exact ⟨v', heq, hVT⟩

/-! ## Inversion for binop typing. -/

theorem binop_apply_sound
    {ht : HeapTy} {op : BinOp} {v1 v2 v : Value} {T1 T2 T : Ty}
    (hOp : op.typeOf T1 T2 = some T)
    (hV1 : HasTypeV ht v1 T1) (hV2 : HasTypeV ht v2 T2)
    (hApp : op.apply v1 v2 = some v) :
    HasTypeV ht v T := by
  cases op <;> cases T1 <;> cases T2 <;> simp [BinOp.typeOf] at hOp <;>
    (try (cases hOp; cases hV1; cases hV2; simp [BinOp.apply] at hApp; cases hApp; constructor))

/-! ## Preservation. -/

theorem preservation :
    ∀ {h env e v h'} (D : BigStep h env e v h')
      {ht Γ T} (hT : HasType Γ e T)
      (hH : HasTypeH ht h) (hE : HasTypeEnv ht env Γ),
      ∃ ht', HeapTy.extends ht ht' ∧ HasTypeH ht' h' ∧
              HasTypeV ht' v T ∧ HasTypeEnv ht' env Γ := by
  intros h env e v h' D
  induction D with
  | unitR =>
    intros ht Γ T hT hH hE
    cases hT
    exact ⟨ht, HeapTy.extends_refl _, hH, .vUnit, hE⟩
  | intLitR n =>
    intros ht Γ T hT hH hE
    cases hT
    exact ⟨ht, HeapTy.extends_refl _, hH, .vInt n, hE⟩
  | boolLitR b =>
    intros ht Γ T hT hH hE
    cases hT
    exact ⟨ht, HeapTy.extends_refl _, hH, .vBool b, hE⟩
  | varR hLook =>
    intros ht Γ T hT hH hE
    cases hT with
    | var hLookT =>
      have ⟨v', hLook', hTV⟩ := hE.lookup_correspondence hLookT
      rw [hLook] at hLook'
      cases hLook'
      exact ⟨ht, HeapTy.extends_refl _, hH, hTV, hE⟩
  | lamR x body =>
    intros ht Γ T hT hH hE
    cases hT with
    | lam hBody => exact ⟨ht, HeapTy.extends_refl _, hH, .vClos hE hBody, hE⟩
  | appR _ _ _ ih1 ih2 ihb =>
    intros ht Γ T hT hH hE
    cases hT with
    | app hT1 hT2 =>
      have ⟨ht1, ext01, hH1, hClosT, hE1⟩ := ih1 hT1 hH hE
      cases hClosT with
      | vClos hEnv' hBody =>
        have ⟨ht2, ext12, hH2, hArgT, _⟩ := ih2 hT2 hH1 hE1
        have hEnv2 := hEnv'.weaken ext12
        have ⟨ht3, ext23, hH3, hValT, _⟩ :=
          ihb hBody hH2 (HasTypeEnv.cons hArgT hEnv2)
        let totalExt := HeapTy.extends_trans (HeapTy.extends_trans ext01 ext12) ext23
        exact ⟨ht3, totalExt, hH3, hValT, hE.weaken totalExt⟩
  | letInR _ _ ih1 ih2 =>
    intros ht Γ T hT hH hE
    cases hT with
    | letIn hT1 hT2 =>
      have ⟨ht1, ext01, hH1, hV1, hE1⟩ := ih1 hT1 hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ :=
        ih2 hT2 hH1 (HasTypeEnv.cons hV1 hE1)
      let totalExt := HeapTy.extends_trans ext01 ext12
      exact ⟨ht2, totalExt, hH2, hV2, hE.weaken totalExt⟩
  | ifTR _ _ ihc iht =>
    intros ht Γ T hT hH hE
    cases hT with
    | ifte hTc hT1 _ =>
      have ⟨ht1, ext01, hH1, _, hE1⟩ := ihc hTc hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ := iht hT1 hH1 hE1
      let totalExt := HeapTy.extends_trans ext01 ext12
      exact ⟨ht2, totalExt, hH2, hV2, hE.weaken totalExt⟩
  | ifFR _ _ ihc ihf =>
    intros ht Γ T hT hH hE
    cases hT with
    | ifte hTc _ hT2 =>
      have ⟨ht1, ext01, hH1, _, hE1⟩ := ihc hTc hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ := ihf hT2 hH1 hE1
      let totalExt := HeapTy.extends_trans ext01 ext12
      exact ⟨ht2, totalExt, hH2, hV2, hE.weaken totalExt⟩
  | binopR _ _ hOp ih1 ih2 =>
    intros ht Γ T hT hH hE
    cases hT with
    | binop hT1 hT2 hOpT =>
      have ⟨ht1, ext01, hH1, hV1, hE1⟩ := ih1 hT1 hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ := ih2 hT2 hH1 hE1
      let totalExt := HeapTy.extends_trans ext01 ext12
      have hV1ext := hV1.weaken ext12
      have hVT := binop_apply_sound hOpT hV1ext hV2 hOp
      exact ⟨ht2, totalExt, hH2, hVT, hE.weaken totalExt⟩
  | refMkR _ ih =>
    intros ht Γ T hT hH hE
    cases hT with
    | @refMk _ _ T_inner hT1 =>
      have ⟨ht1, ext01, hH1, hV, _⟩ := ih hT1 hH hE
      have extPush := HeapTy.extends_push ht1 T_inner
      have newExt := HeapTy.extends_trans ext01 extPush
      have hHnew := hH1.push hV
      have hVnew : HasTypeV (ht1.push T_inner) (.vLoc ht1.size) (.ref T_inner) := by
        apply HasTypeV.vLoc
        rw [Array.getElem?_push]
        simp
      rw [hH1.1] at hVnew
      exact ⟨ht1.push T_inner, newExt, hHnew, hVnew, hE.weaken newExt⟩
  | derefR _ Hget ih =>
    intros ht Γ T hT hH hE
    cases hT with
    | @deref _ _ T_inner hT1 =>
      have ⟨ht1, ext01, hH1, hVLoc, _⟩ := ih hT1 hH hE
      cases hVLoc with
      | vLoc hLocT =>
        have ⟨v', hgv, hVT⟩ := hH1.2 _ _ hLocT
        rw [Hget] at hgv
        cases hgv
        exact ⟨ht1, ext01, hH1, hVT, hE.weaken ext01⟩
  | assignR _ _ Hin ih1 ih2 =>
    intros ht Γ T hT hH hE
    cases hT with
    | @assign _ _ _ T_target hT1 hT2 =>
      have ⟨ht1, ext01, hH1, hVLoc, hE1⟩ := ih1 hT1 hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ := ih2 hT2 hH1 hE1
      cases hVLoc with
      | vLoc hLocT =>
        have hLocT2 := ext12.2 _ _ hLocT
        have hH2new := hH2.setIfInBounds hV2 hLocT2
        have totalExt := HeapTy.extends_trans ext01 ext12
        exact ⟨ht2, totalExt, hH2new, .vUnit, hE.weaken totalExt⟩
  | seqR _ _ ih1 ih2 =>
    intros ht Γ T hT hH hE
    cases hT with
    | seq hT1 hT2 =>
      have ⟨ht1, ext01, hH1, _, hE1⟩ := ih1 hT1 hH hE
      have ⟨ht2, ext12, hH2, hV2, _⟩ := ih2 hT2 hH1 hE1
      let totalExt := HeapTy.extends_trans ext01 ext12
      exact ⟨ht2, totalExt, hH2, hV2, hE.weaken totalExt⟩

end GolangLean.Core