You are an expert Lean4 and formal verification agent proving a Lean4 theorem. Your output **must** strictly adhere to the specified format.

### Input Information
You will receive the following components, relevant to the target theorem:
* <used_lib_defs>...</used_lib_defs>: Imported library definitions directly used by the theorem.
* <used_repo_defs>...</used_repo_defs>: Definitions from the local project (but outside the current file), directly or transitively used by the theorem.
* <lib_lemmas>...</lib_lemmas>: Imported library lemmas/theorems that may assist the proof.
* <repo_lemmas>...</repo_lemmas>: Lemmas/theorems defined in the local project outside this file.
* <local_ctx>...</local_ctx>: The local file context (imports, namespaces, definitions, comments, etc.) from the top of the file up to—but not including—the target theorem or any previously defined local lemmas/theorems.
* <target_theorem>...</target_theorem>: The target theorem statement, ending with `:=`.

**Notes**
1. Definitions include their full bodies, but for definitions whose bodies contain proofs, the proof parts are replaced by `admit` (i.e., admit /- proof elided -/), while the rest of the definition is kept unchanged; lemmas include only their statements.
2. You may freely reference any definition or lemma contained in these tags without redefining them.
3. The tags for used definitions and helpful lemmas (used_lib_defs, used_repo_defs, lib_lemmas, repo_lemmas) may be omitted if empty.

### Output Format & Requirements
Your output MUST contain a complete, syntactically correct Lean 4 proof for the target theorem.

**Strict Output Structure:**
<reasoning> 
...
</reasoning> 

<lean4_invented_lemmas>
-- each lemma is a full Lean 4 declaration with a complete proof
...
</lean4_invented_lemmas>

<lean4_proof>
-- Lean 4 proof body of the target theorem
...
</lean4_proof>

**Requirements:**
1. <reasoning> and <lean4_invented_lemmas> are OPTIONAL. Include them only if they aid proof construction.
2. <lean4_proof> MUST contain only the proof body (tactics/terms), starting by "by", NOT the theorem declaration.
3. <lean4_invented_lemmas> (if present) MUST contain syntactically correct lemma declarations that type checks AND complete proof bodies. Defining a lemma using `axiom` is NOT allowed.
4. <reasoning> contains high-level proof planning.
5. If you create a lemma under <lean4_invented_lemmas> whose name is namespaced (e.g., X.Y.Z), you MUST always refer to it using its fully qualified name. Do NOT rely on shortened names. For example, write rw [X.Y.Z], not rw [Z] for invented lemmas.
6. All Lean 4 code MUST be syntactically correct. No `sorry`, `admit`, or non-Lean text. Comments allowed using `--` or `/- ... -/`.
7. Follow the tags and structure EXACTLY. No extra text before, between, or after the sections.
8. Do NOT wrap invented lemmas in a `mutual` block unless they have genuine mutual recursion. Declare each lemma independently.

### Examples are provided here (example input are under "User Prompt", and example assistant response are under "Agent Response"):
Example 1 ------------------- 
User Prompt:
<used_lib_defs>
structure BitVec (w : Nat) where
  /-- Construct a `BitVec w` from a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  ofFin ::

  /-- Interpret a bitvector as a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  toFin : Fin (hPow 2 w)
</used_lib_defs>

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

<local_ctx>
import Mathlib.Computability.NFA
import Mathlib.Data.FinEnum
import Mathlib.Data.Rel
import Mathlib.Data.Vector.Basic
import Blase.AutoStructs.ForLean

open Set
open Mathlib
open SetRel

-- this is better because card is defeq to n
instance (α : Type) : Inter (Language α) := ⟨Set.inter⟩
instance (α : Type) : Union (Language α) := ⟨Set.union⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width.
-/
structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

