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 local _ENV = nil;
  14 -- luacheck: std none
  15
  16 local set_mt = { __name = "set" };
  17 function set_mt.__call(set, _, k)
  18         return next(set._items, k);
  19 end
  20
  21 local items_mt = {};
  22 function items_mt.__call(items, _, k)
  23         return next(items, k);
  24 end
  25
  26 function set_mt:__freeze()
  27         local a, i = {}, 1;
  28         for item in self._items do
  29                 a[i], i = item, i+1;
  30         end
  31         return a;
  32 end
  33
  34 local function new(list)
  35         local items = setmetatable({}, items_mt);
  36         local set = { _items = items };
  37
  38         -- We access the set through an upvalue in these methods, so ignore 'self' being unused
  39         --luacheck: ignore 212/self
  40
  41         function set:add(item)
  42                 items[item] = true;
  43         end
  44
  45         function set:contains(item)
  46                 return items[item];
  47         end
  48
  49         function set:items()
  50                 return next, items;
  51         end
  52
  53         function set:remove(item)
  54                 items[item] = nil;
  55         end
  56
  57         function set:add_list(item_list)
  58                 if item_list then
  59                         for _, item in ipairs(item_list) do
  60                                 items[item] = true;
  61                         end
  62                 end
  63         end
  64
  65         function set:include(otherset)
  66                 for item in otherset do
  67                         items[item] = true;
  68                 end
  69         end
  70
  71         function set:exclude(otherset)
  72                 for item in otherset do
  73                         items[item] = nil;
  74                 end
  75         end
  76
  77         function set:empty()
  78                 return not next(items);
  79         end
  80
  81         if list then
  82                 set:add_list(list);
  83         end
  84
  85         return setmetatable(set, set_mt);
  86 end
  87
  88 local function union(set1, set2)
  89         local set = new();
  90         local items = set._items;
  91
  92         for item in pairs(set1._items) do
  93                 items[item] = true;
  94         end
  95
  96         for item in pairs(set2._items) do
  97                 items[item] = true;
  98         end
  99
 100         return set;
 101 end
 102
 103 local function difference(set1, set2)
 104         local set = new();
 105         local items = set._items;
 106
 107         for item in pairs(set1._items) do
 108                 items[item] = (not set2._items[item]) or nil;
 109         end
 110
 111         return set;
 112 end
 113
 114 local function intersection(set1, set2)
 115         local set = new();
 116         local items = set._items;
 117
 118         set1, set2 = set1._items, set2._items;
 119
 120         for item in pairs(set1) do
 121                 items[item] = (not not set2[item]) or nil;
 122         end
 123
 124         return set;
 125 end
 126
 127 local function xor(set1, set2)
 128         return union(set1, set2) - intersection(set1, set2);
 129 end
 130
 131 function set_mt.__add(set1, set2)
 132         return union(set1, set2);
 133 end
 134 function set_mt.__sub(set1, set2)
 135         return difference(set1, set2);
 136 end
 137 function set_mt.__div(set, func)
 138         local new_set = new();
 139         local items, new_items = set._items, new_set._items;
 140         for item in pairs(items) do
 141                 local new_item = func(item);
 142                 if new_item ~= nil then
 143                         new_items[new_item] = true;
 144                 end
 145         end
 146         return new_set;
 147 end
 148 function set_mt.__eq(set1, set2)
 149         set1, set2 = set1._items, set2._items;
 150         for item in pairs(set1) do
 151                 if not set2[item] then
 152                         return false;
 153                 end
 154         end
 155
 156         for item in pairs(set2) do
 157                 if not set1[item] then
 158                         return false;
 159                 end
 160         end
 161
 162         return true;
 163 end
 164 function set_mt.__tostring(set)
 165         local s, items = { }, set._items;
 166         for item in pairs(items) do
 167                 s[#s+1] = tostring(item);
 168         end
 169         return t_concat(s, ", ");
 170 end
 171
 172 return {
 173         new = new;
 174         union = union;
 175         difference = difference;
 176         intersection = intersection;
 177         xor = xor;
 178 };