feat(query): enhance query caching and reference management

- Implemented caching for queries based on component types and filters.
- Added reference counting to manage query instances efficiently.
- Updated the Query class to support a unique key for cached queries.
- Enhanced the World class to handle query creation and release with proper reference counting.
- Improved serialization of QueryFilter for cache key generation.

Author: codehz

src/query-filter.ts

@@ -1,6 +1,6 @@
 import { Archetype } from "./archetype";
 import type { EntityId } from "./entity";
-import { getDetailedIdType, decodeRelationId, isRelationId } from "./entity";
+import { decodeRelationId, getDetailedIdType, isRelationId } from "./entity";
 
 /**
  * Filter options for queries
@@ -9,6 +9,16 @@ export interface QueryFilter {
   negativeComponentTypes?: EntityId<any>[];
 }
 
+/**
+ * Serialize a QueryFilter into a deterministic string suitable for cache keys.
+ * Currently only serializes `negativeComponentTypes`.
+ */
+export function serializeQueryFilter(filter: QueryFilter = {}): string {
+  const negative = (filter.negativeComponentTypes || []).slice().sort((a, b) => a - b);
+  if (negative.length === 0) return "";
+  return `neg:${negative.join(",")}`;
+}
+
 /**
  * Check if an archetype matches the given component types
  */

src/query.test.ts

@@ -304,11 +304,113 @@ describe("Query", () => {
       // entity3 doesn't have position
 
       world.flushCommands();
+    });
+  });
 
-      const entities = query.getEntities();
-      expect(entities).toContain(entity1);
-      expect(entities).not.toContain(entity2);
-      expect(entities).not.toContain(entity3);
+  describe("Query Caching and Reference Counting", () => {
+    type Position = { x: number; y: number };
+    type Velocity = { x: number; y: number };
+
+    const positionComponent = component<Position>();
+    const velocityComponent = component<Velocity>();
+
+    it("should cache queries and return the same instance for identical queries", () => {
+      const world = new World();
+
+      // Create two queries with the same component types
+      const query1 = world.createQuery([positionComponent]);
+      const query2 = world.createQuery([positionComponent]);
+
+      // Should return the same cached instance
+      expect(query1).toBe(query2);
+    });
+
+    it("should cache queries with different component orders as the same query", () => {
+      const world = new World();
+
+      // Create queries with same components but different order
+      const query1 = world.createQuery([positionComponent, velocityComponent]);
+      const query2 = world.createQuery([velocityComponent, positionComponent]);
+
+      // Should return the same cached instance (sorted internally)
+      expect(query1).toBe(query2);
+    });
+
+    it("should create different queries for different component combinations", () => {
+      const world = new World();
+
+      const query1 = world.createQuery([positionComponent]);
+      const query2 = world.createQuery([velocityComponent]);
+      const query3 = world.createQuery([positionComponent, velocityComponent]);
+
+      // All should be different instances
+      expect(query1).not.toBe(query2);
+      expect(query1).not.toBe(query3);
+      expect(query2).not.toBe(query3);
+    });
+
+    it("should properly handle reference counting", () => {
+      const world = new World();
+
+      // Create multiple references to the same query
+      const query1 = world.createQuery([positionComponent]);
+      const query2 = world.createQuery([positionComponent]);
+      const query3 = world.createQuery([positionComponent]);
+
+      // All should be the same instance
+      expect(query1).toBe(query2);
+      expect(query2).toBe(query3);
+
+      // Release all three references
+      world.releaseQuery(query1);
+      world.releaseQuery(query2);
+      world.releaseQuery(query3);
+
+      // Now create a new query - should be a new instance since cache was cleared
+      const query4 = world.createQuery([positionComponent]);
+      expect(query4).not.toBe(query1); // Should be a new instance
+    });
+
+    it("should handle releaseQuery on non-cached queries gracefully", () => {
+      const world = new World();
+
+      // Create a query and immediately release it
+      const query = world.createQuery([positionComponent]);
+      world.releaseQuery(query);
+
+      // Should not throw and should create a new instance next time
+      const query2 = world.createQuery([positionComponent]);
+      expect(query2).not.toBe(query);
+    });
+
+    it("should cache queries with filters separately", () => {
+      const world = new World();
+      type Health = { value: number };
+      const healthComponent = component<Health>();
+
+      // Create queries with and without filters
+      const query1 = world.createQuery([positionComponent]);
+      const query2 = world.createQuery([positionComponent], { negativeComponentTypes: [healthComponent] });
+
+      // Should be different instances due to different filters
+      expect(query1).not.toBe(query2);
+    });
+
+    it("should maintain separate caches for queries with different filters", () => {
+      const world = new World();
+      type Health = { value: number };
+      const healthComponent = component<Health>();
+
+      // Create multiple queries with the same filter
+      const query1 = world.createQuery([positionComponent], { negativeComponentTypes: [healthComponent] });
+      const query2 = world.createQuery([positionComponent], { negativeComponentTypes: [healthComponent] });
+
+      // Should return the same cached instance
+      expect(query1).toBe(query2);
+
+      // Create queries with different filters
+      const query3 = world.createQuery([positionComponent], { negativeComponentTypes: [velocityComponent] });
+      expect(query1).not.toBe(query3);
     });
   });
 });

src/query.ts

@@ -1,26 +1,30 @@
 import { Archetype } from "./archetype";
 import type { EntityId } from "./entity";
-import type { ComponentTuple } from "./types";
 import { matchesComponentTypes, matchesFilter, type QueryFilter } from "./query-filter";
