util/iterators.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 --[[ Iterators ]]--
  10
  11 local it = {};
  12
  13 local t_insert = table.insert;
  14 local select, unpack, next = select, unpack, next;
  15 local function pack(...) return { n = select("#", ...), ... }; end
  16
  17 -- Reverse an iterator
  18 function it.reverse(f, s, var)
  19         local results = {};
  20
  21         -- First call the normal iterator
  22         while true do
  23                 local ret = { f(s, var) };
  24                 var = ret[1];
  25                 if var == nil then break; end
  26                 t_insert(results, 1, ret);
  27         end
  28
  29         -- Then return our reverse one
  30         local i,max = 0, #results;
  31         return function (results)
  32                         if i<max then
  33                                 i = i + 1;
  34                                 return unpack(results[i]);
  35                         end
  36                 end, results;
  37 end
  38
  39 -- Iterate only over keys in a table
  40 local function _keys_it(t, key)
  41         return (next(t, key));
  42 end
  43 function it.keys(t)
  44         return _keys_it, t;
  45 end
  46
  47 -- Iterate only over values in a table
  48 function it.values(t)
  49         local key, val;
  50         return function (t)
  51                 key, val = next(t, key);
  52                 return val;
  53         end, t;
  54 end
  55
  56 -- Given an iterator, iterate only over unique items
  57 function it.unique(f, s, var)
  58         local set = {};
  59
  60         return function ()
  61                 while true do
  62                         local ret = pack(f(s, var));
  63                         var = ret[1];
  64                         if var == nil then break; end
  65                         if not set[var] then
  66                                 set[var] = true;
  67                                 return unpack(ret, 1, ret.n);
  68                         end
  69                 end
  70         end;
  71 end
  72
  73 --[[ Return the number of items an iterator returns ]]--
  74 function it.count(f, s, var)
  75         local x = 0;
  76
  77         while true do
  78                 var = f(s, var);
  79                 if var == nil then break; end
  80                 x = x + 1;
  81         end
  82
  83         return x;
  84 end
  85
  86 -- Return the first n items an iterator returns
  87 function it.head(n, f, s, var)
  88         local c = 0;
  89         return function (s, var)
  90                 if c >= n then
  91                         return nil;
  92                 end
  93                 c = c + 1;
  94                 return f(s, var);
  95         end, s;
  96 end
  97
  98 -- Skip the first n items an iterator returns
  99 function it.skip(n, f, s, var)
 100         for i=1,n do
 101                 var = f(s, var);
 102         end
 103         return f, s, var;
 104 end
 105
 106 -- Return the last n items an iterator returns
 107 function it.tail(n, f, s, var)
 108         local results, count = {}, 0;
 109         while true do
 110                 local ret = pack(f(s, var));
 111                 var = ret[1];
 112                 if var == nil then break; end
 113                 results[(count%n)+1] = ret;
 114                 count = count + 1;
 115         end
 116
 117         if n > count then n = count; end
 118
 119         local pos = 0;
 120         return function ()
 121                 pos = pos + 1;
 122                 if pos > n then return nil; end
 123                 local ret = results[((count-1+pos)%n)+1];
 124                 return unpack(ret, 1, ret.n);
 125         end
 126         --return reverse(head(n, reverse(f, s, var))); -- !
 127 end
 128
 129 function it.filter(filter, f, s, var)
 130         if type(filter) ~= "function" then
 131                 local filter_value = filter;
 132                 function filter(x) return x ~= filter_value; end
 133         end
 134         return function (s, var)
 135                 local ret;
 136                 repeat ret = pack(f(s, var));
 137                         var = ret[1];
 138                 until var == nil or filter(unpack(ret, 1, ret.n));
 139                 return unpack(ret, 1, ret.n);
 140         end, s, var;
 141 end
 142
 143 local function _ripairs_iter(t, key) if key > 1 then return key-1, t[key-1]; end end
 144 function it.ripairs(t)
 145         return _ripairs_iter, t, #t+1;
 146 end
 147
 148 local function _range_iter(max, curr) if curr < max then return curr + 1; end end
 149 function it.range(x, y)
 150         if not y then x, y = 1, x; end -- Default to 1..x if y not given
 151         return _range_iter, y, x-1;
 152 end
 153
 154 -- Convert the values returned by an iterator to an array
 155 function it.to_array(f, s, var)
 156         local t, var = {};
 157         while true do
 158                 var = f(s, var);
 159                 if var == nil then break; end
 160                 t_insert(t, var);
 161         end
 162         return t;
 163 end
 164
 165 -- Treat the return of an iterator as key,value pairs,
 166 -- and build a table
 167 function it.to_table(f, s, var)
 168         local t, var2 = {};
 169         while true do
 170                 var, var2 = f(s, var);
 171                 if var == nil then break; end
 172                 t[var] = var2;
 173         end
 174         return t;
 175 end
 176
 177 return it;