feat(recs): add support for custom type in component (#158)

* feat(recs): add ability to specify custom component type * feat(recs): add OptionalT to allow optional custom type
latticexyz · Sep 26, 2022 · fdc781d · fdc781d
1 parent 34790d1
commit fdc781d
Show file tree

Hide file tree

Showing 7 changed files with 111 additions and 82 deletions.
diff --git a/packages/recs/src/Component.ts b/packages/recs/src/Component.ts
@@ -38,7 +38,7 @@ import { isFullComponentValue, isIndexer } from "./utils";
  * const Position = defineComponent(world, { x: Type.Number, y: Type.Number }, { id: "Position" });
  * ```
  */
-export function defineComponent<S extends Schema, M extends Metadata>(
+export function defineComponent<S extends Schema, M extends Metadata, T = undefined>(
   world: World,
   schema: S,
   options?: { id?: string; metadata?: M; indexed?: boolean }
@@ -49,9 +49,9 @@ export function defineComponent<S extends Schema, M extends Metadata>(
   const update$ = new Subject();
   const metadata = options?.metadata;
   const entities = () => (Object.values(values)[0] as Map<EntityIndex, unknown>).keys();
-  let component = { values, schema, id, update$, metadata, entities, world } as Component<S, M>;
+  let component = { values, schema, id, update$, metadata, entities, world } as Component<S, M, T>;
   if (options?.indexed) component = createIndexer(component);
-  world.registerComponent(component);
+  world.registerComponent(component as Component);
   return component;
 }
 
@@ -67,7 +67,11 @@ export function defineComponent<S extends Schema, M extends Metadata>(
  * setComponent(Position, entity, { x: 1, y: 2 });
  * ```
  */
-export function setComponent<S extends Schema>(component: Component<S>, entity: EntityIndex, value: ComponentValue<S>) {
+export function setComponent<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
+  entity: EntityIndex,
+  value: ComponentValue<S, T>
+) {
   const prevValue = getComponentValue(component, entity);
   for (const [key, val] of Object.entries(value)) {
     component.values[key].set(entity, val);
@@ -91,10 +95,10 @@ export function setComponent<S extends Schema>(component: Component<S>, entity:
  * updateComponent(Position, entity, { x: 1 });
  * ```
  */
-export function updateComponent<T extends Schema>(
-  component: Component<T>,
+export function updateComponent<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
   entity: EntityIndex,
-  value: Partial<ComponentValue<T>>
+  value: Partial<ComponentValue<S, T>>
 ) {
   const currentValue = getComponentValueStrict(component, entity);
   setComponent(component, entity, { ...currentValue, ...value });
@@ -106,7 +110,10 @@ export function updateComponent<T extends Schema>(
  * @param component {@link defineComponent Component} to be updated.
  * @param entity {@link EntityIndex} of the entity whose value should be removed from this component.
  */
-export function removeComponent(component: Component, entity: EntityIndex) {
+export function removeComponent<S extends Schema, M extends Metadata, T>(
+  component: Component<S, M, T>,
+  entity: EntityIndex
+) {
   const prevValue = getComponentValue(component, entity);
   for (const key of Object.keys(component.values)) {
     component.values[key].delete(entity);
@@ -121,7 +128,10 @@ export function removeComponent(component: Component, entity: EntityIndex) {
  * @param entity {@link EntityIndex} of the entity to check whether it has a value in the given component.
  * @returns true if the component contains a value for the given entity, else false.
  */
-export function hasComponent<T extends Schema>(component: Component<T>, entity: EntityIndex): boolean {
+export function hasComponent<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
+  entity: EntityIndex
+): boolean {
   const map = Object.values(component.values)[0];
   return map.has(entity);
 }
@@ -134,10 +144,10 @@ export function hasComponent<T extends Schema>(component: Component<T>, entity:
  * @param entity {@link EntityIndex} of the entity to get the value for from the given component.
  * @returns Value of the given entity in the given component or undefined if no value exists.
  */
-export function getComponentValue<S extends Schema>(
-  component: Component<S>,
+export function getComponentValue<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
   entity: EntityIndex
-): ComponentValue<S> | undefined {
+): ComponentValue<S, T> | undefined {
   const value: Record<string, unknown> = {};
 
   // Get the value of each schema key
@@ -148,7 +158,7 @@ export function getComponentValue<S extends Schema>(
     value[key] = val;
   }
 
-  return value as ComponentValue<S>;
+  return value as ComponentValue<S, T>;
 }
 
 /**
@@ -162,10 +172,10 @@ export function getComponentValue<S extends Schema>(
  * @remarks
  * Throws an error if no value exists in the component for the given entity.
  */
-export function getComponentValueStrict<T extends Schema>(
-  component: Component<T>,
+export function getComponentValueStrict<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
   entity: EntityIndex
-): ComponentValue<T> {
+): ComponentValue<S, T> {
   const value = getComponentValue(component, entity);
   if (!value) throw new Error(`No value for component ${component.id} on entity ${component.world.entities[entity]}`);
   return value;
@@ -185,7 +195,10 @@ export function getComponentValueStrict<T extends Schema>(
  * componentValueEquals({ x: 1 }, { x: 1, y: 3 }) // returns true because x is equal and y is not present in a
  * ```
  */
-export function componentValueEquals<T extends Schema>(a?: Partial<ComponentValue<T>>, b?: ComponentValue<T>): boolean {
+export function componentValueEquals<S extends Schema, T = undefined>(
+  a?: Partial<ComponentValue<S, T>>,
+  b?: ComponentValue<S, T>
+): boolean {
   if (!a && !b) return true;
   if (!a || !b) return false;
 
@@ -205,10 +218,10 @@ export function componentValueEquals<T extends Schema>(a?: Partial<ComponentValu
  * @param value {@link ComponentValue} with {@link ComponentSchema} `S`
  * @returns Tuple `[component, value]`
  */
-export function withValue<S extends Schema>(
-  component: Component<S>,
-  value: ComponentValue<S>
-): [Component<S>, ComponentValue<S>] {
+export function withValue<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>,
+  value: ComponentValue<S, T>
+): [Component<S, Metadata, T>, ComponentValue<S, T>] {
   return [component, value];
 }
 
@@ -219,9 +232,9 @@ export function withValue<S extends Schema>(
  * @param value look for entities with this {@link ComponentValue}.
  * @returns Set with {@link EntityIndex EntityIndices} of the entities with the given component value.
  */
-export function getEntitiesWithValue<T extends Schema>(
-  component: Component<T> | Indexer<T>,
-  value: Partial<ComponentValue<T>>
+export function getEntitiesWithValue<S extends Schema>(
+  component: Component<S> | Indexer<S>,
+  value: Partial<ComponentValue<S>>
 ): Set<EntityIndex> {
   // Shortcut for indexers
   if (isIndexer(component) && isFullComponentValue(component, value)) {
@@ -245,7 +258,9 @@ export function getEntitiesWithValue<T extends Schema>(
  * @param component {@link defineComponent Component} to get all entities from
  * @returns Set of all entities in the given component.
  */
-export function getComponentEntities(component: Component): IterableIterator<EntityIndex> {
+export function getComponentEntities<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>
+): IterableIterator<EntityIndex> {
   return component.entities();
 }
 
@@ -263,27 +278,29 @@ export function getComponentEntities(component: Component): IterableIterator<Ent
  * @param component {@link defineComponent Component} to use as underlying source for the overridable component
  * @returns overridable component
  */
-export function overridableComponent<S extends Schema>(component: Component<S>): OverridableComponent<S> {
+export function overridableComponent<S extends Schema, T = undefined>(
+  component: Component<S, Metadata, T>
+): OverridableComponent<S, T> {
   let nonce = 0;
 
   // Map from OverrideId to Override (to be able to add multiple overrides to the same Entity)
-  const overrides = new Map<string, { update: Override<S>; nonce: number }>();
+  const overrides = new Map<string, { update: Override<S, T>; nonce: number }>();
 
   // Map from EntityIndex to current overridden component value
-  const overriddenEntityValues = new Map<EntityIndex, Partial<ComponentValue<S>> | null>();
+  const overriddenEntityValues = new Map<EntityIndex, Partial<ComponentValue<S, T>> | null>();
 
   // Update event stream that takes into account overridden entity values
   const update$ = new Subject<{
     entity: EntityIndex;
-    value: [ComponentValue<S> | undefined, ComponentValue<S> | undefined];
-    component: Component;
+    value: [ComponentValue<S, T> | undefined, ComponentValue<S, T> | undefined];
+    component: Component<S, Metadata, T>;
   }>();
 
   // Channel through update events from the original component if there are no overrides
   component.update$.pipe(filter((e) => !overriddenEntityValues.get(e.entity))).subscribe(update$);
 
   // Add a new override to some entity
-  function addOverride(id: string, update: Override<S>) {
+  function addOverride(id: string, update: Override<S, T>) {
     overrides.set(id, { update, nonce: nonce++ });
     setOverriddenComponentValue(update.entity, update.value);
   }
@@ -310,11 +327,11 @@ export function overridableComponent<S extends Schema>(component: Component<S>):
   }
 
   // Internal function to get the current overridden value or value of the source component
-  function getOverriddenComponentValue(entity: EntityIndex): ComponentValue<S> | undefined {
+  function getOverriddenComponentValue(entity: EntityIndex): ComponentValue<S, T> | undefined {
     const originalValue = getComponentValue(component, entity);
     const overriddenValue = overriddenEntityValues.get(entity);
     return (originalValue || overriddenValue) && overriddenValue !== null // null is a valid override, in this case return undefined
-      ? ({ ...originalValue, ...overriddenValue } as ComponentValue<S>)
+      ? ({ ...originalValue, ...overriddenValue } as ComponentValue<S, T>)
       : undefined;
   }
 
@@ -344,11 +361,11 @@ export function overridableComponent<S extends Schema>(component: Component<S>):
     },
   });
 
-  const partialValues: Partial<Component<S>["values"]> = {};
+  const partialValues: Partial<Component<S, Metadata, T>["values"]> = {};
   for (const key of Object.keys(component.values) as (keyof S)[]) {
     partialValues[key] = new Proxy(component.values[key], valueProxyHandler(key));
   }
-  const valuesProxy = partialValues as Component<S>["values"];
+  const valuesProxy = partialValues as Component<S, Metadata, T>["values"];
 
   const overriddenComponent = new Proxy(component, {
     get(target, prop) {
@@ -363,10 +380,10 @@ export function overridableComponent<S extends Schema>(component: Component<S>):
       if (prop === "addOverride" || prop === "removeOverride") return true;
       return prop in target;
     },
-  }) as OverridableComponent<S>;
+  }) as OverridableComponent<S, T>;
 
   // Internal function to set the current overridden component value and emit the update event
-  function setOverriddenComponentValue(entity: EntityIndex, value?: Partial<ComponentValue<S>> | null) {
+  function setOverriddenComponentValue(entity: EntityIndex, value?: Partial<ComponentValue<S, T>> | null) {
     // Check specifically for undefined - null is a valid override
     const prevValue = getOverriddenComponentValue(entity);
     if (value !== undefined) overriddenEntityValues.set(entity, value);

diff --git a/packages/recs/src/Indexer.ts b/packages/recs/src/Indexer.ts
@@ -15,19 +15,21 @@ import { Component, ComponentValue, EntityIndex, Indexer, Metadata, Schema } fro
  * @param component {@link defineComponent Component} to index.
  * @returns Indexed version of the component.
  */
-export function createIndexer<S extends Schema, M extends Metadata>(component: Component<S, M>): Indexer<S, M> {
+export function createIndexer<S extends Schema, M extends Metadata, T = undefined>(
+  component: Component<S, M, T>
+): Indexer<S, M, T> {
   const valueToEntities = new Map<string, Set<EntityIndex>>();
 
-  function getEntitiesWithValue(value: ComponentValue<S>) {
+  function getEntitiesWithValue(value: ComponentValue<S, T>) {
     const entities = valueToEntities.get(getValueKey(value));
     return entities ? new Set([...entities]) : new Set<EntityIndex>();
   }
 
-  function getValueKey(value: ComponentValue<S>): string {
+  function getValueKey(value: ComponentValue<S, T>): string {
     return Object.values(value).join("/");
   }
 
-  function add(entity: EntityIndex, value: ComponentValue<S> | undefined) {
+  function add(entity: EntityIndex, value: ComponentValue<S, T> | undefined) {
     if (!value) return;
     const valueKey = getValueKey(value);
     let entitiesWithValue = valueToEntities.get(valueKey);
@@ -38,7 +40,7 @@ export function createIndexer<S extends Schema, M extends Metadata>(component: C
     entitiesWithValue.add(entity);
   }
 
-  function remove(entity: EntityIndex, value: ComponentValue<S> | undefined) {
+  function remove(entity: EntityIndex, value: ComponentValue<S, T> | undefined) {
     if (!value) return;
     const valueKey = getValueKey(value);
     const entitiesWithValue = valueToEntities.get(valueKey);

diff --git a/packages/recs/src/constants.ts b/packages/recs/src/constants.ts
@@ -16,6 +16,8 @@ export enum Type {
   OptionalEntity,
   EntityArray,
   OptionalEntityArray,
+  T,
+  OptionalT,
 }
 
 /**
@@ -42,4 +44,5 @@ export const OptionalTypes = [
   Type.OptionalNumberArray,
   Type.OptionalString,
   Type.OptionalStringArray,
+  Type.OptionalT,
 ];
diff --git a/packages/recs/src/types.ts b/packages/recs/src/types.ts
@@ -33,7 +33,7 @@ export type Metadata =
 /**
  * Mapping between JavaScript {@link Type} enum and corresponding TypeScript type.
  */
-export type ValueType = {
+export type ValueType<T = undefined> = {
   [Type.Boolean]: boolean;
   [Type.Number]: number;
   [Type.String]: string;
@@ -47,51 +47,53 @@ export type ValueType = {
   [Type.OptionalStringArray]: string[] | undefined;
   [Type.OptionalEntity]: EntityID | undefined;
   [Type.OptionalEntityArray]: EntityID[] | undefined;
+  [Type.T]: T;
+  [Type.OptionalT]: T | undefined;
 };
 
 /**
  * Used to infer the TypeScript type of a component value corresponding to a given {@link Schema}.
  */
-export type ComponentValue<S extends Schema = Schema> = {
-  [key in keyof S]: ValueType[S[key]];
+export type ComponentValue<S extends Schema = Schema, T = undefined> = {
+  [key in keyof S]: ValueType<T>[S[key]];
 };
 
 /**
  * Type of a component update corresponding to a given {@link Schema}.
  */
-export type ComponentUpdate<S extends Schema = Schema> = {
+export type ComponentUpdate<S extends Schema = Schema, T = undefined> = {
   entity: EntityIndex;
-  value: [ComponentValue<S> | undefined, ComponentValue<S> | undefined];
-  component: Component<S>;
+  value: [ComponentValue<S, T> | undefined, ComponentValue<S, T> | undefined];
+  component: Component<S, Metadata, T>;
 };
 
 /**
  * Type of component returned by {@link defineComponent}.
  */
-export interface Component<S extends Schema = Schema, M extends Metadata = Metadata> {
+export interface Component<S extends Schema = Schema, M extends Metadata = Metadata, T = undefined> {
   id: string;
-  values: { [key in keyof S]: Map<EntityIndex, ValueType[S[key]]> };
+  values: { [key in keyof S]: Map<EntityIndex, ValueType<T>[S[key]]> };
   schema: S;
   metadata: M;
   entities: () => IterableIterator<EntityIndex>;
   world: World;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  update$: Subject<ComponentUpdate<S>> & { observers: any };
+  update$: Subject<ComponentUpdate<S, T>> & { observers: any };
 }
 
 /**
  * Type of indexer returned by {@link createIndexer}.
  */
-export type Indexer<S extends Schema, M extends Metadata = Metadata> = Component<S, M> & {
-  getEntitiesWithValue: (value: ComponentValue<S>) => Set<EntityIndex>;
+export type Indexer<S extends Schema, M extends Metadata = Metadata, T = undefined> = Component<S, M, T> & {
+  getEntitiesWithValue: (value: ComponentValue<S, T>) => Set<EntityIndex>;
 };
 
 export type Components = {
-  [key: string]: Component<Schema>;
+  [key: string]: Component;
 };
 
-export interface ComponentWithStream<T extends Schema> extends Component<T> {
-  stream$: Subject<{ entity: EntityIndex; value: ComponentValue<T> | undefined }>;
+export interface ComponentWithStream<S extends Schema, T = undefined> extends Component<S, Metadata, T> {
+  stream$: Subject<{ entity: EntityIndex; value: ComponentValue<S, T> | undefined }>;
 }
 
 export type AnyComponentValue = ComponentValue<Schema>;
@@ -176,16 +178,16 @@ export type QueryFragments = QueryFragment<Schema>[];
 
 export type SchemaOf<C extends Component<Schema>> = C extends Component<infer S> ? S : never;
 
-export type Override<T extends Schema> = {
+export type Override<S extends Schema, T = undefined> = {
   entity: EntityIndex;
-  value: Partial<ComponentValue<T>> | null;
+  value: Partial<ComponentValue<S, T>> | null;
 };
 
 /**
  * Type of overridable component returned by {@link overridableComponent}.
  */
-export type OverridableComponent<T extends Schema = Schema> = Component<T> & {
-  addOverride: (actionEntityId: EntityID, update: Override<T>) => void;
+export type OverridableComponent<S extends Schema = Schema, T = undefined> = Component<S, Metadata, T> & {
+  addOverride: (actionEntityId: EntityID, update: Override<S, T>) => void;
   removeOverride: (actionEntityId: EntityID) => void;
 };
 

diff --git a/packages/std-client/src/components/ActionComponent.ts b/packages/std-client/src/components/ActionComponent.ts
@@ -1,5 +1,10 @@
-import { defineComponent, World, Type } from "@latticexyz/recs";
+import { defineComponent, World, Type, Component, Metadata, SchemaOf } from "@latticexyz/recs";
 
-export function defineActionComponent(world: World) {
-  return defineComponent(world, { state: Type.Number, on: Type.OptionalEntity }, { id: "Action" });
+export function defineActionComponent<T = undefined>(world: World) {
+  const Action = defineComponent(
+    world,
+    { state: Type.Number, on: Type.OptionalEntity, metadata: Type.OptionalT },
+    { id: "Action" }
+  );
+  return Action as Component<SchemaOf<typeof Action>, Metadata, T>;
 }