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