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