diff --git a/zig-ecs/src/ecs/entity.zig b/zig-ecs/src/ecs/entity.zig
index 92b986a..e8847ec 100644
--- a/zig-ecs/src/ecs/entity.zig
+++ b/zig-ecs/src/ecs/entity.zig
@@ -7,29 +7,37 @@ pub const EntityTraitsSize = enum { small, medium, large };
 
 pub fn EntityTraitsType(comptime size: EntityTraitsSize) type {
     return switch (size) {
-        .small => EntityTraitsDefinition(u16, u12, u4, 0xFFFF, 0xFFF, 10),
-        .medium => EntityTraitsDefinition(u32, u20, u12, 0xFFFFF, 0xFFF, 20),
-        .large => EntityTraitsDefinition(u64, u32, u32, 0xFFFFFFFF, 0xFFFFFFFF, 32),
+        .small => EntityTraitsDefinition(u16, u12, u4),
+        .medium => EntityTraitsDefinition(u32, u20, u12),
+        .large => EntityTraitsDefinition(u64, u32, u32),
     };
 }
 
-fn EntityTraitsDefinition(comptime EntityType: type, comptime IndexType: type, comptime VersionType: type, comptime EntityMask: EntityType, comptime VersionMask: EntityType, comptime EntityShift: EntityType) type {
-    std.debug.assert(@typeInfo(EntityType) == .Int and std.meta.Int(.unsigned, @bitSizeOf(EntityType)) == EntityType);
-    std.debug.assert(@typeInfo(IndexType) == .Int and std.meta.Int(.unsigned, @bitSizeOf(IndexType)) == IndexType);
-    std.debug.assert(@typeInfo(VersionType) == .Int and std.meta.Int(.unsigned, @bitSizeOf(VersionType)) == VersionType);
-
-    if (@bitSizeOf(IndexType) + @bitSizeOf(VersionType) != @bitSizeOf(EntityType))
+fn EntityTraitsDefinition(comptime EntityType: type, comptime IndexType: type, comptime VersionType: type) type {
+    std.debug.assert(std.meta.trait.isUnsignedInt(EntityType));
+    std.debug.assert(std.meta.trait.isUnsignedInt(IndexType));
+    std.debug.assert(std.meta.trait.isUnsignedInt(VersionType));
+    
+    const sizeOfIndexType = @bitSizeOf(IndexType);
+    const sizeOfVersionType = @bitSizeOf(VersionType);
+    const entityShift = sizeOfIndexType;
+    
+    if (sizeOfIndexType + sizeOfVersionType != @bitSizeOf(EntityType))
         @compileError("IndexType and VersionType must sum to EntityType's bit count");
 
+    const entityMask = std.math.maxInt(IndexType);
+    const versionMask = std.math.maxInt(VersionType);
+    
     return struct {
         entity_type: type = EntityType,
         index_type: type = IndexType,
         version_type: type = VersionType,
         /// Mask to use to get the entity index number out of an identifier
-        entity_mask: EntityType = EntityMask,
+        entity_mask: EntityType = entityMask,
         /// Mask to use to get the version out of an identifier
-        version_mask: EntityType = VersionMask,
-        entity_shift: EntityType = EntityShift,
+        version_mask: EntityType = versionMask,
+        /// Bit size of entity in entity_type
+        entity_shift: EntityType = entityShift,
 
         pub fn init() @This() {
             return @This(){};