Add lc module

10 files changed, 78 insertions, 9 deletions

diff --git a/src/args.nim b/src/args.nim
index fd03ee0..9aca7aa 100644
--- a/src/args.nim
+++ b/src/args.nim
@@ -1,6 +1,6 @@
 import
   options, os, posix, sequtils, sets, strutils, sugar,
-  utils
+  utils, "listcomp"
 
 type
   ArgumentType* {.pure.} = enum
diff --git a/src/aur.nim b/src/aur.nim
index 399183c..9b0628f 100644
--- a/src/aur.nim
+++ b/src/aur.nim
@@ -1,7 +1,7 @@
 import
   json, lists, options, re, sequtils, sets, strutils, sugar, tables,
   package, utils,
-  "wrapper/curl"
+  "wrapper/curl", "listcomp"
 
 type
   AurComment* = tuple[
diff --git a/src/common.nim b/src/common.nim
index 8db3a79..055474e 100644
--- a/src/common.nim
+++ b/src/common.nim
@@ -1,7 +1,7 @@
 import
   options, os, osproc, posix, sequtils, sets, strutils, sugar, tables,
   args, config, format, lists, package, pacman, utils,
-  "wrapper/alpm"
+  "wrapper/alpm", "listcomp"
 
 type
   CacheKind* {.pure.} = enum
diff --git a/src/feature/syncinfo.nim b/src/feature/syncinfo.nim
index 4f758fe..a06c88d 100644
--- a/src/feature/syncinfo.nim
+++ b/src/feature/syncinfo.nim
@@ -2,7 +2,7 @@ import
   options, posix, sequtils, strutils, sugar, tables,
   "../args", "../aur", "../common", "../config", "../format", "../package",
     "../pacman", "../utils",
-  "../wrapper/alpm"
+  "../wrapper/alpm", "../listcomp"
 
 const
   pacmanInfoStrings = [
diff --git a/src/feature/syncinstall.nim b/src/feature/syncinstall.nim
index 7aba043..bb97178 100644
--- a/src/feature/syncinstall.nim
+++ b/src/feature/syncinstall.nim
@@ -2,7 +2,7 @@ import
   algorithm, options, os, posix, sequtils, sets, strutils, sugar, tables,
   "../args", "../aur", "../config", "../common", "../format", "../lists", "../package",
     "../pacman", "../utils",
-  "../wrapper/alpm"
+  "../wrapper/alpm", "../listcomp"
 
 type
   Installed = tuple[
diff --git a/src/format.nim b/src/format.nim
index 2ac39dc..5c8e336 100644
--- a/src/format.nim
+++ b/src/format.nim
@@ -1,6 +1,6 @@
 import
   macros, options, posix, sequtils, strutils, sugar, times, unicode,
-  utils
+  utils, "listcomp"
 
 type
   PackageLineFormat* = tuple[
diff --git a/src/listcomp.nim b/src/listcomp.nim
new file mode 100644
index 0000000..c9f8dc0
--- /dev/null
+++ b/src/listcomp.nim
@@ -0,0 +1,69 @@
+import macros
+
+type ListComprehension = object
+var lc* : ListComprehension
+
+template `|`*(lc: ListComprehension, comp: untyped): untyped = lc
+
+macro `[]`*(lc: ListComprehension, comp, typ: untyped): untyped =
+  ## List comprehension, returns a sequence. `comp` is the actual list
+  ## comprehension, for example ``x | (x <- 1..10, x mod 2 == 0)``. `typ` is
+  ## the type that will be stored inside the result seq.
+  ##
+  ## .. code-block:: nim
+  ##
+  ##   echo lc[x | (x <- 1..10, x mod 2 == 0), int]
+  ##
+  ##   const n = 20
+  ##   echo lc[(x,y,z) | (x <- 1..n, y <- x..n, z <- y..n, x*x + y*y == z*z),
+  ##           tuple[a,b,c: int]]
+
+  expectLen(comp, 3)
+  expectKind(comp, nnkInfix)
+  assert($comp[0] == "|")
+
+  result = newCall(
+    newDotExpr(
+      newIdentNode("result"),
+      newIdentNode("add")),
+    comp[1])
+
+  for i in countdown(comp[2].len-1, 0):
+    let x = comp[2][i]
+    expectMinLen(x, 1)
+    if x[0].kind == nnkIdent and x[0].strVal == "<-":
+      expectLen(x, 3)
+      result = newNimNode(nnkForStmt).add(x[1], x[2], result)
+    else:
+      result = newIfStmt((x, result))
+
+  result = newNimNode(nnkCall).add(
+    newNimNode(nnkPar).add(
+      newNimNode(nnkLambda).add(
+        newEmptyNode(),
+        newEmptyNode(),
+        newEmptyNode(),
+        newNimNode(nnkFormalParams).add(
+          newNimNode(nnkBracketExpr).add(
+            newIdentNode("seq"),
+            typ)),
+        newEmptyNode(),
+        newEmptyNode(),
+        newStmtList(
+          newAssignment(
+            newIdentNode("result"),
+            newNimNode(nnkPrefix).add(
+              newIdentNode("@"),
+              newNimNode(nnkBracket))),
+          result))))
+
+
+when isMainModule:
+  var a = lc[x | (x <- 1..10, x mod 2 == 0), int]
+  assert a == @[2, 4, 6, 8, 10]
+
+  const n = 20
+  var b = lc[(x,y,z) | (x <- 1..n, y <- x..n, z <- y..n, x*x + y*y == z*z),
+             tuple[a,b,c: int]]
+  assert b == @[(a: 3, b: 4, c: 5), (a: 5, b: 12, c: 13), (a: 6, b: 8, c: 10),
+(a: 8, b: 15, c: 17), (a: 9, b: 12, c: 15), (a: 12, b: 16, c: 20)]
diff --git a/src/main.nim b/src/main.nim
index 92e8887..3de4a65 100644
--- a/src/main.nim
+++ b/src/main.nim
@@ -1,6 +1,6 @@
 import
   options, os, posix, re, sequtils, strutils, sugar,
-  args, config, format, pacman, utils
+  args, config, format, pacman, utils, "listcomp"
 
 import
   "feature/localquery",
diff --git a/src/package.nim b/src/package.nim
index 5b87184..612f98f 100644
--- a/src/package.nim
+++ b/src/package.nim
@@ -1,6 +1,6 @@
 import
   options, os, re, sequtils, sets, strutils, sugar, tables, utils,
-  "wrapper/alpm"
+  "wrapper/alpm", "listcomp"
 
 type
   ConstraintOperation* {.pure.} = enum
diff --git a/src/pacman.nim b/src/pacman.nim
index 857581f..520a67e 100644
--- a/src/pacman.nim
+++ b/src/pacman.nim
@@ -1,6 +1,6 @@
 import
   macros, options, posix, sequtils, sets, strutils, sugar, tables,
-  args, config, utils
+  args, config, utils, "listcomp"
 
 type
   OpGroup* {.pure.} = enum