From 10aaeb69b87da204424b4bf0526b10c489ff1dc5 Mon Sep 17 00:00:00 2001
From: BlackMATov <blackmat@mail.ru>
Date: Sat, 23 Nov 2024 06:17:47 +0700
Subject: [PATCH] chunk tree first impl

---
 ROADMAP.md                           |   7 +
 develop/untests.lua                  |   1 +
 develop/untests/registry_untests.lua | 245 ++++++++++++++++
 evolved/idpools.lua                  |  14 +-
 evolved/registry.lua                 | 411 ++++++++++++++++++++++++++-
 5 files changed, 667 insertions(+), 11 deletions(-)
 create mode 100644 ROADMAP.md
 create mode 100644 develop/untests/registry_untests.lua

diff --git a/ROADMAP.md b/ROADMAP.md
new file mode 100644
index 0000000..4d26aa8
--- /dev/null
+++ b/ROADMAP.md
@@ -0,0 +1,7 @@
+# Roadmap
+
+## Backlog
+
+- cache matched chunks in queries
+- cache transitions between chunks
+- chunk's children should be sorted by id and stored in an array instead of a table
\ No newline at end of file
diff --git a/develop/untests.lua b/develop/untests.lua
index 391df24..dc23eae 100644
--- a/develop/untests.lua
+++ b/develop/untests.lua
@@ -1 +1,2 @@
 require 'develop.untests.idpools_untests'
