List Comprehensions

wiki

List comprehensions [1] provide concise syntax for building lists in mathematical set-builder notation. A number of programming languages (e.g. Haskell and Python) provide built-in support for list comprehensions. Lua does not; however, there are are ways to implement it in Lua.

List Comprehensions in Pure Lua with Code Generation

Unlike some other approaches, the following approach implements list comprehensions in pure Lua as a library (without patching or token filters). It uses a trick with dynamic code generation (loadstring) and caching somewhat analogous to ShortAnonymousFunctions.

This library has been incorporated into [Penlight].

Examples

local comp = require 'comprehension' . new()
comp 'x^2 for x' {2,3} --> {2^2,3^2}
comp 'x^2 for _,x in ipairs(_1)' {2,3} --> {2^2,3^2}
comp 'x^2 for x=_1,_2' (2,3) --> {2^2,3^2}

comp 'sum(x^2 for x)' {2,3} --> 2^2+3^2
comp 'max(x*y for x for y if x<4 if y<6)' ({2,3,4}, {4,5,6}) --> 3*5
comp 'table(v,k for k,v in pairs(_1))' {[3]=5, [5]=7} --> {[5]=3, [7]=5}

Examples/Tests

-- comprehension_test.lua
-- test of comprehension.lua

-- Utility function for the test suite.
-- Returns true/false whether arrays a and b
-- are equal.  This does a deep by-value comparison (supports nested arrays).
local function eqarray(a, b)
  if #a ~= #b then return false end
  for i=1,#a do
    if type(a[i]) == 'table' and type(b[i]) == 'table' then
      if not eqarray(a[i], b[i]) then return false end
    else
      if a[i] ~= b[i] then return false end
    end
  end
  return true
end

local comp = require 'comprehension' . new()

-- test of list building
assert(eqarray(comp 'x for x' {}, {}))
assert(eqarray(comp 'x for x' {2,3}, {2,3}))
assert(eqarray(comp 'x^2 for x' {2,3}, {2^2,3^2}))
assert(eqarray(comp 'x for x if x % 2 == 0' {4,5,6,7}, {4,6}))
assert(eqarray(comp '{x,y} for x for y if x>2 if y>4' ({2,3},{4,5}), {{3,5}}))

-- test of table building
local t = comp 'table(x,x+1 for x)' {3,4}
assert(t[3] == 3+1 and t[4] == 4+1)
local t = comp 'table(x,x+y for x for y)' ({3,4}, {2})
assert(t[3] == 3+2 and t[4] == 4+2)
local t = comp 'table(v,k for k,v in pairs(_1))' {[3]=5, [5]=7}
assert(t[5] == 3 and t[7] == 5)

-- test of sum
assert(comp 'sum(x for x)' {} == 0)
assert(comp 'sum(x for x)' {2,3} == 2+3)
assert(comp 'sum(x^2 for x)' {2,3} == 2^2+3^2)
assert(comp 'sum(x*y for x for y)' ({2,3}, {4,5}) == 2*4+2*5+3*4+3*5)
assert(comp 'sum(x^2 for x if x % 2 == 0)' {4,5,6,7} == 4^2+6^2)
assert(comp 'sum(x*y for x for y if x>2 if y>4)' ({2,3}, {4,5}) == 3*5)

-- test of min/max
assert(comp 'min(x for x)' {3,5,2,4} == 2)
assert(comp 'max(x for x)' {3,5,2,4} == 5)

-- test of placeholder parameters --
assert(comp 'sum(x^_1 + _3 for x if x >= _4)' (2, nil, 3, 4, {3,4,5})
       == 4^2+3 + 5^2+3)

-- test of for =
assert(comp 'sum(x^2 for x=2,3)' () == 2^2+3^2)
assert(comp 'sum(x^2 for x=2,6,1+1)' () == 2^2+4^2+6^2)
assert(comp 'sum(x*y*z for x=1,2 for y=3,3 for z)' {5,6} ==
  1*3*5 + 2*3*5 + 1*3*6 + 2*3*6)
assert(comp 'sum(x*y*z for z for x=1,2 for y=3,3)' {5,6} ==
  1*3*5 + 2*3*5 + 1*3*6 + 2*3*6)

