/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Sebastian Ullrich
-/
import Lean.Parser.Term
import Lean.Parser.Do

namespace Lean
namespace Parser

/--
  Syntax quotation for terms and (lists of) commands. We prefer terms, so ambiguous quotations like
  `` `($x $y) `` will be parsed as an application, not two commands. Use `` `($x:command $y:command) `` instead.
  Multiple command will be put in a `` `null `` node, but a single command will not (so that you can directly
  match against a quotation in a command kind's elaborator). -/
-- TODO: use two separate quotation parsers with parser priorities instead
@[builtinTermParser] def Term.quot := leading_parser "`(" >> incQuotDepth (termParser <|> many1Unbox commandParser) >> ")"
@[builtinTermParser] def Term.precheckedQuot := leading_parser "`" >> Term.quot

namespace Command

def terminationBy := leading_parser "termination_by " >> sepBy1 termParser ", "

@[builtinCommandParser]
def moduleDoc := leading_parser ppDedent $ "/-!" >> commentBody >> ppLine

def namedPrio := leading_parser (atomic ("(" >> nonReservedSymbol "priority") >> " := " >> priorityParser >> ")")
def optNamedPrio := optional namedPrio

def «private»        := leading_parser "private "
def «protected»      := leading_parser "protected "
def visibility       := «private» <|> «protected»
def «noncomputable»  := leading_parser "noncomputable "
def «unsafe»         := leading_parser "unsafe "
def «partial»        := leading_parser "partial "
def «nonrec»         := leading_parser "nonrec "
def declModifiers (inline : Bool) := leading_parser optional docComment >> optional (Term.«attributes» >> if inline then skip else ppDedent ppLine) >> optional visibility >> optional «noncomputable» >> optional «unsafe» >> optional («partial» <|> «nonrec»)
def declId           := leading_parser ident >> optional (".{" >> sepBy1 ident ", " >> "}")
def declSig          := leading_parser many (ppSpace >> (Term.simpleBinderWithoutType <|> Term.bracketedBinder)) >> Term.typeSpec
def optDeclSig       := leading_parser many (ppSpace >> (Term.simpleBinderWithoutType <|> Term.bracketedBinder)) >> Term.optType
def declValSimple    := leading_parser " :=\n" >> termParser >> optional Term.whereDecls
def declValEqns      := leading_parser Term.matchAltsWhereDecls
def declVal          := declValSimple <|> declValEqns <|> Term.whereDecls
def «abbrev»         := leading_parser "abbrev " >> declId >> optDeclSig >> declVal
def optDefDeriving   := optional (atomic ("deriving " >> notSymbol "instance") >> sepBy1 ident ", ")
def «def»            := leading_parser "def " >> declId >> optDeclSig >> declVal >> optDefDeriving >> optional terminationBy
def «theorem»        := leading_parser "theorem " >> declId >> declSig >> declVal >> optional terminationBy
def «constant»       := leading_parser "constant " >> declId >> declSig >> optional declValSimple
def «instance»       := leading_parser Term.attrKind >> "instance " >> optNamedPrio >> optional declId >> declSig >> declVal >> optional terminationBy
def «axiom»          := leading_parser "axiom " >> declId >> declSig
def «example»        := leading_parser "example " >> declSig >> declVal
def inferMod         := leading_parser atomic (symbol "{" >> "}")
def ctor             := leading_parser "\n| " >> declModifiers true >> ident >> optional inferMod >> optDeclSig
def derivingClasses  := sepBy1 (group (ident >> optional (" with " >> Term.structInst))) ", "
def optDeriving      := leading_parser optional (atomic ("deriving " >> notSymbol "instance") >> derivingClasses)
def «inductive»      := leading_parser "inductive " >> declId >> optDeclSig >> optional (symbol ":=" <|> "where") >> many ctor >> optDeriving
def classInductive   := leading_parser atomic (group (symbol "class " >> "inductive ")) >> declId >> optDeclSig >> optional (symbol ":=" <|> "where") >> many ctor >> optDeriving
def structExplicitBinder := leading_parser atomic (declModifiers true >> "(") >> many1 ident >> optional inferMod >> optDeclSig >> optional (Term.binderTactic <|> Term.binderDefault) >> ")"
def structImplicitBinder := leading_parser atomic (declModifiers true >> "{") >> many1 ident >> optional inferMod >> declSig >> "}"
def structInstBinder     := leading_parser atomic (declModifiers true >> "[") >> many1 ident >> optional inferMod >> declSig >> "]"
def structSimpleBinder   := leading_parser atomic (declModifiers true >> ident) >> optional inferMod >> optDeclSig >> optional (Term.binderTactic <|> Term.binderDefault)
def structFields         := leading_parser manyIndent (ppLine >> checkColGe >>(structExplicitBinder <|> structImplicitBinder <|> structInstBinder <|> structSimpleBinder))
def structCtor           := leading_parser atomic (declModifiers true >> ident >> optional inferMod >> " :: ")
def structureTk          := leading_parser "structure "
def classTk              := leading_parser "class "
def «extends»            := leading_parser " extends " >> sepBy1 termParser ", "
def «structure»          := leading_parser
    (structureTk <|> classTk) >> declId >> many Term.bracketedBinder >> optional «extends» >> Term.optType
    >> optional ((symbol " := " <|> " where ") >> optional structCtor >> structFields)
    >> optDeriving
@[builtinCommandParser] def declaration := leading_parser
declModifiers false >> («abbrev» <|> «def» <|> «theorem» <|> «constant» <|> «instance» <|> «axiom» <|> «example» <|> «inductive» <|> classInductive <|> «structure»)
@[builtinCommandParser] def «deriving»     := leading_parser "deriving " >> "instance " >> derivingClasses >> " for " >> sepBy1 ident ", "
@[builtinCommandParser] def «section»      := leading_parser "section " >> optional ident
@[builtinCommandParser] def «namespace»    := leading_parser "namespace " >> ident
@[builtinCommandParser] def «end»          := leading_parser "end " >> optional ident
@[builtinCommandParser] def «variable»     := leading_parser "variable" >> many1 Term.bracketedBinder
@[builtinCommandParser] def «universe»     := leading_parser "universe " >> many1 ident
@[builtinCommandParser] def check          := leading_parser "#check " >> termParser
@[builtinCommandParser] def check_failure  := leading_parser "#check_failure " >> termParser -- Like `#check`, but succeeds only if term does not type check
@[builtinCommandParser] def reduce         := leading_parser "#reduce " >> termParser
@[builtinCommandParser] def eval           := leading_parser "#eval " >> termParser
@[builtinCommandParser] def synth          := leading_parser "#synth " >> termParser
@[builtinCommandParser] def exit           := leading_parser "#exit"
@[builtinCommandParser] def print          := leading_parser "#print " >> (ident <|> strLit)
@[builtinCommandParser] def printAxioms    := leading_parser "#print " >> nonReservedSymbol "axioms " >> ident
@[builtinCommandParser] def «resolve_name» := leading_parser "#resolve_name " >> ident
@[builtinCommandParser] def «init_quot»    := leading_parser "init_quot"
def optionValue := nonReservedSymbol "true" <|> nonReservedSymbol "false" <|> strLit <|> numLit
@[builtinCommandParser] def «set_option»   := leading_parser "set_option " >> ident >> ppSpace >> optionValue
def eraseAttr := leading_parser "-" >> rawIdent
@[builtinCommandParser] def «attribute»    := leading_parser "attribute " >> "[" >> sepBy1 (eraseAttr <|> Term.attrInstance) ", " >> "] " >> many1 ident
@[builtinCommandParser] def «export»       := leading_parser "export " >> ident >> "(" >> many1 ident >> ")"
def openHiding       := leading_parser atomic (ident >> "hiding") >> many1 (checkColGt >> ident)
def openRenamingItem := leading_parser ident >> unicodeSymbol "→" "->" >> checkColGt >> ident
def openRenaming     := leading_parser atomic (ident >> "renaming") >> sepBy1 openRenamingItem ", "
def openOnly         := leading_parser atomic (ident >> "(") >> many1 ident >> ")"
def openSimple       := leading_parser many1 (checkColGt >> ident)
def openScoped       := leading_parser "scoped " >> many1 (checkColGt >> ident)
def openDecl         := openHiding <|> openRenaming <|> openOnly <|> openSimple <|> openScoped
@[builtinCommandParser] def «open»    := leading_parser withPosition ("open " >> openDecl)

@[builtinCommandParser] def «mutual» := leading_parser "mutual " >> many1 (ppLine >> notSymbol "end" >> commandParser) >> ppDedent (ppLine >> "end") >> optional terminationBy
@[builtinCommandParser] def «initialize» := leading_parser optional visibility >> "initialize " >> optional (atomic (ident >> Term.typeSpec >> Term.leftArrow)) >> Term.doSeq
@[builtinCommandParser] def «builtin_initialize» := leading_parser optional visibility >> "builtin_initialize " >> optional (atomic (ident >> Term.typeSpec >> Term.leftArrow)) >> Term.doSeq

@[builtinCommandParser] def «in»  := trailing_parser withOpen (" in " >> commandParser)

/-
  This is an auxiliary command for generation constructor injectivity theorems for inductive types defined at `Prelude.lean`.
  It is meant for bootstrapping purposes only. -/
@[builtinCommandParser] def genInjectiveTheorems := leading_parser "gen_injective_theorems% " >> ident

@[runBuiltinParserAttributeHooks] abbrev declModifiersF := declModifiers false
@[runBuiltinParserAttributeHooks] abbrev declModifiersT := declModifiers true

builtin_initialize
  register_parser_alias "declModifiers"       declModifiersF
  register_parser_alias "nestedDeclModifiers" declModifiersT
  register_parser_alias                       declId
  register_parser_alias                       declSig
  register_parser_alias                       declVal
  register_parser_alias                       optDeclSig
  register_parser_alias                       openDecl

end Command

