zipmacro.nim

import macros
{.experimental: "forLoopMacros".}

macro zip*(x: ForLoopStmt): untyped =
   expectKind x, nnkForStmt
   result = newStmtList()
   echo x.repr
   var args: seq[NimNode]
   for i in 1 ..< x[^2].len:
      args.add x[^2][i]
   var iter: seq[NimNode]
   if x.len == 3:
      if x[0].kind == nnkVarTuple:
         for i in 0 .. x[0].len-2:
            iter.add x[0][i]
         if args.len != iter.len:
            error("Not as many iterator variables and zip arguments")
      else:
         iter.add x[0]
   else:
      for i in 0 .. x.len-3:
         iter.add x[i]
      if args.len != iter.len:
         error("Not as many iterator variables and zip arguments")
   let mLen = genSym(nskLet, "m")
   let minC = newCall(bindSym"min", newNimNode(nnkBracket))
   let indVar = genSym(nskForVar, "i")
   for i in 0 ..< args.len:
      minC[1].add newCall("len", args[i])
   result.add newLetStmt(mLen, minC)
   let letSec = newNimNode(nnkLetSection)
   for i in 0 ..< iter.len:
      var value = nnkBracketExpr.newTree(args[i], indVar)
      if iter.len != args.len:
         value = nnkTupleConstr.newNimNode
         for j in 0 ..< args.len:
            value.add nnkBracketExpr.newTree(args[j], indVar)
      letSec.add newIdentDefs(iter[i], newEmptyNode(), value)
   var body = x[^1]
   if body.kind != nnkStmtList:
      body = newTree(nnkStmtList, body)
   body.insert(0, letSec)
   let newFor = nnkForStmt.newTree(indVar,
      nnkInfix.newTree(bindSym"..<", newLit(0), mLen))
   newFor.add body
   result.add newFor

macro enumerate*(x: ForLoopStmt): untyped =
   expectKind x, nnkForStmt
   # we strip off the first for loop variable and use
   # it as an integer counter:
   result = newStmtList()
   result.add newVarStmt(x[0], newLit(0))
   var body = x[^1]
   if body.kind != nnkStmtList:
      body = newTree(nnkStmtList, body)
   body.add newCall(bindSym"inc", x[0])
   var newFor = newTree(nnkForStmt)
   for i in 1..x.len-3:
      newFor.add x[i]
   # transform enumerate(X) to 'X'
   newFor.add x[^2][1]
   newFor.add body
   result.add newFor
   # now wrap the whole macro in a block to create a new scope
   result = newBlockStmt(result)

when isMainModule:
   let
      a = [1, 2, 3, 5]
      b = @[1, 2, 3, 5]

   for x, y in zip(a, b):
      echo x, " ", y

   for (x, y, z) in zip(a, b, [1, 2, 3, 5, 4]):
      echo x, " ", y, " ", z

   for i, (x, y) in enumerate(zip(a, b)):
      echo x, " ", y

   for i, a in enumerate(zip(a, b)):
      echo i, " ", a