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 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         local inplace, start_length = ina == outa, #ina;
  37         local write = 1;
  38         for read=1,start_length do
  39                 local v = ina[read];
  40                 if func(v) then
  41                         outa[write] = v;
  42                         write = write + 1;
  43                 end
  44         end
  45
  46         if inplace and write <= start_length then
  47                 for i=write,start_length do
  48                         outa[i] = nil;
  49                 end
  50         end
  51
  52         return outa;
  53 end
  54
  55 function array_base.sort(outa, ina, ...)
  56         if ina ~= outa then
  57                 outa:append(ina);
  58         end
  59         t_sort(outa, ...);
  60         return outa;
  61 end
  62
  63 --- These methods only mutate
  64 function array_methods:random()
  65         return self[math.random(1,#self)];
  66 end
  67
  68 function array_methods:shuffle(outa, ina)
  69         local len = #self;
  70         for i=1,#self do
  71                 local r = math.random(i,len);
  72                 self[i], self[r] = self[r], self[i];
  73         end
  74         return self;
  75 end
  76
  77 function array_methods:reverse()
  78         local len = #self-1;
  79         for i=len,1,-1 do
  80                 self:push(self[i]);
  81                 self:pop(i);
  82         end
  83         return self;
  84 end
  85
  86 function array_methods:append(array)
  87         local len,len2  = #self, #array;
  88         for i=1,len2 do
  89                 self[len+i] = array[i];
  90         end
  91         return self;
  92 end
  93
  94 array_methods.push = table.insert;
  95 array_methods.pop = table.remove;
  96 array_methods.concat = table.concat;
  97 array_methods.length = function (t) return #t; end
  98
  99 --- These methods always create a new array
 100 function array.collect(f, s, var)
 101         local t, var = {};
 102         while true do
 103                 var = f(s, var);
 104                 if var == nil then break; end
 105                 table.insert(t, var);
 106         end
 107         return setmetatable(t, array_mt);
 108 end
 109
 110 ---
 111
 112 -- Setup methods from array_base
 113 for method, f in pairs(array_base) do
 114         local base_method = f;
 115         -- Setup global array method which makes new array
 116         array[method] = function (old_a, ...)
 117                 local a = new_array();
 118                 return base_method(a, old_a, ...);
 119         end
 120         -- Setup per-array (mutating) method
 121         array_methods[method] = function (self, ...)
 122                 return base_method(self, self, ...);
 123         end
 124 end
 125
 126 _G.array = array;
 127 module("array");
 128
 129 return array;