Stashing work, blocked by nim-lang/Nim#7737

src/tensor/parallel_loop_fusion/loopfusion_forEach.nim

@@ -0,0 +1,334 @@
+# Copyright (c) 2018 Mamy André-Ratsimbazafy and the Arraymancer contributors
+# Distributed under the Apache v2 License (license terms are at http://www.apache.org/licenses/LICENSE-2.0).
+# This file may not be copied, modified, or distributed except according to those terms.
+
+import
+  macros,
+  ../data_structure, ../init_cpu
+
+proc injectParam(param: NimNode): NimNode =
+  nnkPragmaExpr.newTree(
+      newIdentNode($param),
+      nnkPragma.newTree(ident("inject"))
+    )
+
+proc getSubType*(T: NimNode): NimNode =
+
+  echo getTypeInst(T).treerepr
+  result = getTypeInst(T)[1] # Unfortunetely I get a NimSym ...
+
+proc pop*(tree: var NimNode): NimNode {. compileTime .}=
+  ## varargs[untyped] consumes all arguments so the actual value should be popped
+  ## https://github.com/nim-lang/Nim/issues/5855
+  result = tree[tree.len-1]
+  tree.del(tree.len-1)
+
+macro getOutputSubtype(values: untyped, containers: varargs[typed], loopBody: untyped): untyped =
+  # Get the type of an expression on items of sevaral containers
+
+  # This is the macro equivalent to the following template.
+  # There is no simpler way as `type` cannot return void: https://github.com/nim-lang/Nim/issues/7397
+  #
+  # template test(a, b, ...: seq[int], loopBody: untyped): typedesc =
+  #
+  #   template test_type(): untyped =
+  #     type((
+  #       block:
+  #         var
+  #           x{.inject.}: type(items(a))
+  #           y{.inject.}: type(items(b));
+  #           ...
+  #         loopBody
+  #       ))
+  #
+  #   when compiles(test_type()):
+  #     test_type()
+  #   else:
+  #     void
+
+
+  let N = values.len
+  assert containers.len == N
+
+  echo containers.treerepr
+  echo containers[1].symbol.getImpl.treerepr
+
+  var inject_params = nnkVarSection.newTree()
+
+  for i in 0 ..< N:
+    inject_params.add nnkIdentDefs.newTree(
+      injectParam(values[i]),
+      getSubType(containers[i]),
+      newEmptyNode()
+    )
+
+  let loopBodyType = nnkTypeOfExpr.newTree(
+      nnkBlockStmt.newTree(
+        newEmptyNode(),
+        nnkStmtList.newTree(
+          inject_params,
+          loopBody
+        )
+      )
+  )
+
+  result = quote do:
+    when compiles(`loopBodyType`):
+      # If it's void it fails to compile
+      # Not sure what happens if it fails due to user error.
+      `loopBodyType`
+    else:
+      void
+
+macro generateZip(
+                  zipName: NimNode,
+                  index: untyped,
+                  enumerate: static[bool],
+                  containers: varargs[typed],
+                  mutables: static[seq[int]] # Those are a seq[bool]: https://github.com/nim-lang/Nim/issues/7375
+                  ): typed =
+
+  let N = containers.len
+  assert mutables.len == N
+  assert N > 1, "Error: only 0 or 1 argument passed." &
+    "\nThe zip macros should be called directly " &
+    "with all input sequences like so: zip(s1, s2, s3)."
+
+  # 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()
+  if enumerate:
+    zipParams[0].add getType(int)
+  for i in 0 ..< N:
+    let subt = getSubType(containers[i])
+    if mutables[i] == 1:
+      zipParams[0].add nnkVarTy.newtree(subt)
+    else:
+      zipParams[0].add subt
+
+  # 2.2 Parameters
+  for i in 0 ..< N:
+    let s = newIdentDefs(ident("loopfusion_container" & $i), containers[i].getTypeInst)
+    zipParams[i+1] = s
+
+  # 3. Body
+
+  # 3.1 Check that the length of the seqs are the same
+  let container0 = containers[0]
+  let size0 = newIdentNode("loopfusion_size0")
+  # NOTE: using quote do will fail with an `illegal capture` error, hence build AST manually
+
+  var zipBody =  nnkStmtList.newTree(
+    nnkLetSection.newTree(
+      nnkIdentDefs.newTree(
+        newIdentNode($size0),
+        newEmptyNode(),
+        nnkDotExpr.newTree(
+          newIdentNode($container0),
+          newIdentNode("size")
+        )
+      )
+    )
+  )
+
+  block:
+    var i = 1 # we don't check the size0 of the first container
+    while i < containers.len:
+      let t = containers[i]
+      let check = quote do:
+        assert(`size0` == `t`.size, "forEach macro: tensor " &
+          `t`.astTostr &
+          " in position #" & $(`i`) &
+          " has a different size.")
+      zipBody.add check
+      inc i
+
+  # 3.2 We setup the loop index
+  let iter = if enumerate:
+                newIdentNode($index)
+              else:
+                newIdentNode("loopfusion_index_")
+  let iter_inject = if enumerate: injectParam(iter)
+                    else: iter
+
+  # 3.2 We create the innermost (s0[i], s1[i], s2[i], ..., s100[i])
+  var inner = newPar()
+  if enumerate:
+    inner.add newIdentNode($index)
+  for arg in containers:
+    inner.add nnkBracketExpr.newTree(arg, iter)
+
+  # 3.3 We create the for loop
+  var forLoop = nnkForStmt.newTree()
+  # for i in 0 ..< size0:
+  #   yield (s0[i], s1[i], s2[i], ..., s100[i])
+  #   OR
+  #   yield (i, s0[i], s1[i], s2[i], ..., s100[i])
+
+  forLoop.add iter_inject
+
+  forLoop.add nnkInfix.newTree(
+      ident("..<"),
+      newIntLitNode(0),
+      size0
+    )
+
+  forLoop.add nnkYieldStmt.newTree(
+      inner
+    )
+
+  zipBody.add forLoop
+
+  # 3.4 Construct the iterator
+  var zipImpl = newProc(
+    name = zipName,
+    params = zipParams,
+    body = zipBody,
+    procType = nnkIteratorDef
+  )
+
+  # 4. Make it visible
+  result.add zipImpl
+
+macro forEachImpl[N: static[int]](
+  index: untyped,
+  enumerate: static[bool],
+  values: untyped,
+  containers: varargs[typed],
+  mutables: static[array[N, int]], # Those are a seq[bool]: https://github.com/nim-lang/Nim/issues/7375
+  loopBody: untyped
+  ): untyped =
+
+  assert values.len == N
+  assert containers.len == N
+
+  # 1. Initialization
+  result = newStmtList()
+
+  # 2. Create the idents injected
+  var idents_inject = newPar()
+  for ident in values:
+    idents_inject.add injectParam(ident)
+
+  # 3. Create the index injected if applicable
+  let idx_inject = injectParam(index)
+
+  # 4. Checking the result type and creating put it in temp space if it's
+  #    not void
+  let outType = getAST(getOutputSubtype(values, containers, loopBody))
+  let loopResult = genSym(nskVar, "loopfusion_result_")
+
+  result.add quote do:
+    when not (`outType` is void):
+      var `loopResult`: Tensor[`outType`] = @[]
+
+  # 5. Generate the matching zip iterator
+  let zipName = if enumerate:
+                  genSym(nskIterator,"loopfusion_enumerateZipImpl_" & $N & "_")
+                else:
+                  genSym(nskIterator,"loopfusion_zipImpl_" & $N & "_")
+
+  result.add getAST(generateZip(zipName, index, enumerate, containers, mutables))
+
+  # 6. Creating the call
+  var zipCall = newCall(zipName)
+  containers.copyChildrenTo zipCall
+
+  # 7. For statement/expression
+  var forLoop = nnkForStmt.newTree()
+  if enumerate:
+    forLoop.add idx_inject
+  idents_inject.copyChildrenTo forLoop
+  forLoop.add zipCall
+
+  forLoop.add quote do:
+    when `outType` is void:
+      `loopbody`
+    else:
+      `loopResult`.add `loopBody`
+
+  # 8. Finalize
+  result.add forLoop
+  result.add quote do:
+    when not (`outType` is void):
+      `loopResult`
+
+
+macro forEach*(args: varargs[untyped]): untyped =
+  ## Iterates over a variadic number of sequences or arrays
+
+  ## Example:
+  ##
+  ## let a = [1, 2, 3]
+  ## let b = @[11, 12, 13]
+  ## let c = @[10, 10, 10]
+  ##
+  ## forEach [x, y, z], [a, b, c]:
+  ##   echo (x + y) * z
+
+  # In an untyped context, we can't deal with types at all so we reformat the args
+  # and then pass the new argument to a typed macro
+
+  var params = args
+  var loopBody = params.pop
+
+  var index = getType(int)              # to be replaced with the index variable if applicable
+  var values = nnkBracket.newTree()
+  var containers = nnkArgList.newTree()
+  var mutables: seq[bool] = @[]
+  var N = 0
+  var enumerate = false
+
+  for arg in params:
+    case arg.kind:
+    of nnkIdent:
+      if N == 0:
+        index = arg
+        enumerate = true
+      else:
+        error "Syntax error: argument " & ($arg.kind).substr(3) & " in position #" & $N & " was unexpected."
+    of nnkInfix:
+      if eqIdent(arg[0], "in"):
+        values.add arg[1]
+        if arg[2].kind == nnkVarTy:
+          containers.add arg[2][0]
+          mutables.add true
+        else:
+          containers.add arg[2] # TODO: use an intermediate assignation if it's a result of a proc to avoid calling it multiple time
+          mutables.add false
+      else:
+        error "Syntax error: argument " & ($arg.kind).substr(3) & " in position #" & $N & " was unexpected."
+    else:
+      error "Syntax error: argument " & ($arg.kind).substr(3) & " in position #" & $N & " was unexpected."
+    inc N
+
+  if enumerate:
+    result = quote do:
+      forEachImpl(`index`, true, `values`, `containers`, `mutables`,`loopBody`)
+  else:
+    result = quote do:
+      forEachImpl(`index`, false, `values`, `containers`, `mutables`,`loopBody`)
+
+
+when isMainModule:
+
+  type
+    CustomSeq*[T] = object
+      data*: seq[T]
+
+  # let a = [1, 2, 3].toTensor
+  # let b = [10, 11, 12].toTensor
+
+  # var c = [0, 0, 0].toTensor
+
+  var a, b, c: CustomSeq[int]
+
+  forEach z in var c, x in a, y in b:
+    echo x