Cache computed 'base' keys for entity classes.

.gitignore

@@ -5,4 +5,7 @@
 tsconfig.tsbuildinfo
 .DS_Store
 /esm
-!.gitkeep
+!.gitkeep
+
+benchmarks/master
+benchmarks/working

package.json

@@ -48,8 +48,8 @@
     "dev:lint": "eslint \"src/**/*.{js,ts}\" --fix --config \".eslintrc.cjs\"",
     "dev:pretty": "prettier \"src/**/*.{js,ts}\" --write",
     "bench": "npm-run-all build bench:working bench:compare",
-    "bench:working": "tsx benchmarks > \"coverage/working\"",
-    "bench:master": "tsx benchmarks > \"coverage/master\"",
-    "bench:compare": "nanobench-compare \"coverage/working\" \"coverage/master\""
+    "bench:working": "tsx benchmarks > \"benchmarks/working\"",
+    "bench:master": "tsx benchmarks > \"benchmarks/master\"",
+    "bench:compare": "nanobench-compare \"benchmarks/working\" \"benchmarks/master\""
   }
 }

src/ecs/Entity.ts

@@ -137,24 +137,35 @@ export class Entity<T extends BaseType = {}> {
 
   protected addComponent<C extends ComponentClass>(ComponentConstructor: C, data?: PartialValueByType<C>): void {
     const type = ComponentConstructor.type as string & keyof T;
+    let dirty = false;
     if (!(type in this.$)) {
       // get the component in question
       this.$[type] = Object.assign(new ComponentConstructor(), data ?? {}) as T[string & keyof T];
-
       this[Components].push(ComponentConstructor);
       // turns out indexing repeatedly is faster than doing a bool set/check
       Manager[ToIndex][this.mid].push([this, this.key ?? null]);
+      // we're opting out of an optimization that uses a precomputed component key for each entity class.
+      dirty = true;
+    }
+    if (dirty) {
+      // reset the parent constructor to "Entity" — the component key now needs to be recomputed
+      this.constructor = Entity;
     }
   }
 
   protected removeComponents(...components: ComponentClass[]): void {
+    let dirty = false;
     for (const C of components) {
       if (C.type in this.$) {
         this[Components].splice(this[Components].indexOf(C), 1);
         delete this.$[C.type];
         Manager[ToIndex][this.mid].push([this, this.key ?? null]);
+        dirty = true;
       }
     }
+    if (dirty) {
+      this.constructor = Entity;
+    }
   }
 
   public constructor(context: Context, data: BaseDataType<T> & { id?: Identifier } = {}, tags: string[] = []) {

src/lib/ChangeSet.ts

@@ -5,6 +5,7 @@
 export class ChangeSet {
   protected onChange?: () => void;
   protected items: Record<string, unknown> = {};
+  protected _size: number = 0;
 
   public *[Symbol.iterator](): Iterator<string> {
     yield* this.all();
@@ -32,13 +33,15 @@ export class ChangeSet {
       }
     }
     if (changed) {
+      this._size = Object.keys(this.items).length;
       this.onChange?.();
     }
     return this;
   }
 
   public clear(): void {
     this.items = {};
+    this._size = 0;
     this.onChange?.();
   }
 
@@ -53,13 +56,19 @@ export class ChangeSet {
     }
 
     if (changed) {
+      this._size = Object.keys(this.items).length;
       this.onChange?.();
     }
     return true;
   }
 
+  public get size() {
+    return this._size;
+  }
+
   public constructor(items: string[], onChange?: () => void) {
     this.items = items.reduce((a, b) => ({ ...a, [b]: true }), {});
+    this._size = Object.keys(this.items).length;
     this.onChange = onChange;
   }
 }

src/lib/Manager.test.ts

@@ -1,7 +1,8 @@
 import { describe, expect, test } from 'vitest';
 
 import { getContext } from '../test/helpers';
