listcomprehensions.nim

import macros

macro lc(init, body): untyped =
  # analyse the body, find the deepest expression 'it' and replace it via
  # 'result.add it'
  let res = genSym(nskVar, "lcResult")

  proc detectForLoopVar(n: NimNode): NimNode =
    if n.kind == nnkForStmt:
      result = n[0]
    else:
      for x in n:
        result = detectForLoopVar(x)
        if result != nil: return result
      return nil

  proc t(n, it, res: NimNode): NimNode =
    case n.kind
    of nnkStmtList, nnkStmtListExpr, nnkBlockStmt, nnkBlockExpr,
       nnkWhileStmt,
       nnkForStmt, nnkIfExpr, nnkIfStmt, nnkTryStmt, nnkCaseStmt,
       nnkElifBranch, nnkElse, nnkElifExpr:
      result = n
      if n.len >= 1:
        result[^1] = t(n[^1], it, res)
    else:
      if n == it:
        template adder(res, it) =
          res.add it
        result = getAst adder(res, it)
      else:
        result = n

  let it = detectForLoopVar(body)
  if it == nil: error("no for loop in body", body)
  result = newTree(nnkStmtListExpr, newVarStmt(res, init), t(body, it, res), res)
  echo repr result

let x = lc(newSeq[int]()):
  for i in [1, 2, 3]:
    if i mod 2 == 0: i

echo x