namespace Term
@[builtinTermParser] def «open» := leading_parser:leadPrec "open " >> Command.openDecl >> withOpenDecl (" in " >> termParser)
@[builtinTermParser] def «set_option» := leading_parser:leadPrec "set_option " >> ident >> ppSpace >> Command.optionValue >> " in " >> termParser
end Term

namespace Tactic
@[builtinTacticParser] def «open» := leading_parser:leadPrec "open " >> Command.openDecl >> withOpenDecl (" in " >> tacticSeq)
@[builtinTacticParser] def «set_option» := leading_parser:leadPrec "set_option " >> ident >> ppSpace >> Command.optionValue >> " in " >> tacticSeq
end Tactic

end Parser
end Lean


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Sebastian Ullrich
-/
import Lean.Parser.Term
import Lean.Parser.Do

namespace Lean
namespace Parser

/--
  Syntax quotation for terms and (lists of) commands. We prefer terms, so ambiguous quotations like
  `($x $y) will be parsed as an application, not two commands. Use `($x:command $y:command) instead.
  Multiple command will be put in a `null node, but a single command will not (so that you can directly
  match against a quotation in a command kind's elaborator). -/
-- TODO: use two separate quotation parsers with parser priorities instead
@[builtinTermParser] def Term.quot := leading_parser "`(" >> incQuotDepth (termParser <|> many1Unbox commandParser) >> ")"
@[builtinTermParser] def Term.precheckedQuot := leading_parser "`" >> Term.quot

namespace Command

def namedPrio := leading_parser (atomic ("(" >> nonReservedSymbol "priority") >> " := " >> priorityParser >> ")")
def optNamedPrio := optional namedPrio

def «private»        := leading_parser "private "
def «protected»      := leading_parser "protected "
def visibility       := «private» <|> «protected»
def «noncomputable»  := leading_parser "noncomputable "
def «unsafe»         := leading_parser "unsafe "
def «partial»        := leading_parser "partial "
def declModifiers (inline : Bool) := leading_parser optional docComment >> optional (Term.«attributes» >> if inline then skip else ppDedent ppLine) >> optional visibility >> optional «noncomputable» >> optional «unsafe» >> optional «partial»
def declId           := leading_parser ident >> optional (".{" >> sepBy1 ident ", " >> "}")
def declSig          := leading_parser many (ppSpace >> (Term.simpleBinderWithoutType <|> Term.bracketedBinder)) >> Term.typeSpec
def optDeclSig       := leading_parser many (ppSpace >> (Term.simpleBinderWithoutType <|> Term.bracketedBinder)) >> Term.optType
def declValSimple    := leading_parser " :=\n" >> termParser >> optional Term.whereDecls
def declValEqns      := leading_parser Term.matchAltsWhereDecls
def declVal          := declValSimple <|> declValEqns <|> Term.whereDecls
def «abbrev»         := leading_parser "abbrev " >> declId >> optDeclSig >> declVal
def «def»            := leading_parser "def " >> declId >> optDeclSig >> declVal
def «theorem»        := leading_parser "theorem " >> declId >> declSig >> declVal
def «constant»       := leading_parser "constant " >> declId >> declSig >> optional declValSimple
def «instance»       := leading_parser Term.attrKind >> "instance " >> optNamedPrio >> optional declId >> declSig >> declVal
def «axiom»          := leading_parser "axiom " >> declId >> declSig
def «example»        := leading_parser "example " >> declSig >> declVal
def inferMod         := leading_parser atomic (symbol "{" >> "}")
def ctor             := leading_parser "\n| " >> declModifiers true >> ident >> optional inferMod >> optDeclSig
def optDeriving      := leading_parser optional (atomic ("deriving " >> notSymbol "instance") >> sepBy1 ident ", ")
def «inductive»      := leading_parser "inductive " >> declId >> optDeclSig >> optional (symbol ":=" <|> "where") >> many ctor >> optDeriving
def classInductive   := leading_parser atomic (group (symbol "class " >> "inductive ")) >> declId >> optDeclSig >> optional (symbol ":=" <|> "where") >> many ctor >> optDeriving
def structExplicitBinder := leading_parser atomic (declModifiers true >> "(") >> many1 ident >> optional inferMod >> optDeclSig >> optional Term.binderDefault >> ")"
def structImplicitBinder := leading_parser atomic (declModifiers true >> "{") >> many1 ident >> optional inferMod >> declSig >> "}"
def structInstBinder     := leading_parser atomic (declModifiers true >> "[") >> many1 ident >> optional inferMod >> declSig >> "]"
def structSimpleBinder   := leading_parser atomic (declModifiers true >> ident) >> optional inferMod >> optDeclSig >> optional Term.binderDefault
def structFields         := leading_parser manyIndent (ppLine >> checkColGe >>(structExplicitBinder <|> structImplicitBinder <|> structInstBinder <|> structSimpleBinder))
def structCtor           := leading_parser atomic (declModifiers true >> ident >> optional inferMod >> " :: ")
def structureTk          := leading_parser "structure "
def classTk              := leading_parser "class "
def «extends»            := leading_parser " extends " >> sepBy1 termParser ", "
def «structure»          := leading_parser
    (structureTk <|> classTk) >> declId >> many Term.bracketedBinder >> optional «extends» >> Term.optType
    >> optional ((symbol " := " <|> " where ") >> optional structCtor >> structFields)
    >> optDeriving
@[builtinCommandParser] def declaration := leading_parser
declModifiers false >> («abbrev» <|> «def» <|> «theorem» <|> «constant» <|> «instance» <|> «axiom» <|> «example» <|> «inductive» <|> classInductive <|> «structure»)
@[builtinCommandParser] def «deriving»     := leading_parser "deriving " >> "instance " >> sepBy1 ident ", " >> " for " >> sepBy1 ident ", "
@[builtinCommandParser] def «section»      := leading_parser "section " >> optional ident
@[builtinCommandParser] def «namespace»    := leading_parser "namespace " >> ident
@[builtinCommandParser] def «end»          := leading_parser "end " >> optional ident
@[builtinCommandParser] def «variable»     := leading_parser "variable" >> many1 Term.bracketedBinder
@[builtinCommandParser] def «universe»     := leading_parser "universe " >> ident
@[builtinCommandParser] def «universes»    := leading_parser "universes " >> many1 ident
@[builtinCommandParser] def check          := leading_parser "#check " >> termParser
@[builtinCommandParser] def check_failure  := leading_parser "#check_failure " >> termParser -- Like `#check`, but succeeds only if term does not type check
@[builtinCommandParser] def reduce         := leading_parser "#reduce " >> termParser
@[builtinCommandParser] def eval           := leading_parser "#eval " >> termParser
@[builtinCommandParser] def synth          := leading_parser "#synth " >> termParser
@[builtinCommandParser] def exit           := leading_parser "#exit"
@[builtinCommandParser] def print          := leading_parser "#print " >> (ident <|> strLit)
@[builtinCommandParser] def printAxioms    := leading_parser "#print " >> nonReservedSymbol "axioms " >> ident
@[builtinCommandParser] def «resolve_name» := leading_parser "#resolve_name " >> ident
@[builtinCommandParser] def «init_quot»    := leading_parser "init_quot"
def optionValue := nonReservedSymbol "true" <|> nonReservedSymbol "false" <|> strLit <|> numLit
@[builtinCommandParser] def «set_option»   := leading_parser "set_option " >> ident >> ppSpace >> optionValue
def eraseAttr := leading_parser "-" >> ident
@[builtinCommandParser] def «attribute»    := leading_parser "attribute " >> "[" >> sepBy1 (eraseAttr <|> Term.attrInstance) ", " >> "] " >> many1 ident
@[builtinCommandParser] def «export»       := leading_parser "export " >> ident >> "(" >> many1 ident >> ")"
def openHiding       := leading_parser atomic (ident >> "hiding") >> many1 ident
def openRenamingItem := leading_parser ident >> unicodeSymbol "→" "->" >> ident
def openRenaming     := leading_parser atomic (ident >> "renaming") >> sepBy1 openRenamingItem ", "
def openOnly         := leading_parser atomic (ident >> "(") >> many1 ident >> ")"
def openSimple       := leading_parser many1 ident
def openDecl         := openHiding <|> openRenaming <|> openOnly <|> openSimple
@[builtinCommandParser] def «open»    := leading_parser "open " >> openDecl

@[builtinCommandParser] def «mutual» := leading_parser "mutual " >> many1 (ppLine >> notSymbol "end" >> commandParser) >> ppDedent (ppLine >> "end")
@[builtinCommandParser] def «initialize» := leading_parser "initialize " >> optional (atomic (ident >> Term.typeSpec >> Term.leftArrow)) >> Term.doSeq
@[builtinCommandParser] def «builtin_initialize» := leading_parser "builtin_initialize " >> optional (atomic (ident >> Term.typeSpec >> Term.leftArrow)) >> Term.doSeq

@[builtinCommandParser] def «in»  := trailing_parser " in " >> commandParser

/-
  This is an auxiliary command for generation constructor injectivity theorems for inductive types defined at `Prelude.lean`.
  It is meant for bootstrapping purposes only. -/
@[builtinCommandParser] def genInjectiveTheorems := leading_parser "gen_injective_theorems% " >> ident

@[runBuiltinParserAttributeHooks] abbrev declModifiersF := declModifiers false
@[runBuiltinParserAttributeHooks] abbrev declModifiersT := declModifiers true

builtin_initialize
  register_parser_alias "declModifiers"       declModifiersF
  register_parser_alias "nestedDeclModifiers" declModifiersT
  register_parser_alias "declId"              declId
  register_parser_alias "declSig"             declSig
  register_parser_alias "declVal"             declVal
  register_parser_alias "optDeclSig"          optDeclSig
  register_parser_alias "openDecl"            openDecl

end Command

namespace Term
@[builtinTermParser] def «open» := leading_parser:leadPrec "open " >> Command.openDecl >> " in " >> termParser
@[builtinTermParser] def «set_option» := leading_parser:leadPrec "set_option " >> ident >> ppSpace >> Command.optionValue >> " in " >> termParser
end Term

