owning group iterator monster

zig-ecs/src/ecs/groups.zig

@@ -13,9 +13,10 @@ pub const BasicGroup = struct {
     registry: *Registry,
     group_data: *Registry.GroupData,
 
+    // TODO: do we even need an iterator for a group?
     pub const Iterator = struct {
         index: usize = 0,
-        entities: []const Entity,
+        entities: []const Entity, // TODO: should this be a pointer to the slice?
 
         pub fn init(entities: []const Entity) Iterator {
             return .{ .entities = entities };
@@ -68,6 +69,58 @@ pub const OwningGroup = struct {
     group_data: *Registry.GroupData,
     super: *usize,
 
+    fn Iterator(comptime Components: var) type {
+        return struct {
+            index: usize = 0,
+            group: OwningGroup,
+            storage: *Storage(u1),
+            component_ptrs: [@typeInfo(Components).Struct.fields.len][*]u8,
+
+            pub fn init(group: OwningGroup) @This() {
+                const component_info = @typeInfo(Components).Struct;
+
+                // get the data pointers for the chunks
+                var component_ptrs: [component_info.fields.len][*]u8 = undefined;
+                inline for (component_info.fields) |field, i| {
+                    const storage = group.registry.assure(field.field_type.Child);
+                    component_ptrs[i] = @ptrCast([*]u8, storage.instances.items.ptr);
+                }
+
+                return .{
+                    .group = group,
+                    .storage = group.firstOwnedStorage(),
+                    .component_ptrs = component_ptrs,
+                };
+            }
+
+            pub fn next(it: *@This()) ?Components {
+                if (it.index >= it.group.group_data.current) return null;
+
+                const ent = it.storage.set.dense.items[it.index];
+                const entity_index = it.storage.set.index(ent);
+                it.index += 1;
+
+                // fill and return the struct
+                var comps: Components = undefined;
+                inline for (@typeInfo(Components).Struct.fields) |field, i| {
+                    const typed_ptr = @ptrCast([*]field.field_type.Child, @alignCast(@alignOf(field.field_type.Child), it.component_ptrs[i]));
+                    @field(comps, field.name) = &typed_ptr[entity_index];
+                }
+                return comps;
+            }
+
+            pub fn entity(it: @This()) Entity {
+                std.debug.assert(it.index > 0 and it.index <= it.group.group_data.current);
+                return it.storage.set.dense.items[it.index - 1];
+            }
+
+            // Reset the iterator to the initial index
+            pub fn reset(it: *@This()) void {
+                it.index = 0;
+            }
+        };
+    }
+
     pub fn init(registry: *Registry, group_data: *Registry.GroupData, super: *usize) OwningGroup {
         return .{
             .registry = registry,
@@ -76,13 +129,69 @@ pub const OwningGroup = struct {
         };
     }
 
+    fn firstOwnedStorage(self: OwningGroup) *Storage(u1) {
+        const ptr = self.registry.components.getValue(self.group_data.owned[0]).?;
+        return @intToPtr(*Storage(u1), ptr);
+    }
+
     pub fn len(self: OwningGroup) usize {
         return self.group_data.current;
     }
 
+    /// direct access to the array of entities of the first owning group
+    pub fn data(self: OwningGroup) []const Entity {
+        return self.firstOwnedStorage().data();
+    }
+
+    pub fn contains(self: OwningGroup, entity: Entity) bool {
+        var storage = self.firstOwnedStorage();
+        return storage.contains(entity) and storage.set.index(entity) < self.len();
+    }
+
+    pub fn getOwned(self: OwningGroup, entity: Entity, comptime Components: var) Components {
+        // TODO: validate that we have a struct
+        // TODO: validate that all fields are pointers
+        // TODO: validate that all fields are owned
+        const component_info = @typeInfo(Components).Struct;
+
+        // get the data pointers for the chunks
+        var component_ptrs: [component_info.fields.len][*]u8 = undefined;
+        inline for (component_info.fields) |field, i| {
+            const storage = self.registry.assure(field.field_type.Child);
+            component_ptrs[i] = @ptrCast([*]u8, storage.instances.items.ptr);
+        }
+
+        // fill the struct
+        const index = self.firstOwnedStorage().set.index(entity);
+        var comps: Components = undefined;
+        inline for (component_info.fields) |field, i| {
+            const typed_ptr = @ptrCast([*]field.field_type.Child, @alignCast(@alignOf(field.field_type.Child), component_ptrs[i]));
+            @field(comps, field.name) = &typed_ptr[index];
+        }
+
+        return comps;
+    }
+
+    /// returns the component storage for the given type for direct access
+    pub fn getStorage(self: *OwningGroup, comptime T: type) *Storage(T) {
+        return self.registry.assure(T);
+    }
+
+    pub fn get(self: *OwningGroup, comptime T: type, entity: Entity) *T {
+        return self.registry.assure(T).get(entity);
+    }
+
+    pub fn getConst(self: *OwningGroup, comptime T: type, entity: Entity) T {
+        return self.registry.assure(T).getConst(entity);
+    }
+
     pub fn sortable(self: OwningGroup, comptime T: type) bool {
         return self.group_data.super == self.group_data.size;
     }
+
+    pub fn iterator(self: OwningGroup, comptime Components: var) Iterator(Components) {
+        return Iterator(Components).init(self);
+    }
 };
 
 test "BasicGroup creation/iteration" {
@@ -160,6 +269,18 @@ test "OwningGroup" {
     reg.add(e0, @as(i32, 44));
     reg.add(e0, @as(u32, 55));
     std.testing.expectEqual(group.len(), 1);
+    std.testing.expect(group.contains(e0));
+
+    std.testing.expectEqual(group.get(i32, e0).*, 44);
+    std.testing.expectEqual(group.getConst(u32, e0), 55);
+
+    var vals = group.getOwned(e0, struct { int: *i32, uint: *u32 });
+    std.testing.expectEqual(vals.int.*, 44);
+    std.testing.expectEqual(vals.uint.*, 55);
+
+    vals.int.* = 666;
+    var vals2 = group.getOwned(e0, struct { int: *i32, uint: *u32 });
+    std.testing.expectEqual(vals2.int.*, 666);
 }
 
 test "OwningGroup add/remove" {
@@ -177,6 +298,31 @@ test "OwningGroup add/remove" {
     std.testing.expectEqual(group.len(), 0);
 }
 
+test "OwningGroup iterate" {
+    var reg = Registry.init(std.testing.allocator);
+    defer reg.deinit();
+
+    var e0 = reg.create();
+    reg.add(e0, @as(i32, 44));
+    reg.add(e0, @as(u32, 55));
+
+    var e1 = reg.create();
+    reg.add(e1, @as(i32, 666));
+    reg.add(e1, @as(u32, 999));
+
+    var group = reg.group(.{ i32, u32 }, .{}, .{});
+    var iter = group.iterator(struct { int: *i32, uint: *u32 });
+    while (iter.next()) |item| {
+        if (iter.entity() == 0) {
+            std.testing.expectEqual(item.int.*, 44);
+            std.testing.expectEqual(item.uint.*, 55);
+        } else {
+            std.testing.expectEqual(item.int.*, 666);
+            std.testing.expectEqual(item.uint.*, 999);
+        }
+    }
+}
+
 test "multiple OwningGroups" {
     const Sprite = struct { x: f32 };
     const Transform = struct { x: f32 };