+import type { ComponentTuple } from "./types";
 import type { World } from "./world";
 
 /**
  * Query class for efficient entity queries with cached archetypes
  */
 export class Query {
+  // Public key used by World to identify cached queries
+  public readonly key: string;
+
   private world: World<any[]>;
   private componentTypes: EntityId<any>[];
   private filter: QueryFilter;
   private cachedArchetypes: Archetype[] = [];
   private isDisposed = false;
 
-  constructor(world: World<any[]>, componentTypes: EntityId<any>[], filter: QueryFilter = {}) {
+  constructor(world: World<any[]>, componentTypes: EntityId<any>[], filter: QueryFilter = {}, key?: string) {
     this.world = world;
     this.componentTypes = [...componentTypes].sort((a, b) => a - b);
     this.filter = filter;
+    this.key = key ?? `${this.componentTypes.join(",")}|`;
     this.updateCache();
     // Register with world for archetype updates
-    world.registerQuery(this);
+    world._registerQuery(this);
   }
 
   /**
@@ -129,9 +133,23 @@ export class Query {
   /**
    * Dispose the query and disconnect from world
    */
+  /**
+   * Request disposal of this query.
+   * This will decrement the world's reference count for the query.
+   * The query will only be fully disposed when the ref count reaches zero.
+   */
   dispose(): void {
+    // Ask the world to release this query (decrement refcount and fully dispose when zero)
+    this.world.releaseQuery(this);
+  }
+
+  /**
+   * Internal full dispose called by World when refCount reaches zero.
+   */
+  _disposeInternal(): void {
     if (!this.isDisposed) {
-      this.world.unregisterQuery(this);
+      // Unregister from world (remove from notification list)
+      this.world._unregisterQuery(this);
       this.cachedArchetypes = [];
       this.isDisposed = true;
     }

src/world.ts

@@ -1,9 +1,9 @@
 import { Archetype } from "./archetype";
 import { CommandBuffer, type Command } from "./command-buffer";
 import type { EntityId, WildcardRelationId } from "./entity";
-import { EntityIdManager, relation, getDetailedIdType, getIdType, isWildcardRelationId } from "./entity";
+import { EntityIdManager, getDetailedIdType, getIdType, relation } from "./entity";
 import { Query } from "./query";
-import type { QueryFilter } from "./query-filter";
+import { serializeQueryFilter, type QueryFilter } from "./query-filter";
 import type { System } from "./system";
 import { SystemScheduler } from "./system-scheduler";
 import type { ComponentTuple, LifecycleHook } from "./types";
@@ -20,6 +20,8 @@ export class World<ExtraParams extends any[] = [deltaTime: number]> {
   private entityToArchetype = new Map<EntityId, Archetype>();
   private systemScheduler = new SystemScheduler<ExtraParams>();
   private queries: Query[] = [];
+  // Cache for queries keyed by component types + filter signature
+  private queryCache = new Map<string, { query: Query; refCount: number }>();
   private commandBuffer: CommandBuffer;
   private componentToArchetypes = new Map<EntityId<any>, Archetype[]>();
 
@@ -257,26 +259,73 @@ export class World<ExtraParams extends any[] = [deltaTime: number]> {
    * Create a cached query for efficient entity lookups
    */
   createQuery(componentTypes: EntityId<any>[], filter: QueryFilter = {}): Query {
-    return new Query(this, componentTypes, filter);
+    // Build a deterministic key for the query (component types sorted + filter negative components sorted)
+    const sortedTypes = [...componentTypes].sort((a, b) => a - b);
+    const filterKey = serializeQueryFilter(filter);
+    const key = `${this.getComponentTypesHash(sortedTypes)}${filterKey ? `|${filterKey}` : ""}`;
+
+    const cached = this.queryCache.get(key);
+    if (cached) {
+      cached.refCount++;
+      return cached.query;
+    }
+
+    const query = new Query(this, sortedTypes, filter, key);
+    this.queryCache.set(key, { query, refCount: 1 });
+    return query;
   }
 
   /**
    * @internal Register a query for archetype update notifications
    */
-  registerQuery(query: Query): void {
+  _registerQuery(query: Query): void {
     this.queries.push(query);
   }
 
   /**
    * @internal Unregister a query
    */
-  unregisterQuery(query: Query): void {
+  _unregisterQuery(query: Query): void {
     const index = this.queries.indexOf(query);
     if (index !== -1) {
       this.queries.splice(index, 1);
     }
   }
 
+  /**
+   * Release a query reference obtained from createQuery.
+   * Decrements the refCount and fully disposes the query when it reaches zero.
+   */
+  releaseQuery(query: Query): void {
+    const key = query.key;
+    // Fallback: try to find by identity
+    for (const [k, v] of this.queryCache.entries()) {
+      if (v.query === query) {
+        v.refCount--;
+        if (v.refCount <= 0) {
+          this.queryCache.delete(k);
+          // Fully dispose the query (will unregister it from notification list)
+          v.query._disposeInternal();
+        }
+        return;
+      }
+    }
+
+    const entry = this.queryCache.get(key);
+    if (!entry) {
+      // Nothing cached, ensure it's unregistered
+      this._unregisterQuery(query);
+      return;
+    }
+
+    entry.refCount--;
+    if (entry.refCount <= 0) {
+      this.queryCache.delete(key);
+      // Fully dispose the query (will unregister it from notification list)
+      entry.query._disposeInternal();
+    }
+  }
+
   /**
    * @internal Get archetypes that match specific component types (for internal use by queries)
    */