namespace Tactic
@[builtinTacticParser] def «open» := leading_parser:leadPrec "open " >> Command.openDecl >> " in " >> tacticSeq
@[builtinTacticParser] def «set_option» := leading_parser:leadPrec "set_option " >> ident >> ppSpace >> Command.optionValue >> " in " >> tacticSeq
end Tactic

end Parser
end Lean


import Lean.Elab

structure Bar where

structure Foo where
  foo₁ : Nat
  foo₂ : Nat
  bar  : Bar

def mkFoo₁ : Foo := {
--v textDocument/definition
  foo₁ := 1
--v textDocument/declaration
  foo₂ := 2
--v textDocument/typeDefinition
  bar := ⟨⟩
}

structure HandWrittenStruct where
  n : Nat

-- def HandWrittenStruct.n := fun | mk n => n

          --v textDocument/definition
def hws : HandWrittenStruct := {
--v textDocument/definition
  n := 3
}

            --v textDocument/declaration
def mkFoo₂ := mkFoo₁

syntax (name := elabTest) "test" : term

@[termElab elabTest] def elabElabTest : Lean.Elab.Term.TermElab := fun _ _ => do
  let stx ← `(2)
  Lean.Elab.Term.elabTerm stx none

     --v textDocument/declaration
#check test
     --^ textDocument/definition

def Baz (α : Type) := α

#check fun (b : Baz Nat) => b
                          --^ textDocument/typeDefinition

example : Nat :=
  let a := 1
--v textDocument/definition
  a + b
    --^ textDocument/definition
where
  b := 2

macro_rules | `(test) => `(3)
#check test
     --^ textDocument/definition


import Lean.Elab.Do
import Lean.Util.CollectLevelParams

/-
## HBind definition
-/

class HBind (m : Type u1 → Type v1) (n : Type u2 → Type v2) where
  hBind {α : Type u1} {β : Type u2} : m α → (α → n β) → n β

@[defaultInstance] instance [Bind m] : HBind m m := ⟨bind⟩

instance: HBind Id.{u} Id.{v} where
  hBind {α: Type u} {β: Type v} (x: Id α) (f: α → Id β): Id β := f x

/-
## `hdo` notation: parsing
-/

syntax (name := hdo) "hdo" doSeq: term

-- Variation where we specify the type of the monad (to avoid inferring it as
-- [Type → Type])
syntax (name := hdo_2) "hdo " atomic("(" &"monad" " := " term ")") doSeq: term

/-
## `hdo` notation: elaboration

Most of this is copied from src/Lean/Elab/Do.lean. Look for "HDO:" to find the
differences with the implementation of "do".
-/

namespace Lean.Elab.Term
open Lean.Parser.Term
open Lean.Elab.Term
open Meta

namespace Lean.Elab.Term
open Lean.Parser.Term
open Meta

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

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

@[builtinTermElab liftMethod] def elabLiftMethod : TermElab := fun stx _ =>
  throwErrorAt stx "invalid use of `(<- ...)`, must be nested inside a 'do' expression"

/-- Return true if we should not lift `(<- ...)` actions nested in the syntax nodes with the given kind. -/
private def liftMethodDelimiter (k : SyntaxNodeKind) : Bool :=
  k == ``Lean.Parser.Term.do ||
  k == ``Lean.Parser.Term.doSeqIndent ||
  k == ``Lean.Parser.Term.doSeqBracketed ||
  k == ``Lean.Parser.Term.termReturn ||
  k == ``Lean.Parser.Term.termUnless ||
  k == ``Lean.Parser.Term.termTry ||
  k == ``Lean.Parser.Term.termFor

/-- Given `stx` which is a `letPatDecl`, `letEqnsDecl`, or `letIdDecl`, return true if it has binders. -/
private def letDeclArgHasBinders (letDeclArg : Syntax) : Bool :=
  let k := letDeclArg.getKind
  if k == ``Lean.Parser.Term.letPatDecl then
    false
  else if k == ``Lean.Parser.Term.letEqnsDecl then
    true
  else if k == ``Lean.Parser.Term.letIdDecl then
    -- letIdLhs := ident >> checkWsBefore "expected space before binders" >> many (ppSpace >> (simpleBinderWithoutType <|> bracketedBinder)) >> optType
    let binders := letDeclArg[1]
    binders.getNumArgs > 0
  else
    false

/-- Return `true` if the given `letDecl` contains binders. -/
private def letDeclHasBinders (letDecl : Syntax) : Bool :=
  letDeclArgHasBinders letDecl[0]

/-- Return true if we should generate an error message when lifting a method over this kind of syntax. -/
private def liftMethodForbiddenBinder (stx : Syntax) : Bool :=
  let k := stx.getKind
  if k == ``Lean.Parser.Term.fun || k == ``Lean.Parser.Term.matchAlts ||
     k == ``Lean.Parser.Term.doLetRec || k == ``Lean.Parser.Term.letrec  then
     -- It is never ok to lift over this kind of binder
    true
  -- The following kinds of `let`-expressions require extra checks to decide whether they contain binders or not
  else if k == ``Lean.Parser.Term.let then
    letDeclHasBinders stx[1]
  else if k == ``Lean.Parser.Term.doLet then
    letDeclHasBinders stx[2]
  else if k == ``Lean.Parser.Term.doLetArrow then
    letDeclArgHasBinders stx[2]
  else
    false

private partial def hasLiftMethod : Syntax → Bool
  | Syntax.node _ k args =>
    if liftMethodDelimiter k then false
    -- NOTE: We don't check for lifts in quotations here, which doesn't break anything but merely makes this rare case a
    -- bit slower
    else if k == ``Lean.Parser.Term.liftMethod then true
    else args.any hasLiftMethod
  | _ => false

structure ExtractMonadResult where
  m            : Expr
  α            : Expr
  expectedType : Expr

private partial def extractBind (expectedType? : Option Expr) : TermElabM ExtractMonadResult := do
  match expectedType? with
  | none => throwError "invalid 'do' notation, expected type is not available"
  | some expectedType =>
    let extractStep? (type : Expr) : MetaM (Option ExtractMonadResult) := do
      match type with
      | Expr.app m α _ =>
        try
          let bindInstType ← mkAppM ``Bind #[m]
          let _  ← Meta.synthInstance bindInstType
          return some { m := m, α := α, expectedType := expectedType }
        catch _ =>
          return none
      | _ =>
        return none
    let rec extract? (type : Expr) : MetaM (Option ExtractMonadResult) := do
      match (← extractStep? type) with
      | some r => return r
      | none =>
        let typeNew ← whnfCore type
        if typeNew != type then
          extract? typeNew
        else
          if typeNew.getAppFn.isMVar then throwError "invalid 'do' notation, expected type is not available"
          match (← unfoldDefinition? typeNew) with
          | some typeNew => extract? typeNew
          | none => return none
    match (← extract? expectedType) with
    | some r => return r
    | none   => throwError "invalid 'do' notation, expected type is not a monad application{indentExpr expectedType}\nYou can use the `do` notation in pure code by writing `Id.run do` instead of `do`, where `Id` is the identity monad."

namespace HDo

abbrev Var := Syntax  -- TODO: should be `TSyntax identKind`

/- A `doMatch` alternative. `vars` is the array of variables declared by `patterns`. -/
structure Alt (σ : Type) where
  ref : Syntax
  vars : Array Var
  patterns : Syntax
  rhs : σ
  deriving Inhabited

/-
  Auxiliary datastructure for representing a `do` code block, and compiling "reassignments" (e.g., `x := x + 1`).
  We convert `Code` into a `Syntax` term representing the:
  - `do`-block, or
  - the visitor argument for the `forIn` combinator.

  We say the following constructors are terminals:
  - `break`:    for interrupting a `for x in s`
  - `continue`: for interrupting the current iteration of a `for x in s`
  - `return e`: for returning `e` as the result for the whole `do` computation block
  - `action a`: for executing action `a` as a terminal
  - `ite`:      if-then-else
  - `match`:    pattern matching
  - `jmp`       a goto to a join-point

  We say the terminals `break`, `continue`, `action`, and `return` are "exit points"

  Note that, `return e` is not equivalent to `action (pure e)`. Here is an example:
  ```
  def f (x : Nat) : IO Unit := do
  if x == 0 then
     return ()
  IO.println "hello"
  ```
  Executing `#eval f 0` will not print "hello". Now, consider
  ```
  def g (x : Nat) : IO Unit := do
  if x == 0 then
     pure ()
  IO.println "hello"
  ```
  The `if` statement is essentially a noop, and "hello" is printed when we execute `g 0`.

  - `decl` represents all declaration-like `doElem`s (e.g., `let`, `have`, `let rec`).
    The field `stx` is the actual `doElem`,
    `vars` is the array of variables declared by it, and `cont` is the next instruction in the `do` code block.
    `vars` is an array since we have declarations such as `let (a, b) := s`.

  - `reassign` is an reassignment-like `doElem` (e.g., `x := x + 1`).

  - `joinpoint` is a join point declaration: an auxiliary `let`-declaration used to represent the control-flow.

  - `seq a k` executes action `a`, ignores its result, and then executes `k`.
    We also store the do-elements `dbg_trace` and `assert!` as actions in a `seq`.

  A code block `C` is well-formed if
  - For every `jmp ref j as` in `C`, there is a `joinpoint j ps b k` and `jmp ref j as` is in `k`, and
    `ps.size == as.size` -/
inductive Code where
  | decl         (xs : Array Var) (doElem : Syntax) (k : Code)
  | reassign     (xs : Array Var) (doElem : Syntax) (k : Code)
  /- The Boolean value in `params` indicates whether we should use `(x : typeof! x)` when generating term Syntax or not -/
  | 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)
  /- Recall that an if-then-else may declare a variable using `optIdent` for the branches `thenBranch` and `elseBranch`. We store the variable name at `var?`. -/
  | 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))
  | jmp          (ref : Syntax) (jpName : Name) (args : Array Syntax)
  deriving Inhabited

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

