Improve enumerate (#16053)

* Improve enumerate

* updates

* add tests

1 files changed, 43 insertions, 10 deletions

diff --git a/lib/std/enumerate.nim b/lib/std/enumerate.nim
index 4620af94d..3def653c8 100644
--- a/lib/std/enumerate.nim
+++ b/lib/std/enumerate.nim
@@ -28,29 +28,62 @@ macro enumerate*(x: ForLoopStmt): untyped {.since: (1, 3).} =
 
     let c = "abcd"
     var d: seq[(int, char)]
-    for i, x in enumerate(97, c):
+    for (i, x) in enumerate(97, c):
       d.add((i, x))
     assert d == @[(97, 'a'), (98, 'b'), (99, 'c'), (100, 'd')]
 
+  template genCounter(x): untyped =
+    # 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.
+    newVarStmt(x, infix(countStart, "-", newLit(1)))
+
+  template genInc(x): untyped =
+    newCall(bindSym"inc", x)
+
   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]
+  if x.len == 3: # single iteration variable
+    if x[0].kind == nnkVarTuple: # for (x, y, ...) in iter
+      result.add genCounter(x[0][0])
+      body.insert(0, genInc(x[0][0]))
+      for i in 1 .. x[0].len-2:
+        newFor.add x[0][i]
+    else:
+      error("Missing second for loop variable") # for x in iter
+  else: # for x, y, ... in iter
+    result.add genCounter(x[0])
+    body.insert(0, genInc(x[0]))
+    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`
+  result = newBlockStmt(result)
+
+when isMainModule:
+  let a = @[1, 3, 5, 7]
+
+  block:
+    var res: seq[(int, int)]
+    for i, x in enumerate(a):
+      res.add (i, x)
+    assert res == @[(0, 1), (1, 3), (2, 5), (3, 7)]
+  block:
+    var res: seq[(int, int)]
+    for (i, x) in enumerate(a.items):
+      res.add (i, x)
+    assert res == @[(0, 1), (1, 3), (2, 5), (3, 7)]
+  block:
+    var res: seq[(int, int)]
+    for i, x in enumerate(3, a):
+      res.add (i, x)
+    assert res == @[(3, 1), (4, 3), (5, 5), (6, 7)]