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