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 function tail(n, f, s, var)
  94         local results, count = {}, 0;
  95         while true do
  96                 local ret = { f(s, var) };
  97                 var = ret[1];
  98                 if var == nil then break; end
  99                 results[(count%n)+1] = ret;
 100                 count = count + 1;
 101         end
 102
 103         if n > count then n = count; end
 104
 105         local pos = 0;
 106         return function ()
 107                 pos = pos + 1;
 108                 if pos > n then return nil; end
 109                 return unpack(results[((count-1+pos)%n)+1]);
 110         end
 111         --return reverse(head(n, reverse(f, s, var)));
 112 end
 113
 114 -- Convert the values returned by an iterator to an array
 115 function it2array(f, s, var)
 116         local t, var = {};
 117         while true do
 118                 var = f(s, var);
 119                 if var == nil then break; end
 120                 table.insert(t, var);
 121         end
 122         return t;
 123 end
 124
 125 -- Treat the return of an iterator as key,value pairs,
 126 -- and build a table
 127 function it2table(f, s, var)
 128         local t, var = {};
 129         while true do
 130                 var, var2 = f(s, var);
 131                 if var == nil then break; end
 132                 t[var] = var2;
 133         end
 134         return t;
 135 end