fix(query): support specific relation queries for dontFragment components

Fixed an issue where specific target relation queries didn't work properly with dontFragment: true relation components. For dontFragment relations, specific relation IDs are not stored in archetype.componentTypes, only the wildcard marker is added to the archetype. The query matching function matchesComponentTypes only checked archetype.componentTypes.includes(type), causing specific dontFragment relation queries to fail.

Key changes:

1. Modified matchesComponentTypes function in filter.ts to recognize specific dontFragment relations, checking if the archetype contains the corresponding wildcard marker when querying specific dontFragment relations

2. Enhanced Query class with specificDontFragmentTypes cache to store specific dontFragment relations requiring entity-level filtering

3. Modified getEntities() method to enable slow path for queries containing specific dontFragment relations

4. Renamed entityHasAllWildcards() to entityMatchesQuery(), checking both wildcard relations and specific dontFragment relations

All 261 tests pass, including the previously failing 2 tests:

- "should support specific relation queries for dontFragment components"

- "should support specific relation queries with exclusive dontFragment when target changes"

Author: codehz

src/__tests__/dont-fragment-query-notification.test.ts

@@ -464,4 +464,73 @@ describe("DontFragment Query Notification Issue", () => {
       expect(archetype.componentTypes).toContain(wildcardChildOf);
     }
   });
+
+  it("should support specific relation queries for dontFragment components", () => {
+    const world = new World();
+
+    const PositionId = component();
+    const ChildOf = component({ dontFragment: true });
+
+    const parent1 = world.new();
+    const parent2 = world.new();
+
+    // Create specific relation queries (not wildcard)
+    const queryParent1 = world.createQuery([relation(ChildOf, parent1), PositionId]);
+    const queryParent2 = world.createQuery([relation(ChildOf, parent2), PositionId]);
+
+    // Create entities with different parents
+    const entity1 = world.new();
+    world.set(entity1, PositionId);
+    world.set(entity1, relation(ChildOf, parent1));
+
+    const entity2 = world.new();
+    world.set(entity2, PositionId);
+    world.set(entity2, relation(ChildOf, parent2));
+
+    world.sync();
+
+    // Specific queries should find their respective entities
+    expect(queryParent1.getEntities()).toContain(entity1);
+    expect(queryParent1.getEntities()).not.toContain(entity2);
+
+    expect(queryParent2.getEntities()).not.toContain(entity1);
+    expect(queryParent2.getEntities()).toContain(entity2);
+  });
+
+  it("should support specific relation queries with exclusive dontFragment when target changes", () => {
+    const world = new World();
+
+    const PositionId = component();
+    const ChildOf = component({ dontFragment: true, exclusive: true });
+
+    const parent1 = world.new();
+    const parent2 = world.new();
+    const entity = world.new();
+    world.set(entity, PositionId);
+
+    // Create specific queries for each parent
+    const queryParent1 = world.createQuery([relation(ChildOf, parent1), PositionId]);
+    const queryParent2 = world.createQuery([relation(ChildOf, parent2), PositionId]);
+
+    // Set relation to parent1
+    world.set(entity, relation(ChildOf, parent1));
+    world.sync();
+
+    expect(queryParent1.getEntities()).toContain(entity);
+    expect(queryParent2.getEntities()).not.toContain(entity);
+
+    // Change to parent2
+    world.set(entity, relation(ChildOf, parent2));
+    world.sync();
+
+    expect(queryParent1.getEntities()).not.toContain(entity);
+    expect(queryParent2.getEntities()).toContain(entity);
+
+    // Change back to parent1
+    world.set(entity, relation(ChildOf, parent1));
+    world.sync();
+
+    expect(queryParent1.getEntities()).toContain(entity);
+    expect(queryParent2.getEntities()).not.toContain(entity);
+  });
 });

src/query/filter.ts

@@ -1,6 +1,12 @@
 import type { Archetype } from "../core/archetype";
 import type { EntityId } from "../core/entity";
