zig-raylib-engine/zig-ecs/src/ecs/registry.zig

const std = @import("std");
const assert = std.debug.assert;
const utils = @import("utils.zig");

const Handles = @import("handles.zig").Handles;
const SparseSet = @import("sparse_set.zig").SparseSet;
const ComponentStorage = @import("component_storage.zig").ComponentStorage;
const Sink = @import("../signals/sink.zig").Sink;

// allow overriding EntityTraits by setting in root via: EntityTraits = EntityTraitsType(.medium);
const root = @import("root");
const entity_traits = if (@hasDecl(root, "EntityTraits")) root.EntityTraits.init() else @import("entity.zig").EntityTraits.init();

// setup the Handles type based on the type set in EntityTraits
const EntityHandles = Handles(entity_traits.entity_type, entity_traits.index_type, entity_traits.version_type);
pub const Entity = entity_traits.entity_type;

pub const BasicView = @import("views.zig").BasicView;
pub const MultiView = @import("views.zig").MultiView;
pub const BasicGroup = @import("groups.zig").BasicGroup;
pub const OwningGroup = @import("groups.zig").OwningGroup;

/// Stores an ArrayList of components. The max amount that can be stored is based on the type below
pub fn Storage(comptime CompT: type) type {
    return ComponentStorage(CompT, Entity, u16); // 65,535 components
}

