exteneded list comprehension

moonscript/compile.lua

@@ -233,25 +233,37 @@ local compiler_index = {
 	end,
 
 	comprehension = function(self, node)
-		local _, exp, names, iter, when = unpack(node)
-		local insert = ("table.insert(tmp, %s)"):format(self:value(exp))
+		local _, exp, clauses = unpack(node)
+		local insert = { ("table.insert(tmp, %s)"):format(self:value(exp)) }
 
-		if when then
-			insert = {
-				("if %s then"):format(self:value(when)), {
-					insert
-				}, "end"
-			}
+		for i = #clauses,1,-1 do
+			local c = clauses[i]
+
+			if "for" == c[1] then
+				local _, names, iter = unpack(c)
+				insert = {
+					("for %s in %s do"):format(self:name_list(names), self:value(iter)),
+					insert,
+					"end"
+				}
+			elseif "when" == c[1] then
+				local _, when = unpack(c)
+				insert = {
+					("if %s then"):format(self:value(when)),
+					insert,
+					"end"
+				}
+			else
+				error("Unknown clause type :"..tostring(c[1]))
+			end
 		end
 
 		return self:pretty{
-			"(function()", {
-				"local tmp = {}",
-				("for %s in %s do"):format(self:name_list(names), self:value(iter)),
-					type(insert) == "table" and insert or {insert},
-				"end",
-				"return tmp"
-			}, "end)()"
+			"(function()",
+				{ "local tmp = {}", },
+				insert,
+				{ "return tmp" },
+			"end)()"
 		}
 	end,
 

moonscript/parse.lua

@@ -208,11 +208,12 @@ local build_grammar = wrap(function()
 
 		While = key"while" * Exp * key"do"^-1 * Body / mark"while",
 
-		Comprehension = sym"[" * Exp * key"for" * Ct(NameList) *
-			key"in" * Exp *
-			(key"when" * Exp)^-1 *
+		Comprehension = sym"[" * Exp *
+			Ct((key"for" * Ct(NameList) * key"in" * Exp / mark"for") * CompClause^0) *
 			sym"]" / mark"comprehension",
 
+		CompClause = key"for" * Ct(NameList) * key"in" * Exp / mark"for" + key"when" * Exp / mark"when",
+
 		Assign = Ct(AssignableList) * sym"=" * Ct(TableBlock + ExpList) / mark"assign",
 
 		-- we can ignore precedence for now

tests/inputs/lists.moon

@@ -11,3 +11,22 @@ rad = [{a} for a in ipairs {
    1,2,3,4,5,6,
 } when good_number a]
 
+
+[z for z in items for j in list when z > 4]
+
+require "util"
+
+dump = (x) -> print util.dump x
+
+range = (count) ->
+  i = 0
+  return coroutine.wrap ->
+    while i < count
+      coroutine.yield i
+      i = i + 1
+
+dump [x for x in range 10]
+dump [{x, y} for x in range 5 when x > 2 for y in range 5]
+
+things = [x + y for x in range 10 when x > 5 for y in range 10 when y > 7]
+

tests/outputs/lists.lua

@@ -49,4 +49,55 @@ local rad = (function()
     end
   end
   return tmp
+end)();
+(function()
+  local tmp = {}
+  for z in items do
+    for j in list do
+      if z > 4 then
+        table.insert(tmp, z)
+      end
+    end
+  end
+  return tmp
+end)()
+require("util")
+local dump = function(x) print(util.dump(x)) end
+local range = function(count)
+  local i = 0
+  return(coroutine.wrap(function() while i < count do
+    coroutine.yield(i)
+    i = i + 1
+  end end))
+end
+dump((function()
+  local tmp = {}
+  for x in range(10) do
+    table.insert(tmp, x)
+  end
+  return tmp
+end)())
+dump((function()
+  local tmp = {}
+  for x in range(5) do
+    if x > 2 then
+      for y in range(5) do
+        table.insert(tmp, { x, y })
+      end
+    end
+  end
+  return tmp
+end)())
+local things = (function()
+  local tmp = {}
+  for x in range(10) do
+    if x > 5 then
+      for y in range(10) do
+        if y > 7 then
+          table.insert(tmp, x + y)
+        end
+      end
+    end
+  end
+  return tmp
 end)()