notation:10 l:10 " ;; " r:11 => com.SEQ l r

ident

def ident := String deriving BEq, Repr, Hashable

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

com

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)

cont

@[grind] inductive cont where
| Kstop
| Kseq (c : com) (k : cont)
| Kwhile (b : bexp) (c : com) (k : cont)

store

def store : Type := ident → Int

update

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

star

@[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

step

inductive step : com × cont × store -> com × cont × store -> Prop where
  | step_assign : ∀ x a k s,
      step (.ASSIGN x a, k, s) (.SKIP, k, update x (aeval s a) s)

  | step_seq : ∀ c1 c2 s k,
      step (.SEQ c1 c2, k, s) (c1, .Kseq c2 k, s)

  | step_ifthenelse : ∀ b c1 c2 k s,
      step (.IFTHENELSE b c1 c2, k, s) ((if beval s b then c1 else c2), k, s)

  | step_while_done : ∀ b c k s,
      beval s b = false ->
      step (.WHILE b c, k, s) (.SKIP, k, s)

  | step_while_true : ∀ b c k s,
      beval s b = true ->
      step (.WHILE b c, k, s) (c, .Kwhile b c k, s)

  | step_skip_seq : ∀ c k s,
      step (.SKIP, .Kseq c k, s) (c, k, s)

  | step_skip_while : ∀ b c k s,
      step (.SKIP, .Kwhile b c k, s) (.WHILE b c, k, s)

aeval

@[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

beval

@[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

plus

@[grind cases]
inductive plus (R : α → α → Prop) : α → α → Prop where
| plus_left : ∀ {a b c}, R a b → star R b c → plus R a c

grind_pattern plus.plus_left => star R b c, plus R a c

@[grind] theorem plus_star {a b} (h : plus R a b) : star R a b

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

theorem plus_right : star R a b -> R b c -> plus R a c

instr

@[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

codelen

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

stack

def stack : Type := List Int

config

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)

transition

@[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)

transitions

@[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

inductive compile_cont (C : List instr) : cont -> Int -> Prop where
  | ccont_stop : ∀ pc,
      instr_at C pc = .some .Ihalt ->
      compile_cont C .Kstop pc
  | ccont_seq : ∀ c k pc pc',
      code_at C pc (compile_com c) ->
      pc' = pc + codelen (compile_com c) ->
      compile_cont C k pc' ->
      compile_cont C (.Kseq c k) pc
  | ccont_while : ∀ b c k pc d pc' pc'',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      code_at C pc' (compile_com (.WHILE b c)) ->
      pc'' = pc' + codelen (compile_com (.WHILE b c)) ->
      compile_cont C k pc'' ->
      compile_cont C (.Kwhile b c k) pc
  | ccont_branch : ∀ d k pc pc',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      compile_cont C k pc' ->
      compile_cont C k pc

inductive match_config (C : List instr) : com × cont × store -> config -> Prop where
  | match_config_intro : ∀ c k st pc,
      code_at C pc (compile_com c) ->
      compile_cont C k (pc + codelen (compile_com c)) ->
      match_config C (c, k, st) (pc, [], st)

def com_size (c : com) : Nat :=
  match c with
  | .SKIP => 1
  | .ASSIGN _ _ => 1
  | (c1 ;; c2) => (com_size c1 + com_size c2 + 1)
  | .IFTHENELSE _ c1 c2 => (com_size c1 + com_size c2 + 1)
  | .WHILE _ c1 => (com_size c1 + 1)

def cont_size (k : cont) : Nat :=
  match k with
  | .Kstop => 0
  | .Kseq c k' => (com_size c + cont_size k')
  | .Kwhile _ _ k' => cont_size k'

measure'

def measure' (impconf : com × cont × store) : Nat :=
  match impconf with
  | (c, k, _) => (com_size c + cont_size k)

@[grind =] theorem codelen_cons :
  ∀ i c, codelen (i :: c) = codelen c + 1

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

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

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

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

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)

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)

theorem compile_cont_Kseq_inv (C : List instr) (c : com) (k :cont) (pc : Int) (s : store) (H : compile_cont C (.Kseq c k) pc) :
  ∃ pc',
  star (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com c)
  ∧ compile_cont C k (pc' + codelen (compile_com c))

theorem compile_cont_Kwhile_inv (C : List instr) (b : bexp) (c : com) (k : cont) (pc : Int) (s : store) (H : compile_cont C (.Kwhile b c k) pc) :
  ∃ pc',
  plus (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com (.WHILE b c))
  ∧ compile_cont C k (pc' + codelen (compile_com (.WHILE b c)))

theorem match_config_skip (C : List instr) (k : cont) (s : store) (pc : Int) (H : compile_cont C k pc) :
 match_config C (.SKIP, k, s) (pc, [], s)

theorem simulation_step :
  ∀ C impconf1 impconf2 machconf1,
  step impconf1 impconf2 ->
  match_config C impconf1 machconf1 ->
  ∃ machconf2,
      (plus (transition C) machconf1 machconf2
       \/ (star (transition C) machconf1 machconf2
           /\ (measure' impconf2 < measure' impconf1)))
   /\ match_config C impconf2 machconf2

theorem simulation_steps :
  ∀ C impconf1 impconf2, star step impconf1 impconf2 ->
  ∀ machconf1, match_config C impconf1 machconf1 ->
  ∃ machconf2,
     star (transition C) machconf1 machconf2
  /\ match_config C impconf2 machconf2 :=

:= by
      intro C impconf1 impconf2 STAR machconf1 MATCH
      induction STAR generalizing machconf1
      case star_refl x =>
        exists machconf1
        constructor
        · apply star.star_refl
        · exact MATCH
      case star_step x y z STEP STAR ih =>
        have ⟨ machconf2, steps2, match2 ⟩ := simulation_step C x y machconf1 STEP MATCH
        specialize ih machconf2 match2
        rcases ih with ⟨ machconf3, steps3, match3⟩
        exists machconf3
        have w : star (transition C) machconf1 machconf2 := by
          cases steps2
          case inl h =>
            apply plus_star
            exact h
          case inr h =>
            exact h.1
        constructor
        · apply star_trans
          · exact w
          · exact steps3
        · exact match3


Join

@[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

Store

def Store := Std.HashMap ident Int

Le

@[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

Equal

def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

com

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)

Beval

@[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

@[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

Aeval

@[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

@[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

ident

def ident := String deriving BEq, Repr, Hashable

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

Update

@[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

notation:10 l:10 " ;; " r:11 => com.SEQ l r

theorem Equal_Le : ∀ S1 S2, Equal S1 S2 -> Le S1 S2

theorem Le_Join_l : ∀ S1 S2, Le S1 (Join S1 S2)

theorem  Le_Join_r : ∀ S1 S2, Le S2 (Join S1 S2)

Eq'

@[grind] theorem iterate_correct (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) (heq : y = iterate _ F x PRE SMALL ) : eq y (F y) ∧ ∀ z, le (F z) z → le y z

@[grind] theorem Eq_Le : ∀ S1 S2, Eq' S1 S2 → Le S1 S2

@[grind] theorem Le_trans : ∀ S1 S2 S3, Le S1 S2 → Le S2 S3 → Le S1 S3

theorem fixpoint_join_eq : Eq' (Join Init (F (fixpoint_join Init F) )) (fixpoint_join Init F)

import LeroyCompilerVerificationCourse.Imp

import LeroyCompilerVerificationCourse.Constprop

import Batteries.Data.List.Perm

@[grind] class OrderStruct (α : Sort u) where
  eq : α → α → Prop
  le : α → α → Prop
  beq : α → α → Bool
  le_trans : ∀ x y z, le x y -> le y z -> le x z
  beq_true' : ∀ x y : α, beq x y = true → eq x y := by admit /- proof elided -/

open OrderStruct

class Monotone (α : Sort u) (F : α → α) [OrderStruct α] where
  F_mon : ∀ {x y : α}, le x y → le (F x) (F y)

open Monotone

section FixpointExistence

variable (α : Sort u) (F : α → α) [OrderWithBot α]

open OrderStruct OrderWithBot

end FixpointExistence

section Iterate

variable (α : Sort u) [inst : OrderStruct α] (F : α → α) [Monotone α F]

open OrderStruct

@[grind] def iterate (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) : α :=
  if beq x (F x) then x else iterate (F x) (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )
  termination_by x
  decreasing_by
    grind [beq_false']

end Iterate

section Fixpoint

open OrderWithBot

variable {α : Sort u} [i : OrderWithBot α] (F : α → α) [Monotone α F]

end Fixpoint

section Constprop

open Std.HashMap

@[grind] def Eq' (S1 S2 : Store) : Prop := Equal S1 S2

def Eq'_sym : ∀ S1 S2, Eq' S1 S2 → Eq' S2 S1 :=

open OrderStruct

noncomputable instance : OrderStruct Store where
  eq := Equal
  le := Le
  beq (S1 S2 : Store) :=  Decidable.decide (Equal S1 S2)
  le_trans := Le_trans
  gt_wf := Gt_wf

end Constprop

section FixpointJoin

variable (Init : Store)

variable (F : Store → Store) [Monotone Store F]

instance : Monotone Store (fun x => Join Init (F x)) where
  F_mon := by admit /- proof elided -/

noncomputable def fixpoint_join : Store :=

theorem fixpoint_join_sound : Le Init (fixpoint_join Init F) /\ Le (F (fixpoint_join Init F)) (fixpoint_join Init F) :=

:= by
  have LE : Le (Join Init (F (fixpoint_join Init F))) (fixpoint_join Init F) := by
    apply Eq_Le
    apply fixpoint_join_eq
  constructor
  · apply Le_trans
    rotate_left
    · exact LE
    · apply Le_Join_l
  · apply Le_trans
    rotate_left
    · exact LE
    · apply Le_Join_r


ident

def ident := String deriving BEq, Repr, Hashable

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

com

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)

cont

@[grind] inductive cont where
| Kstop
| Kseq (c : com) (k : cont)
| Kwhile (b : bexp) (c : com) (k : cont)

store

def store : Type := ident → Int

update

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

infseq

def infseq {α} (R : α → α → Prop) : α → Prop :=
  λ x : α => ∃ y, R x y ∧ infseq R y
  coinductive_fixpoint

step

inductive step : com × cont × store -> com × cont × store -> Prop where
  | step_assign : ∀ x a k s,
      step (.ASSIGN x a, k, s) (.SKIP, k, update x (aeval s a) s)

  | step_seq : ∀ c1 c2 s k,
      step (.SEQ c1 c2, k, s) (c1, .Kseq c2 k, s)

  | step_ifthenelse : ∀ b c1 c2 k s,
      step (.IFTHENELSE b c1 c2, k, s) ((if beval s b then c1 else c2), k, s)

  | step_while_done : ∀ b c k s,
      beval s b = false ->
      step (.WHILE b c, k, s) (.SKIP, k, s)

  | step_while_true : ∀ b c k s,
      beval s b = true ->
      step (.WHILE b c, k, s) (c, .Kwhile b c k, s)

  | step_skip_seq : ∀ c k s,
      step (.SKIP, .Kseq c k, s) (c, k, s)

  | step_skip_while : ∀ b c k s,
      step (.SKIP, .Kwhile b c k, s) (.WHILE b c, k, s)

aeval

@[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

beval

@[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

plus

@[grind cases]
inductive plus (R : α → α → Prop) : α → α → Prop where
| plus_left : ∀ {a b c}, R a b → star R b c → plus R a c

grind_pattern plus.plus_left => star R b c, plus R a c

star

@[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

notation:10 l:10 " ;; " r:11 => com.SEQ l r

@[grind] theorem plus_star {a b} (h : plus R a b) : star R a b

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

theorem plus_right : star R a b -> R b c -> plus R a c

theorem star_plus_trans :
  ∀ a b c, star R a b -> plus R b c -> plus R a c

theorem infseq_coinduction_principle_2
    (X : α → Prop) (h₁ : ∀ (a : α), X a → ∃ b, plus R a b ∧ X b) (a : α) (rel : X a) : infseq R a

stack

def stack : Type := List Int

measure'

@[grind =] theorem codelen_cons :
  ∀ i c, codelen (i :: c) = codelen c + 1

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

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

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

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

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)

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)

theorem compile_cont_Kseq_inv (C : List instr) (c : com) (k :cont) (pc : Int) (s : store) (H : compile_cont C (.Kseq c k) pc) :
  ∃ pc',
  star (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com c)
  ∧ compile_cont C k (pc' + codelen (compile_com c))

theorem compile_cont_Kwhile_inv (C : List instr) (b : bexp) (c : com) (k : cont) (pc : Int) (s : store) (H : compile_cont C (.Kwhile b c k) pc) :
  ∃ pc',
  plus (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com (.WHILE b c))
  ∧ compile_cont C k (pc' + codelen (compile_com (.WHILE b c)))

theorem match_config_skip (C : List instr) (k : cont) (s : store) (pc : Int) (H : compile_cont C k pc) :
 match_config C (.SKIP, k, s) (pc, [], s)

theorem simulation_step :
  ∀ C impconf1 impconf2 machconf1,
  step impconf1 impconf2 ->
  match_config C impconf1 machconf1 ->
  ∃ machconf2,
      (plus (transition C) machconf1 machconf2
       \/ (star (transition C) machconf1 machconf2
           /\ (measure' impconf2 < measure' impconf1)))
   /\ match_config C impconf2 machconf2

theorem match_initial_configs :
  ∀ c s,
  match_config (compile_program c) (c, .Kstop, s) (0, [], s)

theorem simulation_infseq_inv :
  ∀ C n impconf1 machconf1,
  infseq step impconf1 -> match_config C impconf1 machconf1 ->
  (measure' impconf1 < n) ->
  ∃ impconf2 machconf2,
      infseq step impconf2
   /\ plus (transition C) machconf1 machconf2
   /\ match_config C impconf2 machconf2

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)

def machine_diverges (C : List instr) (s_init : store) : Prop :=
  infseq (transition C) (0, [], s_init)

@[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)) :: []

def compile_program (p : com) : List instr :=
  compile_com p ++ .Ihalt :: []

@[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

inductive compile_cont (C : List instr) : cont -> Int -> Prop where
  | ccont_stop : ∀ pc,
      instr_at C pc = .some .Ihalt ->
      compile_cont C .Kstop pc
  | ccont_seq : ∀ c k pc pc',
      code_at C pc (compile_com c) ->
      pc' = pc + codelen (compile_com c) ->
      compile_cont C k pc' ->
      compile_cont C (.Kseq c k) pc
  | ccont_while : ∀ b c k pc d pc' pc'',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      code_at C pc' (compile_com (.WHILE b c)) ->
      pc'' = pc' + codelen (compile_com (.WHILE b c)) ->
      compile_cont C k pc'' ->
      compile_cont C (.Kwhile b c k) pc
  | ccont_branch : ∀ d k pc pc',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      compile_cont C k pc' ->
      compile_cont C k pc

inductive match_config (C : List instr) : com × cont × store -> config -> Prop where
  | match_config_intro : ∀ c k st pc,
      code_at C pc (compile_com c) ->
      compile_cont C k (pc + codelen (compile_com c)) ->
      match_config C (c, k, st) (pc, [], st)

def com_size (c : com) : Nat :=
  match c with
  | .SKIP => 1
  | .ASSIGN _ _ => 1
  | (c1 ;; c2) => (com_size c1 + com_size c2 + 1)
  | .IFTHENELSE _ c1 c2 => (com_size c1 + com_size c2 + 1)
  | .WHILE _ c1 => (com_size c1 + 1)

def cont_size (k : cont) : Nat :=
  match k with
  | .Kstop => 0
  | .Kseq c k' => (com_size c + cont_size k')
  | .Kwhile _ _ k' => cont_size k'

def measure' (impconf : com × cont × store) : Nat :=
  match impconf with
  | (c, k, _) => (com_size c + cont_size k)

theorem compile_program_correct_diverging :
  ∀ c s,
  infseq step (c, .Kstop, s) ->
  machine_diverges (compile_program c) s :=

:= by
    intro c s H
    generalize heq : compile_program c = C
    unfold machine_diverges
    apply infseq_coinduction_principle_2 (fun machconf => ∃ impconf, infseq step impconf /\ match_config C impconf machconf)
    rotate_left
    · exists (c, .Kstop, s)
      constructor
      · exact H
      · have := match_initial_configs c s
        grind
    · intro machconf ⟨ impconf , ⟨INFSEQ, MATCH ⟩⟩
      have ⟨impconf2 , machconf2, INFSEQ2 , PLUS , MATCH2⟩  := simulation_infseq_inv C (measure' impconf +1) impconf machconf INFSEQ MATCH (by omega)
      exists machconf2
      constructor
      · exact PLUS
      · exists impconf2


notation:10 l:10 " ;; " r:11 => com.SEQ l r

com

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)

cont

@[grind] inductive cont where
| Kstop
| Kseq (c : com) (k : cont)
| Kwhile (b : bexp) (c : com) (k : cont)

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

ident

def ident := String deriving BEq, Repr, Hashable

store

def store : Type := ident → Int

update

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

plus

@[grind cases]
inductive plus (R : α → α → Prop) : α → α → Prop where
| plus_left : ∀ {a b c}, R a b → star R b c → plus R a c

grind_pattern plus.plus_left => star R b c, plus R a c

star

@[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

infseq

def infseq {α} (R : α → α → Prop) : α → Prop :=
  λ x : α => ∃ y, R x y ∧ infseq R y
  coinductive_fixpoint

step

inductive step : com × cont × store -> com × cont × store -> Prop where
  | step_assign : ∀ x a k s,
      step (.ASSIGN x a, k, s) (.SKIP, k, update x (aeval s a) s)

  | step_seq : ∀ c1 c2 s k,
      step (.SEQ c1 c2, k, s) (c1, .Kseq c2 k, s)

  | step_ifthenelse : ∀ b c1 c2 k s,
      step (.IFTHENELSE b c1 c2, k, s) ((if beval s b then c1 else c2), k, s)

  | step_while_done : ∀ b c k s,
      beval s b = false ->
      step (.WHILE b c, k, s) (.SKIP, k, s)

  | step_while_true : ∀ b c k s,
      beval s b = true ->
      step (.WHILE b c, k, s) (c, .Kwhile b c k, s)

  | step_skip_seq : ∀ c k s,
      step (.SKIP, .Kseq c k, s) (c, k, s)

  | step_skip_while : ∀ b c k s,
      step (.SKIP, .Kwhile b c k, s) (.WHILE b c, k, s)

aeval

@[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

beval

@[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

@[grind] theorem plus_star {a b} (h : plus R a b) : star R a b

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

theorem plus_right : star R a b -> R b c -> plus R a c

theorem star_plus_trans :
  ∀ a b c, star R a b -> plus R b c -> plus R a c

stack

def stack : Type := List Int

measure'

@[grind =] theorem codelen_cons :
  ∀ i c, codelen (i :: c) = codelen c + 1

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

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

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

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

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)

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)

theorem compile_cont_Kseq_inv (C : List instr) (c : com) (k :cont) (pc : Int) (s : store) (H : compile_cont C (.Kseq c k) pc) :
  ∃ pc',
  star (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com c)
  ∧ compile_cont C k (pc' + codelen (compile_com c))

theorem compile_cont_Kwhile_inv (C : List instr) (b : bexp) (c : com) (k : cont) (pc : Int) (s : store) (H : compile_cont C (.Kwhile b c k) pc) :
  ∃ pc',
  plus (transition C) (pc, [], s) (pc', [], s)
  ∧ code_at C pc' (compile_com (.WHILE b c))
  ∧ compile_cont C k (pc' + codelen (compile_com (.WHILE b c)))

theorem match_config_skip (C : List instr) (k : cont) (s : store) (pc : Int) (H : compile_cont C k pc) :
 match_config C (.SKIP, k, s) (pc, [], s)

theorem simulation_step :
  ∀ C impconf1 impconf2 machconf1,
  step impconf1 impconf2 ->
  match_config C impconf1 machconf1 ->
  ∃ machconf2,
      (plus (transition C) machconf1 machconf2
       \/ (star (transition C) machconf1 machconf2
           /\ (measure' impconf2 < measure' impconf1)))
   /\ match_config C impconf2 machconf2

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

inductive compile_cont (C : List instr) : cont -> Int -> Prop where
  | ccont_stop : ∀ pc,
      instr_at C pc = .some .Ihalt ->
      compile_cont C .Kstop pc
  | ccont_seq : ∀ c k pc pc',
      code_at C pc (compile_com c) ->
      pc' = pc + codelen (compile_com c) ->
      compile_cont C k pc' ->
      compile_cont C (.Kseq c k) pc
  | ccont_while : ∀ b c k pc d pc' pc'',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      code_at C pc' (compile_com (.WHILE b c)) ->
      pc'' = pc' + codelen (compile_com (.WHILE b c)) ->
      compile_cont C k pc'' ->
      compile_cont C (.Kwhile b c k) pc
  | ccont_branch : ∀ d k pc pc',
      instr_at C pc = .some (.Ibranch d) ->
      pc' = pc + 1 + d ->
      compile_cont C k pc' ->
      compile_cont C k pc

inductive match_config (C : List instr) : com × cont × store -> config -> Prop where
  | match_config_intro : ∀ c k st pc,
      code_at C pc (compile_com c) ->
      compile_cont C k (pc + codelen (compile_com c)) ->
      match_config C (c, k, st) (pc, [], st)

def com_size (c : com) : Nat :=
  match c with
  | .SKIP => 1
  | .ASSIGN _ _ => 1
  | (c1 ;; c2) => (com_size c1 + com_size c2 + 1)
  | .IFTHENELSE _ c1 c2 => (com_size c1 + com_size c2 + 1)
  | .WHILE _ c1 => (com_size c1 + 1)

def cont_size (k : cont) : Nat :=
  match k with
  | .Kstop => 0
  | .Kseq c k' => (com_size c + cont_size k')
  | .Kwhile _ _ k' => cont_size k'

def measure' (impconf : com × cont × store) : Nat :=
  match impconf with
  | (c, k, _) => (com_size c + cont_size k)

theorem simulation_infseq_inv :
  ∀ C n impconf1 machconf1,
  infseq step impconf1 -> match_config C impconf1 machconf1 ->
  (measure' impconf1 < n) ->
  ∃ impconf2 machconf2,
      infseq step impconf2
   /\ plus (transition C) machconf1 machconf2
   /\ match_config C impconf2 machconf2 :=

:= by
    intro C n impconf1 h1 h2 h3 h4
    induction n generalizing impconf1 h1
    case zero => contradiction
    case succ n' ih =>
      rw [infseq] at h2
      rcases h2 with ⟨impconf2 , STEP, INFSEQ⟩
      have ⟨ machconf2, h5 , h6 ⟩ := simulation_step C impconf1 impconf2 h1 STEP h3
      cases h5
      next PLUS =>
        exists impconf2
        exists machconf2
      next w =>
        rcases w with ⟨ STAR, MEASURE ⟩
        specialize ih impconf2 machconf2 INFSEQ h6 (by omega)
        rcases ih with ⟨ c1, m1, w⟩
        exists c1
        exists m1
        constructor
        · exact w.1
        · constructor
          · apply star_plus_trans
            · exact STAR
            · exact w.2.1
          · exact w.2.2


ident

def ident := String deriving BEq, Repr, Hashable

theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

Store

def Store := Std.HashMap ident Int

Equal

@[grind] def num_iter : Nat := 20

fixpoint

theorem Equal_Le : ∀ S1 S2, Equal S1 S2 -> Le S1 S2

import LeroyCompilerVerificationCourse.Imp

import Std.Data.HashMap

open Classical in

def Store := Std.HashMap ident Int

@[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

@[grind] def Top : Store := Std.HashMap.emptyWithCapacity

def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

@[grind] noncomputable def fixpoint_rec (F : Store -> Store) (fuel : Nat) (S : Store) : Store :=
  match fuel with
  | 0 => Top
  | fuel + 1 =>
      let S' := F S
      if Equal S' S then S else fixpoint_rec F fuel S'

@[grind] def num_iter : Nat := 20

@[grind] noncomputable def fixpoint (F : Store -> Store) (init_S : Store) : Store :=
  fixpoint_rec F num_iter init_S

theorem fixpoint_sound (F : Store → Store) (init_S : Store) (h : S = fixpoint F init_S) :
  Le (F S) S :=

:= by
    have A : ∀ fuel S,
             fixpoint_rec F fuel S = Top
             \/ Equal (F (fixpoint_rec F fuel S)) (fixpoint_rec F fuel S) := by
      intro fuel
      induction fuel
      case zero => grind
      case succ fuel' ih  =>
        grind
    have E : S = Top \/ Equal (F S) S = true := by grind
    cases E <;> grind [Equal_Le]


Store

def Store := Std.HashMap ident Int

Join

@[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

com

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)

Beval

@[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

@[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

Aeval

@[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

@[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

Update

@[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

notation:10 l:10 " ;; " r:11 => com.SEQ l r

theorem Le_Join_l : ∀ S1 S2, Le S1 (Join S1 S2)

OrderStruct

@[grind] class OrderStruct (α : Sort u) where
  eq : α → α → Prop
  le : α → α → Prop
  beq : α → α → Bool
  le_trans : ∀ x y z, le x y -> le y z -> le x z
  beq_true' : ∀ x y : α, beq x y = true → eq x y := by admit /- proof elided -/

Monotone

class Monotone (α : Sort u) (F : α → α) [OrderStruct α] where
  F_mon : ∀ {x y : α}, le x y → le (F x) (F y)

iterate

@[grind] def iterate (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) : α :=
  if beq x (F x) then x else iterate (F x) (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )
  termination_by x
  decreasing_by
    grind [beq_false']

Eq'

@[grind] def Eq' (S1 S2 : Store) : Prop := Equal S1 S2

def Eq'_sym : ∀ S1 S2, Eq' S1 S2 → Eq' S2 S1 :=

noncomputable instance : OrderStruct Store where
  eq := Equal
  le := Le
  beq (S1 S2 : Store) :=  Decidable.decide (Equal S1 S2)
  le_trans := Le_trans
  gt_wf := Gt_wf

instance : Monotone Store (fun x => Join Init (F x)) where
  F_mon := by admit /- proof elided -/

noncomputable def fixpoint_join : Store :=

@[grind] theorem iterate_correct (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) (heq : y = iterate _ F x PRE SMALL ) : eq y (F y) ∧ ∀ z, le (F z) z → le y z

theorem fixpoint_join_eq : Eq' (Join Init (F (fixpoint_join Init F) )) (fixpoint_join Init F) :=

:= by
  generalize heq1 : fixpoint_join Init F = t
  apply Eq'_sym
  simp [fixpoint_join] at *
  have := (@iterate_correct Store _ (fun x => Join Init (F x)) _ ?_ ?_ ?_ ?_ ?_ ).1
  unfold Eq'
  · exact this
  · exact Init
  · apply Le_Join_l
  · intro z hyp x
    specialize hyp x
    grind
  · rw [heq1]


Equal

def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

Store

def Store := Std.HashMap ident Int

Le

@[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

Gt

theorem hash_set_incl_size_leq (S1 S2 : Store) : Le S2 S1 → List.Subperm (S1.toList) (S2.toList)

@[grind] theorem Le_cardinal :
  ∀ S T : Store,
  Le T S ->
  S.size <= T.size ∧ (S.size = T.size → Equal S T)

@[grind] theorem Gt_cardinal :
  ∀ S S', Gt S S' -> S.size < S'.size

import LeroyCompilerVerificationCourse.Imp

import LeroyCompilerVerificationCourse.Constprop

import Batteries.Data.List.Perm

open OrderStruct

open Monotone

section FixpointExistence

variable (α : Sort u) (F : α → α) [OrderWithBot α]

open OrderStruct OrderWithBot

end FixpointExistence

section Iterate

variable (α : Sort u) [inst : OrderStruct α] (F : α → α) [Monotone α F]

open OrderStruct

end Iterate

section Fixpoint

open OrderWithBot

variable {α : Sort u} [i : OrderWithBot α] (F : α → α) [Monotone α F]

end Fixpoint

section Constprop

open Std.HashMap

@[grind] def Eq' (S1 S2 : Store) : Prop := Equal S1 S2

@[grind] def Gt (S1 S2 : Store) := Le S2 S1 ∧ ¬ Eq' S2 S1

:= by
  have := @InvImage Store Nat Nat.lt fun x => x.size
  have : ∀ (x y : Store), Gt x y → @InvImage Store Nat Nat.lt (fun x => x.size) x y := by
    intro x y heq
    unfold InvImage
    simp
    apply Gt_cardinal
    exact heq
  have subrel : Subrelation Gt (InvImage Nat.lt (fun x : Store => x.size)) := by
    intro x y gt; grind
  apply @Subrelation.wf Store (InvImage Nat.lt (fun x : Store => x.size)) Gt subrel
  exact InvImage.wf (fun x : Store => x.size) (Nat.lt_wfRel.wf)


Store

def Store := Std.HashMap ident Int

Join

@[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

com

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)

Beval

@[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

@[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

Aeval

@[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

@[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

ident

def ident := String deriving BEq, Repr, Hashable

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

Update

@[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

notation:10 l:10 " ;; " r:11 => com.SEQ l r

fixpoint'

@[grind] theorem iterate_correct (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) (heq : y = iterate _ F x PRE SMALL ) : eq y (F y) ∧ ∀ z, le (F z) z → le y z

import LeroyCompilerVerificationCourse.Imp

import LeroyCompilerVerificationCourse.Constprop

import Batteries.Data.List.Perm

@[grind] class OrderStruct (α : Sort u) where
  eq : α → α → Prop
  le : α → α → Prop
  beq : α → α → Bool
  le_trans : ∀ x y z, le x y -> le y z -> le x z
  beq_true' : ∀ x y : α, beq x y = true → eq x y := by admit /- proof elided -/

open OrderStruct

class OrderWithBot (α : Sort u) extends OrderStruct α where
  bot : α
  bot_smallest : ∀ x, le bot x

class Monotone (α : Sort u) (F : α → α) [OrderStruct α] where
  F_mon : ∀ {x y : α}, le x y → le (F x) (F y)

open Monotone

section FixpointExistence

variable (α : Sort u) (F : α → α) [OrderWithBot α]

open OrderStruct OrderWithBot

end FixpointExistence

section Iterate

variable (α : Sort u) [inst : OrderStruct α] (F : α → α) [Monotone α F]

open OrderStruct

@[grind] def iterate (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) : α :=
  if beq x (F x) then x else iterate (F x) (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )
  termination_by x
  decreasing_by
    grind [beq_false']

end Iterate

section Fixpoint

open OrderWithBot

variable {α : Sort u} [i : OrderWithBot α] (F : α → α) [Monotone α F]

@[grind] def fixpoint' : α := iterate α F bot (by admit /- proof elided -/
) (by admit /- proof elided -/
)

theorem fixpoint_correct :
  eq (fixpoint' F) (F (fixpoint' F)) ∧ ∀ z : α, le (F z) z → le (fixpoint' F) z :=

:= by
    unfold fixpoint'
    apply iterate_correct
    rotate_left
    · exact bot
    · apply bot_smallest
    · grind [bot_smallest]
    · rfl


Join

@[grind] def Join (S1 S2 : Store) : Store :=
  S1.filter (fun key _ => S2.get? key == S1.get? key)

Store

def Store := Std.HashMap ident Int

Le

@[grind] def Le (S1 S2 : Store) : Prop :=
  ∀ x n, S2.get? x = .some n -> S1.get? x = .some n

Equal

def Equal (S1 S2 : Store) := Std.HashMap.Equiv S1 S2

com

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)

Beval

@[grind] def Beval (S : Store) (b : bexp) : Option Bool :=
  match b with
  | .TRUE => .some true
  | .FALSE => .some false
  | .EQUAL a1 a2 => lift2 (fun m n => m == n) (Aeval S a1) (Aeval S a2)
  | .LESSEQUAL a1 a2 => lift2 (fun m n => m <= n) (Aeval S a1) (Aeval S a2)
  | .NOT b1 => lift1 (fun m => !m) (Beval S b1)
  | .AND b1 b2 => lift2 (fun m n => m && n) (Beval S b1) (Beval S b2)

@[grind] def lift1 {A B : Type} (f : A -> B) (o : Option A) : Option B :=
  match o with
  | .some x => .some (f x)
  | .none => .none

Aeval

@[grind] def Aeval (S : Store) (a : aexp) : Option Int :=
  match a with
  | .CONST n => .some n
  | .VAR x => S.get? x
  | .PLUS a1 a2 => lift2 (Int.add) (Aeval S a1) (Aeval S a2)
  | .MINUS a1 a2 => lift2 (Int.sub) (Aeval S a1) (Aeval S a2)

@[grind] def lift2 {A B C : Type} (f : A -> B -> C) (o1 : Option A) (o2 : Option B) : Option C :=
  match o1, o2 with
  | .some x1, .some x2 => .some (f x1 x2) | _, _ => .none

aexp

inductive aexp : Type where
  | CONST (n : Int)
  | VAR (x : ident)
  | PLUS (a1 : aexp) (a2 : aexp)
  | MINUS (a1 : aexp) (s2 : aexp)

ident

def ident := String deriving BEq, Repr, Hashable

bexp

inductive bexp : Type where
  | TRUE
  | FALSE
  | EQUAL (a1 : aexp) (a2 : aexp)
  | LESSEQUAL (a1 : aexp) (a2 : aexp)
  | NOT (b1 : bexp)
  | AND (b1 : bexp) (b2 : bexp)

Update

@[grind] def Update (x : ident) (N : Option Int) (S : Store) : Store :=
  match N with
  | .none => S.erase x
  | .some n => S.insert x n

notation:10 l:10 " ;; " r:11 => com.SEQ l r

theorem Le_Join_l : ∀ S1 S2, Le S1 (Join S1 S2)

iterate

@[grind] theorem iterate_correct (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) (heq : y = iterate _ F x PRE SMALL ) : eq y (F y) ∧ ∀ z, le (F z) z → le y z

import LeroyCompilerVerificationCourse.Imp

import LeroyCompilerVerificationCourse.Constprop

import Batteries.Data.List.Perm

@[grind] class OrderStruct (α : Sort u) where
  eq : α → α → Prop
  le : α → α → Prop
  beq : α → α → Bool
  le_trans : ∀ x y z, le x y -> le y z -> le x z
  beq_true' : ∀ x y : α, beq x y = true → eq x y := by admit /- proof elided -/

open OrderStruct

class Monotone (α : Sort u) (F : α → α) [OrderStruct α] where
  F_mon : ∀ {x y : α}, le x y → le (F x) (F y)

open Monotone

section FixpointExistence

variable (α : Sort u) (F : α → α) [OrderWithBot α]

open OrderStruct OrderWithBot

end FixpointExistence

section Iterate

variable (α : Sort u) [inst : OrderStruct α] (F : α → α) [Monotone α F]

open OrderStruct

@[grind] def iterate (x : α) (PRE : le x (F x)) (SMALL : ∀ z, le (F z) z -> le x z) : α :=
  if beq x (F x) then x else iterate (F x) (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )
  termination_by x
  decreasing_by
    grind [beq_false']

end Iterate

section Fixpoint

open OrderWithBot

variable {α : Sort u} [i : OrderWithBot α] (F : α → α) [Monotone α F]

end Fixpoint

section Constprop

open Std.HashMap

open OrderStruct

noncomputable instance : OrderStruct Store where
  eq := Equal
  le := Le
  beq (S1 S2 : Store) :=  Decidable.decide (Equal S1 S2)
  le_trans := Le_trans
  gt_wf := Gt_wf

end Constprop

section FixpointJoin

variable (Init : Store)

variable (F : Store → Store) [Monotone Store F]

instance : Monotone Store (fun x => Join Init (F x)) where
  F_mon := by admit /- proof elided -/

noncomputable def fixpoint_join : Store :=

theorem fixpoint_join_smallest :
  ∀ S, Le (Join Init (F S)) S -> Le (fixpoint_join Init F) S :=

:= by
    intro S LE
    unfold fixpoint_join
    have := (@iterate_correct Store _ (fun x => Join Init (F x)) _ (fixpoint_join Init F) Init (?_) ?_ ?_).2 S LE
    exact this
    · apply Le_Join_l
    · intro z hyp x
      specialize hyp x
      grind
    · unfold fixpoint_join
      dsimp


syntax "method" ident leafny_binder* "return" "(" ident ":" term ")"
  (require_caluse )*
  (ensures_caluse)* "do" doSeq
  Termination.suffix : command

syntax "ensures" termBeforeReqEnsDo : ensures_caluse

syntax "while_some" term ":|" termBeforeDo "do" doSeq : doElem

syntax "while_some" term ":|" term
  (invariantClause)+
  (doneWith)?
  "do" doSeq : doElem

syntax "let" term ":|" term : doElem

syntax "done_with" termBeforeDo : doneWith

syntax "invariant" termBeforeDo linebreak : invariantClause

syntax "while" term
  (invariantClause)*
  (doneWith)?
  (decreasingTerm)?
  "do" doSeq : doElem

syntax "(mut" ident ":" term ")" : leafny_binder

macro_rules
  | `(doElem| let $x:term :| $t) => `(doElem| let $x:term <- pickSuchThat _ (fun $x => type_with_name_prefix `choice $t))

macro_rules
  | `(doElem| while $t do $seq:doSeq) => do
    let decr <- withRef (<- getRef) `(decreasing none)
    let invs <- withRef (<- getRef) `(invariants [])
    `(doElem|
      for _ in Lean.Loop.mk do
        $invs:term
        onDoneGadget (with_name_prefix `done ¬$t:term)
        $decr:term
        if $t then
          $seq:doSeq
        else break)
  | `(doElem| while $t
              $[invariant $inv:term
              ]*
              $[done_with $inv_done]?
              $[decreasing $measure]?
              do $seq:doSeq) => do
      let invs <- `(invariants [ $[(with_name_prefix `invariant $inv:term)],* ])
      let invd_some ← match inv_done with
      | some invd_some => withRef invd_some ``($invd_some)
      | none => ``(¬$t:term)
      match measure with
      | some measure_some =>
        let decr <- withRef measure_some `(decreasing type_with_name_prefix `decreasing $measure_some)
        `(doElem|
          for _ in Lean.Loop.mk do
            $invs:term
            onDoneGadget (with_name_prefix `done $invd_some:term)
            $decr:term
            if $t then
              $seq:doSeq
            else break)
      | none => do
        let decr <- withRef (<- getRef) `(decreasing none)
        let invs <- `(invariants [ $[(with_name_prefix `invariant $inv:term)],* ])
        `(doElem|
          for _ in Lean.Loop.mk do
            $invs:term
            onDoneGadget (with_name_prefix `done $invd_some:term)
            $decr:term
            if $t then
              $seq:doSeq
            else break)
  | `(doElem| while_some $x:ident :| $t do $seq:doSeq) =>
    match seq with
    | `(doSeq| $[$seq:doElem]*)
    | `(doSeq| $[$seq:doElem;]*)
    | `(doSeq| { $[$seq:doElem]* }) =>
      `(doElem|
        while ∃ $x:ident, $t do
          let $x :| $t
          $[$seq:doElem]*)
    | _ => Lean.Macro.throwError "while_some expects a sequence of do-elements"
  | `(doElem| while_some $x:ident :| $t
              $[invariant $inv:term
              ]*
              $[done_with $inv_done]? do
                $seq:doSeq) => do
    let invs <- `(invariants [ $[(with_name_prefix `invariant $inv:term)],* ])
    let invd_some ← match inv_done with
    | some invd_some => withRef invd_some ``($invd_some)
    | none => ``(¬$t:term)
    match seq with
    | `(doSeq| $[$seq:doElem]*)
    | `(doSeq| $[$seq:doElem;]*)
    | `(doSeq| { $[$seq:doElem]* }) =>
      let decr <- withRef (<- getRef) `(decreasing none)
      `(doElem|
        for _ in Lean.Loop.mk do
          $invs:term
          onDoneGadget (with_name_prefix `done $invd_some:term)
          $decr:term
          if ∃ $x:ident, $t then
            let $x :| $t
            $[$seq:doElem]*
          else break)
    | _ => Lean.Macro.throwError "while_some expects a sequence of do-elements"
  | `(doElem| for $x:ident in $t
            $[invariant $inv:term
            ]*
            do $seq:doSeq) => do
      let invs <- `(invariants [ $[(with_name_prefix `invariant $inv:term)],* ])
      match seq with
      | `(doSeq| $[$seq:doElem]*)
      | `(doSeq| $[$seq:doElem;]*)
      | `(doSeq| { $[$seq:doElem]* }) =>
        `(doElem|
          for $x:ident in $t do
            $invs:term
            $[$seq:doElem]*)
      | _ => Lean.Macro.throwError "for expects a sequence of do-elements"

macro_rules
  | `(tactic|loom_solver) =>
    `(tactic|(
      try simp at *
      try aesop))

macro_rules
  | `(doElem| while $t
              $[invariant $inv:term
              ]*
              $[done_with $inv_done]?
              $[decreasing $measure]?
              do $seq:doSeq) => do
      let balance := mkIdent `balance_name
      let balanceType <- `(term| Bal)
      let inv : Array Term <- inv.mapM fun (inv : Term) => withRef inv ``(fun ($(balance):ident : $balanceType)=> with_name_prefix `inv $inv)
      let invd_some <- match inv_done with
      | some invd_some => withRef invd_some ``(fun ($(balance):ident : $balanceType) => with_name_prefix `done $invd_some)
      | none => ``(fun ($(balance):ident : $balanceType) => with_name_prefix `done ¬$t:term)
      match measure with
      | some measure_some =>
        let measure_some ← withRef measure_some ``(type_with_name_prefix `decreasing ($measure_some:term))
        do
        `(doElem|
          for _ in Lean.Loop.mk do
            invariantGadget [ $[$inv:term],* ]
            onDoneGadget ($invd_some:term)
            decreasingGadget ($measure_some:term)
            if $t then
              $seq:doSeq
            else break)
      | none => do
        `(doElem|
          for _ in Lean.Loop.mk do
            invariantGadget [ $[$inv:term],* ]
            onDoneGadget ($invd_some:term)
            decreasingGadget none
            if $t then
              $seq:doSeq
            else break)

macro_rules
| `(doElem|balance_set $t) => do
  let balId := mkIdent `balance
  `(doElem|do
    $balId:ident := $t
    set $balId:ident
    $balId:ident ← get)

macro_rules
  | `(doElem|$id:ident[$idx:term] := $val:term) =>
    `(doElem| $id:term := ($id:term).modify $idx (fun _ => $val))
  | `(doElem|$id:ident[$idx:term] += $val:term) =>
    `(doElem| $id:term := ($id:term).modify $idx (· + $val))

theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

SpV

theorem getValSpV_eq (spv: SpV Int) (j: ℕ) (h_ind: j < spv.size): spv[spv.ind[j]!] = (spv.val)[j]!

theorem getValSpV_empty (spv: SpV Int) (j: ℕ) (h_empty: ∀ i < spv.size, spv.ind[i]! ≠ j): spv[j] = 0

import Auto

import Lean

import Mathlib.Algebra.BigOperators.Intervals

import Mathlib.Algebra.Ring.Int.Defs

import Loom.MonadAlgebras.NonDetT.Extract

import Loom.MonadAlgebras.WP.Tactic

import Loom.MonadAlgebras.WP.DoNames'

import CaseStudies.Velvet.Std

section SpMV

structure SpV (valTyp : Type) where
  ind: Array Nat
  val: Array valTyp
  size: ℕ
  size_eq: ind.size = size ∧ val.size = size
  inc: ∀ (i j: Nat), i < size → j < size →  i < j → ind[i]! < ind[j]!

def spv_dot (spv1 spv2: SpV Int) (pnt1 pnt2: ℕ): Int :=
  if (spv1.size) ≤ pnt1 ∨ (spv2.size) ≤ pnt2 then
    0
  else
    if (spv1.ind)[pnt1]! = (spv2.ind)[pnt2]! then
      (spv1.val)[pnt1]! * (spv2.val)[pnt2]! + spv_dot spv1 spv2 (pnt1 + 1) (pnt2 + 1)
    else
      if (spv1.ind)[pnt1]! < (spv2.ind)[pnt2]! then
        spv_dot spv1 spv2 (pnt1 + 1) pnt2
      else
        spv_dot spv1 spv2 pnt1 (pnt2 + 1)
  termination_by ((spv1.size) + (spv2.size) - pnt1 - pnt2)

method SpVSpV
  (mut out: Array Int)
  (spv1: SpV Int)
  (spv2: SpV Int) return (u: Unit)
  ensures out.size = 1
  ensures out[0]! = spv_dot spv1 spv2 0 0
  do
    out := Array.replicate 1 0
    let mut pnt1 := 0
    let mut pnt2 := 0
    while pnt1 ≠ spv1.size ∧ pnt2 ≠ spv2.size
    invariant out.size = 1
    invariant pnt1 ≤ spv1.size ∧ pnt2 ≤ spv2.size
    invariant out[0]! + spv_dot spv1 spv2 pnt1 pnt2 = spv_dot spv1 spv2 0 0
    done_with pnt1 = spv1.size ∨ pnt2 = spv2.size
    do
      if (spv1.ind)[pnt1]! = (spv2.ind)[pnt2]! then
        out[0] += (spv1.val)[pnt1]! * (spv2.val)[pnt2]!
        pnt1 := pnt1 + 1
        pnt2 := pnt2 + 1
      else
        if (spv1.ind)[pnt1]! < (spv2.ind)[pnt2]! then
          pnt1 := pnt1 + 1
        else
          pnt2 := pnt2 + 1
    return

method SpMSpV
  (mut out: Array Int)
  (spm: Array (SpV Int))
  (spv: SpV Int) return (u: Unit)
  ensures out.size = spm.size
  ensures ∀ i < spm.size, out[i]! = spv_dot spm[i]! spv 0 0
  do
    out := Array.replicate spm.size 0
    let mut spmInd := Array.replicate spm.size 0
    let mut spvInd := Array.replicate spm.size 0
    while_some i :| i < spm.size ∧ spmInd[i]! < spm[i]!.size ∧ spvInd[i]! < spv.size
    invariant spvInd.size = spm.size
    invariant spmInd.size = spm.size
    invariant out.size = spm.size
    invariant ∀ i < spmInd.size, spmInd[i]! <= spm[i]!.size
    invariant ∀ i < spvInd.size, spvInd[i]! <= spv.size
    invariant ∀ i < spm.size, out[i]! + spv_dot spm[i]! spv spmInd[i]! spvInd[i]! = spv_dot spm[i]! spv 0 0
    done_with ∀ i < spm.size, spmInd[i]! = spm[i]!.size ∨ spvInd[i]! = spv.size
    do
      let ind_m := spmInd[i]!
      let ind_v := spvInd[i]!
      if spm[i]!.ind[ind_m]! = spv.ind[ind_v]! then
        out[i] += spm[i]!.val[ind_m]! * spv.val[ind_v]!
        spmInd[i] += 1
        spvInd[i] += 1
      else
        if spm[i]!.ind[ind_m]! < spv.ind[ind_v]! then
          spmInd[i] += 1
        else
          spvInd[i] += 1
    return

theorem spv_dot_pure_gen (spv1: SpV Int) (spv2: SpV Int) (n pnt1 pnt2: ℕ)
  (sz1: ∀ i < spv1.size, spv1.ind[i]! < n)
  (sz2: ∀ i < spv2.size, spv2.ind[i]! < n):
  spv_dot spv1 spv2 pnt1 pnt2 =
    ∑ i ∈ Finset.range n,
      if max
        (if spv1.size ≤ pnt1 then n else spv1.ind[pnt1]!)
        (if spv2.size ≤ pnt2 then n else spv2.ind[pnt2]!) ≤ i then
        spv1[i] * spv2[i]
      else
        0 :=

:= by
  fun_induction spv_dot spv1 spv2 pnt1 pnt2 with
  | case1 p1 p2 h =>
    have all_zero: (∀ x ∈ Finset.range n, (if max (if spv1.size ≤ p1 then n else spv1.ind[p1]!) (if spv2.size ≤ p2 then n else spv2.ind[p2]!) ≤ x then spv1[x] * spv2[x] else 0) = 0) := by
      intro x hx
      simp
      rcases h with ob1 | ob2
      { simp [ob1]
        simp at hx
        simp [hx] }
      simp [ob2]
      simp at hx
      simp [hx]
    rw [Finset.sum_eq_zero all_zero]
  | case2 p1 p2 h1 eq ih =>
    rw [ih, ←getValSpV_eq spv1 p1 (by omega), ←getValSpV_eq spv2 p2 (by omega)]
    have sum_eq_single:
      ∑ i ∈ Finset.range n, (if i = spv1.ind[p1]! then spv1[spv1.ind[p1]!] * spv2[spv1.ind[p1]!] else 0) = (if spv1.ind[p1]! = spv1.ind[p1]! then spv1[spv1.ind[p1]!] * spv2[spv1.ind[p1]!] else 0) := by
      have hb: ∀ i ∈ Finset.range n, i ≠ spv1.ind[p1]! → ((if i = spv1.ind[p1]! then spv1[spv1.ind[p1]!]! * spv2[spv1.ind[p1]!]! else 0) = 0) := by
        intro i hi iq
        simp at hi
        simp [iq]
      have hc: spv1.ind[p1]! ∉ Finset.range n → ((if spv1.ind[p1]! = spv1.ind[p1]! then spv1[spv1.ind[p1]!]! * spv2[spv1.ind[p1]!]! else 0) = 0) := by
        intro nin
        simp at nin
        have bnd := sz1 p1 (by omega)
        omega
      apply Finset.sum_eq_single spv1.ind[p1]! hb hc
    rw [if_pos] at sum_eq_single
    rw [←eq, ←sum_eq_single]
    rw [←Finset.sum_add_distrib]
    apply Finset.sum_congr
    { rfl }
    intro x hx
    simp at hx
    have h2: ¬spv1.size ≤ p1 ∧ ¬spv2.size ≤ p2 := by omega
    simp [h2]
    have hnx: ¬(n ≤ x) := by simp [hx]
    by_cases xeq: x = spv1.ind[p1]! <;> simp [xeq]
    { intro in1 in2
      by_cases edg1: spv1.size ≤ p1 + 1 <;> simp [edg1] at in1
      { omega }
      have p1lt := (spv1.inc) p1 (p1 + 1) (by omega) (by omega) (by simp)
      -- simp [getElem] at in1
      omega }
    have outb_lemma (spv: SpV Int) (pos: ℕ) (hsz: 1 ≤ spv.size) (hb: spv.ind[spv.size - 1]! < pos): spv[pos] = 0 := by
      by_cases ex: ∃ i < spv.size, spv.ind[i]! = pos
      { rcases ex with ⟨i, hbi, hi⟩
        by_cases heq: i = spv.size - 1
        { simp [heq] at hi
          omega }
        have contra := (spv.inc) i (spv.size - 1) (by omega) (by omega) (by omega)
        simp at hb
        -- simp [getElem] at hi
        omega }
      simp at ex
      simp [getValSpV_empty spv pos ex]
    by_cases edg1: spv1.size ≤ p1 + 1 <;> simp [edg1]
    { simp [hnx]
      have p1eq : p1 = spv1.size - 1 := by omega
      by_cases lex: spv1.ind[p1]! ≤ x <;> simp [lex]
      have ltx: spv1.ind[p1]! < x := by omega
      simp [p1eq] at ltx
      simp [outb_lemma spv1 x (by omega) ltx] }
    by_cases edg2: spv2.size ≤ p2 + 1 <;> simp [edg2]
    { simp [hnx]
      by_cases lex: spv1.ind[p1]! ≤ x <;> simp [lex]
      have ltx: spv1.ind[p1]! < x := by omega
      have p2eq: p2 = spv2.size - 1 := by omega
      simp [eq, p2eq] at ltx
      simp [outb_lemma spv2 x (by omega) ltx] }
    by_cases val: spv1.ind[p1 + 1]! ≤ x ∧ spv2.ind[p2 + 1]! ≤ x <;> simp [val]
    { have inc: spv1.ind[p1]! < spv1.ind[p1 + 1]! := (spv1.inc) p1 (p1 + 1) (by omega) (by omega) (by simp)
      simp [le_of_lt (lt_of_lt_of_le inc val.left)] }
    by_cases xb: spv1.ind[p1]! ≤ x <;> simp [xb]
    have ltx: spv1.ind[p1]! < x := by omega
    simp at val
    have interm_lemma (spv: SpV Int) (pos idx: ℕ) (hsz: idx + 1 < spv.size) (inter: spv.ind[idx]! < pos ∧ pos < spv.ind[idx + 1]!): spv[pos] = 0 := by
      by_cases ex: ∃ i < spv.size, spv.ind[i]! = pos
      { rcases ex with ⟨i, hbi, hi⟩
        have lt_lemma (i1 i2: ℕ) (hi1: i1 < spv.size) (hi2: i2 < spv.size): spv.ind[i1]! < spv.ind[i2]! → i1 < i2 := by
          intro hlt
          by_cases contra_lt: i2 < i1
          { have inc := spv.inc i2 i1 hi2 hi1 contra_lt
            simp at hlt
            omega }
          by_cases contra_eq: i1 = i2
          { simp [contra_eq] at hlt }
          omega
        have left_x := lt_lemma idx i (by omega) (by omega) (by omega)
        have right_x := lt_lemma i (idx + 1) (by omega) (by omega) (by omega)
        omega }
      simp at ex
      simp [getValSpV_empty spv pos ex]
    rcases (em (spv1.ind[p1 + 1]! ≤ x)) with c1 | c2
    { simp [c1] at val
      simp [interm_lemma spv2 x p2 (by omega) (by omega)] }
    simp at c2
    simp [interm_lemma spv1 x p1 (by omega) (by omega)]
    rfl
  | case3 p1 p2 h1 neq le ih =>
    rw [ih]
    apply Finset.sum_congr
    { rfl }
    intro x hx
    by_cases edg: spv1.size = p1 + 1
    { have sz_concl := sz1 p1 (by omega)
      simp [edg]
      simp at hx
      have nx : ¬(n ≤ x) := by omega
      by_cases b2: spv2.size ≤ p2 <;> simp [b2] <;> simp [nx]
      intro ifc1 ifc2
      by_cases ex: ∃ i < spv1.size, spv1.ind[i]! = x
      { rcases ex with ⟨i, ib, hi⟩
        by_cases il: p1 < i
        { omega }
        simp at il
        rcases (le_iff_eq_or_lt.mp il) with ieq | ilt
        { simp [ieq] at hi
          omega }
        have contra := spv1.inc i p1 ib (by omega) ilt
        omega }
      simp at ex
      simp [getValSpV_empty spv1 x ex] }
    have inb: ¬(spv1.size ≤ p1 + 1) ∧ ¬(spv1.size ≤ p1):= by omega
    simp [inb]
    by_cases edg2: spv2.size ≤ p2 <;> simp [edg2]
    { simp at hx
      have neg: ¬ (n ≤ x) := by omega
      simp [neg] }
    by_cases le2: spv2.ind[p2]! ≤ x <;> simp [le2]
    by_cases upper: spv1.ind[p1 + 1]! ≤ x <;> simp [upper]
    { simp at inb
      simp [getElem]
      simp [le_trans (le_of_lt (spv1.inc p1 (p1 + 1) inb.right inb.left (by simp))) upper] }
    by_cases lower: spv1.ind[p1]! ≤ x <;> simp [lower]
    by_cases ex: ∃ i < spv1.size, spv1.ind[i]! = x
    { rcases ex with ⟨ind, indb, hind⟩
      by_cases ind_x: ind = p1
      { simp [ind_x] at hind
        omega }
      rcases (lt_or_gt_of_ne ind_x) with ltc | gtc
      { have contra := spv1.inc ind p1 indb (by omega) ltc
        -- simp [getElem] at lower
        -- simp [getElem] at hind
        omega }
      by_cases indp1: ind = p1 + 1
      { simp [indp1] at hind
        omega }
      have indlt: p1 + 1 < ind := by omega
      have contra := spv1.inc (p1 + 1) ind (by omega) indb indlt
      -- simp [getElem] at upper
      -- simp [getElem] at hind
      omega }
    simp at ex
    simp [getValSpV_empty spv1 x ex]
  | case4 p1 p2 h1 neq nle ih =>
    rw [ih]
    apply Finset.sum_congr
    { rfl }
    intro x hx
    by_cases edg: spv2.size = p2 + 1
    { have sz_concl := sz2 p2 (by omega)
      simp [edg]
      simp at hx
      have nx : ¬(n ≤ x) := by omega
      by_cases b1: spv1.size ≤ p1 <;> simp [b1] <;> simp [nx]
      intro ifc1 ifc2
      by_cases ex: ∃ i < spv2.size, spv2.ind[i]! = x
      { rcases ex with ⟨i, ib, hi⟩
        by_cases il: p2 < i
        { omega }
        simp at il
        rcases (le_iff_eq_or_lt.mp il) with ieq | ilt
        { simp [ieq] at hi
          omega }
        have contra := (spv2.inc) i p2 ib (by omega) ilt
        -- simp [getElem] at hi
        -- simp [getElem] at ifc2
        omega }
      simp at ex
      simp [getValSpV_empty spv2 x ex] }
    have inb: ¬(spv2.size ≤ p2 + 1) ∧ ¬(spv2.size ≤ p2):= by omega
    simp [inb]
    by_cases edg1: spv1.size ≤ p1 <;> simp [edg1]
    { simp at hx
      have neg: ¬ (n ≤ x) := by omega
      simp [neg] }
    by_cases le1: spv1.ind[p1]! ≤ x <;> simp [le1]
    by_cases upper: spv2.ind[p2 + 1]! ≤ x <;> simp [upper]
    { simp at inb
      simp [getElem]
      simp [le_trans (le_of_lt ((spv2.inc) p2 (p2 + 1) inb.right inb.left (by simp))) upper] }
    by_cases lower: spv2.ind[p2]! ≤ x <;> simp [lower]
    by_cases ex: ∃ i < spv2.size, spv2.ind[i]! = x
    { rcases ex with ⟨ind, indb, hind⟩
      by_cases ind_x: ind = p2
      { simp [ind_x] at hind
        omega }
      rcases (lt_or_gt_of_ne ind_x) with ltc | gtc
      { have contra := (spv2.inc) ind p2 indb (by omega) ltc
        -- simp [getElem] at lower
        -- simp [getElem] at hind
        omega }
      by_cases indp2: ind = p2 + 1
      { simp [indp2] at hind
        omega }
      have indlt: p2 + 1 < ind := by omega
      have contra := (spv2.inc) (p2 + 1) ind (by omega) indb indlt
      -- simp [getElem] at upper
      -- simp [getElem] at hind
      omega }
    simp at ex
    simp [getValSpV_empty spv2 x ex]


macro "⌜" p:term "⌝" : term => `(LE.pure $p)

open

notation "[part|" t "]" => open ExceptionAsSuccess in t

notation "[totl|" t "]" => open ExceptionAsFailure in t

wp

def wp (c : m α) (post : α -> l) : l := liftM (n := Cont l) c post

NonDetT.wp

def NonDetT.wp {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : {α : Type u} -> NonDetT m α -> Cont l α
  | _, .pure ret => pure ret
  | _, .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | _, @NonDetT.pickCont _ _ τ p _ f => fun post => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post
  | _, .repeatCont init f cont => fun post => ⨆ (inv : ForInStep _ -> l),
    ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) inv⌝ ⊓
    spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

NonDetT

inductive NonDetT (m : Type u -> Type v) : (α : Type u) -> Type _ where
  | pure {α} (ret : α) : NonDetT m α
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) : NonDetT m α
  | pickCont {α} (τ : Type u) (p : τ -> Prop) [Findable p] (f : τ → NonDetT m α) : NonDetT m α
  | repeatCont {α} {β} (init : β) (f : β -> NonDetT m (ForInStep β)) (cont : β -> NonDetT m α) : NonDetT m α

spec

def spec (pre : l) (post : α -> l) : Cont l α :=
  fun p => pre ⊓ ⌜post ≤ p⌝

MAlgOrdered

class MAlgOrdered (l : outParam (Type v)) [Monad m] [CompleteLattice l] where
  μ : m l -> l
  μ_ord_pure : ∀ l, μ (pure l) = l
  μ_ord_bind {α : Type v} :
    ∀ (f g : α -> m l), μ ∘ f ≤ μ ∘ g ->
      ∀ x : m α, μ (x >>= f) ≤ μ (x >>= g)

NonDetT.wp

def   NonDetT.wp {l : Type u} {α : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m α -> Cont l α
  | .pure ret => pure ret
  | .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | .pickCont _ p f => fun post => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post

LE.pure

noncomputable def LE.pure {l : Type u} [inst: LE l] [OrderTop l] [OrderBot l] : Prop -> l := fun p =>
  if p then ⊤ else ⊥

NonDetT.wp

def NonDetT.wp {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : {α : Type u} -> NonDetT m α -> Cont l α
  | _, .pure ret => pure ret
  | _, .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | _, @NonDetT.pickCont _ _ τ p _ f =>
  fun post =>
    let p : Set τ := p;
    ⨅ a ∈ (p : Set τ), wp (f a) post
  | _, @NonDetT.repeatCont _ _ β init f cont => fun post => ⨆ (inv : ForInStep β -> l),
      ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) inv⌝ ⊓
      spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

def NonDetT.wp {l : Type u} {α : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m α -> Cont l α
  | .pure ret => pure ret
  | .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | .pickCont τ p f => fun post => let p : Set τ := p; ⨅ a ∈ (p : Set τ), wp (f a) post

W

structure W (t : Type v) [Preorder t] (α : Type u) where
  wp : Cont t α
  wp_montone : wp.monotone

NonDetT.wp

def NonDetT.wp {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : {α : Type u} -> NonDetT m α -> Cont l α
  | _, .pure ret => pure ret
  | _, .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | _, @NonDetT.pickCont _ _ τ p _ f => fun post => let p : Set τ := p; ⌜∃ h, p h⌝ ⊓  ⨅ a ∈ (p : Set τ), wp (f a) post
  | _, @NonDetT.repeatCont _ _ β init f cont => fun post => ⨆ (inv : ForInStep β -> l) (measure : β -> Nat),
      ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))⌝ ⊓
      spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

def   NonDetT.wp {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : {α : Type u} -> NonDetT m α -> Cont l α
  | _, .pure ret => pure ret
  | _, .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | _, @NonDetT.pickCont _ _ _ p _ f => fun post => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post
  | _, .repeatCont init f cont => fun post => ⨆ (inv : ForInStep _ -> l) (measure : _ -> Nat),
    ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))⌝ ⊓
    spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

DivM

inductive DivM (α : Type u) where
  | res (x : α)
  | div

LogicLift

class LogicLift (l : outParam (Type u)) ( k : Type u) [CompleteLattice l] [CompleteLattice k] where
  [lift : MonadLift (Cont l) (Cont k)]
  lift_top {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊤) = ⊤
  lift_bot {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊥) = ⊥

MAlg.lift

abbrev MAlg.lift {m : Type u -> Type v} {l : Type u} [Monad m] [MAlg m l] :
  {α : Type u} -> m α -> Cont l α := fun x f => μ $ f <$> x

MAlg

class MAlg [Monad m] (l : outParam (Type v)) where
  μ : m l -> l
  pure : ∀ l, μ (pure l) = l
  bind : ∀ {α : Type v} (x : m α) (f g : α -> m l),
    μ ∘ f = μ ∘ g ->
    μ (x >>= f) = μ (x >>= g)

NonDetT.μ

def NonDetT.μ  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m l -> l := fun x => NonDetT.wp x id

NonDetT.bind

def NonDetT.bind (x : NonDetT m α) (f : α → NonDetT m β) : NonDetT m β :=
  match x with
  | pure ret => f ret
  | vis x f' => vis x fun y => bind (f' y) f
  | pickCont τ p f' => pickCont τ p fun t => bind (f' t) f
  | repeatCont init f' cont => repeatCont init f' fun t => bind (cont t) f

WPGen.bind

def WPGen.bind {x : m α} {f : α -> m β} (wpg : WPGen x) (wpgf : ∀ a, WPGen (f a)) :
  WPGen (x >>= f) where
  get := fun post => wpg.get (fun a => (wpgf a).get post)
  prop := by admit /- proof elided -/

private def _root_.Lean.SimpleScopedEnvExtension.get [Inhabited σ] (ext : SimpleScopedEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

Context

structure Context where
  ref         : Syntax

  m           : Syntax

  returnType  : Syntax
  mutableVars : VarSet := {}
  insideFor   : Bool := false

private def _root_.Lean.SimplePersistentEnvExtension.get [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

WPGen

structure WPGen (x : m α) where
  get : Cont l α

  prop : ∀ post, get post <= wp x post

private def _root_.Lean.EnvExtension.get [Inhabited σ] (ext : EnvExtension σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

CCPOBotLawful

class CCPOBotLawful (m : Type u -> Type v) [∀ α, Lean.Order.CCPO (m α)] [CCPOBot m] where
  prop {α} : CCPOBot.compBot (m := m) (α := α) = Lean.Order.bot

lemma TotalCorrectness.DivM.wp_eq (α : Type) (x : DivM α) (post : α -> Prop) :
  wp x post =
    match x with

lemma PartialCorrectness.DivM.wp_eq (α : Type) (x : DivM α) (post : α -> Prop) :
  wp x post =
    match x with

lemma NonDetT.wp_eq_wp {α : Type u} (x : NonDetT m α) (post : α -> l) :
  _root_.wp x post = NonDetT.wp x post

lemma NonDetT.wp_bind  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α β : Type u} (x : NonDetT m α) (f : α -> NonDetT m β)
  (post : β -> l):
  NonDetT.wp (x.bind f) post = NonDetT.wp x (fun x => NonDetT.wp (f x) post)

lemma NonDetT.wp_mono [LawfulMonad m] {α : Type u} {l : Type u} [CompleteLattice l] [MAlgOrdered m l] (x : NonDetT m α) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  NonDetT.wp x f <= NonDetT.wp x g

omit [MAlgOrdered m l] in
lemma spec_mono {α : Type u}  {l : Type u} [CompleteLattice l] (pre : l) (post : α -> l) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  spec pre post f <= spec pre post g

lemma NonDetT.wp_mono  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α : Type u} (x : NonDetT m α) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  NonDetT.wp x f <= NonDetT.wp x g

@[local simp]
lemma DivM.total_decompose (α : Type) (x : DivM α) (post₁ post₂ : α -> Prop) :
  ([totl| wp x post₁] ∧ [part| wp x post₂]) = [totl| wp x (post₁ ⊓ post₂)]

@[local simp]
lemma wp_mono_part (x : NonDetT DivM α) (post₁ post₂ : α -> Prop) :
  (post₁ ≤ post₂) → ([totl|wp x post₁]) ≤ ([part| wp x post₂])

import Loom.MonadAlgebras.NonDetT.Extract

notation "[totl|" t "]" => open TotalCorrectness TotalCorrectness.DemonicChoice in t

notation "[part|" t "]" => open PartialCorrectness PartialCorrectness.DemonicChoice in t

lemma VelvetM.total_decompose {α : Type} (x : VelvetM α) (post₁ post₂ : α -> Prop):
  [totl| wp x post₁] ⊓ [part| wp x post₂] = [totl| wp x (post₁ ⊓ post₂)] :=

:= by
    unhygienic induction x <;> try simp [loomLogicSimp]
    { simp [DivM.total_decompose]
      simp [[totl|DivM.wp_eq]]
      split
      { simp }
      rename_i arg
      have ind := f_ih arg post₁ post₂
      simp at ind
      rw [ind]
      trivial }
    { constructor <;> rintro hyp
      { constructor; aesop
        intro i hi
        have hl := hyp.left.right i hi
        have hr := hyp.right i hi
        have ind := f_ih i post₁ post₂
        simp [hl, hr] at ind
        exact ind }
      constructor; constructor; aesop
      all_goals
        intro i hi
        have conj := hyp.right i hi
        have ind := f_ih i post₁ post₂
        simp [loomLogicSimp] at ind
        rw [←ind] at conj
        simp [conj] }
    constructor
    { intro conj
      rcases conj with ⟨h, inv_spec, hspec⟩
      rcases h with ⟨inv, x1, hinv⟩
      rcases x1 with ⟨x, hx⟩
      exists inv ⊓ inv_spec
      constructor
      { exists x
        intro b hb
        simp [←[totl| NonDetT.wp_eq_wp]]
        simp [loomLogicSimp] at hb
        have hxb := hx b hb.left
        simp [←[totl| NonDetT.wp_eq_wp]] at hxb
        simp [spec, LE.pure, loomLogicSimp, ←[part| NonDetT.wp_eq_wp]] at hspec
        have hspecb := hspec.left b hb.right
        have ind := f_ih b
          (fun x_1 ↦
          match x_1 with
          | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
          | ForInStep.done b' => inv (ForInStep.done b'))
          inv_spec
        simp [loomLogicSimp] at ind
        have indr := ind.mp (And.intro hxb hspecb)
        have v1 := [totl| NonDetT.wp_mono
          (f b)
          (fun x_1 ↦
            (match x_1 with
              | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
              | ForInStep.done b' => inv (ForInStep.done b')) ∧
              inv_spec x_1)
          (fun x_1 ↦
            match x_1 with
            | ForInStep.yield b' => (inv (ForInStep.yield b') ∧ inv_spec (ForInStep.yield b')) ∧ x b' < x b
            | ForInStep.done b' => inv (ForInStep.done b') ∧ inv_spec (ForInStep.done b'))
          ]
        simp [loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]] at v1
        exact v1
          (fun x => by
            match x with
            | .yield b' => intro hb hspec; simp [hb, hspec]
            | .done b' => intro hb hspec; simp [hb, hspec] )
          indr }
      simp [spec, LE.pure, loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]] at hinv
      simp [spec, LE.pure, loomLogicSimp, ←[part| NonDetT.wp_eq_wp]] at hspec
      simp [spec, LE.pure, loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]]
      simp [hinv, hspec]
      intro x inv_x inv_spec
      have h₁ := hinv.right x inv_x
      have h₂ := hspec.right.right x inv_spec
      have cont_ind := cont_ih x post₁ post₂
      simp [loomLogicSimp] at cont_ind
      exact cont_ind.mp (And.intro h₁ h₂) }
    intro hyp
    rcases hyp with ⟨inv, x_ex, h_inv⟩
    rcases x_ex with ⟨x, hx⟩
    simp [spec]
    simp [spec, LE.pure] at h_inv
    constructor <;> exists inv <;> constructor
    { exists x }
    { simp [h_inv, LE.pure]
      exact le_trans h_inv.right (by
        simp [← [totl|NonDetT.wp_eq_wp]]
        simp [loomLogicSimp]
        intro x and_wp
        have cont_ind := cont_ih x post₁ post₂
        simp [loomLogicSimp, and_wp] at cont_ind
        simp [cont_ind] ) }
    { intro b hb
      have hbx := hx b hb
      simp [←[totl| NonDetT.wp_eq_wp]] at hbx
      have hb_triv: True ≤ ([totl| wp (f b) fun x_1 ↦
        match x_1 with
        | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
        | ForInStep.done b' => inv (ForInStep.done b')]) := by
        simp [loomLogicSimp]
        exact hbx
      have tr_intro: True ≤ ([part| NonDetT.wp (f b) inv]) := le_trans hb_triv (by
        simp [loomLogicSimp]
        intro wp_x
        simp [←[part| NonDetT.wp_eq_wp]]
        apply wp_mono_part (f b) (fun x_1 ↦
          match x_1 with
          | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
          | ForInStep.done b' => inv (ForInStep.done b')) (fun x => inv x)
        { simp [loomLogicSimp]
          intro x1
          match x1 with
          | ForInStep.yield b' => simp; intros; simp [*]
          | ForInStep.done b' => simp }
        exact wp_x)
      simp at tr_intro
      simp [tr_intro] }
    simp [LE.pure, h_inv]
    exact le_trans h_inv.right (by
      simp [←[totl| NonDetT.wp_eq_wp], ←[part| NonDetT.wp_eq_wp]]
      simp [loomLogicSimp])


theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

sumUpTo

theorem getValSpV_eq (spv: SpV Int) (j: ℕ) (h_ind: j < spv.size): spv[spv.ind[j]!] = (spv.val)[j]!

theorem getValSpV_empty (spv: SpV Int) (j: ℕ) (h_empty: ∀ i < spv.size, spv.ind[i]! ≠ j): spv[j] = 0

import Auto

import Lean

import Mathlib.Algebra.BigOperators.Intervals

import Mathlib.Algebra.Ring.Int.Defs

import Loom.MonadAlgebras.NonDetT.Extract

import Loom.MonadAlgebras.WP.Tactic

import Loom.MonadAlgebras.WP.DoNames'

import CaseStudies.Velvet.Std

section SpMV

structure SpV (valTyp : Type) where
  ind: Array Nat
  val: Array valTyp
  size: ℕ
  size_eq: ind.size = size ∧ val.size = size
  inc: ∀ (i j: Nat), i < size → j < size →  i < j → ind[i]! < ind[j]!

def sumUpTo
  (spv : SpV Int)
  (v : Array Int) (bound : ℕ) : Int := ∑ i ∈ Finset.range bound, ((spv.val)[i]! * v[(spv.ind)[i]!]!)

theorem VSpV_correct_pure (out: Array Int) (arr: Array Int)
  (spv: SpV Int)
  (h_b: ∀ i < spv.size, spv.ind[i]! < arr.size):
  out.size = 1 → out[0]! = sumUpTo spv arr spv.size →
    out[0]! = ∑ i ∈ Finset.range (arr.size), spv[i] * arr[i]! :=

:= by
      intro sz sum_eq
      rw [sum_eq, sumUpTo]
      have ind_lemma: ∀ k, k ≤ arr.size →
        ∑ i ∈ Finset.range (spv.size), (if spv.ind[i]! < k then spv.val[i]! * arr[spv.ind[i]!]! else 0) =
        ∑ i ∈ Finset.range k, spv[i] * arr[i]! := by
          intro k
          induction k with
          | zero =>
            simp
          | succ m hm =>
            intro lt_m
            simp [Finset.sum_range_succ]
            rw [←hm (by omega)]
            have splitted_sum:
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! < m then spv.val[i]! * arr[spv.ind[i]!]! else 0) +
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! = m then spv.val[i]! * arr[spv.ind[i]!]! else 0) =
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! < m + 1 then spv.val[i]! * arr[spv.ind[i]!]! else 0) := by
                rw [←Finset.sum_add_distrib]
                rw [Finset.sum_congr (by rfl)]
                intro x hx
                by_cases h_eq_m : spv.ind[x]! = m <;> simp [h_eq_m]
                have miff : spv.ind[x]! < m ↔ spv.ind[x]! < m + 1 := by
                  constructor <;> rintro h_lt <;> omega
                simp [miff]
            rw [←spv.size_eq.1, ←splitted_sum, add_left_cancel_iff.mpr]
            by_cases exists_i: ∃ i < spv.size, spv.ind[i]! = m
            { rcases exists_i with ⟨ind, h_ind⟩
              rw [← h_ind.right]
              have lemma_res := getValSpV_eq spv ind h_ind.left
              simp at lemma_res
              simp [lemma_res]
              have almost_zero : ∀ i ∈ Finset.range (spv.ind.size), i ≠ ind →
                ((if spv.ind[i]! = spv.ind[ind]! then spv.val[i]! * arr[spv.ind[i]!]! else 0) = 0) := by
                  intro i i_inb i_not_ind
                  by_cases vind_eq : spv.ind[i]! = spv.ind[ind]!
                  { simp at i_inb
                    have i_sz: i < spv.size := by
                      rw [spv.size_eq.1] at i_inb
                      exact i_inb
                    rcases lt_or_gt_of_ne i_not_ind with i_lt | inb_lt
                    { simp [getElem, Nat.lt_iff_le_and_ne.mp (spv.inc i ind i_sz h_ind.left i_lt)] }
                    have lt := Nat.lt_iff_le_and_ne.mp (spv.inc ind i h_ind.left i_sz inb_lt)
                    simp [getElem] at vind_eq
                    simp [vind_eq] at lt }
                  simp [vind_eq]
              have ind_inb: ind ∉ Finset.range (spv.size) →
                ((if spv.ind[ind]! = spv.ind[ind]! then spv.val[ind]! * arr[spv.ind[ind]!]! else 0) = 0) := by
                  intro ind_not_inb
                  simp at ind_not_inb
                  omega
              rw [←spv.size_eq.1] at ind_inb
              simp [Finset.sum_eq_single ind almost_zero ind_inb] }
            simp at exists_i
            have h_getVal := getValSpV_empty spv m exists_i
            simp at h_getVal
            simp [h_getVal]
            apply Finset.sum_eq_zero
            rintro x hx
            by_cases h_eq: spv.ind[x]! = m <;> simp [h_eq]
            simp [←h_eq] at h_getVal
            simp at hx
            rw [spv.size_eq.1] at hx
            simp [getValSpV_eq spv x hx] at h_getVal
            simp [h_getVal]
      have fin_lemma := ind_lemma (arr.size) (by rfl)
      rw [←fin_lemma]
      exact Finset.sum_congr (by rfl) fun i h_i => by aesop


macro "⌜" p:term "⌝" : term => `(LE.pure $p)

MonadNonDet

class MonadNonDet (m : Type u → Type v) where
  pick : (τ : Type u) → [Inhabited τ] → m τ

  pickSuchThat : (τ : Type u) → (p : τ → Prop) → [Findable p] → m τ
  assume : (as : Prop) → [Decidable as] → m PUnit.{u+1}

  rep {α : Type u} : α → (α → m (ForInStep α)) → m α

NonDetT

inductive NonDetT (m : Type u -> Type v) : (α : Type u) -> Type _ where
  | pure {α} (ret : α) : NonDetT m α
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) : NonDetT m α
  | pickCont {α} (τ : Type u) (p : τ -> Prop) [Findable p] (f : τ → NonDetT m α) : NonDetT m α
  | repeatCont {α} {β} (init : β) (f : β -> NonDetT m (ForInStep β)) (cont : β -> NonDetT m α) : NonDetT m α

CCPOBot

class CCPOBot (m : Type u -> Type v)  where
  compBot {α} : m α

CCPOBotLawful

class CCPOBotLawful (m : Type u -> Type v) [∀ α, Lean.Order.CCPO (m α)] [CCPOBot m] where
  prop {α} : CCPOBot.compBot (m := m) (α := α) = Lean.Order.bot

wp

def wp (c : m α) (post : α -> l) : l := liftM (n := Cont l) c post

MAlgPartial

class MAlgPartial (m : Type u -> Type v) [Monad m] [∀ α, Lean.Order.CCPO (m α)]
  [CompleteLattice l] [MAlgOrdered m l] where
  csup_lift {α : Type u} (xc : Set (m α)) (post : α -> l) :
    Lean.Order.chain xc ->
    ⨅ x ∈ xc, MAlg.lift x post <= MAlg.lift (Lean.Order.CCPO.csup xc) post

MAlg.lift

abbrev MAlg.lift {m : Type u -> Type v} {l : Type u} [Monad m] [MAlg m l] :
  {α : Type u} -> m α -> Cont l α := fun x f => μ $ f <$> x

MAlg

class MAlg [Monad m] (l : outParam (Type v)) where
  μ : m l -> l
  pure : ∀ l, μ (pure l) = l
  bind : ∀ {α : Type v} (x : m α) (f g : α -> m l),
    μ ∘ f = μ ∘ g ->
    μ (x >>= f) = μ (x >>= g)

MAlgOrdered

class MAlgOrdered (l : outParam (Type v)) [Monad m] [CompleteLattice l] where
  μ : m l -> l
  μ_ord_pure : ∀ l, μ (pure l) = l
  μ_ord_bind {α : Type v} :
    ∀ (f g : α -> m l), μ ∘ f ≤ μ ∘ g ->
      ∀ x : m α, μ (x >>= f) ≤ μ (x >>= g)

MAlgDet

class MAlgDet (l : outParam (Type v)) [Monad m] [CompleteLattice l] [MAlgOrdered m l] where

  demonic {α ι : Type v} (c : m α) (p : ι -> α -> l) [Nonempty ι] :
    ⨅ i, MAlg.lift c (p i) ≤ MAlg.lift c (fun x => ⨅ i, p i x)

  angelic {α ι : Type v} (c : m α) (p : ι -> α -> l) [Nonempty ι] :
    ⨆ i, MAlg.lift c (p i) ≥ MAlg.lift c (fun x => ⨆ i, p i x)

LogicLift

class LogicLift (l : outParam (Type u)) ( k : Type u) [CompleteLattice l] [CompleteLattice k] where
  [lift : MonadLift (Cont l) (Cont k)]
  lift_top {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊤) = ⊤
  lift_bot {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊥) = ⊥

NonDetT.wp

def NonDetT.wp {l : Type u} {α : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m α -> Cont l α
  | .pure ret => pure ret
  | .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | .pickCont τ p f => fun post => let p : Set τ := p; ⨅ a ∈ (p : Set τ), wp (f a) post

NonDetT.μ

def NonDetT.μ  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m l -> l := fun x => NonDetT.wp x id

NonDetT.bind

def NonDetT.bind (x : NonDetT m α) (f : α → NonDetT m β) : NonDetT m β :=
  match x with
  | pure ret => f ret
  | vis x f' => vis x fun y => bind (f' y) f
  | pickCont τ p f' => pickCont τ p fun t => bind (f' t) f
  | repeatCont init f' cont => repeatCont init f' fun t => bind (cont t) f

WPGen.bind

def WPGen.bind {x : m α} {f : α -> m β} (wpg : WPGen x) (wpgf : ∀ a, WPGen (f a)) :
  WPGen (x >>= f) where
  get := fun post => wpg.get (fun a => (wpgf a).get post)
  prop := by admit /- proof elided -/

private def _root_.Lean.SimpleScopedEnvExtension.get [Inhabited σ] (ext : SimpleScopedEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

Context

structure Context where
  ref         : Syntax

  m           : Syntax

  returnType  : Syntax
  mutableVars : VarSet := {}
  insideFor   : Bool := false

private def _root_.Lean.SimplePersistentEnvExtension.get [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

WPGen

structure WPGen (x : m α) where
  get : Cont l α

  prop : ∀ post, get post <= wp x post

private def _root_.Lean.EnvExtension.get [Inhabited σ] (ext : EnvExtension σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

LE.pure

noncomputable def LE.pure {l : Type u} [inst: LE l] [OrderTop l] [OrderBot l] : Prop -> l := fun p =>
  if p then ⊤ else ⊥

lemma wp_pure (x : α) (post : α -> l) : wp (m := m) (pure x) post = post x

lemma wp_cons (x : m α) (post post' : α -> l) :
  (∀ y, post y ≤ post' y) ->
  wp x post ≤ wp x post'

@[simp]
lemma NonDetT.wp_vis {β : Type u} (x : m β) (f : β → NonDetT m α) post :
  _root_.wp (NonDetT.vis x f) post = _root_.wp x fun a => _root_.wp (f a) post

lemma NonDetT.wp_eq_wp {α : Type u} (x : NonDetT m α) (post : α -> l) :
  _root_.wp x post = NonDetT.wp x post

lemma NonDetT.wp_bind  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α β : Type u} (x : NonDetT m α) (f : α -> NonDetT m β)
  (post : β -> l):
  NonDetT.wp (x.bind f) post = NonDetT.wp x (fun x => NonDetT.wp (f x) post)

@[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨅ a, ⌜p a⌝ ⇨ _root_.wp (f a) post

@[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨆ a, ⌜p a⌝ ⊓ _root_.wp (f a) post

lemma wp_bind {β} (x : m α) (f : α -> m β) (post : β -> l) :
    wp (x >>= f) post = wp x (fun x => wp (f x) post)

trueE

@[simp]
lemma trueE (l : Type v) [inst: LE l] [OrderTop l] [OrderBot l] : ⌜True⌝ = (⊤ : l)

NonDetT.extractWeak

omit [CCPOBot m] [MAlgDet m l] [LawfulMonad m] in
lemma wp_csup (xc : Set (m α)) (post : α -> l) [∀ α, CCPO (m α)] [MAlgPartial m]:
  Lean.Order.chain xc ->
  ⨅ c ∈ xc, wp c post ≤ wp (Lean.Order.CCPO.csup xc) post

omit [CCPOBot m] [MAlgDet m l] [LawfulMonad m] in
lemma wp_bot [∀ α, CCPO (m α)] [MAlgPartial m]:
  wp (bot : m α) = fun _ => (⊤ : l)

import Mathlib.Logic.Function.Basic

import Mathlib.Order.CompleteBooleanAlgebra

import Mathlib.Order.Lattice

import Mathlib.Order.Basic

import Mathlib.Data.W.Basic

import Mathlib.Data.FinEnum

import Loom.MonadAlgebras.WP.Gen

import Loom.MonadAlgebras.WP.Liberal

import Loom.MonadAlgebras.NonDetT'.Basic

open Lean.Order

def findNat (p : Nat -> Prop) [DecidablePred p] : Option Nat :=
  let rec aux i :=
    if p i then
      some i
    else
      aux (i + 1)
  partial_fixpoint
  aux 0

def find [Encodable α] (p : α -> Prop) [DecidablePred p] : Option α :=
  findNat (fun x => (Encodable.decode x).any (p ·)) |>.bind Encodable.decode

class WeakFindable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_some_p : find () = some x -> p x

instance WeakFindable.of_Findable {α : Type u} (p : α -> Prop) [Findable p] : WeakFindable p where
  find := Findable.find p
  find_some_p := Findable.find_some_p

inductive ExtractNonDet (findable : {τ : Type u} -> (τ -> Prop) -> Type u) {m} : {α : Type u} -> NonDetT m α -> Type _ where
  | pure {α} : ∀ (x : α), ExtractNonDet findable (NonDetT.pure x)
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) :
      (∀ y, ExtractNonDet findable (f y)) → ExtractNonDet findable (.vis x f)
  | pickSuchThat {α} (τ : Type u) (p : τ -> Prop) (f : τ → NonDetT m α)
    {_ : findable p}
     : (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont τ p f)
  | assume {α} (p : PUnit -> Prop) (f : PUnit → NonDetT m α) {_ : Decidable (p .unit)} :
    (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont PUnit p f)

instance ExtractNonDet.pure' : ExtractNonDet findable (Pure.pure (f := NonDetT m) x) :=

instance ExtractNonDet.liftM (x : m α) :
  ExtractNonDet findable (liftM (n := NonDetT m) x) :=

instance ExtractNonDet.assume' {p : Prop} [Decidable p] : ExtractNonDet findable (MonadNonDet.assume (m :=  NonDetT m) p) :=

instance ExtractNonDet.pickSuchThat' {τ : Type u} (p : τ -> Prop) [Findable p] :
  ExtractNonDet Findable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.pickSuchThat_weak {τ : Type u} (p : τ -> Prop) [WeakFindable p] :
  ExtractNonDet WeakFindable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.if {p : Prop} {dec : Decidable p} {x y : NonDetT m α}
  (_ : ExtractNonDet findable x) (_ : ExtractNonDet findable y) :
  ExtractNonDet findable (if p then x else y) :=

instance ExtractNonDet.ForIn_list {xs : List α} {init : β} {f : α → β → NonDetT m (ForInStep β)}
  (_ : ∀ a b, ExtractNonDet findable (f a b)) :
  ExtractNonDet findable (forIn xs init f) :=

variable [Monad m] [CCPOBot m] [CompleteBooleanAlgebra l] [MAlgOrdered m l] [MAlgDet m l] [LawfulMonad m]

@[simp, inline]
def NonDetT.extractGen {findable : {τ : Type u} -> (τ -> Prop) -> Type u} {α : Type u}
  (findOf : ∀ {τ : Type u} (p : τ -> Prop), findable p -> Unit -> Option τ)
  : (s : NonDetT m α) -> (ex : ExtractNonDet findable s := by admit /- proof elided -/
  ) -> m α
  | .pure x, _ => Pure.pure x
  | .vis x f, .vis _ _ _ => liftM x >>= (fun x => extractGen findOf (f x))
  | .pickCont _ p f, .pickSuchThat _ _ _ _ =>
    match findOf p ‹_› () with
    | none => CCPOBot.compBot
    | some x =>  extractGen findOf (f x)
  | .pickCont _ p f, .assume _ _ _ =>
    if p .unit then
      extractGen findOf (f .unit)
    else CCPOBot.compBot

def NonDetT.extractWeak {α : Type u} (s : NonDetT m α) (ex : ExtractNonDet WeakFindable s := by admit /- proof elided -/
) : m α :=
  NonDetT.extractGen WeakFindable.find s

namespace DemonicChoice

omit [MAlgDet m l] in
lemma ExtractNonDet.extract_refines_wp_weak [∀ α, CCPO (m α)] [MAlgPartial m] [CCPOBotLawful m] (s : NonDetT m α) (inst : ExtractNonDet WeakFindable s) :
  wp s post <= wp s.extractWeak post :=

:= by
  unhygienic induction inst
  { simp [wp_pure, NonDetT.extractWeak] }
  { simp only [NonDetT.wp_vis, NonDetT.extractWeak, NonDetT.extractGen, monadLift_self, wp_bind]
    apply wp_cons; aesop (add norm inf_comm) }
  { simp only [NonDetT.wp_pickCont, NonDetT.extractWeak, NonDetT.extractGen]; split
    simp only [CCPOBotLawful.prop, wp_bot, le_top]
    rename_i y h
    refine iInf_le_of_le y ?_
    have := WeakFindable.find_some_p (p := p) (by assumption)
    simp only [this, trueE, top_himp, ge_iff_le]; apply a_ih }
  simp only [NonDetT.wp_pickCont, NonDetT.extractWeak, NonDetT.extractGen]
  have : ∀ a : PUnit.{u_1 + 1}, a = .unit := by simp
  simp [this, iInf_const]; split_ifs <;> simp [*, CCPOBotLawful.prop, wp_bot]
  apply a_ih