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