util/array.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2008-2010 Matthew Wild
   3 -- Copyright (C) 2008-2010 Waqas Hussain
   4 --
   5 -- This project is MIT/X11 licensed. Please see the
   6 -- COPYING file in the source package for more information.
   7 --
   8
   9 local t_insert, t_sort, t_remove, t_concat
  10     = table.insert, table.sort, table.remove, table.concat;
  11
  12 local setmetatable = setmetatable;
  13 local math_random = math.random;
  14 local math_floor = math.floor;
  15 local pairs, ipairs = pairs, ipairs;
  16 local tostring = tostring;
  17
  18 local array = {};
  19 local array_base = {};
  20 local array_methods = {};
  21 local array_mt = { __index = array_methods, __tostring = function (array) return "{"..array:concat(", ").."}"; end };
  22
  23 local function new_array(self, t, _s, _var)
  24         if type(t) == "function" then -- Assume iterator
  25                 t = self.collect(t, _s, _var);
  26         end
  27         return setmetatable(t or {}, array_mt);
  28 end
  29
  30 function array_mt.__add(a1, a2)
  31         local res = new_array();
  32         return res:append(a1):append(a2);
  33 end
  34
  35 setmetatable(array, { __call = new_array });
  36
  37 -- Read-only methods
  38 function array_methods:random()
  39         return self[math_random(1,#self)];
  40 end
  41
  42 -- These methods can be called two ways:
  43 --   array.method(existing_array, [params [, ...]]) -- Create new array for result
  44 --   existing_array:method([params, ...]) -- Transform existing array into result
  45 --
  46 function array_base.map(outa, ina, func)
  47         for k,v in ipairs(ina) do
  48                 outa[k] = func(v);
  49         end
  50         return outa;
  51 end
  52
  53 function array_base.filter(outa, ina, func)
  54         local inplace, start_length = ina == outa, #ina;
  55         local write = 1;
  56         for read=1,start_length do
  57                 local v = ina[read];
  58                 if func(v) then
  59                         outa[write] = v;
  60                         write = write + 1;
  61                 end
  62         end
  63
  64         if inplace and write <= start_length then
  65                 for i=write,start_length do
  66                         outa[i] = nil;
  67                 end
  68         end
  69
  70         return outa;
  71 end
  72
  73 function array_base.sort(outa, ina, ...)
  74         if ina ~= outa then
  75                 outa:append(ina);
  76         end
  77         t_sort(outa, ...);
  78         return outa;
  79 end
  80
  81 function array_base.pluck(outa, ina, key)
  82         for i=1,#ina do
  83                 outa[i] = ina[i][key];
  84         end
  85         return outa;
  86 end
  87
  88 function array_base.reverse(outa, ina)
  89         local len = #ina;
  90         if ina == outa then
  91                 local middle = math_floor(len/2);
  92                 len = len + 1;
  93                 local o; -- opposite
  94                 for i = 1, middle do
  95                         o = len - i;
  96                         outa[i], outa[o] = outa[o], outa[i];
  97                 end
  98         else
  99                 local off = len + 1;
 100                 for i = 1, len do
 101                         outa[i] = ina[off - i];
 102                 end
 103         end
 104         return outa;
 105 end
 106
 107 --- These methods only mutate the array
 108 function array_methods:shuffle(outa, ina)
 109         local len = #self;
 110         for i=1,#self do
 111                 local r = math_random(i,len);
 112                 self[i], self[r] = self[r], self[i];
 113         end
 114         return self;
 115 end
 116
 117 function array_methods:append(array)
 118         local len,len2  = #self, #array;
 119         for i=1,len2 do
 120                 self[len+i] = array[i];
 121         end
 122         return self;
 123 end
 124
 125 function array_methods:push(x)
 126         t_insert(self, x);
 127         return self;
 128 end
 129
 130 function array_methods:pop(x)
 131         local v = self[x];
 132         t_remove(self, x);
 133         return v;
 134 end
 135
 136 function array_methods:concat(sep)
 137         return t_concat(array.map(self, tostring), sep);
 138 end
 139
 140 function array_methods:length()
 141         return #self;
 142 end
 143
 144 --- These methods always create a new array
 145 function array.collect(f, s, var)
 146         local t = {};
 147         while true do
 148                 var = f(s, var);
 149                 if var == nil then break; end
 150                 t_insert(t, var);
 151         end
 152         return setmetatable(t, array_mt);
 153 end
 154
 155 ---
 156
 157 -- Setup methods from array_base
 158 for method, f in pairs(array_base) do
 159         local base_method = f;
 160         -- Setup global array method which makes new array
 161         array[method] = function (old_a, ...)
 162                 local a = new_array();
 163                 return base_method(a, old_a, ...);
 164         end
 165         -- Setup per-array (mutating) method
 166         array_methods[method] = function (self, ...)
 167                 return base_method(self, self, ...);
 168         end
 169 end
 170
 171 _G.array = array;
 172 module("array");
 173
 174 return array;