Macro in for-loop

forloop_macro.nim

import macros, typetraits

#######################################################

static: echo "########################################"
static: echo "### zip Tree ###"
dumpTree:
  iterator zip(a: seq[int], b: seq[bool]): (int,bool) {.noSideEffect.} =
    let len = min(a.len, b.len)

    for i in 0..<len:
      yield (a[i], b[i])

static: echo "########################################"
static: echo "### zip AST Gen ###"
dumpASTGen:
  iterator zip(a: seq[int], b: seq[bool]): (int,bool) {.noSideEffect.} =
    let len = min(a.len, b.len)

    for i in 0..<len:
      yield (a[i], b[i])

#######################################################

static: echo "########################################"
static: echo "### zip macros ###"


macro getSubType*(T: NimNode): untyped =
  # Get the subtype T of an input
  result = getTypeInst(T)[1]

macro zip(args: varargs[typed]): untyped =
  let N = args.len
  assert N > 1, "Error: only 0 or 1 argument passed." &
    "\nZip for Tensors should be called directly " &
    "with all input tensors like so: zip(t1, t2, t3)."

  echo args.treeRepr

  # 1. Initialization
  result = newStmtList()

  # Now we create a `zipImpl` iterator with N arguments

  # 2. Create the parameters: Return type + N arguments
  var zipParams = newSeq[NimNode](N+1)

  # 2.1 Return type
  zipParams[0] = newPar()
  for i in 0 ..< N:
    zipParams[0].add getAST(getSubType(args[i]))

  # 2.2 Parameters
  for i in 0 ..< N:
    let s = newIdentDefs(ident("s" & $i), args[i].getTypeInst)
    zipParams[i+1] = s

  # 3. Body
  var zipBody = newStmtList()

  # 3.1 Check that the length of the seqs are the same
  let arg0 = args[0]
  let size0 = newIdentNode("size0")
  zipBody.add quote do:
    let `size0` = `arg0`.len

  for i, t in args:
    let size_t = newIdentNode("checksize_" & $i)
    let check = quote do:
      let `size_t` = `t`.len
      assert(`size0` == `size_t`, "Zip Macro: argument in position #" & $(`i`) & " has a different length.")
    zipBody.add check

  # 3.2 We create the innermost (s0[i], s1[i], s2[i], ..., s100[i])
  let iter = newIdentNode("i")
  var inner = newPar()
  for arg in args:
    inner.add nnkBracketExpr.newTree(arg, iter)

  zipBody.add nnkForStmt.newTree(
    # for i in 0 ..< size0:
    #   yield (s0[i], s1[i], s2[i], ..., s100[i])
    iter,
    nnkInfix.newTree(
      ident("..<"),
      newIntLitNode(0),
      size0
    ),
    nnkYieldStmt.newTree(
      inner
    )
  )


  # 3.4 Construct the iterator
  let zipName = newIdentNode("zipImpl_" & $N & "_")
  var zipImpl = newProc(
    name = zipName,
    params = zipParams,
    body = zipBody,
    procType = nnkIteratorDef
  )

  # 4. Make it visible
  result.add zipImpl

  # 5. Call the iterator
  var zipCall = newCall(zipName)
  for arg in args:
    zipCall.add arg
  result.add zipCall

  echo "### Zip Impl variadic ###"
  echo zipImpl.treeRepr

macro zipTest(arguments: varargs[untyped]): untyped =

  result = quote do:
    iterator zipZipZip(a: seq[int], b: seq[bool], c: seq[int], d: seq[float]): (int, bool, int, float) =

      let size0 = a.len

      for i in 0..<size0:
        yield (a[i], b[i], c[i], d[i])

    zipZipZip(`arguments`)

for a, b, c, d in zipTest(@[1,10,15], @[false, false, true], @[3,2,1], @[4.0,5.0,6.0]):
  echo (a,b,c,d)