compr.lean

def List.instMonad : Monad List where
  pure x := [x]
  bind := List.flatMap

def List.guard : Bool -> List Unit
  | false => []
  | true => [()]

declare_syntax_cat comprRule

syntax "let" term "<-" term : comprRule
syntax "let" term ":=" term : comprRule
syntax "if" term : comprRule

inductive ComprRule
  | if (cond : Lean.TSyntax `term)
  | bind (pat term : Lean.TSyntax `term) (trivial : Bool)
  | def (pat term : Lean.TSyntax `term) (trivial : Bool)

def parseRule : Lean.TSyntax `comprRule -> Lean.MacroM ComprRule
  | `(comprRule|if $cond)                => return .if cond
  | `(comprRule|let $pat:ident <- $term) => return .bind pat term true
  | `(comprRule|let $pat:term  <- $term) => return .bind pat term false
  | `(comprRule|let $pat:ident := $term) => return .def pat term true
  | `(comprRule|let $pat:term  := $term) => return .def pat term false
  | _ => Lean.Macro.throwUnsupported

def build (t : Lean.TSyntax `term) : List (Lean.TSyntax `comprRule) -> Lean.MacroM (Lean.TSyntax `term)
  | [] => `(return $t)
  | r :: xs => do
    let r <- parseRule r
    match r with
    | .if cond =>             `(List.guard $cond >>= fun  _    => $(<- build t xs))
    | .bind pat term true  => `(term|$term >>= fun        $pat => $(<- build t xs))
    | .bind pat term false => `(term|$term >>= fun      | $pat => $(<- build t xs) | _ => [])
    | .def pat term true   => `(term|match ($term) with | $pat => $(<- build t xs))
    | .def pat term false  => `(term|match ($term) with | $pat => $(<- build t xs) | _ => [])

macro "[" t:term "|" r:comprRule,* "]" : term => do
  let block <- build t r.getElems.toList
  `(have := List.instMonad; $block)

example : [x * x | let x <- [1, 2, 3]] = [1, 4, 9] := by decide
example : [x | let x <- [1, 2, 3], if x % 2 = 0] = [2] := by decide
def ifNN (n : Int) : Option Nat := if n < 0 then none else some n.natAbs
example : [y | let x <- [-1, -2, -3, 1, 2, 3], let .some y := ifNN x] = [1, 2, 3] := by decide