From 11c377c1149a23657a7f0dd897866cb550ade8d1 Mon Sep 17 00:00:00 2001
From: Miran <narimiran@disroot.org>
Date: Tue, 22 Sep 2020 13:08:36 +0200
Subject: [PATCH] add `enumerate` macro (#15297)

* add `enumerate` macro

* address the comments

* put `enumerate` in its own module
---
 doc/manual.rst        | 11 +++++----
 lib/pure/sugar.nim    |  1 +
 lib/std/enumerate.nim | 56 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 5 deletions(-)
 create mode 100644 lib/std/enumerate.nim

diff --git a/doc/manual.rst b/doc/manual.rst
index ada6261b9..8c167bf62 100644
--- a/doc/manual.rst
+++ b/doc/manual.rst
@@ -5508,10 +5508,11 @@ type ``system.ForLoopStmt`` can rewrite the entirety of a ``for`` loop:
 
   macro enumerate(x: ForLoopStmt): untyped =
     expectKind x, nnkForStmt
-    # we strip off the first for loop variable and use
-    # it as an integer counter:
+    # check if the starting count is specified:
+    var countStart = if x[^2].len == 2: newLit(0) else: x[^2][1]
     result = newStmtList()
-    result.add newVarStmt(x[0], newLit(0))
+    # we strip off the first for loop variable and use it as an integer counter:
+    result.add newVarStmt(x[0], countStart)
     var body = x[^1]
     if body.kind != nnkStmtList:
       body = newTree(nnkStmtList, body)
@@ -5520,7 +5521,7 @@ type ``system.ForLoopStmt`` can rewrite the entirety of a ``for`` loop:
     for i in 1..x.len-3:
       newFor.add x[i]
     # transform enumerate(X) to 'X'
-    newFor.add x[^2][1]
+    newFor.add x[^2][^1]
     newFor.add body
     result.add newFor
     # now wrap the whole macro in a block to create a new scope
@@ -5532,7 +5533,7 @@ type ``system.ForLoopStmt`` can rewrite the entirety of a ``for`` loop:
 
   # without wrapping the macro in a block, we'd need to choose different
   # names for `a` and `b` here to avoid redefinition errors
-  for a, b in enumerate([1, 2, 3, 5]):
+  for a, b in enumerate(10, [1, 2, 3, 5]):
     echo a, " ", b
 
 
diff --git a/lib/pure/sugar.nim b/lib/pure/sugar.nim
index 0edcc5697..8bfc7b7ff 100644
--- a/lib/pure/sugar.nim
+++ b/lib/pure/sugar.nim
@@ -335,6 +335,7 @@ macro collect*(init, body: untyped): untyped {.since: (1, 1).} =
       call.add init[i]
   result = newTree(nnkStmtListExpr, newVarStmt(res, call), resBody, res)
 
+
 when isMainModule:
   since (1, 1):
     block dup_with_field:
diff --git a/lib/std/enumerate.nim b/lib/std/enumerate.nim
new file mode 100644
index 000000000..4620af94d
--- /dev/null
+++ b/lib/std/enumerate.nim
@@ -0,0 +1,56 @@
+#
+#
+#            Nim's Runtime Library
+#        (c) Copyright 2020 Nim contributors
+#
+#    See the file "copying.txt", included in this
+#    distribution, for details about the copyright.
+#
+
+## This module implements `enumerate` syntactic sugar based on Nim's
+## macro system.
+
+import std/private/since
+import macros
+
+
+macro enumerate*(x: ForLoopStmt): untyped {.since: (1, 3).} =
+  ## Enumerating iterator for collections.
+  ##
+  ## It yields `(count, value)` tuples (which must be immediately unpacked).
+  ## The default starting count `0` can be manually overridden if needed.
+  runnableExamples:
+    let a = [10, 20, 30]
+    var b: seq[(int, int)]
+    for i, x in enumerate(a):
+      b.add((i, x))
+    assert b == @[(0, 10), (1, 20), (2, 30)]
+
+    let c = "abcd"
+    var d: seq[(int, char)]
+    for i, x in enumerate(97, c):
+      d.add((i, x))
+    assert d == @[(97, 'a'), (98, 'b'), (99, 'c'), (100, 'd')]
+
+  expectKind x, nnkForStmt
+  # check if the starting count is specified:
+  var countStart = if x[^2].len == 2: newLit(0) else: x[^2][1]
+  result = newStmtList()
+  # We strip off the first for loop variable and use it as an integer counter.
+  # We must immediately decrement it by one, because it gets incremented before
+  # the loop body - to be able to use the final expression in other macros.
+  result.add newVarStmt(x[0], infix(countStart, "-", newLit(1)))
+  var body = x[^1]
+  if body.kind != nnkStmtList:
+    body = newTree(nnkStmtList, body)
+  body.insert(0, 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 = quote do:
+    block: `result`