/- A code block, and the collection of variables updated by it. -/
structure CodeBlock where
  code  : Code
  uvars : VarSet := {} -- set of variables updated by `code`

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

private def varsToMessageData (vars : Array Var) : MessageData :=
  MessageData.joinSep (vars.toList.map fun n => MessageData.ofName (n.getId.simpMacroScopes)) " "

partial def CodeBlocl.toMessageData (codeBlock : CodeBlock) : MessageData :=
  let us := MessageData.ofList <| (varSetToArray codeBlock.uvars).toList.map MessageData.ofSyntax
  let rec loop : Code → MessageData
    | Code.decl xs _ k            => m!"let {varsToMessageData xs} := ...\n{loop k}"
    | Code.reassign xs _ k        => m!"{varsToMessageData xs} := ...\n{loop k}"
    | Code.joinpoint n ps body k  => m!"let {n.simpMacroScopes} {varsToMessageData (ps.map Prod.fst)} := {indentD (loop body)}\n{loop k}"
    | Code.seq e k                => m!"{e}\n{loop k}"
    | Code.action e               => e
    | Code.ite _ _ _ c t e        => m!"if {c} then {indentD (loop t)}\nelse{loop e}"
    | Code.jmp _ j xs             => m!"jmp {j.simpMacroScopes} {xs.toList}"
    | Code.«break» _              => m!"break {us}"
    | Code.«continue» _           => m!"continue {us}"
    | Code.«return» _ v           => m!"return {v} {us}"
    | Code.«match» _ _ ds t alts  =>
      m!"match {ds} with"
      ++ alts.foldl (init := m!"") fun acc alt => acc ++ m!"\n| {alt.patterns} => {loop alt.rhs}"
  loop codeBlock.code

/- Return true if the give code contains an exit point that satisfies `p` -/
partial def hasExitPointPred (c : Code) (p : Code → Bool) : Bool :=
  let rec loop : Code → Bool
    | Code.decl _ _ k           => loop k
    | Code.reassign _ _ k       => loop k
    | Code.joinpoint _ _ b k    => loop b || loop k
    | Code.seq _ k              => loop k
    | Code.ite _ _ _ _ t e      => loop t || loop e
    | Code.«match» _ _ _ _ alts => alts.any (loop ·.rhs)
    | Code.jmp _ _ _            => false
    | c                         => p c
  loop c

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

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

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

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

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

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)
  -- Add elaboration hint for producing sane error message
  let y ← `(ensure_expected_type% "type mismatch, result value" $y)
  let k ← mkCont y
  return Code.decl #[y] doElem k

/- Convert `action _ e` instructions in `c` into `let y ← e; jmp _ jp (xs y)`. -/
partial def convertTerminalActionIntoJmp (code : Code) (jp : Name) (xs : Array Var) : MacroM Code :=
  let rec loop : Code → MacroM Code
    | Code.decl xs stx k           => return Code.decl xs stx (← loop k)
    | Code.reassign xs stx k       => return Code.reassign xs stx (← loop k)
    | Code.joinpoint n ps b k      => return Code.joinpoint n ps (← loop b) (← loop k)
    | Code.seq e k                 => return Code.seq e (← loop k)
    | Code.ite ref x? h c t e      => return Code.ite ref x? h c (← loop t) (← loop e)
    | Code.«match» ref g ds t alts => return Code.«match» ref g ds t (← alts.mapM fun alt => do pure { alt with rhs := (← loop alt.rhs) })
    | Code.action e                => mkAuxDeclFor e fun y =>
      let ref := e
      -- We jump to `jp` with xs **and** y
      let jmpArgs := xs.push y
      return Code.jmp ref jp jmpArgs
    | c                            => return c
  loop code

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

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

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

def mkFreshJP (ps : Array (Var × Bool)) (body : Code) : TermElabM JPDecl := do
  let ps ←
    if ps.isEmpty then
      let y ← `(y)
      pure #[(y, false)]
    else
      pure ps
  -- Remark: the compiler frontend implemented in C++ currently detects jointpoints created by
  -- the "do" notation by testing the name. See hack at method `visit_let` at `lcnf.cpp`
  -- We will remove this hack when we re-implement the compiler frontend in Lean.
  let name ← mkFreshUserName `_do_jp
  pure { name := name, params := ps, body := body }

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 insertVars (rs : VarSet) (xs : Array Var) : VarSet :=
  xs.foldl (fun rs x => rs.insert x.getId x) rs

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

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

/- Create a new jointpoint for `c`, and jump to it with the variables `rs` -/
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

/- Create a new joinpoint that takes `rs` and `val` as arguments. `val` must be syntax representing a pure value.
   The body of the joinpoint is created using `mkJPBody yFresh`, where `yFresh`
   is a fresh variable created by this method. -/
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

/- `pullExitPointsAux rs c` auxiliary method for `pullExitPoints`, `rs` is the set of update variable in the current path.  -/
partial def pullExitPointsAux : VarSet → Code → StateRefT (Array JPDecl) TermElabM Code
  | rs, Code.decl xs stx k           => return Code.decl xs stx (← pullExitPointsAux (eraseVars rs xs) k)
  | rs, Code.reassign xs stx k       => return Code.reassign xs stx (← pullExitPointsAux (insertVars rs xs) k)
  | rs, Code.joinpoint j ps b k      => return Code.joinpoint j ps (← pullExitPointsAux rs b) (← pullExitPointsAux rs k)
  | rs, Code.seq e k                 => return Code.seq e (← pullExitPointsAux rs k)
  | rs, Code.ite ref x? o c t e      => return Code.ite ref x? o c (← pullExitPointsAux (eraseOptVar rs x?) t) (← pullExitPointsAux (eraseOptVar rs x?) e)
  | rs, Code.«match» ref g ds t alts => return Code.«match» ref g ds t (← alts.mapM fun alt => do pure { alt with rhs := (← pullExitPointsAux (eraseVars rs alt.vars) alt.rhs) })
  | rs, c@(Code.jmp _ _ _)           => return  c
  | rs, Code.«break» ref             => mkSimpleJmp ref rs (Code.«break» ref)
  | rs, Code.«continue» ref          => mkSimpleJmp ref rs (Code.«continue» ref)
  | rs, Code.«return» ref val        => mkJmp ref rs val (fun y => return Code.«return» ref y)
  | rs, Code.action e                =>
    -- We use `mkAuxDeclFor` because `e` is not pure.
    mkAuxDeclFor e fun y =>
      let ref := e
      mkJmp ref rs y (fun yFresh => return Code.action (← ``(Pure.pure $yFresh)))

/-
Auxiliary operation for adding new variables to the collection of updated variables in a CodeBlock.
When a new variable is not already in the collection, but is shadowed by some declaration in `c`,
we create auxiliary join points to make sure we preserve the semantics of the code block.
Example: suppose we have the code block `print x; let x := 10; return x`. And we want to extend it
with the reassignment `x := x + 1`. We first use `pullExitPoints` to create
```
let jp (x!1) :=  return x!1;
print x;
let x := 10;
jmp jp x
```
and then we add the reassignment
```
x := x + 1
let jp (x!1) := return x!1;
print x;
let x := 10;
jmp jp x
```
Note that we created a fresh variable `x!1` to avoid accidental name capture.
As another example, consider
```
print x;
let x := 10
y := y + 1;
return x;
```
We transform it into
```
let jp (y x!1) := return x!1;
print x;
let x := 10
y := y + 1;
jmp jp y x
```
and then we add the reassignment as in the previous example.
We need to include `y` in the jump, because each exit point is implicitly returning the set of
update variables.

We implement the method as follows. Let `us` be `c.uvars`, then
1- for each `return _ y` in `c`, we create a join point
  `let j (us y!1) := return y!1`
   and replace the `return _ y` with `jmp us y`
2- for each `break`, we create a join point
  `let j (us) := break`
   and replace the `break` with `jmp us`.
3- Same as 2 for `continue`.
-/
def pullExitPoints (c : Code) : TermElabM Code := do
  if hasExitPoint c then
    let (c, jpDecls) ← (pullExitPointsAux {} c).run #[]
    return attachJPs jpDecls c
  else
    return c

partial def extendUpdatedVarsAux (c : Code) (ws : VarSet) : TermElabM Code :=
  let rec update : Code → TermElabM Code
    | Code.joinpoint j ps b k          => return Code.joinpoint j ps (← update b) (← update k)
    | Code.seq e k                     => return Code.seq e (← update k)
    | c@(Code.«match» ref g ds t alts) => do
      if alts.any fun alt => alt.vars.any fun x => ws.contains x.getId then
        -- If a pattern variable is shadowing a variable in ws, we `pullExitPoints`
        pullExitPoints c
      else
        return Code.«match» ref g ds t (← alts.mapM fun alt => do pure { alt with rhs := (← update alt.rhs) })
    | Code.ite ref none o c t e => return Code.ite ref none o c (← update t) (← update e)
    | c@(Code.ite ref (some h) o cond t e) => do
      if ws.contains h.getId then
        -- if the `h` at `if h:c then t else e` shadows a variable in `ws`, we `pullExitPoints`
        pullExitPoints c
      else
        return Code.ite ref (some h) o cond (← update t) (← update e)
    | Code.reassign xs stx k => return Code.reassign xs stx (← update k)
    | c@(Code.decl xs stx k) => do
      if xs.any fun x => ws.contains x.getId then
        -- One the declared variables is shadowing a variable in `ws`
        pullExitPoints c
      else
        return Code.decl xs stx (← update k)
    | c => return  c
  update c

/-
Extend the set of updated variables. It assumes `ws` is a super set of `c.uvars`.
We **cannot** simply update the field `c.uvars`, because `c` may have shadowed some variable in `ws`.
See discussion at `pullExitPoints`.
-/
partial def extendUpdatedVars (c : CodeBlock) (ws : VarSet) : TermElabM CodeBlock := do
  if ws.any fun x _ => !c.uvars.contains x then
    -- `ws` contains a variable that is not in `c.uvars`, but in `c.dvars` (i.e., it has been shadowed)
    pure { code := (← extendUpdatedVarsAux c.code ws), uvars := ws }
  else
    pure { c with uvars := ws }

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

/-
Given two code blocks `c₁` and `c₂`, make sure they have the same set of updated variables.
Let `ws` the union of the updated variables in `c₁‵ and ‵c₂`.
We use `extendUpdatedVars c₁ ws` and `extendUpdatedVars c₂ ws`
-/
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₂)

/-
Extending code blocks with variable declarations: `let x : t := v` and `let x : t ← v`.
We remove `x` from the collection of updated varibles.
Remark: `stx` is the syntax for the declaration (e.g., `letDecl`), and `xs` are the variables
declared by it. It is an array because we have let-declarations that declare multiple variables.
Example: `let (x, y) := t`
-/
def mkVarDeclCore (xs : Array Var) (stx : Syntax) (c : CodeBlock) : CodeBlock := {
  code := Code.decl xs stx c.code,
  uvars := eraseVars c.uvars xs
}

/-
Extending code blocks with reassignments: `x : t := v` and `x : t ← v`.
Remark: `stx` is the syntax for the declaration (e.g., `letDecl`), and `xs` are the variables
declared by it. It is an array because we have let-declarations that declare multiple variables.
Example: `(x, y) ← t`
-/
def mkReassignCore (xs : Array Var) (stx : Syntax) (c : CodeBlock) : TermElabM CodeBlock := do
  let us := c.uvars
  let ws := insertVars us xs
  -- If `xs` contains a new updated variable, then we must use `extendUpdatedVars`.
  -- See discussion at `pullExitPoints`
  let code ← if xs.any fun x => !us.contains x.getId then extendUpdatedVarsAux c.code ws else pure c.code
  pure { code := Code.reassign xs stx code, uvars := ws }

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

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

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

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

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

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
  pure {
    code  := Code.ite ref x? optIdent cond thenBranch.code elseBranch.code,
    uvars := thenBranch.uvars,
  }

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

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

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

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

def mkMatch (ref : Syntax) (genParam : Syntax) (discrs : Syntax) (optMotive : Syntax) (alts : Array (Alt CodeBlock)) : TermElabM CodeBlock := do
  -- nary version of homogenize
  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 := Code.«match» ref genParam discrs optMotive alts, uvars := ws }

/- Return a code block that executes `terminal` and then `k` with the value produced by `terminal`.
   This method assumes `terminal` is a terminal -/
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 getLetIdDeclVar (letIdDecl : Syntax) : Var :=
  letIdDecl[0]

-- support both regular and syntax match
def getPatternVarsEx (pattern : Syntax) : TermElabM (Array Var) :=
  getPatternVars pattern <|>
  Quotation.getPatternVars pattern

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

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

def getLetEqnsDeclVar (letEqnsDecl : Syntax) : Var :=
  letEqnsDecl[0]

def getLetDeclVars (letDecl : Syntax) : TermElabM (Array Var) := do
  let arg := letDecl[0]
  if arg.getKind == ``Lean.Parser.Term.letIdDecl then
    return #[getLetIdDeclVar arg]
  else if arg.getKind == ``Lean.Parser.Term.letPatDecl then
    getLetPatDeclVars arg
  else if arg.getKind == ``Lean.Parser.Term.letEqnsDecl then
    return #[getLetEqnsDeclVar arg]
  else
    throwError "unexpected kind of let declaration"

def getDoLetVars (doLet : Syntax) : TermElabM (Array Var) :=
  -- leading_parser "let " >> optional "mut " >> letDecl
  getLetDeclVars doLet[2]

def getHaveIdLhsVar (optIdent : Syntax) : TermElabM Var :=
  if optIdent.isNone then
    `(this)
  else
    pure optIdent[0]