-import { getComponentIdFromRelationId, getDetailedIdType, isRelationId } from "../core/entity";
+import {
+  getComponentIdFromRelationId,
+  getDetailedIdType,
+  isDontFragmentComponent,
+  isRelationId,
+  relation,
+} from "../core/entity";
 
 /**
  * Filter options for queries
@@ -32,8 +38,16 @@ export function matchesComponentTypes(archetype: Archetype, componentTypes: Enti
         const componentId = getComponentIdFromRelationId(archetypeType);
         return componentId === detailedType.componentId;
       });
+    } else if (
+      (detailedType.type === "entity-relation" || detailedType.type === "component-relation") &&
+      detailedType.componentId !== undefined &&
+      isDontFragmentComponent(detailedType.componentId)
+    ) {
+      // For specific dontFragment relations, check if archetype has the wildcard marker
+      const wildcardMarker = relation(detailedType.componentId, "*");
+      return archetype.componentTypes.includes(wildcardMarker);
     } else {
-      // For regular components, check direct inclusion
+      // For regular components and non-dontFragment relations, check direct inclusion
       return archetype.componentTypes.includes(type);
     }
   });

src/query/query.ts

@@ -1,6 +1,6 @@
 import type { Archetype } from "../core/archetype";
 import type { EntityId, WildcardRelationId } from "../core/entity";
-import { getDetailedIdType } from "../core/entity";
+import { getDetailedIdType, isDontFragmentComponent } from "../core/entity";
 import type { ComponentTuple, ComponentType } from "../core/types";
 import type { World } from "../core/world";
 import { matchesComponentTypes, matchesFilter, type QueryFilter } from "./filter";
@@ -16,6 +16,8 @@ export class Query {
   private isDisposed = false;
   /** Cached wildcard component types for faster entity filtering */
   private wildcardTypes: WildcardRelationId<any>[];
+  /** Cached specific dontFragment relation types that need entity-level filtering */
+  private specificDontFragmentTypes: EntityId<any>[];
 
   constructor(world: World, componentTypes: EntityId<any>[], filter: QueryFilter = {}) {
     this.world = world;
@@ -25,6 +27,15 @@ export class Query {
     this.wildcardTypes = this.componentTypes.filter(
       (ct) => getDetailedIdType(ct).type === "wildcard-relation",
     ) as WildcardRelationId<any>[];
+    // Pre-compute specific dontFragment relation types that need entity-level filtering
+    this.specificDontFragmentTypes = this.componentTypes.filter((ct) => {
+      const detailedType = getDetailedIdType(ct);
+      return (
+        (detailedType.type === "entity-relation" || detailedType.type === "component-relation") &&
+        detailedType.componentId !== undefined &&
+        isDontFragmentComponent(detailedType.componentId)
+      );
+    });
     this.updateCache();
     // Register with world for archetype updates
     world._registerQuery(this);
@@ -45,22 +56,23 @@ export class Query {
   getEntities(): EntityId[] {
     this.ensureNotDisposed();
 
-    // Fast path: no wildcard relations
-    if (this.wildcardTypes.length === 0) {
+    // Fast path: no wildcard relations and no specific dontFragment relations
+    if (this.wildcardTypes.length === 0 && this.specificDontFragmentTypes.length === 0) {
       const result: EntityId[] = [];
       for (const archetype of this.cachedArchetypes) {
         result.push(...archetype.getEntities());
       }
       return result;
     }
 
-    // Slow path: need to filter entities that actually have wildcard relations
-    // This is necessary for dontFragment components where an archetype can contain
-    // entities with and without the relation
+    // Slow path: need to filter entities that actually have the required relations
+    // This is necessary for:
+    // 1. Wildcard relations where an archetype can contain entities with/without the relation
+    // 2. Specific dontFragment relations where the archetype only has the wildcard marker
     const result: EntityId[] = [];
     for (const archetype of this.cachedArchetypes) {
       for (const entity of archetype.getEntities()) {
-        if (this.entityHasAllWildcards(archetype, entity)) {
+        if (this.entityMatchesQuery(archetype, entity)) {
           result.push(entity);
         }
       }
@@ -69,15 +81,25 @@ export class Query {
   }
 
   /**
-   * Check if entity has all required wildcard relations
+   * Check if entity matches all query requirements (wildcards and specific dontFragment relations)
    */
-  private entityHasAllWildcards(archetype: Archetype, entity: EntityId): boolean {
+  private entityMatchesQuery(archetype: Archetype, entity: EntityId): boolean {
+    // Check wildcard relations
     for (const wildcardType of this.wildcardTypes) {
       const relations = archetype.get(entity, wildcardType);
       if (!relations || relations.length === 0) {
         return false;
       }
     }
+
+    // Check specific dontFragment relations
+    for (const specificType of this.specificDontFragmentTypes) {
+      const result = archetype.getOptional(entity, specificType);
+      if (result === undefined) {
+        return false;
+      }
+    }
+
     return true;
   }