added key,value table comprehensions

Author: leafo

docs/reference.md

@@ -281,10 +281,16 @@ same as the variable names, then the `:` prefix operator can be used:
 
     print_table :hair, :height
 
-## Table Comprehensions
+## Comprehensions
 
-Table comprehensions provide a quick way to iterate over a table's values while
-applying a statement and accumulating the result.
+Compiling provide a convenient syntax for constructing a new table by iterating
+over some existing object and applying an expression to its values. There are
+two kinds of comprehensions: list comprehensions and table comprehensions. They
+both produce Lua tables; _list comprehensions_ accumulate values into an
+array-like table, and _table comprehensions_ let you set both the key and the
+value on each iteration.
+
+### List Comprehensions
 
 The following creates a copy of the `items` table but with all the values
 doubled.
@@ -312,6 +318,33 @@ Using multiple `for` clauses is the same as using nested loops:
 
     points = [{x,y} for x in *x_coords for y in *y_coords]
 
+### Table Comprehensions
+
+The syntax for table comprehensions is very similar, differing by using `{` and
+`}` and taking two values from each iteration.
+
+This example copies the key-value table `thing`:
+
+    thing = {
+      color: "red"
+      name: "fast"
+      width: 123
+    }
+
+    thing_copy = {k,v for k,v in pairs thing}
+
+Table comprehensions, like list comprehensions, also support multiple `for` and
+`when` clauses`. In this example we use a `where` clause to prevent the value
+associated with the `color` key from being copied.
+
+    no_color = {k,v for k,v in pairs thing when k != "color"}
+
+The `*` operator is also supported. Here we create a square root look up table
+for a few numbers.
+
+    numbers = {1,2,3,4}
+    sqrts = {i, math.sqrt i for i in *numbers}
+
 ### Slicing
 
 A special syntax is provided to restrict the items that are iterated over when

moonscript/transform.lua

@@ -206,6 +206,35 @@ Transformer = (function()
   _base_0.__class = _class_0
   return _class_0
 end)()
+local construct_comprehension
+construct_comprehension = function(inner, clauses)
+  local current_stms = inner
+  for _, clause in reversed(clauses) do
+    local t = clause[1]
+    if t == "for" then
+      local _, names, iter = unpack(clause)
+      current_stms = {
+        "foreach",
+        names,
+        iter,
+        current_stms
+      }
+    elseif t == "when" then
+      local _, cond = unpack(clause)
+      current_stms = {
+        "if",
+        cond,
+        current_stms
+      }
+    else
+      current_stms = error("Unknown comprehension clause: " .. t)
+    end
+    current_stms = {
+      current_stms
+    }
+  end
+  return current_stms[1]
+end
 Statement = Transformer({
   assign = function(self, node)
     local _, names, values = unpack(node)
@@ -367,34 +396,7 @@ Statement = Transformer({
         exp
       }
     end
-    local current_stms = action(exp)
-    for _, clause in reversed(clauses) do
-      local t = clause[1]
-      if t == "for" then
-        local names, iter
-        _, names, iter = unpack(clause)
-        current_stms = {
-          "foreach",
-          names,
-          iter,
-          current_stms
-        }
-      elseif t == "when" then
-        local cond
-        _, cond = unpack(clause)
-        current_stms = {
-          "if",
-          cond,
-          current_stms
-        }
-      else
-        current_stms = error("Unknown comprehension clause: " .. t)
-      end
-      current_stms = {
-        current_stms
-      }
-    end
-    return current_stms[1]
+    return construct_comprehension(action(exp), clauses)
   end,
   ["if"] = function(self, node, ret)
     if ret then
@@ -820,6 +822,25 @@ Value = Transformer({
     end)
     return a:wrap(node)
   end,
+  tblcomprehension = function(self, node)
+    local _, key_exp, value_exp, clauses = unpack(node)
+    local accum = NameProxy("tbl")
+    local dest = build.chain({
+      base = accum,
+      {
+        "index",
+        key_exp
+      }
+    })
+    local inner = build.assign_one(dest, value_exp)
+    return build.block_exp({
+      build.assign_one(accum, build.table()),
+      construct_comprehension({
+        inner
+      }, clauses),
+      accum
+    })
+  end,
   fndef = function(self, node)
     smart_node(node)
     node.body = apply_to_last(node.body, implicitly_return(self))

moonscript/transform.moon

@@ -108,6 +108,22 @@ class Transformer
   can_transform: (node) =>
     @transformers[ntype node] != nil
 
+construct_comprehension = (inner, clauses) ->
+  current_stms = inner
+  for _, clause in reversed clauses
+    t = clause[1]
+    current_stms = if t == "for"
+      _, names, iter = unpack clause
+      {"foreach", names, iter, current_stms}
+    elseif t == "when"
+      _, cond = unpack clause
+      {"if", cond, current_stms}
+    else
+      error "Unknown comprehension clause: "..t
+    current_stms = {current_stms}
+
+  current_stms[1]
+
 Statement = Transformer {
   assign: (node) =>
     _, names, values = unpack node
@@ -185,23 +201,8 @@ Statement = Transformer {
 
   comprehension: (node, action) =>
     _, exp, clauses = unpack node
-
     action = action or (exp) -> {exp}
-
-    current_stms = action exp
-    for _, clause in reversed clauses
-      t = clause[1]
-      current_stms = if t == "for"
-        _, names, iter = unpack clause
-        {"foreach", names, iter, current_stms}
-      elseif t == "when"
-        _, cond = unpack clause
-        {"if", cond, current_stms}
-      else
-        error "Unknown comprehension clause: "..t
-      current_stms = {current_stms}
-
-    current_stms[1]
+    construct_comprehension action(exp), clauses
 
   -- handle cascading return decorator
   if: (node, ret) =>
@@ -428,6 +429,7 @@ class Accumulator
 default_accumulator = (node) =>
   Accumulator!\convert node
 
+
 implicitly_return = (scope) ->
   fn = (stm) ->
     t = ntype stm
@@ -451,6 +453,19 @@ Value = Transformer {
       a\mutate_body {exp}, false
     a\wrap node
 
+  tblcomprehension: (node) =>
+    _, key_exp, value_exp, clauses = unpack node
+
+    accum = NameProxy "tbl"
+    dest = build.chain { base: accum, {"index", key_exp} }
+    inner = build.assign_one dest, value_exp
+
+    build.block_exp {
+      build.assign_one accum, build.table!
+      construct_comprehension {inner}, clauses
+      accum
+    }
+
   fndef: (node) =>
     smart_node node
     node.body = apply_to_last node.body, implicitly_return self

moonscript/types.lua

@@ -173,6 +173,9 @@ build = setmetatable({
     })
   end,
   table = function(tbl)
+    if tbl == nil then
+      tbl = { }
+    end
     return {
       "table",
       tbl

moonscript/types.moon

@@ -94,7 +94,7 @@ build = setmetatable {
       names: {name}
       values: {value}
     }
-  table: (tbl) ->
+  table: (tbl={}) ->
     {"table", tbl}
   block_exp: (body) ->
     {"block_exp", body}

tests/inputs/comprehension.moon

@@ -0,0 +1,9 @@
+
+items = {1,2,3,4,5,6}
+out = {k,k*2 for k in items}
+
+
+x = hello: "world", okay: 2323
+
+copy  = {k,v for k,v in pairs x when k != "okay"}
+

tests/outputs/comprehension.lua

@@ -0,0 +1,28 @@
+local items = {
+  1,
+  2,
+  3,
+  4,
+  5,
+  6
+}
+local out = (function()
+  local _tbl_0 = { }
+  for k in items do
+    _tbl_0[k] = k * 2
+  end
+  return _tbl_0
+end)()
+local x = {
+  hello = "world",
+  okay = 2323
+}
+local copy = (function()
+  local _tbl_0 = { }
+  for k, v in pairs(x) do
+    if k ~= "okay" then
+      _tbl_0[k] = v
+    end
+  end
+  return _tbl_0
+end)()