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