util/multitable.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 select = select;
  10 local t_insert = table.insert;
  11 local pairs, next, type = pairs, next, type;
  12 local unpack = table.unpack or unpack; --luacheck: ignore 113
  13
  14 local _ENV = nil;
  15
  16 local function get(self, ...)
  17         local t = self.data;
  18         for n = 1,select('#', ...) do
  19                 t = t[select(n, ...)];
  20                 if not t then break; end
  21         end
  22         return t;
  23 end
  24
  25 local function add(self, ...)
  26         local t = self.data;
  27         local count = select('#', ...);
  28         for n = 1,count-1 do
  29                 local key = select(n, ...);
  30                 local tab = t[key];
  31                 if not tab then tab = {}; t[key] = tab; end
  32                 t = tab;
  33         end
  34         t_insert(t, (select(count, ...)));
  35 end
  36
  37 local function set(self, ...)
  38         local t = self.data;
  39         local count = select('#', ...);
  40         for n = 1,count-2 do
  41                 local key = select(n, ...);
  42                 local tab = t[key];
  43                 if not tab then tab = {}; t[key] = tab; end
  44                 t = tab;
  45         end
  46         t[(select(count-1, ...))] = (select(count, ...));
  47 end
  48
  49 local function r(t, n, _end, ...)
  50         if t == nil then return; end
  51         local k = select(n, ...);
  52         if n == _end then
  53                 t[k] = nil;
  54                 return;
  55         end
  56         if k then
  57                 local v = t[k];
  58                 if v then
  59                         r(v, n+1, _end, ...);
  60                         if not next(v) then
  61                                 t[k] = nil;
  62                         end
  63                 end
  64         else
  65                 for _,b in pairs(t) do
  66                         r(b, n+1, _end, ...);
  67                         if not next(b) then
  68                                 t[_] = nil;
  69                         end
  70                 end
  71         end
  72 end
  73
  74 local function remove(self, ...)
  75         local _end = select('#', ...);
  76         for n = _end,1 do
  77                 if select(n, ...) then _end = n; break; end
  78         end
  79         r(self.data, 1, _end, ...);
  80 end
  81
  82
  83 local function s(t, n, results, _end, ...)
  84         if t == nil then return; end
  85         local k = select(n, ...);
  86         if n == _end then
  87                 if k == nil then
  88                         for _, v in pairs(t) do
  89                                 t_insert(results, v);
  90                         end
  91                 else
  92                         t_insert(results, t[k]);
  93                 end
  94                 return;
  95         end
  96         if k then
  97                 local v = t[k];
  98                 if v then
  99                         s(v, n+1, results, _end, ...);
 100                 end
 101         else
 102                 for _,b in pairs(t) do
 103                         s(b, n+1, results, _end, ...);
 104                 end
 105         end
 106 end
 107
 108 -- Search for keys, nil == wildcard
 109 local function search(self, ...)
 110         local _end = select('#', ...);
 111         for n = _end,1 do
 112                 if select(n, ...) then _end = n; break; end
 113         end
 114         local results = {};
 115         s(self.data, 1, results, _end, ...);
 116         return results;
 117 end
 118
 119 -- Append results to an existing list
 120 local function search_add(self, results, ...)
 121         if not results then results = {}; end
 122         local _end = select('#', ...);
 123         for n = _end,1 do
 124                 if select(n, ...) then _end = n; break; end
 125         end
 126         s(self.data, 1, results, _end, ...);
 127         return results;
 128 end
 129
 130 local function iter(self, ...)
 131         local query = { ... };
 132         local maxdepth = select("#", ...);
 133         local stack = { self.data };
 134         local keys = { };
 135         local function it(self)
 136                 local depth = #stack;
 137                 local key = next(stack[depth], keys[depth]);
 138                 if key == nil then -- Go up the stack
 139                         stack[depth], keys[depth] = nil, nil;
 140                         if depth > 1 then
 141                                 return it(self);
 142                         end
 143                         return; -- The end
 144                 else
 145                         keys[depth] = key;
 146                 end
 147                 local value = stack[depth][key];
 148                 if query[depth] == nil or key == query[depth] then
 149                         if depth == maxdepth then -- Result
 150                                 local result = {}; -- Collect keys forming path to result
 151                                 for i = 1, depth do
 152                                         result[i] = keys[i];
 153                                 end
 154                                 result[depth+1] = value;
 155                                 return unpack(result, 1, depth+1);
 156                         elseif type(value) == "table" then
 157                                 t_insert(stack, value); -- Descend
 158                         end
 159                 end
 160                 return it(self);
 161         end;
 162         return it, self;
 163 end
 164
 165 local function new()
 166         return {
 167                 data = {};
 168                 get = get;
 169                 add = add;
 170                 set = set;
 171                 remove = remove;
 172                 search = search;
 173                 search_add = search_add;
 174                 iter = iter;
 175         };
 176 end
 177
 178 return {
 179         iter = iter;
 180         new = new;
 181 };