Hooks cache & component changes iterator for Jecs

HooksCache.luau

--!strict
--!native
--!optimize 2

local ReplicatedStorage = game:GetService("ReplicatedStorage")

local Jecs = require(ReplicatedStorage.Packages.Jecs)
local World = require(script.Parent.World)

type Entity = Jecs.Entity<nil>
type Component<T = any> = Jecs.Entity<T>

export type Callback = (Entity) -> ()
export type Callbacks = { [Callback]: true? }
export type CallbacksCache = {
	__component: Component,
	Added: Callbacks,
	Changed: Callbacks,
	Removed: Callbacks,
}

export type ChangeSet = { [Entity]: true? }
export type ChangeSets = { [ChangeSet]: true? }
export type ChangeSetsCache = {
	__component: Component,
	Added: ChangeSets,
	Changed: ChangeSets,
	Removed: ChangeSets,
}

local cachedCallbacks: { [Component]: CallbacksCache } = {}
local cachedChangeSets: { [Component]: ChangeSetsCache } = {}

local function globalCacheFor(component: Component): (CallbacksCache, ChangeSetsCache)
	if cachedCallbacks[component] == nil or cachedChangeSets[component] == nil then
		local callbacksAdded: Callbacks = {}
		local callbacksChanged: Callbacks = {}
		local callbacksRemoved: Callbacks = {}

		local changeSetsAdded: ChangeSets = {}
		local changeSetsChanged: ChangeSets = {}
		local changeSetsRemoved: ChangeSets = {}

		World:set(component, Jecs.OnAdd, function(id)
			for set in changeSetsAdded do
				set[id] = true
			end
			for cb in callbacksAdded do
				task.spawn(cb, id)
			end
		end)
		World:set(component, Jecs.OnSet, function(id)
			for set in changeSetsChanged do
				set[id] = true
			end
			for cb in callbacksChanged do
				task.spawn(cb, id)
			end
		end)
		World:set(component, Jecs.OnRemove, function(id)
			for set in changeSetsRemoved do
				set[id] = true
			end
			for cb in callbacksRemoved do
				task.spawn(cb, id)
			end
		end)

		cachedCallbacks[component] = table.freeze({
			__component = component,
			Added = callbacksAdded,
			Changed = callbacksChanged,
			Removed = callbacksRemoved,
		})

		cachedChangeSets[component] = table.freeze({
			__component = component,
			Added = changeSetsAdded,
			Changed = changeSetsChanged,
			Removed = changeSetsRemoved,
		})
	end
	return cachedCallbacks[component], cachedChangeSets[component]
end

local module = {}

function module.GetCallbacks(component: Component): CallbacksCache
	local callbacks, _ = globalCacheFor(component)
	return callbacks
end

function module.GetChangeSets(component: Component): ChangeSetsCache
	local _, changeSets = globalCacheFor(component)
	return changeSets
end

table.freeze(module)

return module

IterComponentChanges.luau

--!strict
--!native
--!optimize 2

local HooksCache = require(script.Parent.HooksCache)
local World = require(script.Parent.World)

type Entity = Jecs.Entity<nil>
type Component<T> = Jecs.Entity<T>

export type Iterator<T> = () -> (Entity, T?, T?)
export type Destructor = () -> ()

type ValuesMap<T> = { [Entity]: T? }

--[=[
	Creates a new observer for the given component.

	### Example Usage

	```luau
	local IterComponentChanges = require(path.To.This.Module)

	local Monster = ... -- The component to observe

	local iterMonsterChanges = IterComponentChanges(Monster)

	local function System()
		for id, old, new in iterMonsterChanges do
			if old == nil and new ~= nil then
				print("Monster added with id", id, "and data", new)
			elseif old ~= nil and new ~= nil then
				print("Monster changed with id", id, "from", old, "to", new)
			elseif old ~= nil and new == nil then
				print("Monster removed with id", id)
			end
		end
	end

	return System
	```
]=]
return function<T>(component: Component<T>): (Iterator<T>, Destructor)
	local values: { [Entity]: T? } = {}
	local changeSet: HooksCache.ChangeSet = {}

	for id in World:query(component) do
		changeSet[id] = true
	end

	local changeSets = HooksCache.GetChangeSets(component)
	changeSets.Added[changeSet] = true
	changeSets.Changed[changeSet] = true
	changeSets.Removed[changeSet] = true

	local id: Entity? = nil
	local iter: Iterator<T> = function()
		id = next(changeSet)
		while id do
			changeSet[id] = nil
			local old: T? = values[id]
			local new: T? = World:get(id, component)
			if old == nil and new == nil then
				-- No old value nor new value, this means that the entity
				-- was both spawned and destroyed before reaching this observer,
				-- in this case we will skip to the next possibly changed entity
				id = next(changeSet)
				continue
			elseif old ~= nil and new == nil then
				-- Old value but no new value = removed, we should
				-- clean up the reference to not leak any memory
				values[id] = nil
			else
				-- Old+new value or just new value = new becomes old
				values[id] = new
			end
			return id, old, new
		end
		return nil :: any, nil, nil
	end

	local destroy: Destructor = function()
		changeSets.Added[changeSet] = nil
		changeSets.Changed[changeSet] = nil
		changeSets.Removed[changeSet] = nil
	end

	return iter, destroy
end

World.luau

local ReplicatedStorage = game:GetService("ReplicatedStorage")

local Jecs = require(ReplicatedStorage.Packages.Jecs) -- Or wherever you installed Jecs

return Jecs.World.new()