Multiple Key Indexing

wiki

--[[

	Indexing values by tuple keys, implemented as a hash tree

	Any array works as a key, even arrays with holes, provided keys.n is set

	or n is passed as parameter to get() and set().



	Procedural interface:

		set(t, keys, e, [n])

		get(t, keys, [n]) -> e



		values(t) -> iterator -> e



		t[k1][k2]...[kn][E] -> e



	Objectual interface:

		([t]) -> idx

		wrap(t) -> idx

		idx.index -> t



		idx[keys] = e			idx:set(keys, e, [n])

		idx[keys] -> e			idx:get(keys, [n]) -> e



		idx:values() -> iterator -> e



]]



local coroutine, pairs, next, select, setmetatable =

	  coroutine, pairs, next, select, setmetatable



setfenv(1, {})



local function const(name)

	return setmetatable({}, {__tostring = function() return name end})

end



local NIL = const'NIL'

local NAN = const'NAN'

local E = const'E'



local function tokey(k)

	return (k == nil and NIL) or (k ~= k and NAN) or k

end



local function fromkey(k)

	return (k == NAN and 0/0) or (k ~= NIL and k) or nil

end



local function add(t, keys, e, n)

	n = n or keys.n or #keys

	for i=1,n do

		local k = tokey(keys[i])

		if not t[k] then

			t[k] = {}

		end

		t = t[k]

	end

	t[E] = e

end



local function many(t)

	return next(t,next(t))

end



local function remove(t, keys, n)

	n = n or keys.n or #keys

	local lastt, cleart, cleark

	for i=1,n do

		local k = tokey(keys[i])

		local tk = t[k]

		if not tk then return end

		if many(tk) then

			cleart, cleark = nil,nil

		elseif not cleart then

			cleart, cleark = t,k

		end

		t = tk

	end

	if not t[E] then return end

	if cleart then

		cleart[cleark] = nil

	else

		t[E] = nil

	end

end



local function set(t, keys, e, n)

	if e ~= nil then

		add(t, keys, e, n)

	else

		remove(t, keys, n)

	end

end



local function get(t, keys, n)

	n = n or keys.n or #keys

	for i=1,n do

		t = t[tokey(keys[i])]

		if not t then return end

	end

	return t[E]

end



local function yield_values(t)

	for k,t in pairs(t) do

		if k == E then

			coroutine.yield(t)

		else

			yield_values(t)

		end

	end

end



local function values(t)

	return coroutine.wrap(yield_values), t

end



--objectual interface



local methods = {}

function methods:set(keys, e, n) set(self.index, keys, e, n) end

function methods:get(keys, n) return get(self.index, keys, n) end

function methods:values() return values(self.index) end



local meta = {__type = 'index'}

function meta:__index(k) return methods[k] or get(self.index, k) end

function meta:__newindex(k, v) return set(self.index, k, v) end



local function wrap(t)

	return setmetatable({index = t}, meta)

end



local M = {

	meta = meta,

	methods = methods,

	set = set,

	get = get,

	values = values,

	wrap = wrap,

}



return setmetatable(M, {__call = function(_,t) return wrap(t or {}) end})

And the test unit (uses TupleModule):

local index = require 'index'

local tuple = require 'tuple'



local idx = index()

local NaN = 0/0



t1 = tuple(1,2,3,4)

t2 = tuple(1,2,3)

idx[t1] = t1

idx[t2] = t2

idx[t1] = nil

idx[t2] = nil

assert(next(idx.index) == nil)



local tuples = {

	tuple(1,NaN,3),

	tuple(1,NaN),

	tuple(1),

	tuple(1,3,NaN),

	tuple(1,nil,3),

	tuple(nil,nil),

	tuple(nil,NaN),

	tuple(nil),

	tuple(NaN),

	tuple(),

	tuple(2,3,4,5,6),

}



for i,e in ipairs(tuples) do

	idx[e] = e

end



for i,e in ipairs(tuples) do

	assert(idx[e] == e)

end



for e in idx:values() do

	print(e)

end



for i,e in ipairs(tuples) do

	idx[e] = nil

end



assert(next(idx.index) == nil)