util/set.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 local ipairs, pairs, setmetatable, next, tostring =
  10       ipairs, pairs, setmetatable, next, tostring;
  11 local t_concat = table.concat;
  12
  13 module "set"
  14
  15 local set_mt = {};
  16 function set_mt.__call(set, _, k)
  17         return next(set._items, k);
  18 end
  19 function set_mt.__add(set1, set2)
  20         return _M.union(set1, set2);
  21 end
  22 function set_mt.__sub(set1, set2)
  23         return _M.difference(set1, set2);
  24 end
  25 function set_mt.__div(set, func)
  26         local new_set, new_items = _M.new();
  27         local items, new_items = set._items, new_set._items;
  28         for item in pairs(items) do
  29                 if func(item) then
  30                         new_items[item] = true;
  31                 end
  32         end
  33         return new_set;
  34 end
  35 function set_mt.__eq(set1, set2)
  36         local set1, set2 = set1._items, set2._items;
  37         for item in pairs(set1) do
  38                 if not set2[item] then
  39                         return false;
  40                 end
  41         end
  42
  43         for item in pairs(set2) do
  44                 if not set1[item] then
  45                         return false;
  46                 end
  47         end
  48
  49         return true;
  50 end
  51 function set_mt.__tostring(set)
  52         local s, items = { }, set._items;
  53         for item in pairs(items) do
  54                 s[#s+1] = tostring(item);
  55         end
  56         return t_concat(s, ", ");
  57 end
  58
  59 local items_mt = {};
  60 function items_mt.__call(items, _, k)
  61         return next(items, k);
  62 end
  63
  64 function new(list)
  65         local items = setmetatable({}, items_mt);
  66         local set = { _items = items };
  67
  68         function set:add(item)
  69                 items[item] = true;
  70         end
  71
  72         function set:contains(item)
  73                 return items[item];
  74         end
  75
  76         function set:items()
  77                 return items;
  78         end
  79
  80         function set:remove(item)
  81                 items[item] = nil;
  82         end
  83
  84         function set:add_list(list)
  85                 for _, item in ipairs(list) do
  86                         items[item] = true;
  87                 end
  88         end
  89
  90         function set:include(otherset)
  91                 for item in pairs(otherset) do
  92                         items[item] = true;
  93                 end
  94         end
  95
  96         function set:exclude(otherset)
  97                 for item in pairs(otherset) do
  98                         items[item] = nil;
  99                 end
 100         end
 101
 102         function set:empty()
 103                 return not next(items);
 104         end
 105
 106         if list then
 107                 set:add_list(list);
 108         end
 109
 110         return setmetatable(set, set_mt);
 111 end
 112
 113 function union(set1, set2)
 114         local set = new();
 115         local items = set._items;
 116
 117         for item in pairs(set1._items) do
 118                 items[item] = true;
 119         end
 120
 121         for item in pairs(set2._items) do
 122                 items[item] = true;
 123         end
 124
 125         return set;
 126 end
 127
 128 function difference(set1, set2)
 129         local set = new();
 130         local items = set._items;
 131
 132         for item in pairs(set1._items) do
 133                 items[item] = (not set2._items[item]) or nil;
 134         end
 135
 136         return set;
 137 end
 138
 139 function intersection(set1, set2)
 140         local set = new();
 141         local items = set._items;
 142
 143         set1, set2 = set1._items, set2._items;
 144
 145         for item in pairs(set1) do
 146                 items[item] = (not not set2[item]) or nil;
 147         end
 148
 149         return set;
 150 end
 151
 152 function xor(set1, set2)
 153         return union(set1, set2) - intersection(set1, set2);
 154 end
 155
 156 return _M;