Saturday, July 11, 2009

n-ary zip() and friends in Lua

I've been implementing more list functions in Lua. Source code here has more comments embedded.

The core of three of them (zip, map, filter) is zip_with_helper():


function zip_with_helper(func, result_helper, ...)
local results_l= {}
local args = {...} -- a table of the argument lists
local args_pos = 1 -- position on each of the individual argument lists
local have_args = true

while have_args do
local arg_list = {}
for i = 1, #args do
local a = nil
a = args[i][args_pos]
if a then
arg_list[i] = a
else
have_args = false
break
end
end
if have_args then
result_helper(func, arg_list, results_l)
end
args_pos = args_pos + 1
end

return results_l
end

function zip_helper(func, arg_list, results_l)
table.insert(results_l, arg_list)
end

function zip(...)
return zip_with_helper(nil, zip_helper, ...)
end


On the one hand, it is kind of distressing to see how many lines of code it takes to implement it in Lua compared to Haskell. However, these implementations of map(), filter() and zip() are n-ary functions. So if you add up all the code for Haskell's map(), zipWith() (which is map() using a two-argument function), zipWith3(), zipWith4(), etc. It doesn't look bad at all. In fact, there's a lot of repeated code on the Haskell side. Maybe Template Haskell can take care of that, I don't know.

So with the hard part done, here's the code for map() and filter():



function map_helper(func, arg_list, results_l)
table.insert(results_l, func(unpack(arg_list)))
end

function map(func, ...)
return zip_with_helper(func, map_helper, ...)
end

function filter_helper(func, arg_list, results_l)
local result = func(unpack(arg_list))
if result then
if #arg_list == 1 then
table.insert(results_l, arg_list[1])
else
table.insert(results_l, arg_list)
end
end
end

function filter(func, ...)
return zip_with_helper(func, filter_helper, ...)
end


In retrospect, I think I'll get rid of the named helper functions and used anonymous ones instead... at least for map() and zip().

I've also figured out a n-ary version of unzip() as well:


function unzip(...)
local tables = {...}
if #tables < 1 then
return nil
end
local multi_return = #tables > 1
if not multi_return then
tables = tables[1] -- Given a table of tables, so unzip that.
end
local result_tables = {}
local length = #tables
if length < 1 then
return nil
end
local width = #tables[1]
if width < 1 then
return nil
end

for i = 1, length do
for j = 1, width do
if i == 1 then
result_tables[j] = {}
end
result_tables[j][i] = tables[i][j]
end
end

if multi_return then
return unpack(result_tables)
else
return result_tables
end
end


The fun thing with this is that you can give it three tables as arguments, or one table that holds the other tables. It doesn't matter... up to a couple thousand list entries. So yeah, I don't know how practical that actually is, but it is cool that it is possible at all.

No comments: