Merge pull request #136 from pmndrs/fix--world-set-types

fix: world set types

Author: krispya

packages/core/src/entity/types.ts

@@ -2,6 +2,7 @@ import type { Relation } from '../relation/types';
 import type {
 	ConfigurableTrait,
 	ExtractSchema,
+	SetTraitCallback,
 	Trait,
 	TraitInstance,
 	TraitValue,
@@ -15,9 +16,7 @@ export type Entity = number & {
 	changed: (trait: Trait) => void;
 	set: <T extends Trait>(
 		trait: T,
-		value:
-			| TraitValue<ExtractSchema<T>>
-			| ((prev: TraitInstance<ExtractSchema<T>>) => TraitValue<ExtractSchema<T>>),
+		value: TraitValue<ExtractSchema<T>> | SetTraitCallback<T>,
 		flagChanged?: boolean
 	) => void;
 	get: <T extends Trait | Relation<Trait>>(trait: T) => TraitInstance<ExtractSchema<T>> | undefined;

packages/core/src/trait/types.ts

@@ -53,6 +53,10 @@ type TraitInstanceFromTrait<T extends Trait> = T['schema'] extends AoSFactory
 				: T['schema'][P];
 	  };
 
+export type SetTraitCallback<T extends Trait> = (
+	prev: TraitInstance<ExtractSchema<T>>
+) => TraitValue<ExtractSchema<T>>;
+
 type TraitInstanceFromSchema<T extends Schema> = T extends AoSFactory
 	? ReturnType<T>
 	: {

packages/core/src/world/world.ts

@@ -13,6 +13,7 @@ import type { TraitData } from '../trait/trait-data';
 import type {
 	ConfigurableTrait,
 	ExtractSchema,
+	SetTraitCallback,
 	Trait,
 	TraitInstance,
 	TraitValue,
@@ -111,7 +112,7 @@ export class World {
 		return getTrait(this, this[$internal].worldEntity, trait);
 	}
 
-	set<T extends Trait>(trait: T, value: TraitValue<ExtractSchema<T>>) {
+	set<T extends Trait>(trait: T, value: TraitValue<ExtractSchema<T>> | SetTraitCallback<T>) {
 		setTrait(this, this[$internal].worldEntity, trait, value);
 	}
 

packages/core/tests/world.test.ts

@@ -92,6 +92,24 @@ describe('World', () => {
 		expect(world.has(Time)).toBe(false);
 	});
 
+	it('should set singletons', () => {
+		const Test = trait({ last: 0, delta: 0 });
+		const world = createWorld(Test);
+
+		world.set(Test, { last: 1, delta: 1 });
+
+		expect(world.get(Test)!.last).toBe(1);
+		expect(world.get(Test)!.delta).toBe(1);
+
+		// Use callbacks to set.
+		world.set(Test, (prev) => {
+			return { last: prev.last + 1, delta: prev.delta + 1 };
+		});
+
+		expect(world.get(Test)!.last).toBe(2);
+		expect(world.get(Test)!.delta).toBe(2);
+	});
+
 	it('should observe traits', () => {
 		const TimeOfDay = trait({ hour: 0 });
 		const world = createWorld(TimeOfDay);