From 708e0dd35f0bafffc6af3a936b86d77668914c0e Mon Sep 17 00:00:00 2001
From: leaf corcoran <leafot@gmail.com>
Date: Fri, 27 May 2011 23:21:29 -0700
Subject: [PATCH] exteneded list comprehension

---
 moonscript/compile.lua  | 42 +++++++++++++++++++++------------
 moonscript/parse.lua    |  7 +++---
 tests/inputs/lists.moon | 19 +++++++++++++++
 tests/outputs/lists.lua | 51 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 101 insertions(+), 18 deletions(-)

diff --git a/moonscript/compile.lua b/moonscript/compile.lua
index ac7a661..13fb86f 100644
--- a/moonscript/compile.lua
+++ b/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,
 
diff --git a/moonscript/parse.lua b/moonscript/parse.lua
index dd7c202..0071147 100644
--- a/moonscript/parse.lua
+++ b/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
diff --git a/tests/inputs/lists.moon b/tests/inputs/lists.moon
index d415692..e67ac40 100644
--- a/tests/inputs/lists.moon
+++ b/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]
+
diff --git a/tests/outputs/lists.lua b/tests/outputs/lists.lua
index f278f9d..8165ea5 100644
--- a/tests/outputs/lists.lua
+++ b/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)()
\ No newline at end of file