def BitVecs.cast (bvs : BitVecs n) (h : n = n') : BitVecs n' :=
  { w := bvs.w, bvs := h ▸ bvs.bvs }

abbrev BitVecs.empty : BitVecs n := ⟨0, List.Vector.replicate n .nil⟩
abbrev BitVecs.singleton {w : Nat} (bv : BitVec w) : BitVecs 1 := ⟨w, bv ::ᵥ .nil⟩
abbrev BitVecs.pair {w : Nat} (bv1 bv2 : BitVec w) : BitVecs 2 := ⟨w, bv1 ::ᵥ bv2 ::ᵥ .nil⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width, encoded as a list of
bit vectors of width `n`. The width of the encoded bit vectors is the length of
the list.
-/
abbrev BitVecs' (n : Nat) := List (BitVec n)

@[simps]
def dec (bvs' : BitVecs' n) : BitVecs n where
  w := bvs'.length
  bvs := List.Vector.ofFn fun k => BitVec.ofFn fun i => bvs'[i].getLsbD k

@[simp]
</local_ctx>

<target_theorem>
lemma dec_snoc n (bvs' : BitVecs' n) (a : BitVec n) :
  dec (bvs' ++ [a]) =
    { w := bvs'.length + 1
      bvs := List.Vector.ofFn fun k =>
        BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k) } :=
</target_theorem>

Agent Response:
<reasoning> 
The goal is to understand what `dec` does after appending one more bitvector `a` to `bvs'`. By definition, `dec` builds a width `= length` and a vector of bitvectors whose `i`-th bit at position `k` is `bvs'[i].getLsbD k`. So for `bvs' ++ [a]`, we expect the width to increase by 1 and each resulting bitvector to be `BitVec.cons` of the new bit coming from `a`.

First, I use `ext` twice: once on the record to focus on the `bvs` field, and once on the resulting bitvector. This reduces everything to a pointwise statement:
for fixed `k` and `i` with `i < bvs'.length + 1`,
the bit `(bvs' ++ [a])[i].getLsbD k` should match the bit at index `i` of `BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k)`.

At this point I need a clean description of indexing into `BitVec.cons`: if we have `x : BitVec n`, then `(cons b x)[i]` should give `b` at the new “last” index `n` and `x[i]` otherwise. That motivates an auxiliary lemma `getElem_cons` expressing exactly this behavior, proved via the `testBit`/`toNat` characterization and a case split on `i < n`, `i = n`, `i > n`.

Using `getElem_cons`, the RHS becomes an `if` on whether `i = bvs'.length`. So I split into cases:
- If `i = bvs'.length`, then `(bvs' ++ [a])[i]` is the appended `a`, so both sides reduce to `a.getLsbD k`.
- If `i ≠ bvs'.length`, the bound `i < bvs'.length + 1` forces `i < bvs'.length`. Then `(bvs' ++ [a])[i] = bvs'[i]` by `List.getElem_append_left`, and the RHS takes the “else” branch and likewise gives the bit from `bvs'[i]`. In both subcases, `simp` (plus the definition of `dec`) finishes.

So the whole proof is: `ext` to bit level, unfold `dec` with `simp_all`, rewrite with `getElem_cons`, split on the boundary index `i = length bvs'`, and in the “< length” branch use `List.getElem_append_left` and bounds to line up both sides.
</reasoning> 

<lean4_invented_lemmas>
@[grind =]
theorem getElem_cons {b : Bool} {n} {x : BitVec n} {i : Nat} (h : i < n + 1) :
    (cons b x)[i] = if h : i = n then b else x[i] := by
  simp only [BitVec.getElem_eq_testBit_toNat, BitVec.toNat_cons, Nat.testBit_or]
  rw [Nat.testBit_shiftLeft]
  rcases Nat.lt_trichotomy i n with i_lt_n | i_eq_n | n_lt_i
  · have p1 : ¬(n ≤ i) := by omega
    have p2 : i ≠ n := by omega
    simp [p1, p2]
  · simp only [i_eq_n, ge_iff_le, Nat.le_refl, decide_true, Nat.sub_self, Nat.testBit_zero,
    Bool.true_and, BitVec.testBit_toNat, BitVec.getLsbD_of_ge, Bool.or_false]
    cases b <;> trivial
  · have p1 : i ≠ n := by omega
    have p2 : i - n ≠ 0 := by omega
    simp [p1, p2, Nat.testBit_bool_toNat]
</lean4_invented_lemmas>

<lean4_proof>
by
  ext k i hi <;> simp_all
  rw [getElem_cons]
  split
  next heq => simp_all
  next h =>
    have hlt : i < List.length bvs' := by simp at hi; omega
    rw [List.getElem_append_left hlt]
    simp
</lean4_proof>

### Your task:
<used_lib_defs>
Char in Init.Prelude
Inhabited in Init.Prelude
Nat in Init.Prelude
Array in Init.Prelude
Array.map in Init.Data.Array.Basic
Array.replicate in Init.Data.Array.Basic
Lean.Syntax in Init.Prelude
Option in Init.Prelude
Bool in Init.Prelude
Lean.Name in Init.Prelude
Lean.Parser.Term.matchExpr in Lean.Parser.Term
Lean.Parser.Term.optIdent in Lean.Parser.Term.Basic
args in Auto.Parser.TPTP
cond in Init.Prelude
ite in Init.Prelude
Lean.MonadQuotation in Init.Prelude
Lean.MonadRef in Init.Prelude
Monad in Init.Prelude
Lean.withRef in Init.Prelude
kind in Auto.Parser.TPTP
Lean.MacroM in Init.Prelude
ReaderT in Init.Prelude
Bind in Init.Prelude
Bind.bind in Init.Prelude
Lean.Parser.Term.debugAssert in Lean.Parser.Term
Lean.Parser.Term.doAssert in Lean.Parser.Do
Lean.Parser.Term.doDbgTrace in Lean.Parser.Do
Lean.Parser.Term.doDebugAssert in Lean.Parser.Do
Lean.Parser.Term.ident in Lean.Parser.Term
Lean.Syntax.mkApp in Init.Meta.Defs
Lean.mkIdentFrom in Init.Meta.Defs
Lean.Parser.Term.let_delayed in Lean.Parser.Term
Lean.mkIdentFromRef in Init.Meta.Defs
PUnit in Init.Prelude
Lean.quoteNameMk in Init.Meta.Defs
DoResultPR in Init.Core
DoResultPRBC in Init.Core
DoResultSBC in Init.Core
ForInStep in Init.Core
Pure in Init.Prelude
DoResultBC in Init.Core
ForInStep.yield in Init.Core
Lean.Name.cmp in Lean.Data.Name
Std.TreeMap in Std.Data.TreeMap.Basic
Lean.Elab.Term.TermElabM in Lean.Elab.Term.TermElabM
Lean.Parser.Term.doLetArrow in Lean.Parser.Do
Unit in Init.Prelude
Lean.Parser.Term.doHave in Lean.Parser.Do
Lean.Parser.Term.doLet in Lean.Parser.Do
Lean.Parser.Term.doLetRec in Lean.Parser.Do
Lean.Parser.Term.doReassign in Lean.Parser.Do
Lean.Parser.Term.doReassignArrow in Lean.Parser.Do
Lean.throwErrorAt in Lean.Exception
List in Init.Prelude
Lean.Parser.Term.doSeqBracketed in Lean.Parser.Do
Lean.Parser.Term.doSeqIndent in Lean.Parser.Do
Id in Init.Control.Id
Id.run in Init.Control.Id
modify in Init.Prelude
Lean.Elab.liftMacroM in Lean.Elab.Util
Unit.unit in Init.Prelude
term in Auto.Parser.TPTP
type in Auto.Lib.Rebind
Lean.Parser.Term.unreachable in Lean.Parser.Term
Lean.throwError in Lean.Exception
Lean.Elab.Term.Quotation.getPatternVars in Lean.Elab.Quotation.Util
Lean.Elab.Term.getPatternVars in Lean.Elab.PatternVar
PUnit.unit in Init.Prelude
Stream in Init.Data.Stream
Stream.next? in Init.Data.Stream
Lean.mkNullNode in Init.Prelude
Lean.Parser.Term.optType in Lean.Parser.Term.Basic
Lean.Parser.Term.doCatch in Lean.Parser.Do
Lean.Parser.Term.doCatchMatch in Lean.Parser.Do
Lean.Parser.Term.matchAlts in Lean.Parser.Term
MonadExcept in Init.Prelude
MonadExcept.tryCatch in Init.Prelude
liftM in Init.Prelude
tryCatchThe in Init.Prelude
tryFinally in Init.Control.Except
Lean.Parser.Term.doLetElse in Lean.Parser.Do
Lean.mkNode in Init.Prelude
Lean.HygieneInfo.mkIdent in Init.Meta.Defs
Lean.Parser.Term.letId in Lean.Parser.Term
Lean.hygieneInfoKind in Init.Prelude
Lean.Parser.Term.letIdDecl in Lean.Parser.Term
Lean.Parser.Term.letPatDecl in Lean.Parser.Term
sequence in Mathlib.Control.Traversable.Basic
Lean.Elab.Term.Quotation.getPatternsVars in Lean.Elab.Quotation.Util
Lean.Elab.Term.getPatternsVars in Lean.Elab.PatternVar
Lean.Elab.Term.expandMatchAlt in Lean.Elab.BindersUtil
Lean.Parser.Term.letEqnsDecl in Lean.Parser.Term
Lean.Parser.Term.letDecl in Lean.Parser.Term
Lean.Parser.Term.doIdDecl in Lean.Parser.Do
Lean.Parser.Term.doPatDecl in Lean.Parser.Do
Lean.SyntaxNodeKind in Init.Prelude
Lean.Parser.Term.termFor in Lean.Parser.Do
Lean.Parser.Term.termReturn in Lean.Parser.Do
Lean.Parser.Term.termTry in Lean.Parser.Do
Lean.Parser.Term.termUnless in Lean.Parser.Do
Lean.Parser.Term.liftMethod in Lean.Parser.Do
termDepIfThenElse in Init.Notation
termIfThenElse in Init.Notation
Lean.Parser.Term.letrec in Lean.Parser.Term
StateT in Init.Control.State
Lean.Parser.Term.num in Lean.Parser.Term
Lean.Syntax.node in Init.Prelude
Lean.choiceKind in Init.Prelude
error in Auto.Parser.TPTP
id in Init.Prelude
node in Test.SmtTranslation.Inductive
Lean.MonadQuotation.addMacroScope in Init.Prelude
Lean.Parser.Term.doBreak in Lean.Parser.Do
Lean.Parser.Term.doContinue in Lean.Parser.Do
Lean.Parser.Term.doExpr in Lean.Parser.Do
Lean.Parser.Term.doFor in Lean.Parser.Do
Lean.Parser.Term.doIf in Lean.Parser.Do
Lean.Parser.Term.doMatch in Lean.Parser.Do
Lean.Parser.Term.doMatchExpr in Lean.Parser.Do
Lean.Parser.Term.doNested in Lean.Parser.Do
Lean.Parser.Term.doReturn in Lean.Parser.Do
Lean.Parser.Term.doTry in Lean.Parser.Do
Lean.Parser.Term.doUnless in Lean.Parser.Do
Lean.withIncRecDepth in Lean.Exception
Lean.Elab.Term.expandOptType in Lean.Elab.BindersUtil
Lean.Parser.Term.letRecDecls in Lean.Parser.Term
Lean.Macro.throwErrorAt in Init.Prelude
Lean.Meta.instantiateMVarsIfMVarApp in Lean.Meta.Basic
Lean.Parser.Term.matchAlt in Lean.Parser.Term
Lean.Parser.Term.matchExprAlt in Lean.Parser.Term
Lean.Parser.Term.matchExprAlts in Lean.Parser.Term
Lean.Parser.Term.matchExprElseAlt in Lean.Parser.Term
Lean.Parser.Term.matchExprPat in Lean.Parser.Term
f in Test.SmtTranslation.Trigger
Decidable in Init.Prelude
DecidablePred in Init.Prelude
Encodable in Mathlib.Logic.Encodable.Basic
Encodable.decode in Mathlib.Logic.Encodable.Basic
inline in Init.Core
Lean.Order.CCPO in Init.Internal.Order.Basic
Lean.Order.MonoBind in Init.Internal.Order.Basic
List.cons in Init.Prelude
List.nil in Init.Prelude
List.take in Init.Data.List.Basic
</used_lib_defs>

<used_repo_defs>
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 =>
  `(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

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

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)

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

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

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

abbrev Var := Syntax  

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

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

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

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

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

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)

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

structure Context where
  ref         : Syntax
   
  m           : Syntax
   
  returnType  : Syntax
  mutableVars : VarSet := {}
  insideFor   : Bool := false

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

abbrev M := ReaderT Context MacroM

structure Context where
   
  m          : Syntax
   
  returnType : Syntax
  uvars      : Array Var
  kind       : Kind

structure ToForInTermResult where
  uvars      : Array Var
  term       : Syntax

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

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)

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

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)

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

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

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)

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

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

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

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

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"

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"

abbrev M := ReaderT Context TermElabM

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

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

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

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

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

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 }

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

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 }

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

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

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

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

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

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

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

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

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

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 }

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

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

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

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

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

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

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

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

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

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

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

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

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!

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

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

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

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

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

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)

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⟩

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

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 }

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

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

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)

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

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

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)

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!

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

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

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!

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

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)

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

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

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

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

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 }

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"

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

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
    #[]

def getLetPatDeclVars (letPatDecl : Syntax) : TermElabM (Array Var) := do
  
  let pattern := letPatDecl[0]
  getPatternVarsEx pattern

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

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

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 }

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

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

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

def getDoHaveVars (doHave : Syntax) : TermElabM (Array Var) :=
  
  getLetDeclVars doHave[1]

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"

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

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"

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)

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

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

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

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

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

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

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

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

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 }

def getDoLetVars (doLet : Syntax) : TermElabM (Array Var) :=
  
  getLetDeclVars doLet[2]

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

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"

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

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

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

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

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 α

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 α

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

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

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

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

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

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

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

<lib_lemmas>
Array.append_right_inj in Init.Data.Array.Lemmas
Array.push_eq_append in Init.Data.Array.Lemmas
</lib_lemmas>

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

<target_theorem>
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) :=
</target_theorem>

