staticfor: compile-time loop unrolling

stew/staticfor.nim

@@ -0,0 +1,43 @@
+# stew
+# Copyright (c) 2024 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
+#    http://www.apache.org/licenses/LICENSE-2.0)
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT) or
+#    http://opensource.org/licenses/MIT)
+# at your option. This file may not be copied, modified, or distributed except
+# according to those terms.
+
+import std/macros
+
+proc replaceNodes(ast: NimNode, what: NimNode, by: NimNode): NimNode =
+  # Replace "what" ident node by "by"
+  proc inspect(node: NimNode): NimNode =
+    case node.kind:
+    of {nnkIdent, nnkSym}:
+      if node.eqIdent(what):
+        by
+      else:
+        node
+    of nnkEmpty, nnkLiterals:
+      node
+    else:
+      let rTree = newNimNode(node.kind, lineInfoFrom = node)
+      for child in node:
+        rTree.add inspect(child)
+      rTree
+  inspect(ast)
+
+macro staticFor*(idx: untyped{nkIdent}, slice: static Slice[int], body: untyped): untyped =
+  ## Unrolled `for` loop over the given range:
+  ##
+  ## ```nim
+  ## staticFor(i, 0..<2):
+  ##   echo default(array[i, byte])
+  ## ```
+  result = newNimNode(nnkStmtList, lineInfoFrom = body)
+  for i in slice:
+    result.add nnkBlockStmt.newTree(
+      ident(":staticFor" & $idx & $i),
+      body.replaceNodes(idx, newLit i)
+    )

tests/all_tests.nim

@@ -32,6 +32,7 @@ import
   test_ptrops,
   test_sequtils2,
   test_sets,
+  test_staticfor,
   test_strformat,
   test_templateutils,
   test_winacl

tests/test_staticfor.nim

@@ -0,0 +1,17 @@
+{.used.}
+
+import unittest2, ../stew/staticfor
+
+suite "staticfor":
+  test "basics":
+    var
+      a = 0
+      b = 0
+
+    for i in 0..10:
+      a += i
+
+    staticFor i, 0..10:
+      b += default(array[i, byte]).len
+
+    check: a == b