--
-- Immutable Maps
--
-- SPDX-FileCopyrightText: Chris Pressey, the original author of this work, has dedicated it to the public domain.
-- For more information, please refer to <https://unlicense.org/>
-- SPDX-License-Identifier: Unlicense
--

new = function()
    return {
        store = {},
        depth = 0,
        next = nil
    }
end

map = function(f, map)
    -- Obtain list of maps we can traverse in appropriate order.
    local maps = {}
    local m = map
    while m ~= nil do
        table.insert(maps, 1, m)
        m = m.next
    end

    -- Traverse maps, merging their stores into a single store.
    local store = {}
    local k, v
    for _, m in ipairs(maps) do
        for k, v in pairs(m.store) do
            store[k] = f(v)
        end
    end

    return {
        store = store,
        depth = 0,
        next = nil
    }
end

flatten = function(m)
    return map(function(x) return x end, m)
end

update = function(k, v, map)
    local m = {
        store = { [k] = v },
        depth = map.depth + 1,
        next = map
    }
    if m.depth > 10 then
        return flatten(m)
    else
        return m
    end
end

remove = function(k, m)
    m = flatten(m)
    m.store[k] = nil
    return m
end

lookup = function(k, map)
    local m = map
    local v
    while m ~= nil do
        v = m.store[k]
        if v ~= nil then
            return v
        else
            m = m.next
        end
    end
    return nil
end

demo = function()
    local m0 = new()
    local m1 = update("k", 3, m0)
    local m2 = update("e", 4, m1)
    local m3 = update("e", 5, m2)

    local k
    print("m1:")
    for _, k in ipairs({"k", "e", "y"}) do
        print(k, lookup(k, m1))
    end
    print("m2:")
    for _, k in ipairs({"k", "e", "y"}) do
        print(k, lookup(k, m2))
    end
    print("m3:")
    for _, k in ipairs({"k", "e", "y"}) do
        print(k, lookup(k, m3))
    end

    local m4 = flatten(m3)
    print("m4 = flatten(m3):")
    for _, k in ipairs({"k", "e", "y"}) do
        print(k, lookup(k, m4))
    end

    local m5 = remove("e", m3)
    print("m5:")
    for _, k in ipairs({"k", "e", "y"}) do
        print(k, lookup(k, m5))
    end

end