-- Hoare/While/Types.lean

import Hoare.While.Syntax

namespace While

inductive Ty
| num : Ty
| bool : Ty
deriving Repr, DecidableEq

inductive WellTyped : Expr → Ty → Prop
| num : WellTyped (Expr.num _) Ty.num
| bool : WellTyped (Expr.bool _) Ty.bool
| add : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.add e1 e2) Ty.num
| sub : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.sub e1 e2) Ty.num
| mul : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.mul e1 e2) Ty.num
| eq : ∀ e1 e2 t, WellTyped e1 t → WellTyped e2 t →
    WellTyped (Expr.eq e1 e2) Ty.bool
| lt : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.lt e1 e2) Ty.bool
| gt : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.gt e1 e2) Ty.bool
| le : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.le e1 e2) Ty.bool
| ge : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.ge e1 e2) Ty.bool
| and : ∀ e1 e2, WellTyped e1 Ty.bool → WellTyped e2 Ty.bool →
    WellTyped (Expr.and e1 e2) Ty.bool
| or : ∀ e1 e2, WellTyped e1 Ty.bool → WellTyped e2 Ty.bool →
    WellTyped (Expr.or e1 e2) Ty.bool

def Expr.ty : Expr → Option Ty
  | Expr.num _ => Ty.num
  | Expr.bool _ => Ty.bool
  | Expr.add e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.sub e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.mul e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.eq e1 e2 => if e1.ty = e2.ty ∧ e1.ty.isSome then some .bool else none
  | Expr.lt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.gt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.le e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.ge e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.and e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.or e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.var _ => none

-- Repository definitions used:
content: syntax num : nexpr

syntax "if " bexpr " then " com " else " com " fi" : com

