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