def getDoHaveVars (doHave : Syntax) : TermElabM (Array Var) := do
  -- doHave := leading_parser "have " >> Term.haveDecl
  -- haveDecl := leading_parser haveIdDecl <|> letPatDecl <|> haveEqnsDecl
  let arg := doHave[1][0]
  if arg.getKind == ``Lean.Parser.Term.haveIdDecl then
    -- haveIdDecl := leading_parser atomic (haveIdLhs >> " := ") >> termParser
    -- haveIdLhs := optional (ident >> many (ppSpace >> (simpleBinderWithoutType <|> bracketedBinder))) >> optType
    return #[← getHaveIdLhsVar arg[0]]
  else if arg.getKind == ``Lean.Parser.Term.letPatDecl then
    getLetPatDeclVars arg
  else if arg.getKind == ``Lean.Parser.Term.haveEqnsDecl then
    -- haveEqnsDecl := leading_parser haveIdLhs >> matchAlts
    return #[← getHaveIdLhsVar arg[0]]
  else
    throwError "unexpected kind of have declaration"

def getDoLetRecVars (doLetRec : Syntax) : TermElabM (Array Var) := do
  -- letRecDecls is an array of `(group (optional attributes >> letDecl))`
  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

-- ident >> optType >> leftArrow >> termParser
def getDoIdDeclVar (doIdDecl : Syntax) : Var :=
  doIdDecl[0]

-- termParser >> leftArrow >> termParser >> optional (" | " >> termParser)
def getDoPatDeclVars (doPatDecl : Syntax) : TermElabM (Array Var) := do
  let pattern := doPatDecl[0]
  getPatternVarsEx pattern

-- leading_parser "let " >> optional "mut " >> (doIdDecl <|> doPatDecl)
def getDoLetArrowVars (doLetArrow : Syntax) : TermElabM (Array Var) := do
  let decl := doLetArrow[2]
  if decl.getKind == ``Lean.Parser.Term.doIdDecl then
    return #[getDoIdDeclVar decl]
  else if decl.getKind == ``Lean.Parser.Term.doPatDecl then
    getDoPatDeclVars decl
  else
    throwError "unexpected kind of 'do' declaration"

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

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

def mkSingletonDoSeq (doElem : Syntax) : Syntax :=
  mkDoSeq #[doElem]

/-
  If the given syntax is a `doIf`, return an equivalente `doIf` that has an `else` but no `else if`s or `if let`s.  -/
