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 local type = type;
  18
  19 local array = {};
  20 local array_base = {};
  21 local array_methods = {};
  22 local array_mt = {
  23         __index = array_methods;
  24         __name = "array";
  25         __tostring = function (self) return "{"..self:concat(", ").."}"; end;
  26 };
  27
  28 function array_mt:__freeze() return self; end
  29
  30 local function new_array(self, t, _s, _var)
  31         if type(t) == "function" then -- Assume iterator
  32                 t = self.collect(t, _s, _var);
  33         end
  34         return setmetatable(t or {}, array_mt);
  35 end
  36
  37 function array_mt.__add(a1, a2)
  38         local res = new_array();
  39         return res:append(a1):append(a2);
  40 end
  41
  42 function array_mt.__eq(a, b)
  43         if #a == #b then
  44                 for i = 1, #a do
  45                         if a[i] ~= b[i] then
  46                                 return false;
  47                         end
  48                 end
  49         else
  50                 return false;
  51         end
  52         return true;
  53 end
  54
  55 function array_mt.__div(a1, func)
  56         local a2 = new_array();
  57         local o = 0;
  58         for i = 1, #a1 do
  59                 local new_value = func(a1[i]);
  60                 if new_value ~= nil then
  61                         o = o + 1;
  62                         a2[o] = new_value;
  63                 end
  64         end
  65         return a2;
  66 end
  67
  68 setmetatable(array, { __call = new_array });
  69
  70 -- Read-only methods
  71 function array_methods:random()
  72         return self[math_random(1, #self)];
  73 end
  74
  75 -- Return a random value excluding the one at idx
  76 function array_methods:random_other(idx)
  77         local max = #self;
  78         return self[((math.random(1, max-1)+(idx-1))%max)+1];
  79 end
  80
  81 -- These methods can be called two ways:
  82 --   array.method(existing_array, [params [, ...]]) -- Create new array for result
  83 --   existing_array:method([params, ...]) -- Transform existing array into result
  84 --
  85 function array_base.map(outa, ina, func)
  86         for k, v in ipairs(ina) do
  87                 outa[k] = func(v);
  88         end
  89         return outa;
  90 end
  91
  92 function array_base.filter(outa, ina, func)
  93         local inplace, start_length = ina == outa, #ina;
  94         local write = 1;
  95         for read = 1, start_length do
  96                 local v = ina[read];
  97                 if func(v) then
  98                         outa[write] = v;
  99                         write = write + 1;
 100                 end
 101         end
 102
 103         if inplace and write <= start_length then
 104                 for i = write, start_length do
 105                         outa[i] = nil;
 106                 end
 107         end
 108
 109         return outa;
 110 end
 111
 112 function array_base.sort(outa, ina, ...)
 113         if ina ~= outa then
 114                 outa:append(ina);
 115         end
 116         t_sort(outa, ...);
 117         return outa;
 118 end
 119
 120 function array_base.unique(outa, ina)
 121         local seen = {};
 122         return array_base.filter(outa, ina, function (item)
 123                 if seen[item] then
 124                         return false;
 125                 else
 126                         seen[item] = true;
 127                         return true;
 128                 end
 129         end);
 130 end
 131
 132 function array_base.pluck(outa, ina, key)
 133         for i = 1, #ina do
 134                 outa[i] = ina[i][key];
 135         end
 136         return outa;
 137 end
 138
 139 function array_base.reverse(outa, ina)
 140         local len = #ina;
 141         if ina == outa then
 142                 local middle = math_floor(len/2);
 143                 len = len + 1;
 144                 local o; -- opposite
 145                 for i = 1, middle do
 146                         o = len - i;
 147                         outa[i], outa[o] = outa[o], outa[i];
 148                 end
 149         else
 150                 local off = len + 1;
 151                 for i = 1, len do
 152                         outa[i] = ina[off - i];
 153                 end
 154         end
 155         return outa;
 156 end
 157
 158 --- These methods only mutate the array
 159 function array_methods:shuffle()
 160         local len = #self;
 161         for i = 1, #self do
 162                 local r = math_random(i, len);
 163                 self[i], self[r] = self[r], self[i];
 164         end
 165         return self;
 166 end
 167
 168 function array_methods:append(ina)
 169         local len, len2 = #self, #ina;
 170         for i = 1, len2 do
 171                 self[len+i] = ina[i];
 172         end
 173         return self;
 174 end
 175
 176 function array_methods:push(x)
 177         t_insert(self, x);
 178         return self;
 179 end
 180
 181 array_methods.pop = t_remove;
 182
 183 function array_methods:concat(sep)
 184         return t_concat(array.map(self, tostring), sep);
 185 end
 186
 187 function array_methods:length()
 188         return #self;
 189 end
 190
 191 --- These methods always create a new array
 192 function array.collect(f, s, var)
 193         local t = {};
 194         while true do
 195                 var = f(s, var);
 196                 if var == nil then break; end
 197                 t_insert(t, var);
 198         end
 199         return setmetatable(t, array_mt);
 200 end
 201
 202 ---
 203
 204 -- Setup methods from array_base
 205 for method, f in pairs(array_base) do
 206         local base_method = f;
 207         -- Setup global array method which makes new array
 208         array[method] = function (old_a, ...)
 209                 local a = new_array();
 210                 return base_method(a, old_a, ...);
 211         end
 212         -- Setup per-array (mutating) method
 213         array_methods[method] = function (self, ...)
 214                 return base_method(self, self, ...);
 215         end
 216 end
 217
 218 return array;