-- test of for in
assert(comp 'sum(i*v for i,v in ipairs(_1))' {2,3} == 1*2+2*3)
assert(comp 'sum(i*v for i,v in _1,_2,_3)' (ipairs{2,3}) == 1*2+2*3)

-- test of difficult syntax
assert(eqarray(comp '" x for x " for x' {2}, {' x for x '}))
assert(eqarray(comp 'x --[=[for x\n\n]=] for x' {2}, {2}))
assert(eqarray(comp '(function() for i = 1,1 do return x*2 end end)() for x'
               {2}, {4}))
assert(comp 'sum(("_5" and x)^_1 --[[_6]] for x)' (2, {4,5}) == 4^2 + 5^2)

-- error checking
assert(({pcall(function() comp 'x for __result' end)})[2]
       :find'not contain __ prefix')

-- environment.
-- Note: generated functions are set to the environment of the 'new' call.
assert(5 == (function()
  local env = {d = 5}
  setfenv(1, env)
  local comp = comp.new()
  return comp 'sum(d for x)' {1}
end)());

print 'DONE'

Run-time behavior

To illustrate the run-time characteristics, consider the following code:

comp 'sum(x^2 for x if x % 2 == 0)'

That gets code generated to this Lua function:

local __in1 = ...
local __result = (  0  )
for __idx1 = 1, #__in1 do
  local x = __in1[__idx1]
  if x % 2 == 0 then
    __result = __result + ( __x^2 )
  end
end
return __result

Note that no intermediate lists are built. The code efficiently avoids memory allocations (apart from the allocation of the function itself, but that is done only on the first invocation due to caching/memoization). Also, no global variables are referenced.

Source

Download from [github].

Note: the implementation uses the module LuaBalanced.

--DavidManura

Status

This module is new and likely still has some bugs.

Possible extensions

A simple extension would be to provide a more mathematical (or more Haskell-like) syntax:

assert(comp 'sum { x^2 | x <- ?, x % 2 == 0 }' {2,3,4} == 2^2+4^2)

A compelling extension, as recommended by Greg Fitzgerald, is to implement the generalized list comprehensions proposed by SPJ and Wadler [2]. This provides some clear directions to take this to the next level, and the related work in Microsoft LINQ [3] shows what this could look like in practice.

The "zip" extension to list comprehensions, using the Haskell-like notation in the paper

[ (x,y,z,w) | (x <- xs | y <- ys), (z <- zs | w <- ws) ] ,

would require only small changes. The corresponding Lua function to generate that would be like this:

local __xs, __ys, __zs, __ws = ...
local __ret = {}   -- i.e. $list_init
for __i=1,__math_min(#__xs, #__ys) do
  local x, y = __xs[__i], __ys[__i]
  for __j=1,__math_min(#__zs, #__ws) do
    local z, w = __zs[__j], __ws[__j]
    __ret[#__ret+1] = {x,y,z,w}   -- i.e. $list_accum(__ret, x, y, z, w)
  end
end
return ret

(The "$" notation here is a short-hand for compile-time macros that were used to expand the source.)

Supporting sort or grouped by, e.g. again using notation in the paper

[ the x.name, sum x.deposit | x <- transactions, group by x.name ] ,

could be achieved by generating functions like this:

local __transactions = ...
local __groups1 = {}
local __groups2 = {}
for __i = 1, #__transactions do
  local x = __transactions[__i]
  local __key = ( x.name )  -- i.e. $group_by_key
  __groups1[__key] = ( x.name )
         -- i.e. $accum_the(__groups1[__key], $val1)
  __groups2[__key] = (__groups2[__key] or 0) + ( x.deposit )
         -- i.e. $accum_sum(__groups2[__key], $val2)
end
local __result = {}   -- i.e. $list_init
for __key in __pairs(__groups1) do
  __result[__result+1] = {__groups1[__key], __groups2[__key]}
         -- i.e. $accum_list(__result, __v)
end
return __result

Taking this to its full fruition seems quite achievable (after some research), though it would be a significant extension to this module (any takers?).

Changes

2009-12-01 - fix __call convenience operator not caching (noted by Joshua Jensen on mail list)

MetaLua

List comprehensions have also been implemented in MetaLua (clist):

LuaMacros

List comprehensions have also been implemented in LuaMacros:

MoonScript

Moonscript has list comprehensions -- http://moonscript.org/reference/#comprehensions . Moonscript gets compiled into Lua code.