OrderedMap.lua

OrderedMap.lua

---@class OrderedMap
local OrderedMap = {}
OrderedMap.__name = 'Map'

local keys_key = setmetatable({}, {
  __tostring = function() return "<___keys___>" end
})

local vals_key = setmetatable({}, {
  __tostring = function() return "<___vals___>" end
})

---@param self OrderedMap
---@param k any
---@param v any
local function insert(self, k, v)
  local keys = rawget(self, keys_key)
  local vals = rawget(self, vals_key)
  if vals[k] == nil then
    table.insert(keys, k)
  end
  vals[k] = v
end

---@param self OrderedMap
---@param k any
local function remove(self, k)
  local keys = rawget(self, keys_key)
  local vals = rawget(self, vals_key)
  if vals[k] == nil then
    return
  end
  for i = 1, #keys do
    if keys[i] == k then
      table.remove(keys, i)
      vals[k] = nil
      return
    end
  end
end

---@return table
function OrderedMap.new(...)
  local self = setmetatable({}, OrderedMap)
  rawset(self, keys_key, {})
  rawset(self, vals_key, {})
  local count = select("#", ...)
  if count > 0 then
    local input = { ... }
    for i = 1, count, 2 do
      insert(self, input[i], input[i + 1])
    end
  end
  return self
end

function OrderedMap:__index(k)
  local vals = rawget(self, vals_key)
  return vals[k]
end

function OrderedMap:__newindex(k, v)
  if v == nil then
    remove(self, k)
  else
    insert(self, k, v)
  end
end

function OrderedMap:__len()
  local keys = rawget(self, keys_key)
  return #keys
end

function OrderedMap:__pairs()
  local keys = rawget(self, keys_key)
  local vals = rawget(self, vals_key)
  local i = 0
  local n = #keys
  return function()
    i = i + 1
    if i <= n then
      local k = keys[i]
      return k, vals[k]
    end
  end
end

return OrderedMap