Cleanup & refactor

lib/utilities/ref.luau

@@ -0,0 +1,85 @@
+--!strict
+--!optimize 2
+local jecs = require("../../jecs")
+type Entity<T = unknown> = jecs.Entity<T>
+
+export type Identity = typeof(setmetatable(
+    {},
+    {} :: {
+        __call: <T>(any, key: unknown) -> (Entity<T>, Cleaner),
+        __index: {
+            reference: <T>(key: unknown) -> (Entity<T>, Cleaner),
+            find: <T>(key: unknown) -> (Entity<T>?, Cleaner?),
+        },
+    }
+))
+
+type Cleaner = () -> ()
+
+local ref_cache: { [jecs.World]: Identity } = {}
+
+local function construct(world: jecs.World, skip_cache: boolean?): Identity
+    if not skip_cache then
+        local hit = ref_cache[world]
+        if hit then
+            return hit
+        end
+    end
+
+    local lookup: { [unknown]: Entity } = {}
+    local cleaner_cache: { [unknown]: Cleaner } = {}
+
+    local function serve_cleaner(key: unknown): () -> ()
+        local hit = cleaner_cache[key]
+        if hit then
+            return hit
+        end
+
+        local function cleaner()
+            lookup[key] = nil
+            cleaner_cache[key] = nil
+        end
+        cleaner_cache[key] = cleaner
+
+        return cleaner
+    end
+
+    local function ref<T>(key: unknown): (Entity<T>, Cleaner)
+        local entity = lookup[key]
+        if not entity then
+            entity = world:entity()
+            lookup[key] = entity
+        end
+
+        return entity, serve_cleaner(key)
+    end
+
+    local function find<T>(key: unknown): (Entity<T>?, Cleaner?)
+        local entity = lookup[key]
+        if not entity then
+            return nil, nil
+        end
+
+        return entity, serve_cleaner(key)
+    end
+
+    local function call<T>(_, key: unknown): (Entity<T>, Cleaner)
+        return ref(key)
+    end
+
+    local self = setmetatable({}, {
+        __call = call,
+        __index = {
+            reference = ref,
+            find = find,
+        },
+    })
+
+    if not skip_cache then
+        ref_cache[world] = self
+    end
+
+    return self
+end
+
+return construct