-import { A } from '../test/helpers/components';
+import { A, B } from '../test/helpers/components';
+import { WithAB } from '../test/helpers/entities';
 
 describe('queries', () => {
   test('identical queries should be cached', () => {
@@ -12,3 +13,18 @@ describe('queries', () => {
     expect(a).toBe(b);
   });
 });
+
+describe('entity class key caching', () => {
+  test('entities with modified components should "lose" their cached base key', () => {
+    const ctx = getContext();
+    const e1 = ctx.create(WithAB);
+    const e2 = ctx.create(WithAB);
+
+    ctx.tick();
+
+    e1.components.remove(B), ctx.tick();
+
+    // @ts-expect-error: accessibility abuse
+    expect(ctx.manager.getEntityKey(e1)).not.toBe(ctx.manager.getEntityKey(e2));
+  });
+});

src/lib/Manager.ts

@@ -2,7 +2,7 @@ import type { ComponentClass, Context, Entity, EntityClass } from '../ecs';
 import type { $AnyEvil, BaseDataType, BaseType, Identifier, QueryStep } from '../types';
 
 import { anonymous, Components, ToDestroy, ToIndex } from '../types';
-import { match } from '../utils';
+import { match, union } from '../utils';
 import { EntityIndex } from './EntityIndex';
 import { Query } from './Query';
 import { Registry } from './Registry';
@@ -39,6 +39,7 @@ export class Manager {
    */
   protected registry = new Registry();
   protected ctx: Context;
+  protected baseEntityKeys: Map<EntityClass, bigint> = new Map();
 
   /**
    * Cached queries, indexed by string key (basically a concatenation of query components).
@@ -61,27 +62,30 @@ export class Manager {
   }
 
   public register(entities: EntityClass[], components: ComponentClass[], tags: string[]): void {
-    const regs = this.registrations;
+    const allComponents = new Set([...components, ...entities.flatMap(e => e.prototype[Components])]);
+
+    for (const component of allComponents) {
+      this.registrations.components[component.type] = component;
+      this.registry.add(component.type);
+    }
 
     for (const entity of entities) {
+      const types = entity.prototype[Components].map((e: ComponentClass) => e.type);
       let key = entity.name;
       if (key === anonymous) {
-        key = entity.prototype[Components].map((e: ComponentClass) => e.type).join('|');
+        key = types.join('|');
       }
-      regs.entities[key] = entity;
-    }
-
-    for (const component of components) {
-      regs.components[component.type] = component;
+      this.registrations.entities[key] = entity;
+      // Each EntityClass has its 'base' (i.e., tag-independent) key. Since tags change more frequently than components, we're going to avoid recomputing the entity's entire key every time it changes.
+      this.baseEntityKeys.set(entity, this.registry.getID(...types)!);
     }
 
     for (const tag of tags) {
-      if (!(tag in regs.tags)) {
-        regs.tags[tag] = this.ctx.ids.id.next();
+      if (!(tag in this.registrations.tags)) {
+        const id = (this.registrations.tags[tag] = this.ctx.ids.id.next());
+        this.registry.add(id);
       }
     }
-
-    this.registry.add(...components.map(c => c.type), ...tags.map(t => regs.tags[t]));
   }
 
   /**
@@ -90,42 +94,39 @@ export class Manager {
   public tick(): void {
     const added: Entity[] = [];
     const removed: Entity[] = Array.from(Manager[ToDestroy][this.id]);
-    const index = Manager[ToIndex][this.id];
-    const hasIndexed = index.length > 0;
+    const reindex = Manager[ToIndex][this.id];
 
-    if (!removed.length && !hasIndexed) {
+    if (!removed.length && !reindex.length) {
       return;
     }
 
     // get the union of every key modified over the course of the last tick
     let modified = 0n;
 
-    if (hasIndexed) {
-      for (const [entity, oldKey] of index) {
-        this.entities[entity.id] = entity;
-
-        entity.key = this.getEntityKey(entity);
-
-        // the entry existed, but has changed
-        if (entity.key !== oldKey) {
-          // we need to update queries touching both the old key and the new key
-          modified |= oldKey | entity.key;
-
-          // the entity has not yet been indexed
-          if (!oldKey) {
-            added.push(entity);
-            this.index.append(entity.key, entity.id);
-            continue;
-          }
-
-          // entities need to be removed from queries and unindexed whether or not they continue to exist
-          removed.push(entity);
-          this.index.remove(oldKey, entity.id);
-          // ...but only bother re-indexing it if it's going to exist next tick.
-          if (!Manager[ToDestroy][this.id].includes(entity)) {
-            this.index.append(entity.key, entity.id);
-            added.push(entity);
-          }
+    for (const [entity, oldKey] of reindex) {
+      this.entities[entity.id] = entity;
+
+      entity.key = this.getEntityKey(entity);
+
+      // the entry existed, but has changed
+      if (entity.key !== oldKey) {
+        // we need to update queries touching both the old key and the new key
+        modified |= oldKey | entity.key;
+
+        // the entity has not yet been indexed
+        if (!oldKey) {
+          added.push(entity);
+          this.index.append(entity.key, entity.id);
+          continue;
+        }
+
+        // entities need to be removed from queries and unindexed whether or not they continue to exist
+        removed.push(entity);
+        this.index.remove(oldKey, entity.id);
+        // ...but only bother re-indexing it if it's going to exist next tick.
+        if (!Manager[ToDestroy][this.id].includes(entity)) {
+          this.index.append(entity.key, entity.id);
+          added.push(entity);
         }
       }
     }
@@ -181,9 +182,17 @@ export class Manager {
    * Given an entity, return its bigint key (union of all components/tags)
    */
   protected getEntityKey(entity: Entity): bigint {
-    const tags = this.registrations.tags;
-    const arr = Array.from(entity.tags).map(t => tags[t]);
-    return this.registry.add(...entity[Components].map(e => e.type), ...arr);
+    const baseKey =
+      // if this entity has a base key, use that instead of recomputing it
+      this.baseEntityKeys.get(entity.constructor as EntityClass) ??
+      this.registry.add(...entity[Components].map(e => e.type));
+
+    if (entity.tags.size > 0) {
+      // if we have tags, then get those too.
+      const arr = Array.from(entity.tags, t => this.registrations.tags[t]);
+      return union(baseKey, this.registry.add(...arr));
+    }
+    return baseKey;
   }
 
   // protected createEntityRefTag(entity: Entity): void {