Merge pull request #145 from pmndrs/refactor/remove-trait-data-class

refactor: remove TraitData class

Author: krispya

packages/core/src/query/query.ts

@@ -2,8 +2,7 @@ import { $internal } from '../common';
 import type { Entity } from '../entity/types';
 import { getEntityId } from '../entity/utils/pack-entity';
 import { registerTrait, trait } from '../trait/trait';
-import type { TraitData } from '../trait/trait-data';
-import type { Trait } from '../trait/types';
+import type { Trait, TraitData } from '../trait/types';
 import { SparseSet } from '../utils/sparse-set';
 import type { World } from '../world/world';
 import { ModifierData } from './modifier';

packages/core/src/trait/trait-data.ts

@@ -5,8 +5,16 @@ import { setChanged } from '../query/modifiers/changed';
 import { getRelationTargets, Pair, Wildcard } from '../relation/relation';
 import { incrementWorldBitflag } from '../world/utils/increment-world-bit-flag';
 import type { World } from '../world/world';
-import { TraitData } from './trait-data';
-import type { ConfigurableTrait, ExtractStore, Norm, Schema, Store, Trait, TraitType } from './types';
+import type {
+	ConfigurableTrait,
+	ExtractStore,
+	Norm,
+	Schema,
+	Store,
+	Trait,
+	TraitData,
+	TraitType,
+} from './types';
 import {
 	createFastSetChangeFunction,
 	createFastSetFunction,
@@ -49,13 +57,29 @@ export const trait = defineTrait;
 
 export function registerTrait(world: World, trait: Trait) {
 	const ctx = world[$internal];
-	const data = new TraitData(world, trait);
+	const traitCtx = trait[$internal];
+
+	const data: TraitData = {
+		generationId: ctx.entityMasks.length - 1,
+		bitflag: ctx.bitflag,
+		trait,
+		store: traitCtx.createStore(),
+		queries: new Set(),
+		notQueries: new Set(),
+		schema: trait.schema,
+		changeSubscriptions: new Set(),
+		addSubscriptions: new Set(),
+		removeSubscriptions: new Set(),
+	};
+
+	// Bind a reference to the store on the trait for direct access in queries.
+	traitCtx.stores[world.id] = data.store;
 
 	// Add trait to the world.
 	ctx.traitData.set(trait, data);
 	world.traits.add(trait);
 
-	// Increment the world bitflag.
+	// Increment the bitflag used for the trait.
 	incrementWorldBitflag(world);
 }
 

packages/core/src/trait/trait.ts

@@ -1,4 +1,6 @@
 import { $internal } from '../common';
+import { Entity } from '../entity/types';
+import { Query } from '../query/query';
 import type { Relation, RelationTarget } from '../relation/types';
 import type { IsEmpty } from '../utils/types';
 
@@ -116,3 +118,16 @@ export type ExtractStore<T extends Trait> = T extends Trait<any, infer S> ? S :
 export type ExtractIsTag<T extends Trait> = T extends Trait<any, any, infer Tag> ? Tag : false;
 
 export type IsTag<T extends Trait> = T extends Trait<any, any, infer Tag> ? Tag : false;
+
+export interface TraitData<T extends Trait = Trait, S extends Schema = ExtractSchema<T>> {
+	generationId: number;
+	bitflag: number;
+	trait: Trait;
+	store: Store<S>;
+	queries: Set<Query>;
+	notQueries: Set<Query>;
+	schema: S;
+	changeSubscriptions: Set<(entity: Entity) => void>;
+	addSubscriptions: Set<(entity: Entity) => void>;
+	removeSubscriptions: Set<(entity: Entity) => void>;
+}

packages/core/src/trait/types.ts

@@ -8,8 +8,8 @@ import type { QueryHash, QueryParameter, QueryResult, QueryUnsubscriber } from '
 import { createQueryHash } from '../query/utils/create-query-hash';
 import { getTrackingCursor, setTrackingMasks } from '../query/utils/tracking-cursor';
 import type { RelationTarget } from '../relation/types';
+import type { TraitData } from '../trait/types';
 import { addTrait, getTrait, hasTrait, registerTrait, removeTrait, setTrait } from '../trait/trait';
-import type { TraitData } from '../trait/trait-data';
 import type {
 	ConfigurableTrait,
 	ExtractSchema,