/// the registry is the main gateway to all ecs functionality. It assumes all internal allocations will succeed and returns
/// no errors to keep the API clean and because if a component array cant be allocated you've got bigger problems.
/// Stores a maximum of u8 (256) component Storage(T).
pub const Registry = struct {
    handles: EntityHandles,
    components: std.AutoHashMap(u32, usize),
    contexts: std.AutoHashMap(u32, usize),
    groups: std.ArrayList(*GroupData),
    allocator: *std.mem.Allocator,

    /// internal, persistant data structure to manage the entities in a group
    const GroupData = struct {
        hash: u32,
        entity_set: SparseSet(Entity, u16) = undefined, // TODO: dont hardcode this. put it in EntityTraits maybe. All SparseSets would need to use the value.
        owned: []u32,
        include: []u32,
        exclude: []u32,
        registry: *Registry,
        current: usize,

        pub fn initPtr(allocator: *std.mem.Allocator, registry: *Registry, hash: u32, owned: []u32, include: []u32, exclude: []u32) *GroupData {
            std.debug.assert(std.mem.indexOfAny(u32, owned, include) == null);
            std.debug.assert(std.mem.indexOfAny(u32, owned, exclude) == null);
            std.debug.assert(std.mem.indexOfAny(u32, include, exclude) == null);

            var group_data = allocator.create(GroupData) catch unreachable;
            group_data.hash = hash;
            if (owned.len == 0) {
                group_data.entity_set = SparseSet(Entity, u16).init(allocator);
            }
            group_data.owned = std.mem.dupe(allocator, u32, owned) catch unreachable;
            group_data.include = std.mem.dupe(allocator, u32, include) catch unreachable;
            group_data.exclude = std.mem.dupe(allocator, u32, exclude) catch unreachable;
            group_data.registry = registry;
            group_data.current = 0;

            return group_data;
        }

        pub fn deinit(self: *GroupData, allocator: *std.mem.Allocator) void {
            // only deinit th SparseSet for non-owning groups
            if (self.owned.len == 0) {
                self.entity_set.deinit();
            }
            allocator.free(self.owned);
            allocator.free(self.include);
            allocator.free(self.exclude);
            allocator.destroy(self);
        }

        fn maybeValidIf(self: *GroupData, entity: Entity) void {
            const isValid: bool = blk: {
                for (self.owned) |tid| {
                    const ptr = self.registry.components.getValue(tid).?;
                    if (!@intToPtr(*Storage(u1), ptr).contains(entity))
                        break :blk false;
                }

                for (self.include) |tid| {
                    const ptr = self.registry.components.getValue(tid).?;
                    if (!@intToPtr(*Storage(u1), ptr).contains(entity))
                        break :blk false;
                }

                for (self.exclude) |tid| {
                    const ptr = self.registry.components.getValue(tid).?;
                    if (@intToPtr(*Storage(u1), ptr).contains(entity))
                        break :blk false;
                }
                break :blk true;
            };

            if (self.owned.len == 0) {
                if (isValid and !self.entity_set.contains(entity))
                    self.entity_set.add(entity);
            } else {
                if (isValid) {
                    const ptr = self.registry.components.getValue(self.owned[0]).?;
                    if (!(@intToPtr(*Storage(u1), ptr).set.index(entity) < self.current)) {
                        for (self.owned) |tid| {
                            const store_ptr = self.registry.components.getValue(tid).?;
                            var store = @intToPtr(*Storage(u1), store_ptr);
                            store.swap(store.data().*[self.current], entity);
                        }
                        self.current += 1;
                    }
                }
                std.debug.assert(self.owned.len >= 0);
            }
        }

        fn discardIf(self: *GroupData, entity: Entity) void {
            if (self.owned.len == 0) {
                if (self.entity_set.contains(entity))
                    self.entity_set.remove(entity);
            } else {
                const ptr = self.registry.components.getValue(self.owned[0]).?;
                var store = @intToPtr(*Storage(u1), ptr);
                if (store.contains(entity) and store.set.index(entity) < self.current) {
                    self.current -= 1;
                    for (self.owned) |tid| {
                        const store_ptr = self.registry.components.getValue(tid).?;
                        store = @intToPtr(*Storage(u1), store_ptr);
                        store.swap(store.data().*[self.current], entity);
                    }
                }
            }
        }
    };

    pub fn init(allocator: *std.mem.Allocator) Registry {
        return Registry{
            .handles = EntityHandles.init(allocator),
            .components = std.AutoHashMap(u32, usize).init(allocator),
            .contexts = std.AutoHashMap(u32, usize).init(allocator),
            .groups = std.ArrayList(*GroupData).init(allocator),
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Registry) void {
        var it = self.components.iterator();
        while (it.next()) |ptr| {
            // HACK: we dont know the Type here but we need to call deinit
            var storage = @intToPtr(*Storage(u1), ptr.value);
            storage.deinit();
        }

        for (self.groups.items) |grp| {
            grp.deinit(self.allocator);
        }

        self.components.deinit();
        self.contexts.deinit();
        self.groups.deinit();
        self.handles.deinit();
    }

    pub fn assure(self: *Registry, comptime T: type) *Storage(T) {
        var type_id = utils.typeId(T);
        if (self.components.get(type_id)) |kv| {
            return @intToPtr(*Storage(T), kv.value);
        }

        var comp_set = Storage(T).initPtr(self.allocator);
        var comp_set_ptr = @ptrToInt(comp_set);
        _ = self.components.put(type_id, comp_set_ptr) catch unreachable;
        return comp_set;
    }

    /// Prepares a pool for the given type if required
    pub fn prepare(self: *Registry, comptime T: type) void {
        _ = self.assure(T);
    }

    /// Returns the number of existing components of the given type
    pub fn len(self: *Registry, comptime T: type) usize {
        self.assure(T).len();
    }

    /// Increases the capacity of the registry or of the pools for the given component
    pub fn reserve(self: *Self, comptime T: type, cap: usize) void {
        self.assure(T).reserve(cap);
    }

    /// Direct access to the list of components of a given pool
    pub fn raw(self: Registry, comptime T: type) []T {
        return self.assure(T).raw();
    }

    /// Direct access to the list of entities of a given pool
    pub fn data(self: Registry, comptime T: type) []Entity {
        return self.assure(T).data();
    }

    pub fn valid(self: *Registry, entity: Entity) bool {
        return self.handles.isAlive(entity);
    }

    /// Returns the entity identifier without the version
    pub fn entityId(self: Registry, entity: Entity) Entity {
        return entity & entity_traits.entity_mask;
    }

    /// Returns the version stored along with an entity identifier
    pub fn version(self: *Registry, entity: Entity) entity_traits.version_type {
        return @truncate(entity_traits.version_type, entity >> @bitSizeOf(entity_traits.index_type));
    }

    /// Creates a new entity and returns it
    pub fn create(self: *Registry) Entity {
        return self.handles.create();
    }

    /// Destroys an entity
    pub fn destroy(self: *Registry, entity: Entity) void {
        assert(self.valid(entity));
        self.removeAll(entity);
        self.handles.remove(entity) catch unreachable;
    }

    pub fn add(self: *Registry, entity: Entity, value: var) void {
        assert(self.valid(entity));
        self.assure(@TypeOf(value)).add(entity, value);
    }

    /// shortcut for adding raw comptime_int/float without having to @as cast
    pub fn addTyped(self: *Registry, comptime T: type, entity: Entity, value: T) void {
        self.add(entity, value);
    }

    /// Replaces the given component for an entity
    pub fn replace(self: *Registry, entity: Entity, value: var) void {
        assert(self.valid(entity));
        self.assure(@TypeOf(value)).replace(entity, value);
    }

    /// shortcut for replacing raw comptime_int/float without having to @as cast
    pub fn replaceTyped(self: *Registry, comptime T: type, entity: Entity, value: T) void {
        self.replace(entity, value);
    }

    pub fn addOrReplace(self: *Registry, entity: Entity, value: var) void {
        assert(self.valid(entity));

        const store = self.assure(@TypeOf(value));
        if (store.tryGet(entity)) |found| {
            found.* = value;
        } else {
            store.add(entity, value);
        }
    }

    /// shortcut for add-or-replace raw comptime_int/float without having to @as cast
    pub fn addOrReplaceTyped(self: *Registry, T: type, entity: Entity, value: T) void {
        self.addOrReplace(entity, value);
    }

    /// Removes the given component from an entity
    pub fn remove(self: *Registry, comptime T: type, entity: Entity) void {
        assert(self.valid(entity));
        self.assure(T).remove(entity);
    }

    pub fn removeIfExists(self: *Registry, comptime T: type, entity: Entity) void {
        assert(self.valid(entity));
        var store = self.assure(T);
        if (store.contains(entity))
            store.remove(entity);
    }

    /// Removes all the components from an entity and makes it orphaned
    pub fn removeAll(self: *Registry, entity: Entity) void {
        assert(self.valid(entity));

        var it = self.components.iterator();
        while (it.next()) |ptr| {
            // HACK: we dont know the Type here but we need to be able to call methods on the Storage(T)
            var store = @intToPtr(*Storage(u128), ptr.value);
            if (store.contains(entity)) store.remove(entity);
        }
    }

    pub fn has(self: *Registry, comptime T: type, entity: Entity) bool {
        assert(self.valid(entity));
        return self.assure(T).set.contains(entity);
    }

    pub fn get(self: *Registry, comptime T: type, entity: Entity) *T {
        assert(self.valid(entity));
        return self.assure(T).get(entity);
    }

    pub fn getConst(self: *Registry, comptime T: type, entity: Entity) T {
        assert(self.valid(entity));
        return self.assure(T).getConst(entity);
    }

    /// Returns a reference to the given component for an entity
    pub fn getOrAdd(self: *Registry, comptime T: type, entity: Entity) *T {
        if (self.has(T, entity)) return self.get(T, entity);
        self.add(T, entity, std.mem.zeros(T));
        return self.get(T, type);
    }

    pub fn tryGet(self: *Registry, comptime T: type, entity: Entity) ?*T {
        return self.assure(T).tryGet(entity);
    }

    /// Returns a Sink object for the given component to add/remove listeners with
    pub fn onConstruct(self: *Registry, comptime T: type) Sink(Entity) {
        return self.assure(T).onConstruct();
    }

    /// Returns a Sink object for the given component to add/remove listeners with
    pub fn onUpdate(self: *Registry, comptime T: type) Sink(Entity) {
        return self.assure(T).onUpdate();
    }

    /// Returns a Sink object for the given component to add/remove listeners with
    pub fn onDestruct(self: *Registry, comptime T: type) Sink(Entity) {
        return self.assure(T).onDestruct();
    }

    /// Binds an object to the context of the registry
    pub fn setContext(self: *Registry, context: var) void {
        std.debug.assert(@typeInfo(@TypeOf(context)) == .Pointer);

        var type_id = utils.typeId(@typeInfo(@TypeOf(context)).Pointer.child);
        _ = self.contexts.put(type_id, @ptrToInt(context)) catch unreachable;
    }

    /// Unsets a context variable if it exists
    pub fn unsetContext(self: *Registry, comptime T: type) void {
        std.debug.assert(@typeInfo(T) != .Pointer);
        _ = self.contexts.put(utils.typeId(T), 0) catch unreachable;
    }

    /// Returns a pointer to an object in the context of the registry
    pub fn getContext(self: *Registry, comptime T: type) ?*T {
        std.debug.assert(@typeInfo(T) != .Pointer);

        return if (self.contexts.get(utils.typeId(T))) |ptr|
            return if (ptr.value > 0) @intToPtr(*T, ptr.value) else null
        else
            null;
    }

    /// Checks whether the given component belongs to any group
    pub fn sortable(self: Registry, comptime T: type) bool {
        return true;
    }

    pub fn view(self: *Registry, comptime includes: var, comptime excludes: var) ViewType(includes, excludes) {
        if (@typeInfo(@TypeOf(includes)) != .Struct)
            @compileError("Expected tuple or struct argument, found " ++ @typeName(@TypeOf(args)));
        if (@typeInfo(@TypeOf(excludes)) != .Struct)
            @compileError("Expected tuple or struct argument, found " ++ @typeName(@TypeOf(excludes)));
        std.debug.assert(includes.len > 0);

        if (includes.len == 1 and excludes.len == 0)
            return BasicView(includes[0]).init(self.assure(includes[0]));

        var includes_arr: [includes.len]u32 = undefined;
        inline for (includes) |t, i| {
            _ = self.assure(t);
            includes_arr[i] = utils.typeId(t);
        }

        var excludes_arr: [excludes.len]u32 = undefined;
        inline for (excludes) |t, i| {
            _ = self.assure(t);
            excludes_arr[i] = utils.typeId(t);
        }

        return MultiView(includes.len, excludes.len).init(self, includes_arr, excludes_arr);
    }

    /// returns the Type that a view will be based on the includes and excludes
    fn ViewType(comptime includes: var, comptime excludes: var) type {
        if (includes.len == 1 and excludes.len == 0) return BasicView(includes[0]);
        return MultiView(includes.len, excludes.len);
    }

    pub fn group(self: *Registry, comptime owned: var, comptime includes: var, comptime excludes: var) GroupType(owned, includes, excludes) {
        if (@typeInfo(@TypeOf(owned)) != .Struct)
            @compileError("Expected tuple or struct argument, found " ++ @typeName(@TypeOf(owned)));
        if (@typeInfo(@TypeOf(includes)) != .Struct)
            @compileError("Expected tuple or struct argument, found " ++ @typeName(@TypeOf(includes)));
        if (@typeInfo(@TypeOf(excludes)) != .Struct)
            @compileError("Expected tuple or struct argument, found " ++ @typeName(@TypeOf(excludes)));
        std.debug.assert(owned.len + includes.len > 0);
        std.debug.assert(owned.len + includes.len + excludes.len > 1);

        var owned_arr: [owned.len]u32 = undefined;
        inline for (owned) |t, i| {
            _ = self.assure(t);
            owned_arr[i] = utils.typeId(t);
        }

        var includes_arr: [includes.len]u32 = undefined;
        inline for (includes) |t, i| {
            _ = self.assure(t);
            includes_arr[i] = utils.typeId(t);
        }

        var excludes_arr: [excludes.len]u32 = undefined;
        inline for (excludes) |t, i| {
            _ = self.assure(t);
            excludes_arr[i] = utils.typeId(t);
        }

        // create a unique hash to identify the group
        var maybe_group_data: ?*GroupData = null;
        comptime const hash = owned.len + (31 * includes.len) + (31 * 31 * excludes.len);

        for (self.groups.items) |grp| {
            if (grp.hash == hash and std.mem.eql(u32, grp.owned, owned_arr[0..]) and std.mem.eql(u32, grp.include, includes_arr[0..]) and std.mem.eql(u32, grp.exclude, excludes_arr[0..])) {
                maybe_group_data = grp;
                break;
            }
        }

        // do we already have the GroupData?
        if (maybe_group_data) |group_data| {
            // non-owning groups
            if (owned.len == 0) {
                return BasicGroup(includes.len, excludes.len).init(&group_data.entity_set, self, includes_arr, excludes_arr);
            } else {
                return OwningGroup(owned.len, includes.len, excludes.len).init(&group_data.current, self, owned_arr, includes_arr, excludes_arr);
            }
        }

        // we need to create a new GroupData
        var new_group_data = GroupData.initPtr(self.allocator, self, hash, owned_arr[0..], includes_arr[0..], excludes_arr[0..]);
        self.groups.append(new_group_data) catch unreachable;

        // wire up our listeners
        inline for (owned) |t| self.onConstruct(t).connectBound(new_group_data, "maybeValidIf");
        inline for (includes) |t| self.onConstruct(t).connectBound(new_group_data, "maybeValidIf");
        inline for (excludes) |t| self.onDestruct(t).connectBound(new_group_data, "maybeValidIf");

        inline for (owned) |t| self.onDestruct(t).connectBound(new_group_data, "discardIf");
        inline for (includes) |t| self.onDestruct(t).connectBound(new_group_data, "discardIf");
        inline for (excludes) |t| self.onConstruct(t).connectBound(new_group_data, "discardIf");

        // pre-fill the GroupData with any existing entitites that match
        if (owned.len == 0) {
            var tmp_view = self.view(owned ++ includes, excludes);
            var view_iter = tmp_view.iterator();
            while (view_iter.next()) |entity| {
                new_group_data.entity_set.add(entity);
            }
        } else {}

        if (owned.len == 0) {
            return BasicGroup(includes.len, excludes.len).init(&new_group_data.entity_set, self, includes_arr, excludes_arr);
        } else {
            return OwningGroup(owned.len, includes.len, excludes.len).init(&new_group_data.current, self, owned_arr, includes_arr, excludes_arr);
        }
    }

    /// returns the Type that a view will be based on the includes and excludes
    fn GroupType(comptime owned: var, comptime includes: var, comptime excludes: var) type {
        if (owned.len == 0) return BasicGroup(includes.len, excludes.len);
        return OwningGroup(owned.len, includes.len, excludes.len);
    }
};