content: macro_rules
| `([nexpr| $n:num]) => `(Expr.num $n)
| `([nexpr| $x:ident]) => `(Expr.var $(Lean.quote x.getId.toString))
| `([bexpr| true]) => `(Expr.bool «true»)
| `([bexpr| false]) => `(Expr.bool «false»)
| `([nexpr| $e1 + $e2]) => `(Expr.add [nexpr| $e1] [nexpr| $e2])
| `([nexpr| $e1 - $e2]) => `(Expr.sub [nexpr| $e1] [nexpr| $e2])
| `([nexpr| $e1 * $e2]) => `(Expr.mul [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:nexpr == $e2]) => `(Expr.eq [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:nexpr < $e2]) => `(Expr.lt [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:nexpr > $e2]) => `(Expr.gt [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:nexpr <= $e2]) => `(Expr.le [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:nexpr >= $e2]) => `(Expr.ge [nexpr| $e1] [nexpr| $e2])
| `([bexpr| $e1:bexpr && $e2]) => `(Expr.and [bexpr| $e1] [bexpr| $e2])
| `([bexpr| $e1:bexpr || $e2]) => `(Expr.or [bexpr| $e1] [bexpr| $e2])
| `([bexpr| ($e)]) => `([bexpr| $e])
| `([nexpr| ($e)]) => `([nexpr| $e])
| `([com| skip]) => `(Com.skip)
| `([com| let $x:ident := $e]) => `(Com.assign $(Lean.quote x.getId.toString) [nexpr| $e])
| `([com| $x:ident := $e]) => `(Com.assign $(Lean.quote x.getId.toString) [nexpr| $e])
| `([com| $c1; $c2]) => `(Com.seq [com| $c1] [com| $c2])
| `([com| if $e then $c1 else $c2 fi]) => `(Com.cond [bexpr| $e] [com| $c1] [com| $c2])
| `([com| while $e do $c od]) => `(Com.while [bexpr| $e] [com| $c])

content: macro_rules
| `([bexpr| $($t:term)]) => `($t)
| `([nexpr| $($t:term)]) => `($t)
| `([com| $($t:term)]) => `($t)

content: def Expr.ty : Expr → Option Ty
  | Expr.num _ => Ty.num
  | Expr.bool _ => Ty.bool
  | Expr.add e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.sub e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.mul e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.eq e1 e2 => if e1.ty = e2.ty ∧ e1.ty.isSome then some .bool else none
  | Expr.lt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.gt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.le e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.ge e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.and e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.or e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.var _ => none 

content: inductive Com
| assign : String → Expr → Com
| seq : Com → Com → Com
| cond : Expr → Com → Com → Com
| while : Expr → Com → Com
| skip : Com

content: inductive Expr
| num : Nat → Expr
| bool : Bool → Expr
| var : String → Expr
| add : Expr → Expr → Expr
| sub : Expr → Expr → Expr
| mul : Expr → Expr → Expr
| eq : Expr → Expr → Expr
| lt : Expr → Expr → Expr
| gt : Expr → Expr → Expr
| le : Expr → Expr → Expr
| ge : Expr → Expr → Expr
| and : Expr → Expr → Expr
| or : Expr → Expr → Expr
deriving Repr, DecidableEq

-- 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: inductive Ty
| num : Ty
| bool : Ty
deriving Repr, DecidableEq

content: inductive WellTyped : Expr → Ty → Prop
| num : WellTyped (Expr.num _) Ty.num
| bool : WellTyped (Expr.bool _) Ty.bool
| add : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.add e1 e2) Ty.num
| sub : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.sub e1 e2) Ty.num
| mul : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.mul e1 e2) Ty.num
| eq : ∀ e1 e2 t, WellTyped e1 t → WellTyped e2 t →
    WellTyped (Expr.eq e1 e2) Ty.bool
| lt : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.lt e1 e2) Ty.bool
| gt : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.gt e1 e2) Ty.bool
| le : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.le e1 e2) Ty.bool
| ge : ∀ e1 e2, WellTyped e1 Ty.num → WellTyped e2 Ty.num →
    WellTyped (Expr.ge e1 e2) Ty.bool
| and : ∀ e1 e2, WellTyped e1 Ty.bool → WellTyped e2 Ty.bool →
    WellTyped (Expr.and e1 e2) Ty.bool
| or : ∀ e1 e2, WellTyped e1 Ty.bool → WellTyped e2 Ty.bool →
    WellTyped (Expr.or e1 e2) Ty.bool

content: def Expr.ty : Expr → Option Ty
  | Expr.num _ => Ty.num
  | Expr.bool _ => Ty.bool
  | Expr.add e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.sub e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.mul e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .num else none
  | Expr.eq e1 e2 => if e1.ty = e2.ty ∧ e1.ty.isSome then some .bool else none
  | Expr.lt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.gt e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.le e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.ge e1 e2 => if e1.ty = some .num ∧ e2.ty = some .num then some .bool else none
  | Expr.and e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.or e1 e2 => if e1.ty = some .bool ∧ e2.ty = some .bool then some .bool else none
  | Expr.var _ => none

-- Local lemmas used:
content: theorem WellTyped.ty_some {e ty} : WellTyped e ty → e.ty = some ty

-- Target theorem and proof:
theorem WellTyped.some_ty {e ty} : e.ty = some ty → WellTyped e ty
:= by
  intro h
  induction e generalizing ty <;> simp_all [Expr.ty]
  case num _ => rw [← h]; exact WellTyped.num
  case bool _ => rw [← h]; exact WellTyped.bool
  case add e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.add e1 e2 ih_e1 ih_e2
  case sub e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.sub e1 e2 ih_e1 ih_e2
  case mul e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.mul e1 e2 ih_e1 ih_e2
  case eq e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'
    match hty : e1.ty with
    | some .num => have hty' := h.1.1.symm; rw [hty] at hty'; exact WellTyped.eq e1 e2 .num (ih_e1 hty') (ih_e2 hty')
    | some .bool => have hty' := h.1.1.symm; rw [hty] at hty'; exact WellTyped.eq e1 e2 .bool (ih_e1 hty') (ih_e2 hty')
    | none => rw [hty] at h; have hc := h.1.2; simp [Option.isSome] at hc
  case lt e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.lt e1 e2 ih_e1 ih_e2
  case gt e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.gt e1 e2 ih_e1 ih_e2
  case le e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.le e1 e2 ih_e1 ih_e2
  case ge e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.ge e1 e2 ih_e1 ih_e2
  case and e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.and e1 e2 ih_e1 ih_e2
  case or e1 e2 ih_e1 ih_e2 =>
    have h' := h.2.symm; subst h'; exact WellTyped.or e1 e2 ih_e1 ih_e2