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