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