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

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

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

theorem spv_dot_pure (spv1 spv2: SpV Int) (n: ℕ)
  (sz1: ∀ i < spv1.size, spv1.ind[i]! < n) (sz2: ∀ i < spv2.size, spv2.ind[i]! < n):
    spv_dot spv1 spv2 0 0 = ∑ i ∈ Finset.range n, spv1[i] * spv2[i] :=

:= by
      simp [spv_dot_pure_gen spv1 spv2 n 0 0 sz1 sz2]
      apply Finset.sum_congr
      { rfl }
      intro x hx
      simp at hx
      by_cases em1: spv1.size = 0 <;> simp [em1]
      { simp [lt_iff_le_not_ge.mp hx]
        simp [getValSpV_empty spv1 x (by intro i hi; omega)] }
      by_cases em2: spv2.size = 0 <;> simp [em2]
      { intros
        simp [getValSpV_empty spv2 x (by intro i hi; omega)] }
      intro zer_ineq
      have zer_lemma (spv: SpV Int) (i: ℕ) (hsz: spv.size ≠ 0): i < spv.ind[0]! → spv[i] = 0 := by
        intro lt
        have all_none: ∀ j < spv.size, spv.ind[j]! ≠ i := by
          intro i1 hi1
          by_cases i10 : i1 = 0
          { simp [←i10] at lt
            omega }
          have contra := spv.inc 0 i1 (by omega) (by omega) (by omega)
          simp at lt
          omega
        simp [getValSpV_empty spv i all_none]
      by_cases sm1: spv1.ind[0]! ≤ x
      { simp [sm1] at zer_ineq
        simp [zer_lemma spv2 x em2 zer_ineq] }
      simp at sm1
      simp [zer_lemma spv1 x em1 sm1]


lemma array_extract_split_i_j_k (arr : Array α) (i j k: Nat) :
    i < j -> j < k -> k ≤ arr.size →
    arr.extract i k = (arr.extract i j) ++ (arr.extract j k)

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

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

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

macro "extract_tactic" : tactic =>

macro "extract_tactic" : tactic =>
  `(tactic| repeat' (intros; extract_step <;> try dsimp))

syntax "loom_solve" : loom_solve_tactic

syntax "ensures" termBeforeReqEnsDo : ensures_caluse

syntax "require" termBeforeReqEnsDo : require_caluse

syntax "prove_correct" : prove_correct_command

syntax "prove_correct" ident Termination.suffix "by" tacticSeq : command

syntax "done_with" termBeforeDo : doneWith

syntax "invariant" termBeforeDo linebreak : invariantClause

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

VelvetM.extract

def VelvetM.extract {α : Type} (x : VelvetM α)
  [∀ α, Lean.Order.CCPO (DivM α)] [Lean.Order.MonoBind DivM] [Inhabited α] : α :=
  x.run.run

run

def run (code : Code) (m : Syntax) (returnType : Syntax) (uvars : Array Var := #[]) (kind := Kind.regular) : MacroM Syntax :=
  toTerm code { m, returnType, kind, uvars }

Kind

inductive Kind where
  | regular
  | forIn
  | forInWithReturn
  | nestedBC
  | nestedPR
  | nestedSBC
  | nestedPRBC

Var

abbrev Var := Syntax

Code

inductive Code where
  | decl         (xs : Array Var) (doElem : Syntax) (k : Code)
  | reassign     (xs : Array Var) (doElem : Syntax) (k : Code)

  | joinpoint    (name : Name) (params : Array (Var × Bool)) (body : Code) (k : Code)
  | seq          (action : Syntax) (k : Code)
  | action       (action : Syntax)
  | break        (ref : Syntax)
  | continue     (ref : Syntax)
  | return       (ref : Syntax) (val : Syntax)

  | ite          (ref : Syntax) (h? : Option Var) (optIdent : Syntax) (cond : Syntax) (thenBranch : Code) (elseBranch : Code)
  | match        (ref : Syntax) (gen : Syntax) (discrs : Syntax) (optMotive : Syntax) (alts : Array (Alt Code))
  | matchExpr    (ref : Syntax) («meta» : Bool) (discr : Syntax) (alts : Array (AltExpr Code)) (elseBranch : Code)
  | jmp          (ref : Syntax) (jpName : Name) (args : Array Syntax)
  deriving Inhabited

AltExpr

structure AltExpr (σ : Type) where
  ref     : Syntax
  var?    : Option Var
  funName : Syntax
  pvars   : Array Syntax
  rhs     : σ
  deriving Inhabited

DoIfView

structure DoIfView where
  ref        : Syntax
  optIdent   : Syntax
  cond       : Syntax
  thenBranch : Syntax
  elseBranch : Syntax

Alt

structure Alt (σ : Type) where
  ref      : Syntax
  vars     : Array Var
  patterns : Syntax
  rhs      : σ
  deriving Inhabited

JPDecl

structure JPDecl where
  name : Name
  params : Array (Var × Bool)
  body : Code

toTerm

partial def toTerm (c : Code) : M Syntax := do
  let term ← go c
  if let some ref := c.getRef? then
    annotate ref term
  else
    return term
where
  go (c : Code) : M Syntax := do
    match c with
    | .return ref val     => withRef ref <| returnToTerm val
    | .continue ref       => withRef ref continueToTerm
    | .break ref          => withRef ref breakToTerm
    | .action e           => actionTerminalToTerm e
    | .joinpoint j ps b k => mkJoinPoint j ps (← toTerm b) (← toTerm k)
    | .jmp ref j args     => return mkJmp ref j args
    | .decl _ stx k       => declToTerm stx (← toTerm k)
    | .reassign _ stx k   => reassignToTerm stx (← toTerm k)
    | .seq stx k          => seqToTerm stx (← toTerm k)
    | .ite ref _ o c t e  => withRef ref <| do mkIte o c (← toTerm t) (← toTerm e)
    | .match ref genParam discrs optMotive alts =>
      let mut termAlts := #[]
      for alt in alts do
        let rhs ← toTerm alt.rhs
        let termAlt := mkNode ``Parser.Term.matchAlt #[mkAtomFrom alt.ref "|", mkNullNode #[alt.patterns], mkAtomFrom alt.ref "=>", rhs]
        termAlts := termAlts.push termAlt
      let termMatchAlts := mkNode ``Parser.Term.matchAlts #[mkNullNode termAlts]
      return mkNode ``Parser.Term.«match» #[mkAtomFrom ref "match", genParam, optMotive, discrs, mkAtomFrom ref "with", termMatchAlts]
    | .matchExpr ref «meta» d alts elseBranch => withFreshMacroScope do
      let d' ← `(discr)
      let mut termAlts := #[]
      for alt in alts do
        let rhs ← `(($(← toTerm alt.rhs) : $((← read).m) _))
        let optVar := if let some var := alt.var? then mkNullNode #[var, mkAtomFrom var "@"] else mkNullNode #[]
        let pat := mkNode ``Parser.Term.matchExprPat #[optVar, alt.funName, mkNullNode alt.pvars]
        let termAlt := mkNode ``Parser.Term.matchExprAlt #[mkAtomFrom alt.ref "|", pat, mkAtomFrom alt.ref "=>", rhs]
        termAlts := termAlts.push termAlt
      let elseBranch := mkNode ``Parser.Term.matchExprElseAlt #[mkAtomFrom ref "|", mkHole ref, mkAtomFrom ref "=>", (← toTerm elseBranch)]
      let termMatchExprAlts := mkNode ``Parser.Term.matchExprAlts #[mkNullNode termAlts, elseBranch]
      let body := mkNode ``Parser.Term.matchExpr #[mkAtomFrom ref "match_expr", d', mkAtomFrom ref "with", termMatchExprAlts]
      if «meta» then
        `(Bind.bind (instantiateMVarsIfMVarApp $d) fun discr => $body)
      else
        `(let discr := $d; $body)

annotate

def annotate [Monad m] [MonadRef m] [MonadQuotation m] (ref : Syntax) (term : Syntax) : m Syntax :=
  withRef term <| `(with_annotate_term $ref $term)

Context

structure Context where
  ref         : Syntax

  m           : Syntax

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

seqToTerm

def seqToTerm (action : Syntax) (k : Syntax) : M Syntax := withRef action <| withFreshMacroScope do
  if action.getKind == ``Parser.Term.doDbgTrace then
    let msg := action[1]
    `(dbg_trace $msg; $k)
  else if action.getKind == ``Parser.Term.doAssert then
    let cond := action[1]
    `(assert! $cond; $k)
  else if action.getKind == ``Parser.Term.doDebugAssert then
    let cond := action[1]
    `(debugAssert| debug_assert! $cond; $k)
  else
    let action ← withRef action ``(($action : $((←read).m) PUnit))
    ``(Bind.bind $action (fun (_ : PUnit) => $k))

M

abbrev M := ReaderT Context MacroM

Context

structure Context where

  m          : Syntax

  returnType : Syntax
  uvars      : Array Var
  kind       : Kind

ToForInTermResult

structure ToForInTermResult where
  uvars      : Array Var
  term       : Syntax

Code.getRef?

def Code.getRef? : Code → Option Syntax
  | .decl _ doElem _     => doElem
  | .reassign _ doElem _ => doElem
  | .joinpoint ..        => none
  | .seq a _             => a
  | .action a            => a
  | .break ref           => ref
  | .continue ref        => ref
  | .return ref _        => ref
  | .ite ref ..          => ref
  | .match ref ..        => ref
  | .matchExpr ref ..    => ref
  | .jmp ref ..          => ref

mkIte

def mkIte (optIdent : Syntax) (cond : Syntax) (thenBranch : Syntax) (elseBranch : Syntax) : MacroM Syntax := do
  if optIdent.isNone then
    ``(if $cond then $thenBranch else $elseBranch)
  else
    let h := optIdent[0]
    ``(if $h:ident : $cond then $thenBranch else $elseBranch)

mkJmp

def mkJmp (ref : Syntax) (j : Name) (args : Array Syntax) : Syntax :=
  Syntax.mkApp (mkIdentFrom ref j) args

mkJoinPoint

def mkJoinPoint (j : Name) (ps : Array (Syntax × Bool)) (body : Syntax) (k : Syntax) : M Syntax := withRef body <| withFreshMacroScope do
  let pTypes ← ps.mapM fun ⟨id, useTypeOf⟩ => do if useTypeOf then `(type_of% $id) else `(_)
  let ps     := ps.map (·.1)

  `(let_delayed $(← mkIdentFromRef j):ident $[($ps : $pTypes)]* : $((← read).m) _ := $body; $k)

