Merge pull request #8949 from nim-lang/araq-for-loop-expressions

For loop expressions

1 files changed, 63 insertions, 0 deletions

diff --git a/tests/macros/tcollect.nim b/tests/macros/tcollect.nim
new file mode 100644
index 000000000..ae28ab61b
--- /dev/null
+++ b/tests/macros/tcollect.nim
@@ -0,0 +1,63 @@
+discard """
+  output: '''@[2, 3, 4, 2, 3, 4, 2, 3, 4, 2, 3, 4, 2, 3, 4]
+@[0, 1, 2, 3]'''
+"""
+
+const data = [1,2,3,4,5,6]
+
+import macros
+
+macro collect(body): untyped =
+  # analyse the body, find the deepest expression 'it' and replace it via
+  # 'result.add it'
+  let res = genSym(nskVar, "collectResult")
+
+  when false:
+    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, res: NimNode): NimNode =
+    case n.kind
+    of nnkStmtList, nnkStmtListExpr, nnkBlockStmt, nnkBlockExpr,
+       nnkWhileStmt,
+       nnkForStmt, nnkIfExpr, nnkIfStmt, nnkTryStmt, nnkCaseStmt,
+       nnkElifBranch, nnkElse, nnkElifExpr:
+      result = copyNimTree(n)
+      if n.len >= 1:
+        result[^1] = t(n[^1], res)
+    else:
+      if true: #n == it:
+        template adder(res, it) =
+          res.add it
+        result = getAst adder(res, n)
+      else:
+        result = n
+
+  when false:
+    let it = detectForLoopVar(body)
+    if it == nil: error("no for loop in body", body)
+
+  let v = newTree(nnkVarSection,
+     newTree(nnkIdentDefs, res, newTree(nnkBracketExpr, bindSym"seq",
+     newCall(bindSym"type", body)), newEmptyNode()))
+
+  result = newTree(nnkStmtListExpr, v, t(body, res), res)
+  #echo repr result
+
+let stuff = collect:
+  var i = -1
+  while i < 4:
+    inc i
+    for it in data:
+      if it < 5 and it > 1:
+        it
+
+echo stuff
+
+echo collect(for i in 0..3: i)