private def expandDoIf? (stx : Syntax) : MacroM (Option Syntax) := match stx with
  | `(doElem|if $p:doIfProp then $t else $e) => pure none
  | `(doElem|if%$i $cond:doIfCond then $t $[else if%$is $conds:doIfCond then $ts]* $[else $e?]?) => withRef stx do
    let mut e      := e?.getD (← `(doSeq|pure PUnit.unit))
    let mut eIsSeq := true
    for (i, cond, t) in Array.zip (is.reverse.push i) (Array.zip (conds.reverse.push cond) (ts.reverse.push t)) do
      e ← if eIsSeq then pure e else `(doSeq|$e:doElem)
      e ← withRef cond <| match cond with
        | `(doIfCond|let $pat := $d) => `(doElem| match%$i $d:term with | $pat:term => $t | _ => $e)
        | `(doIfCond|let $pat ← $d)  => `(doElem| match%$i ← $d    with | $pat:term => $t | _ => $e)
        | `(doIfCond|$cond:doIfProp) => `(doElem| if%$i $cond:doIfProp then $t else $e)
        | _                          => `(doElem| if%$i $(Syntax.missing) then $t else $e)
      eIsSeq := false
    return some e
  | _ => pure none

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

/- This method assumes `expandDoIf?` is not applicable. -/
private def mkDoIfView (doIf : Syntax) : MacroM DoIfView := do
  pure {
    ref        := doIf,
    optIdent   := doIf[1][0],
    cond       := doIf[1][1],
    thenBranch := doIf[3],
    elseBranch := doIf[5][1]
  }

/-
HDO: We use `Prod` instead of `MProd` for exactly the opposite reason
-/
private def mkTuple (elems : Array Syntax) : MacroM Syntax := do
  if elems.size == 0 then
    mkUnit
  else if elems.size == 1 then
    return elems[0]
  else
    elems.extract 0 (elems.size - 1) |>.foldrM (init := elems.back) fun elem tuple =>
      ``(Prod.mk $elem $tuple)

/- Return `some action` if `doElem` is a `doExpr <action>`-/
def isDoExpr? (doElem : Syntax) : Option Syntax :=
  if doElem.getKind == ``Lean.Parser.Term.doExpr then
    some doElem[0]
  else
    none

/--
  Given `uvars := #[a_1, ..., a_n, a_{n+1}]` construct term
  ```
  let a_1     := x.1
  let x       := x.2
  let a_2     := x.1
  let x       := x.2
  ...
  let a_n     := x.1
  let a_{n+1} := x.2
  body
  ```
  Special cases
  - `uvars := #[]` => `body`
  - `uvars := #[a]` => `let a := x; body`

  We use this method when expanding the `for-in` notation.
-/
private def destructTuple (uvars : Array Var) (x : Syntax) (body : Syntax) : MacroM Syntax := do
  if uvars.size == 0 then
    return body
  else if uvars.size == 1 then
    `(let $(uvars[0]):ident := $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 := $x.2; $body)
      | a :: as => withFreshMacroScope do
        let rest ← destruct as (← `(x)) body
        `(let $a:ident := $x.1; let x := $x.2; $rest)
      | _ => unreachable!

/-
The procedure `ToTerm.run` converts a `CodeBlock` into a `Syntax` term.
We use this method to convert
1- The `CodeBlock` for a root `do ...` term into a `Syntax` term. This kind of
   `CodeBlock` never contains `break` nor `continue`. Moreover, the collection
   of updated variables is not packed into the result.
   Thus, we have two kinds of exit points
     - `Code.action e` which is converted into `e`
     - `Code.return _ e` which is converted into `pure e`

   We use `Kind.regular` for this case.

2- The `CodeBlock` for `b` at `for x in xs do b`. In this case, we need to generate
   a `Syntax` term representing a function for the `xs.forIn` combinator.

   a) If `b` contain a `Code.return _ a` exit point. The generated `Syntax` term
      has type `m (ForInStep (Option α × σ))`, where `a : α`, and the `σ` is the type
      of the tuple of variables reassigned by `b`.
      We use `Kind.forInWithReturn` for this case

   b) If `b` does not contain a `Code.return _ a` exit point. Then, the generated
      `Syntax` term has type `m (ForInStep σ)`.
      We use `Kind.forIn` for this case.

3- The `CodeBlock` `c` for a `do` sequence nested in a monadic combinator (e.g., `MonadExcept.tryCatch`).

   The generated `Syntax` term for `c` must inform whether `c` "exited" using `Code.action`, `Code.return`,
   `Code.break` or `Code.continue`. We use the auxiliary types `DoResult`s for storing this information.
   For example, the auxiliary type `DoResultPBC α σ` is used for a code block that exits with `Code.action`,
   **and** `Code.break`/`Code.continue`, `α` is the type of values produced by the exit `action`, and
   `σ` is the type of the tuple of reassigned variables.
   The type `DoResult α β σ` is usedf for code blocks that exit with
   `Code.action`, `Code.return`, **and** `Code.break`/`Code.continue`, `β` is the type of the returned values.
   We don't use `DoResult α β σ` for all cases because:

      a) The elaborator would not be able to infer all type parameters without extra annotations. For example,
         if the code block does not contain `Code.return _ _`, the elaborator will not be able to infer `β`.

      b) We need to pattern match on the result produced by the combinator (e.g., `MonadExcept.tryCatch`),
         but we don't want to consider "unreachable" cases.

   We do not distinguish between cases that contain `break`, but not `continue`, and vice versa.

   When listing all cases, we use `a` to indicate the code block contains `Code.action _`, `r` for `Code.return _ _`,
   and `b/c` for a code block that contains `Code.break _` or `Code.continue _`.

   - `a`: `Kind.regular`, type `m (α × σ)`

   - `r`: `Kind.regular`, type `m (α × σ)`
           Note that the code that pattern matches on the result will behave differently in this case.
           It produces `return a` for this case, and `pure a` for the previous one.

   - `b/c`: `Kind.nestedBC`, type `m (DoResultBC σ)`

   - `a` and `r`:   `Kind.nestedPR`, type `m (DoResultPR α β σ)`

   - `a` and `bc`:  `Kind.nestedSBC`, type `m (DoResultSBC α σ)`

   - `r` and `bc`:  `Kind.nestedSBC`, type `m (DoResultSBC α σ)`
         Again the code that pattern matches on the result will behave differently in this case and
         the previous one. It produces `return a` for the constructor `DoResultSPR.pureReturn a u` for
         this case, and `pure a` for the previous case.

   - `a`, `r`, `b/c`: `Kind.nestedPRBC`, type type `m (DoResultPRBC α β σ)`

Here is the recipe for adding new combinators with nested `do`s.
Example: suppose we want to support `repeat doSeq`. Assuming we have `repeat : m α → m α`
1- Convert `doSeq` into `codeBlock : CodeBlock`
2- Create term `term` using `mkNestedTerm code m uvars a r bc` where
   `code` is `codeBlock.code`, `uvars` is an array containing `codeBlock.uvars`,
   `m` is a `Syntax` representing the Monad, and
   `a` is true if `code` contains `Code.action _`,
   `r` is true if `code` contains `Code.return _ _`,
   `bc` is true if `code` contains `Code.break _` or `Code.continue _`.

   Remark: for combinators such as `repeat` that take a single `doSeq`, all
   arguments, but `m`, are extracted from `codeBlock`.
3- Create the term `repeat $term`
4- and then, convert it into a `doSeq` using `matchNestedTermResult ref (repeat $term) uvsar a r bc`

-/
namespace ToTerm

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

instance : Inhabited Kind := ⟨Kind.regular⟩

def Kind.isRegular : Kind → Bool
  | Kind.regular => true
  | _            => false

structure Context where
  m     : Syntax -- Syntax to reference the monad associated with the do notation.
  uvars : Array Var
  kind  : Kind

abbrev M := ReaderT Context MacroM

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

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

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

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

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

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

def declToTerm (decl : Syntax) (k : Syntax) : M Syntax := withRef decl <| withFreshMacroScope do
  let kind := decl.getKind
  if kind == ``Lean.Parser.Term.doLet then
    let letDecl := decl[2]
    `(let $letDecl:letDecl; $k)
  else if kind == ``Lean.Parser.Term.doLetRec then
    let letRecToken := decl[0]
    let letRecDecls := decl[1]
    return mkNode ``Lean.Parser.Term.letrec #[letRecToken, letRecDecls, mkNullNode, k]
  else if kind == ``Lean.Parser.Term.doLetArrow then
    let arg := decl[2]
    let ref := arg
    if arg.getKind == ``Lean.Parser.Term.doIdDecl then
      let id     := arg[0]
      let type   := expandOptType id arg[1]
      let doElem := arg[3]
      -- `doElem` must be a `doExpr action`. See `doLetArrowToCode`
      match isDoExpr? doElem with
      | some action =>
        let action ← withRef action `(($action : $((← read).m) $type))
        ``(HBind.hBind $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 == ``Lean.Parser.Term.doHave then
    -- The `have` term is of the form  `"have " >> haveDecl >> optSemicolon termParser`
    let args := decl.getArgs
    let args := args ++ #[mkNullNode /- optional ';' -/, k]
    return mkNode `Lean.Parser.Term.«have» args
  else
    Macro.throwErrorAt decl "unexpected kind of 'do' declaration"

def reassignToTerm (reassign : Syntax) (k : Syntax) : MacroM Syntax := withRef reassign <| withFreshMacroScope do
  let kind := reassign.getKind
  if kind == ``Lean.Parser.Term.doReassign then
    -- doReassign := leading_parser (letIdDecl <|> letPatDecl)
    let arg := reassign[0]
    if arg.getKind == ``Lean.Parser.Term.letIdDecl then
      -- letIdDecl := leading_parser ident >> many (ppSpace >> bracketedBinder) >> optType >>  " := " >> termParser
      let x   := arg[0]
      let val := arg[4]
      let newVal ← `(ensure_type_of% $x $(quote "invalid reassignment, value") $val)
      let arg := arg.setArg 4 newVal
      let letDecl := mkNode `Lean.Parser.Term.letDecl #[arg]
      `(let $letDecl:letDecl; $k)
    else
      -- TODO: ensure the types did not change
      let letDecl := mkNode `Lean.Parser.Term.letDecl #[arg]
      `(let $letDecl:letDecl; $k)
  else
    -- Note that `doReassignArrow` is expanded by `doReassignArrowToCode
    Macro.throwErrorAt reassign "unexpected kind of 'do' reassignment"

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 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)
  /-
  We use `let_delayed` instead of `let` for joinpoints to make sure `$k` is elaborated before `$body`.
  By elaborating `$k` first, we "learn" more about `$body`'s type.
  For example, consider the following example `do` expression
  ```
  def f (x : Nat) : IO Unit := do
  if x > 0 then
    IO.println "x is not zero" -- Error is here
  IO.mkRef true
  ```
  it is expanded into
  ```
  def f (x : Nat) : IO Unit := do
  let jp (u : Unit) : IO _ :=
    IO.mkRef true;
  if x > 0 then
    IO.println "not zero"
    jp ()
  else
    jp ()
  ```
  If we use the regular `let` instead of `let_delayed`, the joinpoint `jp` will be elaborated and its type will be inferred to be `Unit → IO (IO.Ref Bool)`.
  Then, we get a typing error at `jp ()`. By using `let_delayed`, we first elaborate `if x > 0 ...` and learn that `jp` has type `Unit → IO Unit`.
  Then, we get the expected type mismatch error at `IO.mkRef true`. -/
  `(let_delayed $(← mkIdentFromRef j):ident $[($ps : $pTypes)]* : $((← read).m) _ := $body; $k)

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

partial def toTerm (c : Code) : M Syntax := do
  match c with
  | Code.return ref val     => withRef ref <| returnToTerm val
  | Code.continue ref       => withRef ref continueToTerm
  | Code.break ref          => withRef ref breakToTerm
  | Code.action e           => actionTerminalToTerm e
  | Code.joinpoint j ps b k => mkJoinPoint j ps (← toTerm b) (← toTerm k)
  | Code.jmp ref j args     => return mkJmp ref j args
  | Code.decl _ stx k       => declToTerm stx (← toTerm k)
  | Code.reassign _ stx k   => reassignToTerm stx (← toTerm k)
  | Code.seq stx k          => seqToTerm stx (← toTerm k)
  | Code.ite ref _ o c t e  => withRef ref <| do mkIte o c (← toTerm t) (← toTerm e)
  | Code.«match» ref genParam discrs optMotive alts =>
    let mut termAlts := #[]
    for alt in alts do
      let rhs ← toTerm alt.rhs
      let termAlt := mkNode `Lean.Parser.Term.matchAlt #[mkAtomFrom alt.ref "|", mkNullNode #[alt.patterns], mkAtomFrom alt.ref "=>", rhs]
      termAlts := termAlts.push termAlt
    let termMatchAlts := mkNode `Lean.Parser.Term.matchAlts #[mkNullNode termAlts]
    return mkNode `Lean.Parser.Term.«match» #[mkAtomFrom ref "match", genParam, optMotive, discrs, mkAtomFrom ref "with", termMatchAlts]

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

/- Given
   - `a` is true if the code block has a `Code.action _` exit point
   - `r` is true if the code block has a `Code.return _ _` exit point
   - `bc` is true if the code block has a `Code.break _` or `Code.continue _` exit point

   generate Kind. See comment at the beginning of the `ToTerm` namespace. -/
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 mkNestedTerm (code : Code) (m : Syntax) (uvars : Array Var) (a r bc : Bool) : MacroM Syntax := do
  ToTerm.run code m uvars (mkNestedKind a r bc)

/- Given a term `term` produced by `ToTerm.run`, pattern match on its result.
   See comment at the beginning of the `ToTerm` namespace.

   - `a` is true if the code block has a `Code.action _` exit point
   - `r` is true if the code block has a `Code.return _ _` exit point
   - `bc` is true if the code block has a `Code.break _` or `Code.continue _` exit point

   The result is a sequence of `doElem` -/
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 ← mkTuple 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
         | DoResultBC.«break» u => $u:term := u; break
         | DoResultBC.«continue» u => $u:term := u; continue)
  | true, true, false => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | DoResultPR.«pure» a u => $u:term := u; pure a
         | DoResultPR.«return» b u => $u:term := u; return b)
  | true, false, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | DoResultSBC.«pureReturn» a u => $u:term := u; pure a
         | DoResultSBC.«break» u => $u:term := u; break
         | DoResultSBC.«continue» u => $u:term := u; continue)
  | false, true, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | DoResultSBC.«pureReturn» a u => $u:term := u; return a
         | DoResultSBC.«break» u => $u:term := u; break
         | DoResultSBC.«continue» u => $u:term := u; continue)
  | true, true, true => toDoElems <$>
    `(do let r ← $term:term;
         match r with
         | DoResultPRBC.«pure» a u => $u:term := u; pure a
         | DoResultPRBC.«return» a u => $u:term := u; return a
         | DoResultPRBC.«break» u => $u:term := u; break
         | DoResultPRBC.«continue» u => $u:term := u; continue)
  | false, false, false => unreachable!

end ToTerm

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

namespace ToCodeBlock

structure Context where
  ref         : Syntax
  m           : Syntax -- Syntax representing the monad associated with the do notation.
  mutableVars : VarSet := {}
  insideFor   : Bool := false

abbrev M := ReaderT Context TermElabM

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 checkReassignable (xs : Array Var) : M Unit := do
  let throwInvalidReassignment (x : Name) : M Unit :=
    throwError "'{x.simpMacroScopes}' cannot be reassigned"
  let ctx ← read
  for x in xs do
    unless ctx.mutableVars.contains x.getId do
      throwInvalidReassignment x.getId

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
      throwInvalidShadowing x.getId

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

structure ToForInTermResult where
  uvars      : Array Var
  term       : Syntax

def mkForInBody  (x : 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 uvars (if hasReturn forInBody.code then ToTerm.Kind.forInWithReturn else ToTerm.Kind.forIn)
  return ⟨uvars, term⟩

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

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

private partial def expandLiftMethodAux (inQuot : Bool) (inBinder : Bool) : Syntax → StateT (List Syntax) M Syntax
  | stx@(Syntax.node i k args) =>
    if liftMethodDelimiter k then
      return stx
    else if k == ``Lean.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 auxDoElem ← `(doElem| let a ← $term:term)
      modify fun s => s ++ [auxDoElem]
      `(a)
    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 expandLiftMethod (doElem : Syntax) : M (List Syntax × Syntax) := do
  if !hasLiftMethod doElem then
    return ([], doElem)
  else
    let (doElem, doElemsNew) ← (expandLiftMethodAux false false doElem).run []
    return (doElemsNew, doElem)

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

/- Generate `CodeBlock` for `doReturn` which is of the form
   ```
   "return " >> optional termParser
   ```
   `doElems` is only used for sanity checking. -/
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

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

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

mutual
  /- "Concatenate" `c` with `doSeqToCode doElems` -/
  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

  /- Generate `CodeBlock` for `doLetArrow; doElems`
     `doLetArrow` is of the form
     ```
     "let " >> optional "mut " >> (doIdDecl <|> doPatDecl)
     ```
     where
     ```
     def doIdDecl   := leading_parser ident >> optType >> leftArrow >> doElemParser
     def doPatDecl  := leading_parser termParser >> leftArrow >> doElemParser >> optional (" | " >> doElemParser)
     ```
  -/
  partial def doLetArrowToCode (doLetArrow : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let ref     := doLetArrow
    let decl    := doLetArrow[2]
    if decl.getKind == ``Lean.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 action => 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 == ``Lean.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 discr ← $doElem; let mut $pattern:term := discr)
          else
            `(do let discr ← $doElem; let $pattern:term := discr)
        doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems
      else
        if isMutableLet doLetArrow then
          throwError "'mut' is currently not supported in let-decls with 'else' case"
        let contSeq := mkDoSeq doElems.toArray
        let elseSeq := mkSingletonDoSeq optElse[1]
        let auxDo ← `(do let discr ← $doElem; match discr with | $pattern:term => $contSeq | _ => $elseSeq)
        doSeqToCode <| getDoSeqElems (getDoSeq auxDo)
    else
      throwError "unexpected kind of 'do' declaration"

  partial def doLetElseToCode (doLetElse : Syntax) (doElems : List Syntax) : M CodeBlock := do
    -- "let " >> termParser >> " := " >> termParser >> checkColGt >> " | " >> doElemParser
    let pattern := doLetElse[1]
    let val     := doLetElse[3]
    let elseSeq := mkSingletonDoSeq doLetElse[5]
    let contSeq := mkDoSeq doElems.toArray
    let auxDo ← `(do let discr := $val; match discr with | $pattern:term => $contSeq | _ => $elseSeq)
    doSeqToCode <| getDoSeqElems (getDoSeq auxDo)

  /- Generate `CodeBlock` for `doReassignArrow; doElems`
     `doReassignArrow` is of the form
     ```
     (doIdDecl <|> doPatDecl)
     ```
  -/
  partial def doReassignArrowToCode (doReassignArrow : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let ref  := doReassignArrow
    let decl := doReassignArrow[0]
    if decl.getKind == ``Lean.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 == ``Lean.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"

  /- Generate `CodeBlock` for `doIf; doElems`
     `doIf` is of the form
     ```
     "if " >> optIdent >> termParser >> " then " >> doSeq
      >> many (group (try (group (" else " >> " if ")) >> optIdent >> termParser >> " then " >> doSeq))
      >> optional (" else " >> doSeq)
     ```  -/
  partial def doIfToCode (doIf : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let view ← liftMacroM $ 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

  /- Generate `CodeBlock` for `doUnless; doElems`
     `doUnless` is of the form
     ```
     "unless " >> termParser >> "do " >> doSeq
     ```  -/
  partial def doUnlessToCode (doUnless : Syntax) (doElems : List Syntax) : M CodeBlock := withRef doUnless do
    let ref   := doUnless
    let cond  := doUnless[1]
    let doSeq := doUnless[3]
    let body ← doSeqToCode (getDoSeqElems doSeq)
    let unlessCode ← liftMacroM <| mkUnless cond body
    concatWith unlessCode doElems

  /- Generate `CodeBlock` for `doFor; doElems`
     `doFor` is of the form
     ```
     def doForDecl := leading_parser termParser >> " in " >> withForbidden "do" termParser
     def doFor := leading_parser "for " >> sepBy1 doForDecl ", " >> "do " >> doSeq
     ```
  -/
  partial def doForToCode (doFor : Syntax) (doElems : List Syntax) : M CodeBlock := do
    let doForDecls := doFor[1].getSepArgs
    if doForDecls.size > 1 then
      /-
        Expand
        ```
        for x in xs, y in ys do
          body
        ```
        into
        ```
        let s := toStream ys
        for x in xs do
          match Stream.next? s with
          | none => break
          | some (y, s') =>
            s := s'
            body
        ```
      -/
      -- Extract second element
      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 toStreamFn ← withRef ys ``(toStream)
        let auxDo ←
          `(do let mut s := $toStreamFn:ident $ys
               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
      -- semantic no-op that replaces the `uvars`' position information (which all point inside the loop)
      -- with that of the respective mutable declarations outside the loop, which allows the language
      -- server to identify them as conceptually identical variables
      let uvars := uvars.map fun v => ctx.mutableVars.findD v.getId v
      let uvarsTuple ← liftMacroM do mkTuple uvars
      if hasReturn forInBodyCodeBlock.code then
        let forInBody ← liftMacroM <| destructTuple uvars (← `(r)) forInBody
        let forInTerm ←
          if let some h := h? then
            `(for_in'% $(xs) (Prod.mk none $uvarsTuple) fun $x $h r => let r := r.2; $forInBody) -- HDO: Prod
          else
            `(for_in% $(xs) (Prod.mk none $uvarsTuple) fun $x r => let r := r.2; $forInBody) -- HDO: Prod
        let auxDo ← `(do let r ← $forInTerm:term;
                         $uvarsTuple: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
            `(for_in'% $(xs) $uvarsTuple fun $x $h r => $forInBody)
          else
            `(for_in% $(xs) $uvarsTuple fun $x r => $forInBody)
        if doElems.isEmpty then
          let auxDo ← `(do let r ← $forInTerm:term;
                           $uvarsTuple:term := r;
                           Pure.pure (ensure_expected_type% "type mismatch, 'for'" PUnit.unit))
          doSeqToCode <| getDoSeqElems (getDoSeq auxDo)
        else
          let auxDo ← `(do let r ← $forInTerm:term; $uvarsTuple:term := r)
          doSeqToCode <| getDoSeqElems (getDoSeq auxDo) ++ doElems

  /-- Generate `CodeBlock` for `doMatch; doElems` -/
  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 -- Array of `doMatchAlt`
    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

  /--
    Generate `CodeBlock` for `doTry; doElems`
    ```
    def doTry := leading_parser "try " >> doSeq >> many (doCatch <|> doCatchMatch) >> optional doFinally
    def doCatch      := leading_parser "catch " >> binderIdent >> optional (":" >> termParser) >> darrow >> doSeq
    def doCatchMatch := leading_parser "catch " >> doMatchAlts
    def doFinally    := leading_parser "finally " >> doSeq
    ```
  -/
  partial def doTryToCode (doTry : Syntax) (doElems: List Syntax) : M CodeBlock := do
    let ref := doTry
    let tryCode ← doSeqToCode (getDoSeqElems doTry[1])
    let optFinally := doTry[3]
    let catches ← doTry[2].getArgs.mapM fun catchStx => do
      if catchStx.getKind == ``Lean.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 == ``Lean.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 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
        ``(MonadExcept.tryCatch $term (fun $(«catch».x):ident => $catchTerm))
      else
        let type := «catch».optType[1]
        ``(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 {} ToTerm.Kind.regular
        ``(tryFinally $term $finallyTerm)
    let doElemsNew ← liftMacroM <| ToTerm.matchNestedTermResult term uvars a r bc
    doSeqToCode (doElemsNew ++ doElems)

  partial def doSeqToCode : List Syntax → M CodeBlock
    | [] => do liftMacroM mkPureUnitAction
    | doElem::doElems => withIncRecDepth <| withRef doElem do
      checkMaxHeartbeats "'do'-expander"
      match (← liftMacroM <| expandMacro? doElem) with
      | some doElem => doSeqToCode (doElem::doElems)
      | none =>
      match (← liftMacroM <| expandDoIf? doElem) with
      | some doElem => doSeqToCode (doElem::doElems)
      | none =>
        let (liftedDoElems, doElem) ← expandLiftMethod doElem
        if !liftedDoElems.isEmpty then
          doSeqToCode (liftedDoElems ++ [doElem] ++ doElems)
        else
          let ref := doElem
          let concatWithRest (c : CodeBlock) : M CodeBlock := concatWith c doElems
          let k := doElem.getKind
          if k == ``Lean.Parser.Term.doLet then
            let vars ← getDoLetVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> withNewMutableVars vars (isMutableLet doElem) (doSeqToCode doElems)
          else if k == ``Lean.Parser.Term.doHave then
            let vars ← getDoHaveVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> (doSeqToCode doElems)
          else if k == ``Lean.Parser.Term.doLetRec then
            let vars ← getDoLetRecVars doElem
            checkNotShadowingMutable vars
            mkVarDeclCore vars doElem <$> (doSeqToCode doElems)
          else if k == ``Lean.Parser.Term.doReassign then
            let vars ← getDoReassignVars doElem
            checkReassignable vars
            let k ← doSeqToCode doElems
            mkReassignCore vars doElem k
          else if k == ``Lean.Parser.Term.doLetArrow then
            doLetArrowToCode doElem doElems
          else if k == ``Lean.Parser.Term.doLetElse then
            doLetElseToCode doElem doElems
          else if k == ``Lean.Parser.Term.doReassignArrow then
            doReassignArrowToCode doElem doElems
          else if k == ``Lean.Parser.Term.doIf then
            doIfToCode doElem doElems
          else if k == ``Lean.Parser.Term.doUnless then
            doUnlessToCode doElem doElems
          else if k == ``Lean.Parser.Term.doFor then withFreshMacroScope do
            doForToCode doElem doElems
          else if k == ``Lean.Parser.Term.doMatch then
            doMatchToCode doElem doElems
          else if k == ``Lean.Parser.Term.doTry then
            doTryToCode doElem doElems
          else if k == ``Lean.Parser.Term.doBreak then
            ensureInsideFor
            ensureEOS doElems
            return mkBreak ref
          else if k == ``Lean.Parser.Term.doContinue then
            ensureInsideFor
            ensureEOS doElems
            return mkContinue ref
          else if k == ``Lean.Parser.Term.doReturn then
            doReturnToCode doElem doElems
          else if k == ``Lean.Parser.Term.doDbgTrace then
            return mkSeq doElem (← doSeqToCode doElems)
          else if k == ``Lean.Parser.Term.doAssert then
            return mkSeq doElem (← doSeqToCode doElems)
          else if k == ``Lean.Parser.Term.doNested then
            let nestedDoSeq := doElem[1]
            doSeqToCode (getDoSeqElems nestedDoSeq ++ doElems)
          else if k == ``Lean.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}"
end

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

end ToCodeBlock

/- Create a synthetic metavariable `?m` and assign `m` to it.
   We use `?m` to refer to `m` when expanding the `do` notation. -/

-- HDO: Metavariables can't be universe polymorphic, so we instead hack and
-- create a quantified external definition, then return the name of that
-- definition as Syntax instead of ?m
private def mkMonadAlias (m : Expr) : TermElabM Syntax := do
  let levelParams := collectLevelParams {} m |>.params
  let mType ← inferType m
  let name ← mkFreshUserName `_hdo
  let decl := Declaration.defnDecl {
      name := name, levelParams := levelParams.toList, type := mType,
      value := m, hints := ReducibilityHints.opaque,
      safety := DefinitionSafety.unsafe
  }
  ensureNoUnassignedMVars decl
  dbg_trace m
  addAndCompile decl
  return mkIdent name

-- HDO: This elaborates like the normal do command (but with Prod/HBind)
@[termElab «hdo»] def elabHDo : TermElab := fun stx expectedType? => do
  tryPostponeIfNoneOrMVar expectedType?
  let bindInfo ← extractBind expectedType?
  let m ← mkMonadAlias bindInfo.m
  let codeBlock ← ToCodeBlock.run stx m
  let stxNew ← liftMacroM <| ToTerm.run codeBlock.code m
  trace[Elab.do] stxNew
  withMacroExpansion stx stxNew <| elabTermEnsuringType stxNew bindInfo.expectedType

-- HDO: Variation with additional information for testing
@[termElab «hdo_2»] def elabHDo2 : TermElab := fun stx expectedType? => do
  match stx with
  | `(hdo (monad := $stx_monad:term) $stx_seq) =>
      let stx ← `(hdo $stx_seq)

      -- We get an expression for the universe polymorphic monad as parameter
      let monad ← elabTerm stx_monad none
      dbg_trace "Monad:"
      dbg_trace monad
      dbg_trace "Inferred monad type:"
      dbg_trace (← inferType monad)

      tryPostponeIfNoneOrMVar expectedType?
      let bindInfo ← extractBind expectedType?
      let m ← mkMonadAlias monad
      let codeBlock ← ToCodeBlock.run stx m
      let stxNew ← liftMacroM <| ToTerm.run codeBlock.code m
      trace[Elab.do] stxNew
      withMacroExpansion stx stxNew <| elabTermEnsuringType stxNew bindInfo.expectedType

  | _ => throwError "unrecognized syntax for hdo_2"

end HDo

builtin_initialize registerTraceClass `Elab.do

private def toDoElem (newKind : SyntaxNodeKind) : Macro := fun stx => do
  let stx := stx.setKind newKind
  withRef stx `(do $stx:doElem)

@[builtinMacro Lean.Parser.Term.termFor]
def expandTermFor : Macro := toDoElem ``Lean.Parser.Term.doFor

@[builtinMacro Lean.Parser.Term.termTry]
def expandTermTry : Macro := toDoElem ``Lean.Parser.Term.doTry

@[builtinMacro Lean.Parser.Term.termUnless]
def expandTermUnless : Macro := toDoElem ``Lean.Parser.Term.doUnless

@[builtinMacro Lean.Parser.Term.termReturn]
def expandTermReturn : Macro := toDoElem ``Lean.Parser.Term.doReturn

end Lean.Elab.Term

/-
## Tests
-/

set_option trace.Elab.do true

-- Elaborating with `hdo` gives us a term absed on `hBind`
def test_IO: IO Unit := hdo
  IO.println "Lean4"
  IO.println "hBind"
  return ()
#print test_IO

def get_0: Id Nat :=
  pure 0
def get_1: Id ((α: Type) → α → α) :=
  pure @id
def get_any: Id PUnit.{u+1} :=
  pure .unit

example: Id Unit :=
  HBind.hBind get_0 fun n =>
  HBind.hBind get_1 fun f => -- Heterogeneous bind succeeds
  pure ()

-- If we specify the monad as a polymorphic one, elaboration succeeds
-- We hacked the universe name so we have to define it
universe u in
example: Id Unit := hdo (monad := Id.{u})
  let _ ← get_0
  let _ ← get_1
  let _ ← get_any.{2}
  get_any.{0}


import data.cpi.species

namespace cpi

variables {ℍ : Type} {ω : context}

/-- A function to look up names within the environment. -/
@[nolint has_inhabited_instance]
def lookup (ℍ : Type) (ω Γ : context) := ∀ n, reference n ω → species.choices ℍ ω (context.extend n Γ)

/-- Rename a lookup function, embedding the returned species into another
    context.-/
def lookup.rename {Γ Δ} (ρ : name Γ → name Δ) : lookup ℍ ω Γ → lookup ℍ ω Δ
| f n r := species.rename (name.ext ρ) (f n r)

/-- Rewrite lemma for when lookups get expanded incorrectly. -/
lemma lookup.rename.def {Γ Δ} (ρ : name Γ → name Δ) (ℓ : lookup ℍ ω Γ)
  : (λ n r, species.rename (name.ext ρ) (ℓ n r)) = lookup.rename ρ ℓ
  := rfl

lemma lookup.rename.inj {Γ Δ} {ρ : name Γ → name Δ} (inj : function.injective ρ)
  : function.injective (@lookup.rename ℍ ω Γ Δ ρ)
| x y eq := funext $ λ n, funext $ λ r, begin
  have : species.rename (name.ext ρ) (x n r) = species.rename (name.ext ρ) (y n r)
    := congr_fun (congr_fun eq n) r,
  from species.rename.inj (name.ext.inj inj) this,
end

lemma lookup.rename_compose {Γ Δ η} (ρ : name Γ → name Δ) (σ : name Δ → name η)
  : ∀ (ℓ : lookup ℍ ω Γ)
  , lookup.rename σ (lookup.rename ρ ℓ) = lookup.rename (σ ∘ ρ) ℓ
| f := funext $ λ n, funext $ λ r, begin
  simp only [lookup.rename, function.comp],
  rw [species.rename_compose (name.ext ρ) (name.ext σ) (f n r), name.ext_comp],
end

end cpi

#lint-