returnToTerm

def returnToTerm (val : Syntax) : M Syntax := do
  let ctx ← read
  let u ← mkUVarTuple
  match ctx.kind with
  | .regular         => if ctx.uvars.isEmpty then ``(Pure.pure $val) else ``(Pure.pure (MProdWithNames.mk $val $u))
  | .forIn           => ``(Pure.pure (ForInStep.done $u))
  | .forInWithReturn => ``(Pure.pure (ForInStep.done (MProdWithNames.mk (some $val) $u)))
  | .nestedBC        => unreachable!
  | .nestedPR        => ``(Pure.pure (DoResultPR.«return» $val $u))
  | .nestedSBC       => ``(Pure.pure (DoResultSBC.«pureReturn» $val $u))
  | .nestedPRBC      => ``(Pure.pure (DoResultPRBC.«return» $val $u))

mkUVarTuple

def mkUVarTuple : M Syntax := do
  let ctx ← read
  mkTuple ctx.uvars

mkTuple

private def mkTuple (elems : Array Syntax) : MacroM Syntax := do
  if elems.size = 0 then
    mkUnit
  else if h : elems.size = 1 then
    ``(WithName.mk' $(elems[0]) $(Lean.quoteNameMk elems[0].getId))
  else
    let init <- ``(WithName.mk' $(elems.back!) $(Lean.quoteNameMk elems.back!.getId))
    elems.extract 0 (elems.size - 1) |>.foldrM (init := init) fun elem tuple => do
      let name := Lean.quoteNameMk elem.getId
      ``(MProdWithNames.mk' $elem $tuple $name)

MProdWithNames.mk'

abbrev MProdWithNames.mk' {α β : Type u} (a : α) (b : β)
  (αName : Lean.Name := default) : MProdWithNames α β αName :=
  @MProdWithNames.mk _ _ αName a b

mkUnit

private def mkUnit : MacroM Syntax :=
  ``((⟨⟩ : PUnit))

WithName.mk'

abbrev WithName.mk' {α : Sort u} (a : α) (name : Lean.Name := default) : WithName α name :=
  a

continueToTerm

def continueToTerm : M Syntax := do
  let ctx ← read
  let u ← mkUVarTuple
  match ctx.kind with
  | .regular         => unreachable!
  | .forIn           => ``(Pure.pure (ForInStep.yield $u))
  | .forInWithReturn => ``(Pure.pure (ForInStep.yield (MProdWithNames.mk none $u)))
  | .nestedBC        => ``(Pure.pure (DoResultBC.«continue» $u))
  | .nestedPR        => unreachable!
  | .nestedSBC       => ``(Pure.pure (DoResultSBC.«continue» $u))
  | .nestedPRBC      => ``(Pure.pure (DoResultPRBC.«continue» $u))

declToTerm

def declToTerm (decl : Syntax) (k : Syntax) : M Syntax := withRef decl <| withFreshMacroScope do
  let kind := decl.getKind
  if kind == ``Parser.Term.doLet then
    let letDecl := decl[2]
    `(let $letDecl:letDecl; $k)
  else if kind == ``Parser.Term.doLetRec then
    let letRecToken := decl[0]
    let letRecDecls := decl[1]
    return mkNode ``Parser.Term.letrec #[letRecToken, letRecDecls, mkNullNode, k]
  else if kind == ``Parser.Term.doLetArrow then
    let arg := decl[2]
    if arg.getKind == ``Parser.Term.doIdDecl then
      let id     := arg[0]
      let type   := expandOptType id arg[1]
      let doElem := arg[3]

      match isDoExpr? doElem with
      | some action =>
        let action ← withRef action `(($action : $((← read).m) $type))
        ``(Bind.bind $action (fun ($id:ident : $type) => $k))
      | none        => Macro.throwErrorAt decl "unexpected kind of `do` declaration"
    else
      Macro.throwErrorAt decl "unexpected kind of `do` declaration"
  else if kind == ``Parser.Term.doHave then

    let args := decl.getArgs
    let args := args ++ #[mkNullNode  , k]
    return mkNode `Lean.Parser.Term.«have» args
  else
    Macro.throwErrorAt decl "unexpected kind of `do` declaration"

  partial def doLetArrowToCode (doLetArrow : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let decl    := doLetArrow[2]
    if decl.getKind == ``Parser.Term.doIdDecl then
      let y := decl[0]
      checkNotShadowingMutable #[y]
      let doElem := decl[3]
      let k ← withNewMutableVars #[y] (isMutableLet doLetArrow) (doSeqToCode doElems)
      match isDoExpr? doElem with
      | some _      => return mkVarDeclCore #[y] doLetArrow k
      | none =>
        checkLetArrowRHS doElem
        let c ← doSeqToCode [doElem]
        match doElems with
        | []       => pure c
        | kRef::_  => concat c kRef y k
    else if decl.getKind == ``Parser.Term.doPatDecl then
      let pattern := decl[0]
      let doElem  := decl[2]
      let optElse := decl[3]
      if optElse.isNone then withFreshMacroScope do
        let auxDo ← if isMutableLet doLetArrow then
          `(do let%$doLetArrow __discr ← $doElem; let%$doLetArrow mut $pattern:term := __discr)
        else
          `(do let%$doLetArrow __discr ← $doElem; let%$doLetArrow $pattern:term := __discr)
        doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems
      else
        let contSeq ← if isMutableLet doLetArrow then
          let vars ← (← getPatternVarsEx pattern).mapM fun var => `(doElem| let mut $var := $var)
          pure (vars ++ doElems.toArray)
        else
          pure doElems.toArray
        let contSeq := mkDoSeq contSeq
        let elseSeq := optElse[1]
        let auxDo ← `(do let%$doLetArrow __discr ← $doElem; match%$doLetArrow __discr with | $pattern:term => $contSeq | _ => $elseSeq)
        doSeqToCode <| getDoSeqElems (getDoSeq auxDo)
    else
      throwError "unexpected kind of `do` declaration"

checkLetArrowRHS

def checkLetArrowRHS (doElem : Syntax) : M Unit := do
  let kind := doElem.getKind
  if kind == ``Parser.Term.doLetArrow ||
     kind == ``Parser.Term.doLet ||
     kind == ``Parser.Term.doLetRec ||
     kind == ``Parser.Term.doHave ||
     kind == ``Parser.Term.doReassign ||
     kind == ``Parser.Term.doReassignArrow then
    throwErrorAt doElem "invalid kind of value `{kind}` in an assignment"

M

abbrev M := ReaderT Context TermElabM

VarSet

abbrev VarSet := Std.TreeMap Name Syntax Name.cmp

getDoSeq

private def getDoSeq (doStx : Syntax) : Syntax :=
  doStx[1]

  partial def doSeqToCode : List Syntax → M CodeBlock
    | [] => do liftMacroM mkPureUnitAction
    | doElem::doElems => withIncRecDepth <| withRef doElem do
      checkSystem "`do`-expander"
      match (← liftMacroM <| expandMacro? doElem) with
      | some doElem => doSeqToCode (doElem::doElems)
      | none =>
      match (← liftMacroM <| expandDoIf? doElem) with
      | some doElem => doSeqToCode (doElem::doElems)
      | none =>
      match (← liftMacroM <| expandDoLetExpr? doElem doElems) with
      | some doElem => doSeqToCode [doElem]
      | none =>
        let (liftedDoElems, doElem) ← expandLiftMethod doElem
        if !liftedDoElems.isEmpty then
          doSeqToCode (liftedDoElems ++ [doElem] ++ doElems)
        else
          let ref := doElem
          let k := doElem.getKind
          if k == ``Parser.Term.doLet then
            let vars ← getDoLetVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> withNewMutableVars vars (isMutableLet doElem) (doSeqToCode doElems)
          else if k == ``Parser.Term.doHave then
            let vars ← getDoHaveVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> (doSeqToCode doElems)
          else if k == ``Parser.Term.doLetRec then
            let vars ← getDoLetRecVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> (doSeqToCode doElems)
          else if k == ``Parser.Term.doReassign then
            let vars ← getDoReassignVars doElem
            checkReassignable vars
            let k ← doSeqToCode doElems
            mkReassignCore vars doElem k
          else if k == ``Parser.Term.doLetArrow then
            doLetArrowToCode doElem doElems
          else if k == ``Parser.Term.doLetElse then
            doLetElseToCode doElem doElems
          else if k == ``Parser.Term.doReassignArrow then
            doReassignArrowToCode doElem doElems
          else if k == ``Parser.Term.doIf then
            doIfToCode doElem doElems
          else if k == ``Parser.Term.doUnless then
            doUnlessToCode doElem doElems
          else if k == ``Parser.Term.doFor then withFreshMacroScope do
            doForToCode doElem doElems
          else if k == ``Parser.Term.doMatch then
            doMatchToCode doElem doElems
          else if k == ``Parser.Term.doMatchExpr then
            doMatchExprToCode doElem doElems
          else if k == ``Parser.Term.doTry then
            doTryToCode doElem doElems
          else if k == ``Parser.Term.doBreak then
            ensureInsideFor
            ensureEOS doElems
            return mkBreak ref
          else if k == ``Parser.Term.doContinue then
            ensureInsideFor
            ensureEOS doElems
            return mkContinue ref
          else if k == ``Parser.Term.doReturn then
            doReturnToCode doElem doElems
          else if k == ``Parser.Term.doDbgTrace then
            return mkSeq doElem (← doSeqToCode doElems)
          else if k == ``Parser.Term.doAssert then
            return mkSeq doElem (← doSeqToCode doElems)
          else if k == ``Parser.Term.doDebugAssert then
            return mkSeq doElem (← doSeqToCode doElems)
          else if k == ``Parser.Term.doNested then
            let nestedDoSeq := doElem[1]
            doSeqToCode (getDoSeqElems nestedDoSeq ++ doElems)
          else if k == ``Parser.Term.doExpr then
            let term := doElem[0]
            if doElems.isEmpty then
              return mkTerminalAction term
            else
              return mkSeq term (← doSeqToCode doElems)
          else
            throwError "unexpected do-element of kind {doElem.getKind}:\n{doElem}"

getDoSeqElems

private def getDoSeqElems (doSeq : Syntax) : List Syntax :=
  if doSeq.getKind == ``Parser.Term.doSeqBracketed then
    doSeq[1].getArgs.toList.map fun arg => arg[0]
  else if doSeq.getKind == ``Parser.Term.doSeqIndent then
    doSeq[0].getArgs.toList.map fun arg => arg[0]
  else
    []

  partial def doIfToCode (doIf : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let view := mkDoIfView doIf
    let thenBranch ← doSeqToCode (getDoSeqElems view.thenBranch)
    let elseBranch ← doSeqToCode (getDoSeqElems view.elseBranch)
    let ite ← mkIte view.ref view.optIdent view.cond thenBranch elseBranch
    concatWith ite doElems

CodeBlock

structure CodeBlock where
  code  : Code
  uvars : VarSet := {}

mkDoIfView

private def mkDoIfView (doIf : Syntax) : DoIfView := {
  ref        := doIf
  optIdent   := doIf[1][0]
  cond       := doIf[1][1]
  thenBranch := doIf[3]
  elseBranch := doIf[5][1]
}

  partial def concatWith (c : CodeBlock) (doElems : List Syntax) : M CodeBlock :=
    match doElems with
    | [] => pure c
    | nextDoElem :: _  => do
      let k ← doSeqToCode doElems
      let ref := nextDoElem
      concat c ref none k

concat

def concat (terminal : CodeBlock) (kRef : Syntax) (y? : Option Var) (k : CodeBlock) : TermElabM CodeBlock := do
  unless hasTerminalAction terminal.code do
    throwErrorAt kRef "`do` element is unreachable"
  let (terminal, k) ← homogenize terminal k
  let xs := varSetToArray k.uvars
  let y ← match y? with | some y => pure y | none => `(y)
  let ps := xs.map fun x => (x, true)
  let ps := ps.push (y, false)
  let jpDecl ← mkFreshJP ps k.code
  let jp := jpDecl.name
  let terminal ← liftMacroM <| convertTerminalActionIntoJmp terminal.code jp xs
  return { code  := attachJP jpDecl terminal, uvars := k.uvars }

homogenize

def homogenize (c₁ c₂ : CodeBlock) : TermElabM (CodeBlock × CodeBlock) := do
  let ws := union c₁.uvars c₂.uvars
  let c₁ ← extendUpdatedVars c₁ ws
  let c₂ ← extendUpdatedVars c₂ ws
  pure (c₁, c₂)

extendUpdatedVars

partial def extendUpdatedVars (c : CodeBlock) (ws : VarSet) : TermElabM CodeBlock := do
  if ws.any fun x _ => !c.uvars.contains x then

    pure { code := (← extendUpdatedVarsAux c.code ws), uvars := ws }
  else
    pure { c with uvars := ws }

extendUpdatedVarsAux

partial def extendUpdatedVarsAux (c : Code) (ws : VarSet) : TermElabM Code :=
  let rec update (c : Code) : TermElabM Code := do
    match c with
    | .joinpoint j ps b k    => return .joinpoint j ps (← update b) (← update k)
    | .seq e k               => return .seq e (← update k)
    | .match ref g ds t alts =>
      if alts.any fun alt => alt.vars.any fun x => ws.contains x.getId then

        pullExitPoints c
      else
        return .match ref g ds t (← alts.mapM fun alt => do pure { alt with rhs := (← update alt.rhs) })
    | .matchExpr ref «meta» d alts e =>
      if alts.any fun alt => alt.vars.any fun x => ws.contains x.getId then

        pullExitPoints c
      else
        let alts ← alts.mapM fun alt => do pure { alt with rhs := (← update alt.rhs) }
        let e ← update e
        return .matchExpr ref «meta» d alts e
    | .ite ref none o c t e => return .ite ref none o c (← update t) (← update e)
    | .ite ref (some h) o cond t e =>
      if ws.contains h.getId then

        pullExitPoints c
      else
        return Code.ite ref (some h) o cond (← update t) (← update e)
    | .reassign xs stx k => return .reassign xs stx (← update k)
    | .decl xs stx k => do
      if xs.any fun x => ws.contains x.getId then

        pullExitPoints c
      else
        return .decl xs stx (← update k)
    | c => return  c
  update c

AltExpr.vars

def AltExpr.vars (alt : AltExpr σ) : Array Var := Id.run do
  let mut vars := #[]
  if let some var := alt.var? then
    vars := vars.push var
  for pvar in alt.pvars do
    match pvar with
    | `(_) => pure ()
    | _ => vars := vars.push pvar
  return vars

pullExitPoints

def pullExitPoints (c : Code) : TermElabM Code := do
  if hasExitPoint c then
    let (c, jpDecls) ← (pullExitPointsAux {} c).run #[]
    return attachJPs jpDecls c
  else
    return c

hasExitPoint

def hasExitPoint (c : Code) : Bool :=
  hasExitPointPred c fun _ => true

hasExitPointPred

partial def hasExitPointPred (c : Code) (p : Code → Bool) : Bool :=
  let rec loop : Code → Bool
    | .decl _ _ k             => loop k
    | .reassign _ _ k         => loop k
    | .joinpoint _ _ b k      => loop b || loop k
    | .seq _ k                => loop k
    | .ite _ _ _ _ t e        => loop t || loop e
    | .match _ _ _ _ alts     => alts.any (loop ·.rhs)
    | .matchExpr _ _ _ alts e => alts.any (loop ·.rhs) || loop e
    | .jmp ..                 => false
    | c                       => p c
  loop c

pullExitPointsAux

partial def pullExitPointsAux (rs : VarSet) (c : Code) : StateRefT (Array JPDecl) TermElabM Code := do
  match c with
  | .decl xs stx k         => return .decl xs stx (← pullExitPointsAux (eraseVars rs xs) k)
  | .reassign xs stx k     => return .reassign xs stx (← pullExitPointsAux (insertVars rs xs) k)
  | .joinpoint j ps b k    => return .joinpoint j ps (← pullExitPointsAux rs b) (← pullExitPointsAux rs k)
  | .seq e k               => return .seq e (← pullExitPointsAux rs k)
  | .ite ref x? o c t e    => return .ite ref x? o c (← pullExitPointsAux (eraseOptVar rs x?) t) (← pullExitPointsAux (eraseOptVar rs x?) e)
  | .jmp ..                => return  c
  | .break ref             => mkSimpleJmp ref rs (.break ref)
  | .continue ref          => mkSimpleJmp ref rs (.continue ref)
  | .return ref val        => mkJmp ref rs val (fun y => return .return ref y)
  | .action e              =>

    mkAuxDeclFor e fun y =>
      let ref := e
      mkJmp ref rs y (fun yFresh => return .action (← ``(Pure.pure $yFresh)))
  | .match ref g ds t alts =>
    let alts ← alts.mapM fun alt => do pure { alt with rhs := (← pullExitPointsAux (eraseVars rs alt.vars) alt.rhs) }
    return .match ref g ds t alts
  | .matchExpr ref «meta» d alts e =>
    let alts ← alts.mapM fun alt => do pure { alt with rhs := (← pullExitPointsAux (eraseVars rs alt.vars) alt.rhs) }
    let e ← pullExitPointsAux rs e
    return .matchExpr ref «meta» d alts e

mkJmp

def mkJmp (ref : Syntax) (rs : VarSet) (val : Syntax) (mkJPBody : Syntax → MacroM Code) : StateRefT (Array JPDecl) TermElabM Code := do
  let xs := varSetToArray rs
  let args := xs.push val
  let yFresh ← withRef ref `(y)
  let ps := xs.map fun x => (x, true)
  let ps := ps.push (yFresh, false)
  let jpBody ← liftMacroM <| mkJPBody yFresh
  let jp ← addFreshJP ps jpBody
  return Code.jmp ref jp args

varSetToArray

private def varSetToArray (s : VarSet) : Array Var :=
  s.foldl (fun xs _ x => xs.push x) #[]

addFreshJP

def addFreshJP (ps : Array (Var × Bool)) (body : Code) : StateRefT (Array JPDecl) TermElabM Name := do
  let jp ← mkFreshJP ps body
  modify fun (jps : Array JPDecl) => jps.push jp
  pure jp.name

mkFreshJP

def mkFreshJP (ps : Array (Var × Bool)) (body : Code) : TermElabM JPDecl := do
  let ps ← if ps.isEmpty then
    let y ← `(y)
    pure #[(y.raw, false)]
  else
    pure ps

  let name ← mkFreshUserName `__do_jp
  pure { name := name, params := ps, body := body }

mkSimpleJmp

def mkSimpleJmp (ref : Syntax) (rs : VarSet) (c : Code) : StateRefT (Array JPDecl) TermElabM Code := do
  let xs := varSetToArray rs
  let jp ← addFreshJP (xs.map fun x => (x, true)) c
  if xs.isEmpty then
    let unit ← ``(Unit.unit)
    return Code.jmp ref jp #[unit]
  else
    return Code.jmp ref jp xs

eraseOptVar

def eraseOptVar (rs : VarSet) (x? : Option Var) : VarSet :=
  match x? with
  | none   => rs
  | some x => rs.insert x.getId x

mkAuxDeclFor

def mkAuxDeclFor {m} [Monad m] [MonadQuotation m] (e : Syntax) (mkCont : Syntax → m Code) : m Code := withRef e <| withFreshMacroScope do
  let y ← `(y)
  let doElem ← `(doElem| let y ← $e:term)

  let y ← `(ensure_expected_type% "type mismatch, result value" $y)
  let k ← mkCont y
  return .decl #[y] doElem k

eraseVars

def eraseVars (rs : VarSet) (xs : Array Var) : VarSet :=
  xs.foldl (·.erase ·.getId) rs

insertVars

def insertVars (rs : VarSet) (xs : Array Var) : VarSet :=
  xs.foldl (fun rs x => rs.insert x.getId x) rs

attachJPs

def attachJPs (jpDecls : Array JPDecl) (k : Code) : Code :=
  jpDecls.foldr attachJP k

attachJP

def attachJP (jpDecl : JPDecl) (k : Code) : Code :=
  Code.joinpoint jpDecl.name jpDecl.params jpDecl.body k

union

private def union (s₁ s₂ : VarSet) : VarSet :=
  s₁.foldl (·.insert ·) s₂

convertTerminalActionIntoJmp

partial def convertTerminalActionIntoJmp (code : Code) (jp : Name) (xs : Array Var) : MacroM Code :=
  let rec loop : Code → MacroM Code
    | .decl xs stx k         => return .decl xs stx (← loop k)
    | .reassign xs stx k     => return .reassign xs stx (← loop k)
    | .joinpoint n ps b k    => return .joinpoint n ps (← loop b) (← loop k)
    | .seq e k               => return .seq e (← loop k)
    | .ite ref x? h c t e    => return .ite ref x? h c (← loop t) (← loop e)
    | .action e              => mkAuxDeclFor e fun y =>
      let ref := e

      let jmpArgs := xs.push y
      return Code.jmp ref jp jmpArgs
    | .match ref g ds t alts =>
      return .match ref g ds t (← alts.mapM fun alt => do pure { alt with rhs := (← loop alt.rhs) })
    | .matchExpr ref «meta» d alts e => do
      let alts ← alts.mapM fun alt => do pure { alt with rhs := (← loop alt.rhs) }
      let e ← loop e
      return .matchExpr ref «meta» d alts e
    | c => return c
  loop code

hasTerminalAction

def hasTerminalAction (c : Code) : Bool :=
  hasExitPointPred c fun
    | .action _ => true
    | _ => false

mkIte

def mkIte (ref : Syntax) (optIdent : Syntax) (cond : Syntax) (thenBranch : CodeBlock) (elseBranch : CodeBlock) : TermElabM CodeBlock := do
  let x? := optIdent.getOptional?
  let (thenBranch, elseBranch) ← homogenize thenBranch elseBranch
  return {
    code  := .ite ref x? optIdent cond thenBranch.code elseBranch.code,
    uvars := thenBranch.uvars,
  }

ensureInsideFor

def ensureInsideFor : M Unit :=
  unless (← read).insideFor do
    throwError "invalid `do` element, it must be inside `for`"

  partial def doForToCode (doFor : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let doForDecls := doFor[1].getSepArgs
    if h : doForDecls.size > 1 then

      let doForDecl := doForDecls[1]!
      unless doForDecl[0].isNone do
        throwErrorAt doForDecl[0] "the proof annotation here has not been implemented yet"
      let y  := doForDecl[1]
      let ys := doForDecl[3]
      let doForDecls := doForDecls.eraseIdx 1
      let body := doFor[3]
      withFreshMacroScope do

        let toStreamApp ← withRef ys `(@toStream _ _ _ $ys)
        let auxDo ←
          `(do let mut s := $toStreamApp:term
               for $doForDecls:doForDecl,* do
                 match @Stream.next? _ _ _ s with
                 | none => break
                 | some ($y, s') =>
                   s := s'
                   do $body)
        doSeqToCode (getDoSeqElems (getDoSeq auxDo) ++ doElems)
    else withRef doFor do
      let h?        := if doForDecls[0]![0].isNone then none else some doForDecls[0]![0][0]
      let x         := doForDecls[0]![1]
      withRef x <| checkNotShadowingMutable (← getPatternVarsEx x)
      let xs        := doForDecls[0]![3]
      let forElems  := getDoSeqElems doFor[3]
      let forInBodyCodeBlock ← withFor (doSeqToCode forElems)
      let ⟨uvars, forInBody⟩ ← mkForInBody x forInBodyCodeBlock
      let ctx ← read

      let uvars := uvars.map fun v => ctx.mutableVars.getD v.getId v
      let uvarsTuple ← liftMacroM do mkTuple uvars
      let uvarsTuplePat ← liftMacroM do mkTuplePat uvars

      if hasReturn forInBodyCodeBlock.code then
        let forInBody ← liftMacroM <| destructTuple uvars (← `(r)) forInBody
        let optType ← `(Option $((← read).returnType))
        let forInTerm ← if let some h := h? then
          annotate doFor
            (← `(for_in'% $(xs) (MProdWithNames.mk (none : $optType) $uvarsTuple) fun $x $h (r : MProdWithNames $optType _) => let r := r.2; $forInBody))
        else
          annotate doFor
            (← `(for_in% $(xs) (MProdWithNames.mk (none : $optType) $uvarsTuple) fun $x (r : MProdWithNames $optType _) => let r := r.2; $forInBody))
        let auxDo ← `(do let r ← $forInTerm:term;
                         $uvarsTuplePat:term := r.2;
                         match r.1 with
                         | none => Pure.pure (ensure_expected_type% "type mismatch, `for`" PUnit.unit)
                         | some a => return ensure_expected_type% "type mismatch, `for`" a)
        doSeqToCode (getDoSeqElems (getDoSeq auxDo) ++ doElems)
      else
        let forInBody ← liftMacroM <| destructTuple uvars (← `(r)) forInBody
        let forInTerm ← if let some h := h? then
          annotate doFor (← `(for_in'% $(xs) $uvarsTuple fun $x $h r => $forInBody))
        else
          annotate doFor (← `(for_in% $(xs) $uvarsTuple fun $x r => $forInBody))
        if doElems.isEmpty then
          let auxDo ← `(do let r ← $forInTerm:term;
                           $uvarsTuplePat:term := r;
                           Pure.pure (ensure_expected_type% "type mismatch, `for`" PUnit.unit))
          doSeqToCode <| getDoSeqElems (getDoSeq auxDo)
        else
          let auxDo ← `(do let r ← $forInTerm:term; $uvarsTuplePat:term := r)
          doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems

destructTuple

private def destructTuple (uvars : Array Var) (x : Syntax) (body : Syntax) : MacroM Syntax := do
  if uvars.size = 0 then
    return body
  else if h : uvars.size = 1 then
    `(let $(uvars[0]):ident := WithName.erase $x; $body)
  else
    destruct uvars.toList x body
where
  destruct (as : List Var) (x : Syntax) (body : Syntax) : MacroM Syntax := do
    match as with
      | [a, b]  => `(let $a:ident := $x.1; let $b:ident := WithName.erase $x.2; $body)
      | a :: as => withFreshMacroScope do
        let rest ← destruct as (← `(x)) body
        `(let $a:ident := $x.1; let x := $x.2; $rest)
      | _ => unreachable!

WithName.erase

abbrev WithName.erase {α : Type u} {name} (a : WithName α name) : α := a

withFor

def withFor {α} (x : M α) : M α :=
  withReader (fun ctx => { ctx with insideFor := true }) x

checkNotShadowingMutable

def checkNotShadowingMutable (xs : Array Var) : M Unit := do
  let throwInvalidShadowing (x : Name) : M Unit :=
    throwError "mutable variable `{x.simpMacroScopes}` cannot be shadowed"
  let ctx ← read
  for x in xs do
    if ctx.mutableVars.contains x.getId then
      withRef x <| throwInvalidShadowing x.getId

hasReturn

def hasReturn (c : Code) : Bool :=
  hasExitPointPred c fun
    | .return .. => true
    | _ => false

getPatternVarsEx

def getPatternVarsEx (pattern : Syntax) : TermElabM (Array Var) :=
  getPatternVars pattern <|>
  Quotation.getPatternVars pattern

mkTuplePat

private def mkTuplePat (elems : Array Syntax) : MacroM Syntax := do
  if elems.size = 0 then
    mkUnit
  else if h : elems.size = 1 then
    return elems[0]
  else
    elems.extract 0 (elems.size - 1) |>.foldrM (init := elems.back!) fun elem tuple => do
        ``(MProdWithNames.mk $elem $tuple)

mkForInBody

def mkForInBody  (_ : Syntax) (forInBody : CodeBlock) : M ToForInTermResult := do
  let ctx ← read
  let uvars := forInBody.uvars
  let uvars := varSetToArray uvars
  let term ← liftMacroM <| ToTerm.run forInBody.code ctx.m ctx.returnType uvars (if hasReturn forInBody.code then ToTerm.Kind.forInWithReturn else ToTerm.Kind.forIn)
  return ⟨uvars, term⟩

mkBreak

def mkBreak (ref : Syntax) : CodeBlock :=
  { code := .break ref }

  partial def doUnlessToCode (doUnless : Syntax) (doElems : List Syntax) : M CodeBlock := withRef doUnless do
    let cond  := doUnless[1]
    let doSeq := doUnless[3]
    let body ← doSeqToCode (getDoSeqElems doSeq)
    let unlessCode ← liftMacroM <| mkUnless cond body
    concatWith unlessCode doElems

mkUnless

def mkUnless (cond : Syntax) (c : CodeBlock) : MacroM CodeBlock := do
  let thenBranch ← mkPureUnitAction
  return { c with code := .ite (← getRef) none mkNullNode cond thenBranch.code c.code }

mkPureUnitAction

def mkPureUnitAction : MacroM CodeBlock := do
  return mkTerminalAction (← mkPureUnit)

mkPureUnit

private def mkPureUnit : MacroM Syntax :=
  ``(pure PUnit.unit)

mkTerminalAction

def mkTerminalAction (action : Syntax) : CodeBlock :=
  { code := .action action }

  partial def doTryToCode (doTry : Syntax) (doElems: List Syntax) : M CodeBlock := do
    let tryCode ← doSeqToCode (getDoSeqElems doTry[1])
    let optFinally := doTry[3]
    let catches ← doTry[2].getArgs.mapM fun catchStx : Syntax => do
      if catchStx.getKind == ``Parser.Term.doCatch then
        let x       := catchStx[1]
        if x.isIdent then
          withRef x <| checkNotShadowingMutable #[x]
        let optType := catchStx[2]
        let c ← doSeqToCode (getDoSeqElems catchStx[4])
        return { x := x, optType := optType, codeBlock := c : Catch }
      else if catchStx.getKind == ``Parser.Term.doCatchMatch then
        let matchAlts := catchStx[1]
        let x ← `(ex)
        let auxDo ← `(do match ex with $matchAlts)
        let c ← doSeqToCode (getDoSeqElems (getDoSeq auxDo))
        return { x := x, codeBlock := c, optType := mkNullNode : Catch }
      else
        throwError "unexpected kind of `catch`"
    let finallyCode? ← if optFinally.isNone then pure none else some <$> doSeqToCode (getDoSeqElems optFinally[0][1])
    if catches.isEmpty && finallyCode?.isNone then
      throwError "invalid `try`, it must have a `catch` or `finally`"
    let ctx ← read
    let ws    := getTryCatchUpdatedVars tryCode catches finallyCode?
    let uvars := varSetToArray ws
    let a     := tryCatchPred tryCode catches finallyCode? hasTerminalAction
    let r     := tryCatchPred tryCode catches finallyCode? hasReturn
    let bc    := tryCatchPred tryCode catches finallyCode? hasBreakContinue
    let toTerm (codeBlock : CodeBlock) : M Syntax := do
      let codeBlock ← liftM $ extendUpdatedVars codeBlock ws
      liftMacroM <| ToTerm.mkNestedTerm codeBlock.code ctx.m ctx.returnType uvars a r bc
    let term ← toTerm tryCode
    let term ← catches.foldlM (init := term) fun term «catch» => do
      let catchTerm ← toTerm «catch».codeBlock
      if catch.optType.isNone then
        annotate doTry (← ``(MonadExcept.tryCatch $term (fun $(«catch».x):ident => $catchTerm)))
      else
        let type := «catch».optType[1]
        annotate doTry (← ``(tryCatchThe $type $term (fun $(«catch».x):ident => $catchTerm)))
    let term ← match finallyCode? with
      | none             => pure term
      | some finallyCode => withRef optFinally do
        unless finallyCode.uvars.isEmpty do
          throwError "`finally` currently does not support reassignments"
        if hasBreakContinueReturn finallyCode.code then
          throwError "`finally` currently does `return`, `break`, nor `continue`"
        let finallyTerm ← liftMacroM <| ToTerm.run finallyCode.code ctx.m ctx.returnType {} ToTerm.Kind.regular
        annotate doTry (← ``(tryFinally $term $finallyTerm))
    let doElemsNew ← liftMacroM <| ToTerm.matchNestedTermResult term uvars a r bc
    doSeqToCode (doElemsNew ++ doElems)

tryCatchPred

def tryCatchPred (tryCode : CodeBlock) (catches : Array Catch) (finallyCode? : Option CodeBlock) (p : Code → Bool) : Bool :=
  p tryCode.code ||
  catches.any (fun «catch» => p «catch».codeBlock.code) ||
  match finallyCode? with
  | none => false
  | some finallyCode => p finallyCode.code

Catch

structure Catch where
  x         : Syntax
  optType   : Syntax
  codeBlock : CodeBlock

mkNestedTerm

def mkNestedTerm (code : Code) (m : Syntax) (returnType : Syntax) (uvars : Array Var) (a r bc : Bool) : MacroM Syntax := do
  ToTerm.run code m returnType uvars (mkNestedKind a r bc)

mkNestedKind

def mkNestedKind (a r bc : Bool) : Kind :=
  match a, r, bc with
  | true,  false, false => .regular
  | false, true,  false => .regular
  | false, false, true  => .nestedBC
  | true,  true,  false => .nestedPR
  | true,  false, true  => .nestedSBC
  | false, true,  true  => .nestedSBC
  | true,  true,  true  => .nestedPRBC
  | false, false, false => unreachable!

getTryCatchUpdatedVars

def getTryCatchUpdatedVars (tryCode : CodeBlock) (catches : Array Catch) (finallyCode? : Option CodeBlock) : VarSet :=
  let ws := tryCode.uvars
  let ws := catches.foldl (init := ws) fun ws alt => union alt.codeBlock.uvars ws
  let ws := match finallyCode? with
    | none   => ws
    | some c => union c.uvars ws
  ws

hasBreakContinueReturn

def hasBreakContinueReturn (c : Code) : Bool :=
  hasExitPointPred c fun
    | .break _    => true
    | .continue _ => true
    | .return _ _ => true
    | _ => false

matchNestedTermResult

def matchNestedTermResult (term : Syntax) (uvars : Array Var) (a r bc : Bool) : MacroM (List Syntax) := do
  let toDoElems (auxDo : Syntax) : List Syntax := getDoSeqElems (getDoSeq auxDo)
  let u ← mkTuplePat uvars
  match a, r, bc with
  | true, false, false =>
    if uvars.isEmpty then
      return toDoElems (← `(do $term:term))
    else
      return toDoElems (← `(do let r ← $term:term; $u:term := r.2; pure r.1))
  | false, true, false =>
    if uvars.isEmpty then
      return toDoElems (← `(do let r ← $term:term; return r))
    else
      return toDoElems (← `(do let r ← $term:term; $u:term := r.2; return r.1))
  | false, false, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | .break u => $u:term := u; break
         | .continue u => $u:term := u; continue)
  | true, true, false => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | .pure a u => $u:term := u; pure a
         | .return b u => $u:term := u; return b)
  | true, false, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | .pureReturn a u => $u:term := u; pure a
         | .break u => $u:term := u; break
         | .continue u => $u:term := u; continue)
  | false, true, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | .pureReturn a u => $u:term := u; return a
         | .break u => $u:term := u; break
         | .continue u => $u:term := u; continue)
  | true, true, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | .pure a u => $u:term := u; pure a
         | .return a u => $u:term := u; return a
         | .break u => $u:term := u; break
         | .continue u => $u:term := u; continue)
  | false, false, false => unreachable!

hasBreakContinue

def hasBreakContinue (c : Code) : Bool :=
  hasExitPointPred c fun
    | .break _    => true
    | .continue _ => true
    | _ => false

mkContinue

def mkContinue (ref : Syntax) : CodeBlock :=
  { code := .continue ref }

  partial def doLetElseToCode (doLetElse : Syntax) (doElems : List Syntax) : M CodeBlock := do

    let pattern := doLetElse[2]
    let val     := doLetElse[4]
    let elseSeq := doLetElse[6]
    let contSeq ← if isMutableLet doLetElse then
      let vars ← (← getPatternVarsEx pattern).mapM fun var => `(doElem| let mut $var := $var)
      pure (vars ++ doElems.toArray)
    else
      pure doElems.toArray
    let contSeq := mkDoSeq contSeq
    let auxDo ← `(do match $val:term with | $pattern:term => $contSeq | _ => $elseSeq)
    doSeqToCode <| getDoSeqElems (getDoSeq auxDo)

isMutableLet

def isMutableLet (doElem : Syntax) : Bool :=
  let kind := doElem.getKind
  (kind == ``doLetArrow || kind == ``doLet || kind == ``doLetElse)
  &&
  !doElem[1].isNone

mkDoSeq

def mkDoSeq (doElems : Array Syntax) : Syntax :=
  mkNode `Lean.Parser.Term.doSeqIndent #[mkNullNode <| doElems.map fun doElem => mkNullNode #[doElem, mkNullNode]]

mkSeq

def mkSeq (action : Syntax) (c : CodeBlock) : CodeBlock :=
  { c with code := .seq action c.code }

withNewMutableVars

def withNewMutableVars {α} (newVars : Array Var) (mutable : Bool) (x : M α) : M α :=
  withReader (fun ctx => if mutable then { ctx with mutableVars := insertVars ctx.mutableVars newVars } else ctx) x

mkReassignCore

def mkReassignCore (xs : Array Var) (stx : Syntax) (c : CodeBlock) : TermElabM CodeBlock := do
  let us := c.uvars
  let ws := insertVars us xs

  let code ← if xs.any fun x => !us.contains x.getId then extendUpdatedVarsAux c.code ws else pure c.code
  pure { code := .reassign xs stx code, uvars := ws }

getDoReassignVars

def getDoReassignVars (doReassign : Syntax) : TermElabM (Array Var) := do
  let arg := doReassign[0]
  if arg.getKind == ``Parser.Term.letIdDecl then
    return getLetIdDeclVars arg
  else if arg.getKind == ``Parser.Term.letPatDecl then
    getLetPatDeclVars arg
  else
    throwError "unexpected kind of reassignment"

getLetIdDeclVars

def getLetIdDeclVars (letIdDecl : Syntax) : Array Var :=
  assert! letIdDecl.isOfKind ``Parser.Term.letIdDecl

  getLetIdVars letIdDecl[0]

getLetIdVars

def getLetIdVars (letId : Syntax) : Array Var :=
  assert! letId.isOfKind ``Parser.Term.letId

  if letId[0].isIdent then
    #[letId[0]]
  else if letId[0].isOfKind hygieneInfoKind then
    #[HygieneInfo.mkIdent letId[0] `this (canonical := true)]
  else
    #[]

getLetPatDeclVars

def getLetPatDeclVars (letPatDecl : Syntax) : TermElabM (Array Var) := do

  let pattern := letPatDecl[0]
  getPatternVarsEx pattern

ensureEOS

def ensureEOS (doElems : List Syntax) : M Unit :=
  unless doElems.isEmpty do
    throwError "must be last element in a `do` sequence"

mkVarDeclCore

def mkVarDeclCore (xs : Array Var) (stx : Syntax) (c : CodeBlock) : CodeBlock := {
  code := Code.decl xs stx c.code,
  uvars := eraseVars c.uvars xs
}

  partial def doMatchToCode (doMatch : Syntax) (doElems: List Syntax) : M CodeBlock := do
    let ref       := doMatch
    let genParam  := doMatch[1]
    let optMotive := doMatch[2]
    let discrs    := doMatch[3]
    let matchAlts := doMatch[5][0].getArgs
    let matchAlts ← matchAlts.foldlM (init := #[]) fun result matchAlt => return result ++ (← liftMacroM <| expandMatchAlt matchAlt)
    let alts ←  matchAlts.mapM fun matchAlt => do
      let patterns := matchAlt[1][0]
      let vars ← getPatternsVarsEx patterns.getSepArgs
      withRef patterns <| checkNotShadowingMutable vars
      let rhs  := matchAlt[3]
      let rhs ← doSeqToCode (getDoSeqElems rhs)
      pure { ref := matchAlt, vars := vars, patterns := patterns, rhs := rhs : Alt CodeBlock }
    let matchCode ← mkMatch ref genParam discrs optMotive alts
    concatWith matchCode doElems

mkMatch

def mkMatch (ref : Syntax) (genParam : Syntax) (discrs : Syntax) (optMotive : Syntax) (alts : Array (Alt CodeBlock)) : TermElabM CodeBlock := do

  let ws := alts.foldl (union · ·.rhs.uvars) {}
  let alts ← alts.mapM fun alt => do
    let rhs ← extendUpdatedVars alt.rhs ws
    return { ref := alt.ref, vars := alt.vars, patterns := alt.patterns, rhs := rhs.code : Alt Code }
  return { code := .match ref genParam discrs optMotive alts, uvars := ws }

getPatternsVarsEx

def getPatternsVarsEx (patterns : Array Syntax) : TermElabM (Array Var) :=
  getPatternsVars patterns <|>
  Quotation.getPatternsVars patterns

doReturnToCode

def doReturnToCode (doReturn : Syntax) (doElems: List Syntax) : M CodeBlock := withRef doReturn do
  ensureEOS doElems
  let argOpt := doReturn[1]
  let arg ← if argOpt.isNone then liftMacroM mkUnit else pure argOpt[0]
  return mkReturn (← getRef) arg

mkReturn

def mkReturn (ref : Syntax) (val : Syntax) : CodeBlock :=
  { code := .return ref val }

getDoHaveVars

def getDoHaveVars (doHave : Syntax) : TermElabM (Array Var) :=

  getLetDeclVars doHave[1]

getLetDeclVars

def getLetDeclVars (letDecl : Syntax) : TermElabM (Array Var) := do

  let arg := letDecl[0]
  if arg.getKind == ``Parser.Term.letIdDecl then
    return getLetIdDeclVars arg
  else if arg.getKind == ``Parser.Term.letPatDecl then
    getLetPatDeclVars arg
  else if arg.getKind == ``Parser.Term.letEqnsDecl then
    return getLetEqnsDeclVars arg
  else
    throwError "unexpected kind of let declaration"

getLetEqnsDeclVars

def getLetEqnsDeclVars (letEqnsDecl : Syntax) : Array Var :=
  assert! letEqnsDecl.isOfKind ``Parser.Term.letEqnsDecl

  getLetIdVars letEqnsDecl[0]

getDoLetRecVars

def getDoLetRecVars (doLetRec : Syntax) : TermElabM (Array Var) := do

  let letRecDecls := doLetRec[1][0].getSepArgs
  let letDecls := letRecDecls.map fun p => p[2]
  let mut allVars := #[]
  for letDecl in letDecls do
    let vars ← getLetDeclVars letDecl
    allVars := allVars ++ vars
  return allVars

  partial def doReassignArrowToCode (doReassignArrow : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let decl := doReassignArrow[0]
    if decl.getKind == ``Parser.Term.doIdDecl then
      let doElem := decl[3]
      let y      := decl[0]
      let auxDo ← `(do let r ← $doElem; $y:ident := r)
      doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems
    else if decl.getKind == ``Parser.Term.doPatDecl then
      let pattern := decl[0]
      let doElem  := decl[2]
      let optElse := decl[3]
      if optElse.isNone then withFreshMacroScope do
        let auxDo ← `(do let __discr ← $doElem; $pattern:term := __discr)
        doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems
      else
        throwError "reassignment with `|` (i.e., \"else clause\") is not currently supported"
    else
      throwError "unexpected kind of `do` reassignment"

expandLiftMethod

def expandLiftMethod (doElem : Syntax) : M (List Syntax × Syntax) := do
  if !hasLiftMethod doElem then
    return ([], doElem)
  else
    let baseId ← withFreshMacroScope (MonadQuotation.addMacroScope `__do_lift)
    let (doElem, doElemsNew) ← (expandLiftMethodAux baseId false false doElem).run []
    return (doElemsNew, doElem)

hasLiftMethod

private partial def hasLiftMethod : Syntax → Bool
  | Syntax.node _ k args =>
    if liftMethodDelimiter k then false

    else if k == ``Parser.Term.liftMethod then true

    else if k == ``termDepIfThenElse || k == ``termIfThenElse then args.size >= 2 && hasLiftMethod args[1]!
    else args.any hasLiftMethod
  | _ => false

liftMethodDelimiter

private def liftMethodDelimiter (k : SyntaxNodeKind) : Bool :=
  k == ``Parser.Term.do ||
  k == ``Parser.Term.doSeqIndent ||
  k == ``Parser.Term.doSeqBracketed ||
  k == ``Parser.Term.termReturn ||
  k == ``Parser.Term.termUnless ||
  k == ``Parser.Term.termTry ||
  k == ``Parser.Term.termFor

expandLiftMethodAux

private partial def expandLiftMethodAux (inQuot : Bool) (inBinder : Bool) : Syntax → StateT (List Syntax) M Syntax
  | stx@(Syntax.node i k args) =>
    if k == choiceKind then do

      let alts ← stx.getArgs.mapM (expandLiftMethodAux inQuot inBinder · |>.run [])
      let (_, lifts) := alts[0]!
      unless alts.all (·.2 == lifts) do
        throwErrorAt stx "cannot lift `(<- ...)` over inconsistent syntax variants, consider lifting out the binding manually"
      modify (· ++ lifts)
      return .node i k (alts.map (·.1))
    else if liftMethodDelimiter k then
      return stx

    else if h : args.size >= 2 ∧ (k == ``termDepIfThenElse || k == ``termIfThenElse) then do
      let inAntiquot := stx.isAntiquot && !stx.isEscapedAntiquot
      let arg1 ← expandLiftMethodAux (inQuot && !inAntiquot || stx.isQuot) inBinder args[1]
      let args := args.set! 1 arg1
      return Syntax.node i k args
    else if k == ``Parser.Term.liftMethod && !inQuot then withFreshMacroScope do
      if inBinder then
        throwErrorAt stx "cannot lift `(<- ...)` over a binder, this error usually happens when you are trying to lift a method nested in a `fun`, `let`, or `match`-alternative, and it can often be fixed by adding a missing `do`"
      let term := args[1]!
      let term ← expandLiftMethodAux inQuot inBinder term

      let id ← mkIdentFromRef (.num baseId (← get).length)
      let auxDoElem : Syntax ← `(doElem| let $id:ident ← $term:term)
      modify fun s => s ++ [auxDoElem]
      return id
    else do
      let inAntiquot := stx.isAntiquot && !stx.isEscapedAntiquot
      let inBinder   := inBinder || (!inQuot && liftMethodForbiddenBinder stx)
      let args ← args.mapM (expandLiftMethodAux (inQuot && !inAntiquot || stx.isQuot) inBinder)
      return Syntax.node i k args
  | stx => return stx

run

def run (doStx : Syntax) (m : Syntax) (returnType : Syntax) : TermElabM CodeBlock :=
  (doSeqToCode <| getDoSeqElems <| getDoSeq doStx).run { ref := doStx, m, returnType }

liftMethodForbiddenBinder

private def liftMethodForbiddenBinder (stx : Syntax) : Bool :=
  let k := stx.getKind

  if k == ``Parser.Term.fun || k == ``Parser.Term.matchAlts ||
     k == ``Parser.Term.doLetRec || k == ``Parser.Term.letrec then

    true

  else if k == ``Parser.Term.let then
    letDeclHasBinders stx[1]
  else if k == ``Parser.Term.doLet then
    letDeclHasBinders stx[2]
  else if k == ``Parser.Term.doLetArrow then
    letDeclArgHasBinders stx[2]
  else
    false

letDeclHasBinders

private def letDeclHasBinders (letDecl : Syntax) : Bool :=
  letDeclArgHasBinders letDecl[0]

letDeclArgHasBinders

private def letDeclArgHasBinders (letDeclArg : Syntax) : Bool :=
  let k := letDeclArg.getKind
  if k == ``Parser.Term.letPatDecl then
    false
  else if k == ``Parser.Term.letEqnsDecl then
    true
  else if k == ``Parser.Term.letIdDecl then

    let binders := letDeclArg[1]
    binders.getNumArgs > 0
  else
    false

checkReassignable

def checkReassignable (xs : Array Var) : M Unit := do
  let throwInvalidReassignment (x : Name) : M Unit :=
    throwError "`{x.simpMacroScopes}` cannot be mutated, only variables declared using `let mut` can be mutated. If you did not intend to mutate but define `{x.simpMacroScopes}`, consider using `let {x.simpMacroScopes}` instead"
  let ctx ← read
  for x in xs do
    unless ctx.mutableVars.contains x.getId do
      throwInvalidReassignment x.getId

  partial def doMatchExprToCode (doMatchExpr : Syntax) (doElems: List Syntax) : M CodeBlock := do
    let ref       := doMatchExpr
    let «meta»    := doMatchExpr[1].isNone
    let discr     := doMatchExpr[2]
    let alts      := doMatchExpr[4][0].getArgs
    let alts ← alts.mapM fun alt => do
      let pat      := alt[1]
      let var?     := if pat[0].isNone then none else some pat[0][0]
      let funName  := pat[1]
      let pvars    := pat[2].getArgs
      let rhs      := alt[3]
      let rhs ← doSeqToCode (getDoSeqElems rhs)
      pure { ref, var?, funName, pvars, rhs }
    let elseBranch ← doSeqToCode (getDoSeqElems doMatchExpr[4][1][3])
    let matchCode ← mkMatchExpr ref «meta» discr alts elseBranch
    concatWith matchCode doElems

mkMatchExpr

def mkMatchExpr (ref : Syntax) («meta» : Bool) (discr : Syntax) (alts : Array (AltExpr CodeBlock)) (elseBranch : CodeBlock) : TermElabM CodeBlock := do

  let ws := alts.foldl (union · ·.rhs.uvars) {}
  let ws := union ws elseBranch.uvars
  let alts ← alts.mapM fun alt => do
    let rhs ← extendUpdatedVars alt.rhs ws
    return { alt with rhs := rhs.code : AltExpr Code }
  let elseBranch ← extendUpdatedVars elseBranch ws
  return { code := .matchExpr ref «meta» discr alts elseBranch.code, uvars := ws }

getDoLetVars

def getDoLetVars (doLet : Syntax) : TermElabM (Array Var) :=

  getLetDeclVars doLet[2]

isDoExpr?

def isDoExpr? (doElem : Syntax) : Option Syntax :=
  if doElem.getKind == ``Parser.Term.doExpr then
    some doElem[0]
  else
    none

reassignToTerm

def reassignToTerm (reassign : Syntax) (k : Syntax) : MacroM Syntax := withRef reassign <| withFreshMacroScope do
  match reassign with
  | `(doElem| $x:ident := $rhs) => `(let $x:ident := ensure_type_of% $x $(quote "invalid reassignment, value") $rhs; $k)
  | `(doElem| $e:term  := $rhs) => `(let $e:term  := ensure_type_of% $e $(quote "invalid reassignment, value") $rhs; $k)
  | _ =>

    Macro.throwErrorAt reassign "unexpected kind of `do` reassignment"

actionTerminalToTerm

def actionTerminalToTerm (action : Syntax) : M Syntax := withRef action <| withFreshMacroScope do
  let ctx ← read
  let u ← mkUVarTuple
  match ctx.kind with
  | .regular         => if ctx.uvars.isEmpty then pure action else ``(Bind.bind $action fun y => Pure.pure (MProdWithNames.mk y $u))
  | .forIn           => ``(Bind.bind $action fun (_ : PUnit) => Pure.pure (ForInStep.yield $u))
  | .forInWithReturn => ``(Bind.bind $action fun (_ : PUnit) => Pure.pure (ForInStep.yield (MProdWithNames.mk none $u)))
  | .nestedBC        => unreachable!
  | .nestedPR        => ``(Bind.bind $action fun y => (Pure.pure (DoResultPR.«pure» y $u)))
  | .nestedSBC       => ``(Bind.bind $action fun y => (Pure.pure (DoResultSBC.«pureReturn» y $u)))
  | .nestedPRBC      => ``(Bind.bind $action fun y => (Pure.pure (DoResultPRBC.«pure» y $u)))

breakToTerm

def breakToTerm : M Syntax := do
  let ctx ← read
  let u ← mkUVarTuple
  match ctx.kind with
  | .regular         => unreachable!
  | .forIn           => ``(Pure.pure (ForInStep.done $u))
  | .forInWithReturn => ``(Pure.pure (ForInStep.done (MProdWithNames.mk none $u)))
  | .nestedBC        => ``(Pure.pure (DoResultBC.«break» $u))
  | .nestedPR        => unreachable!
  | .nestedSBC       => ``(Pure.pure (DoResultSBC.«break» $u))
  | .nestedPRBC      => ``(Pure.pure (DoResultPRBC.«break» $u))

DivM

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

NonDetT.run

def NonDetT.run {α : Type u} :  NonDetT m α -> m α
  | .pure x => Pure.pure x
  | .vis x f => liftM x >>= (fun x => (f x).run)
  | @NonDetT.pickCont _ _ _ p _ f =>
    match Findable.find (p := p) () with
    | none => CCPOBot.compBot
    | some x =>  (f x).run
  | .repeatCont init f cont =>
    forIn Lean.Loop.mk init (fun _ x => (f x).run) >>= (fun x => (cont x).run)

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

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 α

Findable

class Findable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_none : (find ()).isNone -> ∀ x, ¬ p x
  find_some_p : find () = some x -> p x

find

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

findNat

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

NonDetT.extract

@[inline]
def NonDetT.extract {α : Type u} (s : NonDetT m α) (ex : ExtractNonDet Findable s := by admit /- proof elided -/
) : m α :=
  NonDetT.extractGen Findable.find s

NonDetT.extractGen

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

CCPOBot

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

DivM.run

def DivM.run {α : Type u} [Inhabited α] : DivM α -> α
  | DivM.res x => x
  | DivM.div => default

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

Encoding

structure Encoding where
  cnt: Nat
  c: Char
  deriving Inhabited

def get_cnt_sum (l: List Encoding) :=
  match l with
  | List.nil => 0
  | List.cons x xs => x.cnt + get_cnt_sum xs

decodeStrLean

@[grind]
def decodeStrLean (encoded_str: Array Encoding) : Array Char :=
  let mp := Array.map (fun e => Array.replicate e.cnt e.c) encoded_str
  mp.flatten

lemma decodeStrLean_append : forall arr1 arr2,
  decodeStrLean (arr1 ++ arr2) = ( decodeStrLean arr1 ) ++ ( decodeStrLean arr2 )

lemma array_extract_split (arr : Array α) (i j : Nat) :
    i < j → j ≤ arr.size →
    arr.extract 0 j = (arr.extract 0 i) ++ (arr.extract i j)

import Auto

import Loom.MonadAlgebras.NonDetT.Extract

import Loom.MonadAlgebras.WP.Tactic

import Loom.MonadAlgebras.WP.DoNames'

import CaseStudies.Velvet.Std

import CaseStudies.TestingUtil

open PartialCorrectness DemonicChoice

section RunLengthEncoding

structure Encoding where
  cnt: Nat
  c: Char
  deriving Inhabited

def get_cnt_sum (l: List Encoding) :=
  match l with
  | List.nil => 0
  | List.cons x xs => x.cnt + get_cnt_sum xs

@[reducible]
def is_valid_run_sequence (encoded_str: Array Encoding) :=
    forall i, ( h: i < encoded_str.size ) -> (encoded_str[i]'h).cnt > 0

method decodeStr' (encoded_str: Array Encoding)
   return (res: Array Char)
   require is_valid_run_sequence encoded_str
   ensures (res.size = get_cnt_sum encoded_str.toList)
     do
       let mut decoded := Array.replicate 0 'x'
       let mut i := 0
       while i < encoded_str.size
          invariant 0 <= i ∧ i <= encoded_str.size
          invariant decoded.size = get_cnt_sum (encoded_str.extract 0 i).toList
          done_with i = encoded_str.size
          do
            let elem := encoded_str[i]!
            let elem_decoded := Array.replicate elem.cnt elem.c
            decoded :=  decoded ++ elem_decoded
            i := i + 1
       return decoded

prove_correct decodeStr' by
  loom_solve
  · simp[*] at *
    have : decoded.size = get_cnt_sum (List.take i encoded_str.toList) := by admit /- proof elided -/

@[grind]
def decodeStrLean (encoded_str: Array Encoding) : Array Char :=
  let mp := Array.map (fun e => Array.replicate e.cnt e.c) encoded_str
  mp.flatten

lemma array_extract_split_i_j_k (arr : Array α) (i j k: Nat) :
    i < j -> j < k -> k ≤ arr.size →
    arr.extract i k = (arr.extract i j) ++ (arr.extract j k) :=

:= by
  -- 1. Introduce hypotheses into the context
  intro h_lt h_le
  simp[*]
  grind

prove_correct encodeStr by
  loom_solve
  · rw [array_extract_split str i j (by assumption) (by assumption)]
    rw [Array.push_eq_append]
    rw [decodeStrLean_append]
    have : decodeStrLean encoding = str.extract 0 i := by trivial
    rw [this]
    refine (Array.append_right_inj (str.extract 0 i)).mpr ?_
    repeat (unfold decodeStrLean ; simp_all)
    apply Array.ext
    grind
    intros l hl hl2
    have inv : ∀ (k : ℕ), i ≤ k → k < j → str[k]! = str[i] := by trivial
    have inv' := inv (i+l)
    grind


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

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)

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)

Cont.monotone

@[simp]
def Cont.monotone {t : Type v} {α : Type u} [Preorder t] (wp : Cont t α) :=
  ∀ (f f' : α -> t), (∀ a, f a ≤ f' a) → wp f ≤ wp f'

Cont

abbrev Cont (t : Type v) (α : Type u) := (α -> t) -> t

triple

notation "{" P "}" c "{" v "," Q "}" => triple P c (fun v => Q)

lemma Cont.monotone_lift {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteLattice l] [MAlgOrdered m l] :
  ∀ {α : Type u} (x : m α), MAlg.lift x |>.monotone

lemma triple_bind {β} (pre : l) (x : m α) (cut : α -> l)
  (f : α -> m β) (post : β -> l) :
  triple pre x cut ->
  (∀ y, triple (cut y) (f y) post) ->
  triple pre (x >>= f) post

lemma triple_pure (pre : l) (x : α) (post : α -> l) :
  triple pre (pure (f := m) x) post <-> pre ≤ (post x)

Loop.forIn

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

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

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

import Loom.MonadAlgebras.Defs

import Loom.MonadAlgebras.Instances.Basic

import Loom.MonadAlgebras.Instances.ExceptT

import Loom.MonadAlgebras.Instances.StateT

import Loom.MonadAlgebras.Instances.ReaderT

import Loom.MonadAlgebras.Instances.Gen

variable {m : Type u -> Type v} [Monad m] [LawfulMonad m] {α : Type u} {l : Type u}

section

variable [CompleteLattice l]

section

variable [mprop : MAlgOrdered m l]

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

def triple (pre : l) (c : m α) (post : α -> l) : Prop :=
  pre ≤ wp c post

end

variable [MAlgOrdered m l]

end

section

variable [CompleteLattice l] [MAlgOrdered m l]

noncomputable

end

section Determinism

variable [inst: CompleteLattice l] [MAlgOrdered m l]

end Determinism

section Loops

open Lean.Order

@[specialize, inline]
def Loop.forIn.loop {m : Type u -> Type v} [Monad m] [∀ α, CCPO (m α)] [MonoBind m] (f : Unit → β → m (ForInStep β)) (b : β) : m β := do
    match ← f () b with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop f b
  partial_fixpoint

@[inline]
def Loop.forIn {β : Type u} [Monad m] [∀ α, CCPO (m α)] [MonoBind m]
  (_ : Lean.Loop) (init : β) (f : Unit → β → m (ForInStep β)) : m β :=
  Loop.forIn.loop f init

variable [inst: _root_.CompleteLattice l] [MAlgOrdered m l]

namespace PartialCorrectness

variable [∀ α, CCPO (m α)] [MonoBind m] [MAlgPartial m]

end PartialCorrectness

namespace TotalCorrectness

variable [∀ α, CCPO (m α)] [MonoBind m]

lemma repeat_inv (f : Unit -> β -> m (ForInStep β))
  (inv : ForInStep β -> l) (measure : β -> Nat)
  init :
   (∀ b, triple (inv (.yield b)) (f () b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))) ->
   triple (inv (.yield init)) (Loop.forIn.loop f init) (fun b => inv (.done b)) :=

:= by
  intro hstep
  have induc (C : β → Prop) (a : β) (h : ∀ x, (∀ y, measure y < measure x → C y) → C x): C a := by
    have lem: ∀ n: Nat, ∀ x, measure x ≤ n →  C x := by
      intro n
      induction n with
      | zero =>
        intro x hx
        exact h x (fun y => by omega)
      | succ m ih =>
        intro x hx
        by_cases neq: measure x ≤ m
        { exact ih x neq }
        have eq: measure x = m + 1 := by omega
        exact h x (fun y hy => ih y (by omega))
    exact lem (measure a) a (by simp)
  apply induc
    (fun ini => triple (inv (.yield ini)) (Loop.forIn.loop f ( ini)) (fun b => inv (.done b))) init
  intro b ih; unfold Loop.forIn.loop; simp [triple, wp_bind]; apply le_trans
  apply hstep; apply wp_cons; rintro (_|_)
  { simp [wp_pure] }
  rename_i a
  match ForInStep.yield a with
  | .yield b' =>
    simp; intro m_lt
    have ihb := ih b' m_lt
    simp [triple] at ihb
    exact ihb
  | .done b' =>
    simp
    rw [wp_pure b' (fun s: β => inv (ForInStep.done s))]


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

WithName

abbrev WithName (α : Sort u) (name : Lean.Name := default) := α

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) => ⊥) = ⊥

triple

notation "{" P "}" c "{" v "," Q "}" => triple P c (fun v => Q)

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

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

theorem spv_dot_pure (spv1 spv2: SpV Int) (n: ℕ)
  (sz1: ∀ i < spv1.size, spv1.ind[i]! < n) (sz2: ∀ i < spv2.size, spv2.ind[i]! < n):
    spv_dot spv1 spv2 0 0 = ∑ i ∈ Finset.range n, spv1[i] * spv2[i]

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 SpMSpV_correct_triple
  (out: Array Int)
  (spm: Array (SpV Int))
  (spv: SpV Int)
  (n: ℕ):
  triple
    (∀ i < spm.size, (∀ j < spm[i]!.size, spm[i]!.ind[j]! < n) ∧ (∀ j < spv.size, spv.ind[j]! < n))
    (SpMSpV out spm spv)
    fun ⟨_, outNew⟩ =>
      outNew.size = spm.size ∧ ∀ i < spm.size, outNew[i]! = ∑ idx ∈ Finset.range n, spm[i]![idx]! * spv[idx]! :=

:= by
        simp [triple]
        intro inb
        apply wp_cons (SpMSpV out spm spv)
          fun ⟨_, outNew⟩ =>
            (∀ i < spm.size, outNew[i]! = spv_dot spm[i]! spv 0 0) ∧ outNew.size = spm.size
        { rintro outNew; simp
          intro sum_eq sz_eq
          simp [sz_eq]
          intro i ib
          simp [←spv_dot_pure spm[i]! spv n (inb i ib).left (inb i ib).right]
          exact sum_eq i ib }
        simp
        have triple_true := SpMSpV_correct out spm spv
        simp [triple] at triple_true
        simp [WithName] at triple_true
        exact triple_true


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)

Cont.monotone

@[simp]
def Cont.monotone {t : Type v} {α : Type u} [Preorder t] (wp : Cont t α) :=
  ∀ (f f' : α -> t), (∀ a, f a ≤ f' a) → wp f ≤ wp f'

Cont

abbrev Cont (t : Type v) (α : Type u) := (α -> t) -> t

lemma Cont.monotone_lift {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteLattice l] [MAlgOrdered m l] :
  ∀ {α : Type u} (x : m α), MAlg.lift x |>.monotone

Loop.forIn

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

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

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

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

import Loom.MonadAlgebras.Defs

import Loom.MonadAlgebras.Instances.Basic

import Loom.MonadAlgebras.Instances.ExceptT

import Loom.MonadAlgebras.Instances.StateT

import Loom.MonadAlgebras.Instances.ReaderT

import Loom.MonadAlgebras.Instances.Gen

variable {m : Type u -> Type v} [Monad m] [LawfulMonad m] {α : Type u} {l : Type u}

section

variable [CompleteLattice l]

section

variable [mprop : MAlgOrdered m l]

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

def triple (pre : l) (c : m α) (post : α -> l) : Prop :=
  pre ≤ wp c post

end

variable [MAlgOrdered m l]

end

section

variable [CompleteLattice l] [MAlgOrdered m l]

noncomputable

end

section Determinism

variable [inst: CompleteLattice l] [MAlgOrdered m l]

end Determinism

section Loops

open Lean.Order

@[specialize, inline]
def Loop.forIn.loop {m : Type u -> Type v} [Monad m] [∀ α, CCPO (m α)] [MonoBind m] (f : Unit → β → m (ForInStep β)) (b : β) : m β := do
    match ← f () b with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop f b
  partial_fixpoint

@[inline]
def Loop.forIn {β : Type u} [Monad m] [∀ α, CCPO (m α)] [MonoBind m]
  (_ : Lean.Loop) (init : β) (f : Unit → β → m (ForInStep β)) : m β :=
  Loop.forIn.loop f init

variable [inst: _root_.CompleteLattice l] [MAlgOrdered m l]

namespace PartialCorrectness

variable [∀ α, CCPO (m α)] [MonoBind m] [MAlgPartial m]

lemma repeat_inv (f : Unit -> β -> m (ForInStep β))
  (inv : ForInStep β -> l)
  init :
   (∀ b, triple (inv (.yield b)) (f () b) (inv)) ->
   triple (inv (.yield init)) (Loop.forIn.loop f init) (fun b => inv (.done b)) :=

:= by
  intro hstep
  revert init
  apply Loop.forIn.loop.fixpoint_induct (f := f) (motive :=
    fun loop => ∀ init, triple (inv (.yield init)) (loop init) (fun b =>inv (.done b)))
  { apply Lean.Order.admissible_pi_apply
      (P := fun init loop => triple (inv (.yield init)) (loop) (fun b =>inv (.done b)))
    simp [admissible, triple]; intro init loops cl h
    apply le_trans'; apply wp_csup; solve_by_elim
    simp; solve_by_elim }
  intro loop ih init; simp [triple, wp_bind]; apply le_trans; apply hstep
  apply wp_cons; rintro (_|_); simp [wp_pure]
  apply ih