util/array.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2008-2009 Matthew Wild
   3 -- Copyright (C) 2008-2009 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 array = {};
  13 local array_base = {};
  14 local array_methods = {};
  15 local array_mt = { __index = array_methods, __tostring = function (array) return array:concat(", "); end };
  16
  17 local function new_array(_, t)
  18         return setmetatable(t or {}, array_mt);
  19 end
  20
  21 function array_mt.__add(a1, a2)
  22         local res = new_array();
  23         return res:append(a1):append(a2);
  24 end
  25
  26 setmetatable(array, { __call = new_array });
  27
  28 function array_base.map(outa, ina, func)
  29         for k,v in ipairs(ina) do
  30                 outa[k] = func(v);
  31         end
  32         return outa;
  33 end
  34
  35 function array_base.filter(outa, ina, func)
  36         for k,v in ipairs(ina) do
  37                 if func(v) then
  38                         outa:push(v);
  39                 end
  40         end
  41         return outa;
  42 end
  43
  44 function array_base.sort(outa, ina, ...)
  45         if ina ~= outa then
  46                 outa:append(ina);
  47         end
  48         t_sort(outa, ...);
  49         return outa;
  50 end
  51
  52 --- These methods only mutate
  53 function array_methods:random()
  54         return self[math.random(1,#self)];
  55 end
  56
  57 function array_methods:shuffle(outa, ina)
  58         local len = #self;
  59         for i=1,#self do
  60                 local r = math.random(i,len);
  61                 self[i], self[r] = self[r], self[i];
  62         end
  63         return self;
  64 end
  65
  66 function array_methods:reverse()
  67         local len = #self-1;
  68         for i=len,1,-1 do
  69                 self:push(self[i]);
  70                 self:pop(i);
  71         end
  72         return self;
  73 end
  74
  75 function array_methods:append(array)
  76         local len,len2  = #self, #array;
  77         for i=1,len2 do
  78                 self[len+i] = array[i];
  79         end
  80         return self;
  81 end
  82
  83 array_methods.push = table.insert;
  84 array_methods.pop = table.remove;
  85 array_methods.concat = table.concat;
  86 array_methods.length = function (t) return #t; end
  87
  88 --- These methods always create a new array
  89 function array.collect(f, s, var)
  90         local t, var = {};
  91         while true do
  92                 var = f(s, var);
  93                 if var == nil then break; end
  94                 table.insert(t, var);
  95         end
  96         return setmetatable(t, array_mt);
  97 end
  98
  99 ---
 100
 101 -- Setup methods from array_base
 102 for method, f in pairs(array_base) do
 103         local method = method; -- Yes, this is necessary :)
 104         local base_method = f;
 105         -- Setup global array method which makes new array
 106         array[method] = function (old_a, ...)
 107                 local a = new_array();
 108                 return base_method(a, old_a, ...);
 109         end
 110         -- Setup per-array (mutating) method
 111         array_methods[method] = function (self, ...)
 112                 return base_method(self, self, ...);
 113         end
 114 end
 115
 116 _G.array = array;
 117 module("array");
 118
 119 return array;