Syntax Quotation in a Macro (with Antiquotes)

MacroQuoteSyntax.lean

import Lean
open Lean Syntax

instance : Quote SourceInfo :=
  Quote.mk fun info =>
    match info with
    | SourceInfo.original leading pos trailing =>
      mkCApp ``SourceInfo.original #[quote leading, quote pos, quote trailing]
    | SourceInfo.synthetic pos endPos =>
      mkCApp ``SourceInfo.synthetic #[quote pos, quote endPos]
    | SourceInfo.none => 
      mkCIdent ``SourceInfo.none

/-- `C[$(e)]` ~> `let a := e; C[$a]`. Used in the implementation of antiquot splices. -/
partial def floatOutAntiquotTerms : 
  Syntax -> StateT (Syntax -> MacroM Syntax) MacroM Syntax
  | stx@(Syntax.node k args) => do
    if isAntiquot stx && !isEscapedAntiquot stx then
      let e := getAntiquotTerm stx
      if !e.isIdent || !e.getId.isAtomic then
        return ← withFreshMacroScope do
          let a ← `(a)
          modify (fun cont stx => (`(let $a:ident := $e; $stx) : MacroM _))
          stx.setArg 2 a
    Syntax.node k (← args.mapM floatOutAntiquotTerms)
  | stx => pure stx

def getAntiquotationIds (stx : Syntax) : MacroM (Array Syntax) := do
  let mut ids := #[]
  for stx in stx.topDown do
    if (isAntiquot stx || isTokenAntiquot stx) && !isEscapedAntiquot stx then
      let anti := getAntiquotTerm stx
      if anti.isIdent then ids := ids.push anti
      else Macro.throwErrorAt stx "complex antiquotation not allowed here"
  return ids

def getSepFromSplice (splice : Syntax) : Syntax := do
  let Syntax.atom _ sep ← getAntiquotSpliceSuffix splice | unreachable!
  Syntax.mkStrLit (sep.dropRight 1)

partial def mkTuple : Array Syntax -> MacroM Syntax
  | #[]  => `(Unit.unit)
  | #[e] => e
  | es   => do
    let stx ← mkTuple (es.eraseIdx 0)
    `(Prod.mk $(es[0]) $stx)

partial def macroQuoteSyntax : Syntax -> MacroM Syntax
  | Syntax.ident info rawVal val preresolved =>
    mkCApp ``Syntax.ident 
      #[quote info, quote rawVal, quote val, quote preresolved]
  | stx@(Syntax.node k _) =>
    if isAntiquot stx && !isEscapedAntiquot stx then
      getAntiquotTerm stx
    else if isTokenAntiquot stx && !isEscapedAntiquot stx then
      match stx[0] with
      | Syntax.atom info val => 
        `(Syntax.atom 
          (Option.getD (getHeadInfo $(getAntiquotTerm stx)) $(quote info)) 
          $(quote val))
      | _                 => 
        Macro.throwErrorAt stx "expected token"
    else if isAntiquotSuffixSplice stx && !isEscapedAntiquot stx then
      -- splices must occur in a `many` node
      Macro.throwErrorAt stx "unexpected antiquotation splice"
    else if isAntiquotSplice stx && !isEscapedAntiquot stx then
      Macro.throwErrorAt stx "unexpected antiquotation splice"
    else do
      let empty := mkCIdent ``Array.empty
      -- if escaped antiquotation, decrement by one escape level
      let stx := unescapeAntiquot stx
      let args <-
        stx.getArgs.foldlM (fun args arg => do
          if k == nullKind && isAntiquotSuffixSplice arg then
            let antiquot := getAntiquotSuffixSpliceInner arg
            match antiquotSuffixSplice? arg with
            | `optional => 
              `(Array.appendCore $args 
                  (match $(getAntiquotTerm antiquot):term with
                  | some x => Array.empty.push x
                  | none   => Array.empty))
            | `many => 
              `(Array.appendCore $args $(getAntiquotTerm antiquot))
            | `sepBy => 
              `(Array.appendCore $args 
                  (@SepArray.elemsAndSeps 
                    $(getSepFromSplice arg) $(getAntiquotTerm antiquot)))
            | k => 
              Macro.throwErrorAt arg 
                "invalid antiquotation suffix splice kind '{k}'"
          else if k == nullKind && isAntiquotSplice arg then
            let k := antiquotSpliceKind? arg
            let (arg, bindLets) ← floatOutAntiquotTerms arg |>.run pure
            let inner ← (getAntiquotSpliceContents arg).mapM macroQuoteSyntax
            let ids ← getAntiquotationIds arg
            if ids.isEmpty then
              Macro.throwErrorAt stx 
                "antiquotation splice must contain at least one antiquotation"
            let arr ← match k with
              | `optional => `(match $[$ids:ident],* with
                  | $[some $ids:ident],* => $(quote inner)
                  | none                 => Array.empty)
              | _ =>
                let arr ← ids[:ids.size-1].foldrM 
                  (fun id arr => `(Array.zip $id $arr)) ids.back
                `(Array.map (fun $(← mkTuple ids) => $(inner[0])) $arr)
            let arr ←
              if k == `sepBy then
                `(mkSepArray $arr (mkAtom $(getSepFromSplice arg)))
              else arr
            let arr ← bindLets arr
            `(Array.appendCore $args $arr)
          else do
            let arg ← macroQuoteSyntax arg;
            `(Array.push $args $arg)) 
        empty
      mkCApp ``Syntax.node #[quote k, args]
  | Syntax.atom info val =>
    mkCApp ``Syntax.atom #[quote info, quote val]
  | Syntax.missing => 
    Macro.throwUnsupported

macro "static_quote" x:term : term => macroQuoteSyntax x

#check (static_quote (2 + 2 : Nat))