+require 'develop.untests.registry_untests'
diff --git a/develop/untests/registry_untests.lua b/develop/untests/registry_untests.lua
new file mode 100644
index 0000000..d107a27
--- /dev/null
+++ b/develop/untests/registry_untests.lua
@@ -0,0 +1,245 @@
+local evo = require 'evolved.evolved'
+
+do
+    local f1, f2 = evo.registry.entity(), evo.registry.entity()
+
+    local e = evo.registry.entity()
+    assert(e.chunk == nil)
+
+    evo.registry.insert(e, f1)
+    assert(evo.registry.has(e, f1))
+    assert(not evo.registry.has(e, f2))
+
+    evo.registry.insert(e, f2)
+    assert(evo.registry.has(e, f1))
+    assert(evo.registry.has(e, f2))
+
+    evo.registry.remove(e, f1)
+    assert(not evo.registry.has(e, f1))
+    assert(evo.registry.has(e, f2))
+
+    evo.registry.remove(e, f2)
+    assert(not evo.registry.has(e, f1))
+    assert(not evo.registry.has(e, f2))
+end
+
+do
+    local f = evo.registry.entity()
+    local e = evo.registry.entity()
+
+    if not os.getenv("LOCAL_LUA_DEBUGGER_VSCODE") then
+        assert(not pcall(evo.registry.get, e, f))
+        assert(not pcall(evo.registry.assign, e, f, 42))
+    end
+
+    assert(evo.registry.get_or(e, f) == nil)
+    assert(evo.registry.get_or(e, f, 42) == 42)
+
+    evo.registry.insert(e, f, 84)
+
+    assert(evo.registry.get(e, f) == 84)
+    assert(evo.registry.get_or(e, f) == 84)
+    assert(evo.registry.get_or(e, f, 42) == 84)
+
+    if not os.getenv("LOCAL_LUA_DEBUGGER_VSCODE") then
+        assert(not pcall(evo.registry.insert, e, f, 42))
+    end
+
+    evo.registry.assign(e, f)
+    assert(evo.registry.get(e, f) == true)
+
+    evo.registry.assign(e, f, 21)
+    assert(evo.registry.get(e, f) == 21)
+end
+
+do
+    local f1, f2 = evo.registry.entity(), evo.registry.entity()
+
+    local e = evo.registry.entity()
+
+    evo.registry.insert(e, f1, f1.guid)
+    assert(e.chunk == evo.registry.chunk(f1))
+
+    do
+        local chunk_f1 = evo.registry.chunk(f1)
+        assert(#chunk_f1.entities == 1)
+        assert(#chunk_f1.components[f1] == 1)
+    end
+
+    evo.registry.insert(e, f2, f2.guid)
+    assert(e.chunk == evo.registry.chunk(f1, f2))
+
+    do
+        local chunk_f1 = evo.registry.chunk(f1)
+        assert(#chunk_f1.entities == 0)
+        assert(#chunk_f1.components[f1] == 0)
+
+        local chunk_f1_f2 = evo.registry.chunk(f1, f2)
+        assert(#chunk_f1_f2.entities == 1)
+        assert(#chunk_f1_f2.components[f1] == 1)
+        assert(#chunk_f1_f2.components[f2] == 1)
+    end
+
+    evo.registry.remove(e, f1)
+    assert(e.chunk == evo.registry.chunk(f2))
+
+    do
+        local chunk_f1 = evo.registry.chunk(f1)
+        assert(#chunk_f1.entities == 0)
+        assert(#chunk_f1.components[f1] == 0)
+
+        local chunk_f2 = evo.registry.chunk(f2)
+        assert(#chunk_f2.entities == 1)
+        assert(#chunk_f2.components[f2] == 1)
+
+        local chunk_f1_f2 = evo.registry.chunk(f1, f2)
+        assert(#chunk_f1_f2.entities == 0)
+        assert(#chunk_f1_f2.components[f1] == 0)
+        assert(#chunk_f1_f2.components[f2] == 0)
+    end
+end
+
+for _ = 1, 100 do
+    local insert_fragments = {} ---@type evolved.entity[]
+    local insert_fragment_count = math.random(1, 10)
+
+    for _ = 1, insert_fragment_count do
+        local fragment = evo.registry.entity()
+        table.insert(insert_fragments, fragment)
+    end
+
+    local remove_fragments = {} ---@type evolved.entity[]
+    local remove_fragment_count = math.random(1, insert_fragment_count)
+
+    for _ = 1, remove_fragment_count do
+        local fragment = insert_fragments[math.random(1, #insert_fragments)]
+        table.insert(remove_fragments, fragment)
+    end
+
+    ---@param array any[]
+    local function shuffle_array(array)
+        for i = #array, 2, -1 do
+            local j = math.random(i)
+            array[i], array[j] = array[j], array[i]
+        end
+    end
+
+    local entities = {} ---@type evolved.entity[]
+
+    for _ = 1, 100 do
+        local e1, e2 = evo.registry.entity(), evo.registry.entity()
+        table.insert(entities, e1)
+        table.insert(entities, e2)
+
+        shuffle_array(insert_fragments)
+        for _, f in ipairs(insert_fragments) do
+            evo.registry.insert(e1, f, f.guid)
+        end
+
+        shuffle_array(insert_fragments)
+        for _, f in ipairs(insert_fragments) do
+            evo.registry.insert(e2, f, f.guid)
+        end
+
+        assert(e1.chunk == e2.chunk)
+        assert(evo.registry.has_all(e1, unpack(insert_fragments)))
+        assert(evo.registry.has_all(e2, unpack(insert_fragments)))
+
+        shuffle_array(remove_fragments)
+        for _, f in ipairs(remove_fragments) do
+            if evo.registry.has(e1, f) then
+                evo.registry.remove(e1, f)
+            end
+        end
+
+        shuffle_array(remove_fragments)
+        for _, f in ipairs(remove_fragments) do
+            if evo.registry.has(e2, f) then
+                evo.registry.remove(e2, f)
+            end
+        end
+
+        assert(e1.chunk == e2.chunk)
+        assert(not evo.registry.has_any(e1, unpack(remove_fragments)))
+        assert(not evo.registry.has_any(e2, unpack(remove_fragments)))
+
+        if e1.chunk ~= nil then
+            for f, _ in pairs(e1.chunk.components) do
+                assert(evo.registry.get(e1, f) == f.guid)
+                assert(evo.registry.get(e2, f) == f.guid)
+            end
+        end
+    end
+end
+
+do
+    local f1, f2, f3 = evo.registry.entity(), evo.registry.entity(), evo.registry.entity()
+
+    local e1 = evo.registry.entity()
+    evo.registry.insert(e1, f1)
+
+    local e2 = evo.registry.entity()
+    evo.registry.insert(e2, f1)
+    evo.registry.insert(e2, f2)
+
+    local e3 = evo.registry.entity()
+    evo.registry.insert(e3, f1)
+    evo.registry.insert(e3, f2)
+    evo.registry.insert(e3, f3)
+
+    do
+        local e = evo.registry.entity()
+
+        evo.registry.insert(e, f1)
+        evo.registry.remove(e, f1)
+
+        evo.registry.insert(e, f1)
+        evo.registry.insert(e, f2)
+        evo.registry.remove(e, f1)
+        evo.registry.remove(e, f2)
+
+        evo.registry.insert(e, f1)
+        evo.registry.insert(e, f2)
+        evo.registry.insert(e, f3)
+        evo.registry.remove(e, f1)
+        evo.registry.remove(e, f2)
+        evo.registry.remove(e, f3)
+    end
+
+    local q1 = evo.registry.query(f1)
+    local q2 = evo.registry.query(f1, f2)
+    local q3 = evo.registry.query(f1, f2, f3)
+
+    ---@param query evolved.query
+    ---@return evolved.entity[]
+    ---@nodiscard
+    local function collect_query_entities(query)
+        local entities = {} ---@type evolved.entity[]
+        for chunk in evo.registry.execute(query) do
+            for _, e in ipairs(chunk.entities) do
+                table.insert(entities, e)
+            end
+        end
+        return entities
+    end
+
+    ---@param array1 any[]
+    ---@param array2 any[]
+    ---@return boolean
+    ---@nodiscard
+    local function is_array_equal(array1, array2)
+        if #array1 ~= #array2 then
+            return false
+        end
+        for i = 1, #array1 do
+            if array1[i] ~= array2[i] then
+                return false
+            end
+        end
+        return true
+    end
+
+    assert(is_array_equal(collect_query_entities(q1), { e1, e2, e3 }))
+    assert(is_array_equal(collect_query_entities(q2), { e2, e3 }))
+    assert(is_array_equal(collect_query_entities(q3), { e3 }))
+end
diff --git a/evolved/idpools.lua b/evolved/idpools.lua
index e69b388..b410831 100644
--- a/evolved/idpools.lua
+++ b/evolved/idpools.lua
@@ -12,11 +12,11 @@ evolved_idpool_mt.__index = evolved_idpool_mt
 ---@return evolved.idpool
 function idpools.idpool()
     ---@type evolved.idpool
-    local self = {
+    local idpool = {
         acquired_ids = {},
         available_index = 0,
     }
-    return setmetatable(self, evolved_idpool_mt)
+    return setmetatable(idpool, evolved_idpool_mt)
 end
 
 ---@param id integer
@@ -52,7 +52,15 @@ end
 function idpools.release_id(idpool, id)
     local index = bit.band(id, 0xFFFFF)
     local version = bit.band(id, 0x7FF00000)
-    version = version == 0x7FF00000 and 0x100000 or version + 0x100000
+
+    if idpool.acquired_ids[index] ~= id then
+        error('id is not acquired or already released', 2)
+    end
+
+    version = version == 0x7FF00000
+        and 0x100000
+        or version + 0x100000
+
     idpool.acquired_ids[index] = idpool.available_index + version
     idpool.available_index = index
 end
diff --git a/evolved/registry.lua b/evolved/registry.lua
index 4bda1f5..48ed4a1 100644
--- a/evolved/registry.lua
+++ b/evolved/registry.lua
@@ -1,16 +1,41 @@
+local idpools = require 'evolved.idpools'
+
 ---@class evolved.registry
 local registry = {}
 
+---
+---
+---
+---
+---
+
+local __guids = idpools.idpool()
+local __roots = {} ---@type table<evolved.entity, evolved.chunk>
+local __chunks = {} ---@type table<evolved.entity, evolved.chunk[]>
+local __queries = {} ---@type table<evolved.entity, evolved.query[]>
+
+---
+---
+---
+---
+---
+
 ---@class evolved.entity
+---@field guid integer
 ---@field chunk? evolved.chunk
+---@field index_in_chunk integer
 local evolved_entity_mt = {}
 evolved_entity_mt.__index = evolved_entity_mt
 
 ---@class evolved.query
+---@field fragments evolved.entity[]
 local evolved_query_mt = {}
 evolved_query_mt.__index = evolved_query_mt
 
 ---@class evolved.chunk
+---@field parent? evolved.chunk
+---@field fragment evolved.entity
+---@field children table<evolved.entity, evolved.chunk>
 ---@field entities evolved.entity[]
 ---@field components table<evolved.entity, any[]>
 local evolved_chunk_mt = {}
@@ -22,6 +47,54 @@ evolved_chunk_mt.__index = evolved_chunk_mt
 ---
 ---
 
+---@param query evolved.query
+local function __on_new_query(query)
+    local main_fragment = query.fragments[#query.fragments]
+    local main_fragment_queries = __queries[main_fragment] or {}
+    main_fragment_queries[#main_fragment_queries + 1] = query
+    __queries[main_fragment] = main_fragment_queries
+end
+
+---@param chunk evolved.chunk
+local function __on_new_chunk(chunk)
+    local main_fragment = chunk.fragment
+    local main_fragment_chunks = __chunks[main_fragment] or {}
+    main_fragment_chunks[#main_fragment_chunks + 1] = chunk
+    __chunks[main_fragment] = main_fragment_chunks
+end
+
+---
+---
+---
+---
+---
+
+---@param entity evolved.entity
+local function __detach_entity(entity)
+    local chunk = assert(entity.chunk)
+    local index_in_chunk = entity.index_in_chunk
+
+    if index_in_chunk == #chunk.entities then
+        chunk.entities[index_in_chunk] = nil
+
+        for _, cs in pairs(chunk.components) do
+            cs[index_in_chunk] = nil
+        end
+    else
+        chunk.entities[index_in_chunk] = chunk.entities[#chunk.entities]
+        chunk.entities[index_in_chunk].index_in_chunk = index_in_chunk
+        chunk.entities[#chunk.entities] = nil
+
+        for _, cs in pairs(chunk.components) do
+            cs[index_in_chunk] = cs[#cs]
+            cs[#cs] = nil
+        end
+    end
+
+    entity.chunk = nil
+    entity.index_in_chunk = 0
+end
+
 ---@param chunk evolved.chunk
 ---@param fragment evolved.entity
 ---@return boolean
@@ -72,6 +145,101 @@ local function __chunk_has_any_fragments(chunk, fragment, ...)
     return false
 end
 
+---@param fragment evolved.entity
+---@return evolved.chunk
+---@nodiscard
+local function __root_chunk(fragment)
+    do
+        local root_chunk = __roots[fragment]
+        if root_chunk then return root_chunk end
+    end
+
+    ---@type evolved.chunk
+    local root_chunk = {
+        parent = nil,
+        fragment = fragment,
+        children = {},
+        entities = {},
+        components = { [fragment] = {} },
+    }
+
+    setmetatable(root_chunk, evolved_chunk_mt)
+
+    do
+        __roots[fragment] = root_chunk
+    end
+
+    __on_new_chunk(root_chunk)
+    return root_chunk
+end
+
+---@param chunk? evolved.chunk
+---@param fragment evolved.entity
+---@return evolved.chunk
+---@nodiscard
+local function __chunk_with_fragment(chunk, fragment)
+    if chunk == nil then
+        return __root_chunk(fragment)
+    end
+
+    if fragment.guid == chunk.fragment.guid then
+        return chunk
+    end
+
+    if fragment.guid < chunk.fragment.guid then
+        local sibling_chunk = __chunk_with_fragment(chunk.parent, fragment)
+        return __chunk_with_fragment(sibling_chunk, chunk.fragment)
+    end
+
+    do
+        local child_chunk = chunk.children[fragment]
+        if child_chunk then return child_chunk end
+    end
+
+    ---@type evolved.chunk
+    local child_chunk = {
+        parent = chunk,
+        fragment = fragment,
+        children = {},
+        entities = {},
+        components = { [fragment] = {} },
+    }
+
+    for f, _ in pairs(chunk.components) do
+        child_chunk.components[f] = {}
+    end
+
+    setmetatable(child_chunk, evolved_chunk_mt)
+
+    do
+        chunk.children[fragment] = child_chunk
+    end
+
+    __on_new_chunk(child_chunk)
+    return child_chunk
+end
+
+---@param chunk? evolved.chunk
+---@param fragment evolved.entity
+---@return evolved.chunk?
+---@nodiscard
+local function __chunk_without_fragment(chunk, fragment)
+    if chunk == nil then
+        return nil
+    end
+
+    if fragment.guid == chunk.fragment.guid then
+        return chunk.parent
+    end
+
+    if fragment.guid < chunk.fragment.guid then
+        local sibling_chunk = __chunk_without_fragment(chunk.parent, fragment)
+        return __chunk_with_fragment(sibling_chunk, chunk.fragment)
+    end
+
+    return chunk
+end
+
 ---
 ---
 ---
@@ -80,16 +248,67 @@ end
 
 ---@return evolved.entity
 ---@nodiscard
-function registry.entity() end
+function registry.entity()
+    local guid = idpools.acquire_id(__guids)
+
+    ---@type evolved.entity
+    local entity = {
+        guid = guid,
+        chunk = nil,
+        index_in_chunk = 0,
+    }
+
+    return setmetatable(entity, evolved_entity_mt)
+end
 
 ---@param entity evolved.entity
-function registry.destroy(entity) end
+---@return boolean
+---@nodiscard
+function registry.is_alive(entity)
+    return idpools.is_id_alive(__guids, entity.guid)
+end
+
+---@param entity evolved.entity
+function registry.destroy(entity)
+    if not registry.is_alive(entity) then
+        error(string.format('entity %s is not alive', entity), 2)
+    end
+
+    if entity.chunk ~= nil then
+        __detach_entity(entity)
+    end
+
+    idpools.release_id(__guids, entity.guid)
+end
 
 ---@param entity evolved.entity
 ---@param fragment evolved.entity
 ---@return any
 ---@nodiscard
-function registry.get(entity, fragment) end
+function registry.get(entity, fragment)
+    local chunk_components = entity.chunk and entity.chunk.components[fragment]
+
+    if chunk_components == nil then
+        error(string.format('entity %s does not have fragment %s', entity, fragment), 2)
+    end
+
+    return chunk_components[entity.index_in_chunk]
+end
+
+---@param entity evolved.entity
+---@param fragment evolved.entity
+---@param default any
+---@return any
+---@nodiscard
+function registry.get_or(entity, fragment, default)
+    local chunk_components = entity.chunk and entity.chunk.components[fragment]
+
+    if chunk_components == nil then
+        return default
+    end
+
+    return chunk_components[entity.index_in_chunk]
+end
 
 ---@param entity evolved.entity
 ---@param fragment evolved.entity
@@ -120,25 +339,201 @@ end
 ---@param entity evolved.entity
 ---@param fragment evolved.entity
 ---@param component any
-function registry.assign(entity, fragment, component) end
+function registry.assign(entity, fragment, component)
+    component = component == nil and true or component
+
+    local chunk_components = entity.chunk and entity.chunk.components[fragment]
+
+    if chunk_components == nil then
+        error(string.format('entity %s does not have fragment %s', entity, fragment), 2)
+    end
+
+    chunk_components[entity.index_in_chunk] = component
+end
 
 ---@param entity evolved.entity
 ---@param fragment evolved.entity
 ---@param component any
-function registry.insert(entity, fragment, component) end
+function registry.insert(entity, fragment, component)
+    component = component == nil and true or component
+
+    local old_chunk = entity.chunk
+    local new_chunk = __chunk_with_fragment(old_chunk, fragment)
+
+    if old_chunk == new_chunk then
+        error(string.format('entity %s already has fragment %s', entity, fragment), 2)
+    end
+
+    local old_index_in_chunk = entity.index_in_chunk
+    local new_index_in_chunk = new_chunk and #new_chunk.entities + 1 or 0
+
+    if new_chunk ~= nil then
+        new_chunk.entities[new_index_in_chunk] = entity
+        new_chunk.components[fragment][new_index_in_chunk] = component
+
+        if old_chunk ~= nil then
+            for old_f, old_cs in pairs(old_chunk.components) do
+                local new_cs = new_chunk.components[old_f]
+                new_cs[new_index_in_chunk] = old_cs[old_index_in_chunk]
+            end
+        end
+    end
+
+    if old_chunk ~= nil then
+        __detach_entity(entity)
+    end
+
+    entity.chunk = new_chunk
+    entity.index_in_chunk = new_index_in_chunk
+end
 
 ---@param entity evolved.entity
 ---@param fragment evolved.entity
-function registry.remove(entity, fragment) end
+function registry.remove(entity, fragment)
+    local old_chunk = entity.chunk
+    local new_chunk = __chunk_without_fragment(old_chunk, fragment)
 
+    if old_chunk == new_chunk then
+        error(string.format('entity %s does not have fragment %s', entity, fragment), 2)
+    end
+
+    local old_index_in_chunk = entity.index_in_chunk
+    local new_index_in_chunk = new_chunk and #new_chunk.entities + 1 or 0
+
+    if new_chunk ~= nil then
+        new_chunk.entities[new_index_in_chunk] = entity
+
+        if old_chunk ~= nil then
+            for new_f, new_cs in pairs(new_chunk.components) do
+                local old_cs = old_chunk.components[new_f]
+                new_cs[new_index_in_chunk] = old_cs[old_index_in_chunk]
+            end
+        end
+    end
+
+    if old_chunk ~= nil then
+        __detach_entity(entity)
+    end
+
+    entity.chunk = new_chunk
+    entity.index_in_chunk = new_index_in_chunk
+end
+
+---@param fragment evolved.entity
 ---@param ... evolved.entity
 ---@return evolved.query
 ---@nodiscard
-function registry.query(...) end
+function registry.query(fragment, ...)
+    local fragments = { fragment, ... }
+
+    table.sort(fragments, function(a, b)
+        return a.guid < b.guid
+    end)
+
+    ---@type evolved.query
+    local query = {
+        fragments = fragments,
+    }
+
+    setmetatable(query, evolved_query_mt)
+
+    __on_new_query(query)
+    return query
+end
 
 ---@param query evolved.query
 ---@return fun(): evolved.chunk?
 ---@nodiscard
-function registry.execute(query) end
+function registry.execute(query)
+    local main_fragment = query.fragments[#query.fragments]
+    local main_fragment_chunks = __chunks[main_fragment] or {}
+
+    ---@type evolved.chunk[]
+    local matched_chunk_stack = {}
+
+    for _, main_fragment_chunk in ipairs(main_fragment_chunks) do
+        if __chunk_has_all_fragments(main_fragment_chunk, unpack(query.fragments)) then
+            matched_chunk_stack[#matched_chunk_stack + 1] = main_fragment_chunk
+        end
+    end
+
+    return function()
+        while #matched_chunk_stack > 0 do
+            local matched_chunk = matched_chunk_stack[#matched_chunk_stack]
+            matched_chunk_stack[#matched_chunk_stack] = nil
+
+            for _, matched_chunk_child in pairs(matched_chunk.children) do
+                matched_chunk_stack[#matched_chunk_stack + 1] = matched_chunk_child
+            end
+
+            return matched_chunk
+        end
+    end
+end
+
+---
+---
+---
+---
+---
+
+---@param fragment evolved.entity
+---@param ... evolved.entity
+---@return evolved.chunk
+---@nodiscard
+function registry.chunk(fragment, ...)
+    local fragments = { fragment, ... }
+
+    table.sort(fragments, function(a, b)
+        return a.guid < b.guid
+    end)
+
+    local chunk = __root_chunk(fragments[1])
+
+    for i = 2, #fragments do
+        chunk = __chunk_with_fragment(chunk, fragments[i])
+    end
+
+    return chunk
+end
+
+---
+---
+---
+---
+---
+
+function evolved_entity_mt:__tostring()
+    local index, version = idpools.unpack_id(self.guid)
+
+    return string.format('[%d;%d]', index, version)
+end
+
+function evolved_query_mt:__tostring()
+    local fragment_ids = ''
+
+    for _, fragment in ipairs(self.fragments) do
+        fragment_ids = string.format('%s%s', fragment_ids, fragment)
+    end
+
+    return string.format('(%s)', fragment_ids)
+end
+
+function evolved_chunk_mt:__tostring()
+    local fragment_ids = ''
+
+    local chunk_iter = self; while chunk_iter do
+        fragment_ids = string.format('%s%s', chunk_iter.fragment, fragment_ids)
+        chunk_iter = chunk_iter.parent
+    end
+
+    return string.format('{%s}', fragment_ids)
+end
+
+---
+---
+---
+---
+---
 
 return registry