Cleanup & refactor

lib/utilities/tracker.luau

@@ -0,0 +1,273 @@
+--!strict
+--!optimize 2
+local command_buffer = require("./command_buffer")
+local jecs = require("../../jecs")
+type Entity<T = unknown> = jecs.Entity<T>
+type Id<T = unknown> = jecs.Id<T>
+type i53 = number
+
+local OnAdd = jecs.OnAdd
+local OnSet = jecs.OnSet
+local OnRemove = jecs.OnRemove
+
+local construct_ref = require("./ref")
+
+-- The external type differs for better DX
+export type Commands = {
+    added: { [i53]: { i53 } },
+    set: { [i53]: { [i53]: unknown } },
+    removed: { [i53]: { i53 } },
+}
+
+type Added = { [Id]: { Entity } }
+type Set = { [Id]: { [Entity]: unknown } }
+type Removed = { [Id]: { Entity } }
+type Lookup = { [Id]: { [Entity]: number } }
+type InternalCommands = {
+    added: Added,
+    set: Set,
+    removed: Removed,
+}
+
+--- Tracks operations on entities for the provided world.
+export type Identity = {
+    --- Gets the current state.
+    --- A state is a representation of the minimum of commands necessary to produce the current world from a clean slate.
+    state: () -> Commands,
+    --- Gets the currently tracked snapshot.
+    --- A snapshot is a representation of the minimum of commands necessary to produce the current world back from when the last snapshot was taken.
+    snapshot: () -> Commands?,
+    --- Applies a set of commands to the tracked world, optionally doing it through a command buffer.
+    apply: (snapshot: Commands, buf: command_buffer.Identity?) -> (),
+}
+
+local function get_non_nilable<T, K>(container: {} & T, index: K): index<T, K>
+    local data = container[index]
+    if not data then
+        data = {}
+        container[index] = data
+    end
+    return data
+end
+
+local function insert_unique<T, V>(container: T, value: V, lookup: { [V]: number }?)
+    if lookup then
+        if lookup[value] then
+            return
+        end
+
+        local idx = #lookup + 1
+        lookup[value] = idx;
+        (container :: any)[idx] = value
+        return
+    end
+
+    if table.find(container, value) then
+        return
+    end
+    table.insert(container, value)
+end
+
+local function construct(world: jecs.World, ...: Entity<any>): Identity
+    local components = { ... }
+    local ref = construct_ref(world, true)
+
+    local state_added: Added = {}
+    local state_added_lookup: Lookup = {}
+    local state_set: Set = {}
+    local state_removed: Removed = {}
+    local state_removed_lookup: Lookup = {}
+
+    local snapshot_added: Added = {}
+    local snapshot_set: Set = {}
+    local snapshot_removed: Removed = {}
+
+    for _, component in components do
+        world:set(component, OnAdd, function(entity: Entity)
+            local snapshot = get_non_nilable(snapshot_added, component)
+            insert_unique(snapshot, entity)
+
+            local state = get_non_nilable(state_added, component)
+            local lookup = get_non_nilable(state_added_lookup, component)
+            insert_unique(state, entity, lookup)
+
+            -- Clean up previous operations
+            local set_state = state_set[component]
+            if set_state and set_state[entity] then
+                set_state[entity] = nil
+            end
+
+            local removed_lookup = state_removed_lookup[component]
+            if removed_lookup then
+                local idx = removed_lookup[entity]
+                if idx then
+                    removed_lookup[entity] = nil
+                    local removed_state = state_removed[component]
+                    if removed_state then
+                        -- Shifting around the array could be expensive, prefer `tbl[idx] = nil`
+                        removed_state[idx] = nil
+                    end
+                end
+            end
+        end)
+
+        world:set(component, OnSet, function(entity, value)
+            local snapshot = get_non_nilable(snapshot_set, component)
+            snapshot[entity] = value
+
+            local state = get_non_nilable(state_set, component)
+            state[entity] = value
+
+            -- Clean up previous operations
+            local added_lookup = state_added_lookup[component]
+            if added_lookup then
+                local idx = added_lookup[entity]
+                if idx then
+                    added_lookup[entity] = nil
+                    local added_state = state_added[component]
+                    if added_state then
+                        -- Shifting around the array could get expensive, prefer `array[idx] = nil`
+                        added_state[idx] = nil
+                    end
+                end
+            end
+
+            local removed_lookup = state_removed_lookup[component]
+            if removed_lookup then
+                local idx = removed_lookup[entity]
+                if idx then
+                    removed_lookup[entity] = nil
+                    local removed_state = state_removed[component]
+                    if removed_state then
+                        -- Shifting around the array could get expensive, prefer `array[idx] = nil`
+                        removed_state[idx] = nil
+                    end
+                end
+            end
+        end)
+
+        world:set(component, OnRemove, function(entity: Entity)
+            local snapshot = get_non_nilable(snapshot_removed, component)
+            insert_unique(snapshot, entity)
+
+            local state = get_non_nilable(state_removed, component)
+            local lookup = get_non_nilable(state_removed_lookup, component)
+
+            -- Clean up previous operations
+            local added_lookup = state_added_lookup[component]
+            if added_lookup then
+                local idx = added_lookup[entity]
+                if idx then
+                    added_lookup[entity] = nil
+                    local added_state = state_added[component]
+                    if added_state then
+                        -- Shifting around the array could get expensive, prefer `array[idx] = nil`
+                        added_state[idx] = nil
+                    end
+                end
+            end
+
+            local set_state = state_set[component]
+            if set_state and set_state[entity] then
+                set_state[entity] = nil
+            end
+
+            insert_unique(state, entity, lookup)
+        end)
+    end
+
+    -- We cast anything exposing `Commands` as `any` to improve the types for the end user
+    local function get_state(): InternalCommands
+        return {
+            added = state_added,
+            set = state_set,
+            removed = state_removed,
+        }
+    end
+
+    local function get_snapshot(): InternalCommands?
+        local diff_added = snapshot_added
+        local diff_set = snapshot_set
+        local diff_removed = snapshot_removed
+        snapshot_added = {}
+        snapshot_set = {}
+        snapshot_removed = {}
+
+        if next(diff_added) == nil and next(diff_set) == nil and next(diff_removed) == nil then
+            return nil
+        end
+
+        return {
+            added = diff_added,
+            set = diff_set,
+            removed = diff_removed,
+        }
+    end
+
+    local function apply_snapshot(snapshot: InternalCommands, buf: command_buffer.Identity?)
+        local add
+        local set
+        local remove
+        do
+            if buf then
+                add = buf.add
+                set = buf.set
+                remove = buf.remove
+            else
+                function add<T>(entity: Entity, component: Id<T>)
+                    world:add(entity, component)
+                end
+
+                function set<T>(entity: Entity, component: Id<T>, data: T)
+                    world:set(entity, component, data)
+                end
+
+                function remove<T>(entity: Entity, component: Id<T>)
+                    world:remove(entity, component)
+                end
+            end
+        end
+
+        for component, entities in snapshot.added do
+            for _, id in entities do
+                local entity = ref(`foreign-{id}`)
+
+                if world:has(entity, component) then
+                    continue
+                end
+                add(entity, component)
+            end
+        end
+
+        for component, entities in snapshot.set do
+            for id, data in entities do
+                local entity = ref(`foreign-{id}`)
+
+                if world:get(entity, component) == data then
+                    continue
+                end
+                set(entity, component, data)
+            end
+        end
+
+        for component, entities in snapshot.removed do
+            for _, id in entities do
+                local entity = ref(`foreign-{id}`)
+
+                if world:has(entity, component) then
+                    continue
+                end
+                remove(entity, component)
+            end
+        end
+    end
+
+    -- Public types differ for better DX
+    return {
+        state = get_state,
+        snapshot = get_snapshot,
+        apply = apply_snapshot,
+    } :: any
+end
+
+return construct