Array library for lua

array.lua

-- A simple array-manipulation library for Lua.
-- The natural question is: why bother writing this?
-- After all, it's very short, and anyone reasonably familiar with
-- normal array-manipulation libraries could replicate it in a few
-- hours.
-- But despite this, such a library does not exist in the wild, and
-- the existing counterparts are bad. For instance, penlight.List
-- lacks foldr(), and includes a bunch of things that it shouldn't,
-- like t:put(x), which is equivalent to t:insert(1, x), but also
-- prevents anyone reading your code from knowing what it does.
-- Likewise, the other "array.lua" does not provide a metatable
-- and lacks mapv() (aka zipWith)
-- The goal is to provide convenient methods that most people
-- familiar with common data structure manipulations would already
-- recognize and understand, not to build a very sophisticated lib.
-- So without further ado...

local array = {__actions = {}}
array.__actions.__index = array

-- don't mess with perfection!
array.insert = table.insert
array.remove = table.remove
array.sort = table.sort
array.concat = table.concat

-- basic LIFO stack functions
function array:push(_1, ...)
  if _1 then
    self[#self+1] = _1
    return self:push(...)
  else
    return self
  end
end

function array:push_arr(arr)
  for i = 1, #arr do
    self[#self+1] = arr[i]
  end
  return self
end

function array:peek(n)
  n = n or 1 -- zero is true in lua (returns nil)
  local len = #self
  local function do_peek(n, offset)
    if n > 0 and offset < len then
      return self[len - offset], do_peek(n-1, offset+1)
    end
  end
  return do_peek(n, 0)
end

function array:peek_arr(n)
  n = n or 1
  local res = array{}
  local len = #self
  for i = 1, math.min(n, len) do
    res[i] = self[len + 1 - i]
  end
  return res
end

function array:pop(n)
  n = n or 1
  if n > 0 then
    v = self:remove()
    if v then
      return v, self:pop(n-1)
    end
  end
end

function array:pop_arr(n)
  n = n or 1
  local res = array{}
  for i = 1, math.min(n, #self) do
    res[i] = self:remove()
  end
  return res
end

-- one might reasonably inquire: "why not write deep_copy()"?
-- after all, people often want something like deep_copy().
-- but for an array library, it is not clear what this should
-- mean. do we deep copy only internal arrays, or their hash
-- structures too? do we deep copy userdata? how about closures?
function array:copy()
  local res = array{}
  for i = 1, #self do
    res[i] = self[i]
  end
  return res
end

-- there is no good reason to rename "fold" to "reduce",
-- since neither is descriptive enough for anyone who is
-- not already familiar with these functions
function array:foldl(fn, arg)
  local res = self[1]
  if arg then res = fn(arg, res) end
  for i = 2, #self do
    res = fn(res, self[i])
  end
  return res
end

function array:foldr(fn, arg)
  local res = self[#self]
  if arg then res = fn(res, arg) end
  for i = #self - 1, 1, -1 do
    res = fn(self[i], res)
  end
  return res
end

-- a little-known feature of Lua is its support for
-- "true" enums, which can only be matched by actually
-- referencing the values. this prevents confusion.
-- fear not, table equality is fast (compares pointers).
local function enum()
  local res = {}
  setmetatable(res, {__newindex = function() error"don't do that!" end})
  return res
end

array.SEPARATE = false
array.IN_PLACE = enum()
array.SINGLE_COPY = enum()

function array:map(fn, option)
  local res = array{}
  if option == array.IN_PLACE then
    res = self
  elseif option then
    error"invalid option to map()"
  end
  for i = 1, #self do
    res[i] = fn(self[i])
  end
  return res
end

function array:map_indexed(fn, option)
  local res = array{}
  if option == array.IN_PLACE then
    res = self
  elseif option then
    error"invalid option to map_indexed()"
  end
  for i = 1, #self do
    res[i] = fn(i, self[i])
  end
  return res
end

--note that filter() is "stable", while split() may not be
function array:filter(fn, option)
  if option == array.IN_PLACE then
    j = 1
    for i = 1, #self do
      local v = self[i]
      if fn(v) then
        self[j] = v
        j = j + 1
      end
    end
    for i = #self, j+1, -1 do
      self[i] = nil
    end
    return self
  elseif not option then
    local res = array{}
    for i = 1, #self do
      local v = self[i]
      if fn(v) then
        res[#res+1] = v
      end
    end
    return res
  else
    error"invalid option to filter()"
  end
end

function array:split(fn, option)
  if not option then
    local pass, fail = array{}, array{}
    for i = 1, #self do
      local v = self[i]
      if fn(v) then
        pass[#pass+1] = v
      else
        fail[#fail+1] = v
      end
    end
    return pass, fail
  elseif option == array.IN_PLACE then
    local first, last = 1, #self
    repeat
      if not fn(first) then
        self[first], self[last] = self[last], self[first]
        last = last - 1
      else
        first = first + 1
      end
    until first >= last
    return self
  elseif option == array.SINGLE_COPY then
    local first, last = 1, #self
    local res = array{}
    for i = 1, #self do
      local v = self[i]
      if fn(v) then
        res[first] = self[i]
        first = first + 1
      else
        res[last] = self[i]
        last = last - 1
      end
    end
    return res
  else
    error"invalid option to split()"
  end
end

-- these functions mimic the lua string library, which
-- accounts for the weird behavior of negative indices
function array:sub(start, stop)
  stop = stop or #self
  start = start < 0 and #self + 1 - start or start
  stop = stop < 0 and #self + 1 - stop or stop
  local res = {}
  for k = start, stop, 1 do
    res[#res+1] = self[k]
  end
  return res
end

function array:reverse()
  local res = array{}
  for i = #self, 1, -1 do
    res[#res+1] = self[i]
  end
  return res
end

function array:rep(n)
  local res = array{}
  for i = 1, n do
    for j = 1, #self do
      res[#res+1] = self[j]
    end
  end
  return res
end

-- for mapv(), we use an auxiliary zip list of arrays
-- it is not intended to create the zip list yourself.
-- zipping only takes as many arguments as the *shortest*
-- array provided to zip().
-- usage: arr1:zip(arr2, arr3):mapv(fn) [noncommutative]
local ziparr = {__actions = {}}
ziparr.__actions.__index = ziparr

setmetatable(ziparr, 
  {__call = function(self, t)
      setmetatable(t, self.__actions)
      t.min_len = inf
      for i = 1, #t do
        local len = #(t[i])
        if t.min_len > len then t.min_len = len end
      end
      return t
    end})

function ziparr:push(_1, ...)
  if _1 then
    local len = #_1
    if self.min_len > len then self.min_len = len end
    self[#self+1] = _1
    return self:push(...)
  else
    return self
  end
end

function ziparr:push_zip(t)
  if not t.min_len then
    for i = 1, #t do
      local len = #(t[i])
      if self.min_len > len then self.min_len = len end
      self[#self+1] = t[i]
    end
  else
    self.min_len = math.min(self.min_len, t.min_len)
    for i = 1, #t do
      self[#self+1] = t[i]
    end
  end
  return self
end

function ziparr:insert(pos, val)
  table.insert(self, pos, val)
  local len = #(val or pos)
  if self.min_len > len then self.min_len = len end
end

function ziparr:mapv(fn)
  local function get_args(row, col)
    if col < #self then
      return self[row][col], get_args(row, col+1)
    else
      return self[row][col]
    end
  end
  local res = array{}
  for i = 1, self.min_len do
    res[#res+1] = fn(get_args(i, 1))
  end
  return res
end

array.FLATTEN = false
array.TRANSPOSE = enum()
-- array.INTERLEAVE = enum() (maybe? :p)

function ziparr:combine(option)
  local res = array{}
  if not option then
    for i = 1, #self do
      local arr = self[i]
      for j = 1, #arr do
        res[#res+1] = arr[j]
      end
    end
    return res
  elseif option == array.TRANSPOSE then
    for i = 1, self.min_len do
      res[i] = array{}
      for j = 1, #self do
        res[i][j] = self[j][i]
      end
    end
    return res
  else
    error"bad option to combine()"
  end
end

function array:zip(...)
  return ziparr{self, ...}
end

function array:to_zip()
  return ziparr(self)
end

function ziparr:to_array()
  return array(self)
end
    
local function get_concat_fn(colltype, pushfn)
  return function(self, obj)
    local res = colltype{}
    pushfn(res, self)
    return pushfn(res, obj)
  end
end

array.__actions.__concat = get_concat_fn(array, array.push_arr)
ziparr.__actions.__concat = get_concat_fn(ziparr, ziparr.push_zip)

-- initializer
setmetatable(array, 
  {__call = function(self, arr)
      setmetatable(arr, self.__actions)
      return arr
    end})

return array