Basic do notation

open Lean

declare_syntax_cat stmt
syntax "do'" stmt : term

-- Prevent `if ...` from being parsed as a term
syntax (priority := low) term : stmt
syntax "let" ident "←" stmt:1 ";" stmt : stmt
macro "{" 
s: TSyntax `stmt
s:stmt "}" : stmt => `($
s: TSyntax `stmt
s)

/-
  Remark: we annotate `macro`s and `macro_rules` with their corresponding
  translation/abbreviation id from the paper.
-/

syntax "d!" stmt : term  -- corresponds to `D(s)`

macro_rules
  | `(do' $
s: TSyntax `stmt
s) => `(d! $
s: TSyntax `stmt
s)  -- (1)

-- helper function; see usage below
def 
expandStmt: TSyntax `stmt → MacroM (TSyntax `stmt)
expandStmt (
s: TSyntax `stmt
s : 
TSyntax: SyntaxNodeKinds → Type
TSyntax 
`stmt: Name
`stmt) : 
MacroM: Type → Type
MacroM (
TSyntax: SyntaxNodeKinds → Type
TSyntax 
`stmt: Name
`stmt) := do
  let 
s': Syntax
s' ← 
expandMacros: Syntax → (optParam (SyntaxNodeKind → Bool) fun k => k != `Lean.Parser.Term.byTactic) → MacroM Syntax
expandMacros 
s: TSyntax `stmt
s
  if 
s: TSyntax `stmt
s == 
s': Syntax
s' then
    
Macro.throwUnsupported: {α : Type} → MacroM α
Macro.throwUnsupported
  else
    -- There is no static guarantee that `expandMacros` stays in the `stmt` category,
    -- but it is true for all our macros
    return 
TSyntax.mk: {ks : SyntaxNodeKinds} → Syntax → TSyntax ks
TSyntax.mk 
s': Syntax
s'

macro_rules
  | `(d! $
e: TSyntax `term
e:term) => `($
e: TSyntax `term
e)                                       -- (D1)
  | `(d! let $
x: TSyntax `ident
x ← $
s: TSyntax `stmt
s; $
s': TSyntax `stmt
s') => `((d! $
s: TSyntax `stmt
s) >>= fun $
x: TSyntax `ident
x => (d! $
s': TSyntax `stmt
s'))  -- (D2)
  | `(d! $
s: TSyntax `stmt
s) => do
    -- fallback rule: try to expand abbreviation
    let 
s': TSyntax `stmt
s' ← 
expandStmt: TSyntax `stmt → MacroM (TSyntax `stmt)
expandStmt 
s: TSyntax `stmt
s
    `(d! $
s': TSyntax `stmt
s')

macro "let" 
x: TSyntax `ident
x:ident ":=" 
e: TSyntax `term
e:term ";" 
s: TSyntax `stmt
s:stmt : stmt => `(let $
x: TSyntax `ident
x ← pure $
e: TSyntax `term
e; $
s: TSyntax `stmt
s)  -- (A1)
-- priority `0` prevents `;` from being used in trailing contexts without braces (see e.g. `:1` above)
macro:0 
s₁: TSyntax `stmt
s₁:stmt ";" 
s₂: TSyntax `stmt
s₂:stmt : stmt => `(let x ← $
s₁: TSyntax `stmt
s₁; $
s₂: TSyntax `stmt
s₂)                     -- (A2)