Fix wildcard queries for dontFragment relations (#5)

* Initial plan

* Add test cases to reproduce wildcard query bug with dontFragment components

Co-authored-by: codehz <13158903+codehz@users.noreply.github.com>

* Fix wildcard query issue with dontFragment components

- Added hasRelationWithComponentId method to Archetype to check for relations including dontFragment
- Updated getMatchingArchetypes to use new method for wildcard relation filtering
- Modified Query.getEntities to filter entities with wildcard relations at entity level for dontFragment components
- Enhanced buildWildcardRelationValue to include dontFragment relations in iteration
- Added comprehensive tests for wildcard queries with dontFragment components

Co-authored-by: codehz <13158903+codehz@users.noreply.github.com>

* Address code review feedback

- Add null safety check for wildcard relation results in Query.getEntities
- Improve error message with component ID and entity ID for better debugging

Co-authored-by: codehz <13158903+codehz@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: codehz <13158903+codehz@users.noreply.github.com>

Author: Copilot

src/archetype.ts

@@ -457,10 +457,11 @@ export class Archetype {
     componentTypes: T,
     componentDataSources: (any[] | EntityId<any>[] | undefined)[],
     entityIndex: number,
+    entityId: EntityId,
   ): ComponentTuple<T> {
     return componentDataSources.map((dataSource, i) => {
       const compType = componentTypes[i]!;
-      return this.buildSingleComponent(compType, dataSource, entityIndex);
+      return this.buildSingleComponent(compType, dataSource, entityIndex, entityId);
     }) as ComponentTuple<T>;
   }
 
@@ -471,12 +472,19 @@ export class Archetype {
     compType: ComponentType<any>,
     dataSource: any[] | EntityId<any>[] | undefined,
     entityIndex: number,
+    entityId: EntityId,
   ): any {
     const optional = isOptionalEntityId(compType);
     const actualType = optional ? compType.optional : compType;
 
     if (getIdType(actualType) === "wildcard-relation") {
-      return this.buildWildcardRelationValue(dataSource, entityIndex, optional);
+      return this.buildWildcardRelationValue(
+        actualType as WildcardRelationId<any>,
+        dataSource,
+        entityIndex,
+        entityId,
+        optional,
+      );
     } else {
       return this.buildRegularComponentValue(dataSource, entityIndex, optional);
     }
@@ -486,27 +494,54 @@ export class Archetype {
    * Build wildcard relation value from matching relations
    */
   private buildWildcardRelationValue(
+    wildcardRelationType: WildcardRelationId<any>,
     dataSource: any[] | EntityId<any>[] | undefined,
     entityIndex: number,
+    entityId: EntityId,
     optional: boolean,
   ): any {
-    if (dataSource === undefined) {
-      if (optional) {
-        return undefined;
-      }
-      throw new Error(`No matching relations found for mandatory wildcard relation component type`);
-    }
-
-    const matchingRelations = dataSource as EntityId<any>[];
+    const matchingRelations = (dataSource as EntityId<any>[]) || [];
     const relations: [EntityId<unknown>, any][] = [];
 
+    // Add regular archetype relations
     for (const relType of matchingRelations) {
       const dataArray = this.getComponentData(relType);
       const data = dataArray[entityIndex];
       const decodedRel = decodeRelationId(relType as RelationId<any>);
       relations.push([decodedRel.targetId, data === MISSING_COMPONENT ? undefined : data]);
     }
 
+    // Add dontFragment relations for this entity
+    // Get the component ID from the wildcard relation type
+    const wildcardDecoded = decodeRelationId(wildcardRelationType);
+    const targetComponentId = wildcardDecoded.componentId;
+
+    const dontFragmentData = this.dontFragmentRelations.get(entityId);
+    if (dontFragmentData) {
+      // Check dontFragment relations for matching component ID
+      for (const [relType, data] of dontFragmentData) {
+        const relDetailed = getDetailedIdType(relType);
+        if (
+          (relDetailed.type === "entity-relation" || relDetailed.type === "component-relation") &&
+          relDetailed.componentId === targetComponentId
+        ) {
+          relations.push([relDetailed.targetId, data]);
+        }
+      }
+    }
+
+    // If no relations found and not optional, this entity doesn't match
+    if (relations.length === 0) {
+      if (!optional) {
+        const wildcardDecoded = decodeRelationId(wildcardRelationType);
+        throw new Error(
+          `No matching relations found for mandatory wildcard relation component ${wildcardDecoded.componentId} on entity ${entityId}`,
+        );
+      }
+      // For optional, return undefined when there are no relations
+      return undefined;
+    }
+
     return optional ? { value: relations } : relations;
   }
 
@@ -570,7 +605,7 @@ export class Archetype {
     for (let entityIndex = 0; entityIndex < this.entities.length; entityIndex++) {
       const entity = this.entities[entityIndex]!;
 
-      const components = this.buildComponentsForIndex(componentTypes, componentDataSources, entityIndex);
+      const components = this.buildComponentsForIndex(componentTypes, componentDataSources, entityIndex, entity);
 
       yield [entity, ...components];
     }
@@ -593,7 +628,7 @@ export class Archetype {
       const entity = this.entities[entityIndex]!;
 
       // Direct array access for each component type using pre-cached sources
-      const components = this.buildComponentsForIndex(componentTypes, componentDataSources, entityIndex);
+      const components = this.buildComponentsForIndex(componentTypes, componentDataSources, entityIndex, entity);
 
       callback(entity, ...components);
     }
@@ -613,4 +648,41 @@ export class Archetype {
       callback(this.entities[i]!, components);
     }
   }
+
+  /**
+   * Check if any entity in this archetype has a relation matching the given component ID
+   * This includes both regular relations in componentTypes and dontFragment relations
+   * @param componentId The component ID to match
+   * @returns true if any entity has a matching relation
+   */
+  hasRelationWithComponentId(componentId: EntityId<any>): boolean {
+    // Check regular archetype components
+    for (const componentType of this.componentTypes) {
+      const detailedType = getDetailedIdType(componentType);
+      if (
+        (detailedType.type === "entity-relation" || detailedType.type === "component-relation") &&
+        detailedType.componentId === componentId
+      ) {
+        return true;
+      }
+    }
+
+    // Check dontFragment relations for any entity in this archetype
+    for (const entityId of this.entities) {
+      const entityDontFragmentRelations = this.dontFragmentRelations.get(entityId);
+      if (entityDontFragmentRelations) {
+        for (const relationType of entityDontFragmentRelations.keys()) {
+          const detailedType = getDetailedIdType(relationType);
+          if (
+            (detailedType.type === "entity-relation" || detailedType.type === "component-relation") &&
+            detailedType.componentId === componentId
+          ) {
+            return true;
+          }
+        }
+      }
+    }
+
+    return false;
+  }
 }

src/dont-fragment.test.ts

@@ -222,4 +222,73 @@ describe("DontFragment Relations", () => {
     expect(archetypes2.length).toBe(1);
     expect(archetypes2[0].size).toBe(5);
   });
+
+  it("should query entities with wildcard relation on dontFragment component using createQuery", () => {
+    const world = new World();
+
+    const PositionId = component();
+    const ChildOf = component({ dontFragment: true });
+
+    const parent1 = world.new();
+    const parent2 = world.new();
+
+    const child1 = world.new();
+    world.set(child1, PositionId);
+    world.set(child1, relation(ChildOf, parent1));
+
+    const child2 = world.new();
+    world.set(child2, PositionId);
+    world.set(child2, relation(ChildOf, parent2));
+
+    world.sync();
+
+    // Try to query entities with wildcard ChildOf relation
+    const wildcardChildOf = relation(ChildOf, "*");
+    const query = world.createQuery([wildcardChildOf]);
+    const entities = query.getEntities();
+
+    // This should find both child1 and child2
+    expect(entities.length).toBe(2);
+    expect(entities).toContain(child1);
+    expect(entities).toContain(child2);
+  });
+
+  it("should query entities with wildcard relation + other components on dontFragment", () => {
+    const world = new World();
+
+    const PositionId = component();
+    const VelocityId = component();
+    const ChildOf = component({ dontFragment: true });
+
+    const parent1 = world.new();
+    const parent2 = world.new();
+
+    const child1 = world.new();
+    world.set(child1, PositionId);
+    world.set(child1, VelocityId);
+    world.set(child1, relation(ChildOf, parent1));
+
+    const child2 = world.new();
+    world.set(child2, PositionId);
+    world.set(child2, VelocityId);
+    world.set(child2, relation(ChildOf, parent2));
+
+    // Entity without ChildOf relation
+    const child3 = world.new();
+    world.set(child3, PositionId);
+    world.set(child3, VelocityId);
+
+    world.sync();
+
+    // Query for entities with wildcard ChildOf relation AND Position
+    const wildcardChildOf = relation(ChildOf, "*");
+    const query = world.createQuery([wildcardChildOf, PositionId]);
+    const entities = query.getEntities();
+
+    // Should find child1 and child2, but not child3 (no ChildOf relation)
+    expect(entities.length).toBe(2);
+    expect(entities).toContain(child1);
+    expect(entities).toContain(child2);
+    expect(entities).not.toContain(child3);
+  });
 });

src/query.ts

@@ -1,5 +1,6 @@
 import { Archetype } from "./archetype";
-import type { EntityId } from "./entity";
+import type { EntityId, WildcardRelationId } from "./entity";
+import { getDetailedIdType } from "./entity";
 import { matchesComponentTypes, matchesFilter, type QueryFilter } from "./query-filter";
 import type { ComponentTuple, ComponentType } from "./types";
 import type { World } from "./world";
@@ -38,9 +39,44 @@ export class Query {
   getEntities(): EntityId[] {
     this.ensureNotDisposed();
     const result: EntityId[] = [];
-    for (const archetype of this.cachedArchetypes) {
-      result.push(...archetype.getEntities());
+
+    // Check if any component types are wildcard relations
+    const hasWildcardRelations = this.componentTypes.some((ct) => {
+      const detailed = getDetailedIdType(ct);
+      return detailed.type === "wildcard-relation";
+    });
+
+    // If there are wildcard relations, we need to filter entities that actually have them
+    // This is necessary for dontFragment components where an archetype can contain entities
+    // with and without the relation
+    if (hasWildcardRelations) {
+      for (const archetype of this.cachedArchetypes) {
+        for (const entity of archetype.getEntities()) {
+          // Check if entity has all required wildcard relations
+          let hasAllRelations = true;
+          for (const componentType of this.componentTypes) {
+            const detailed = getDetailedIdType(componentType);
+            if (detailed.type === "wildcard-relation") {
+              // Check if entity has at least one relation matching this wildcard
+              const relations = archetype.get(entity, componentType as WildcardRelationId<any>);
+              if (!relations || relations.length === 0) {
+                hasAllRelations = false;
+                break;
+              }
+            }
+          }
+          if (hasAllRelations) {
+            result.push(entity);
+          }
+        }
+      }
+    } else {
+      // No wildcard relations, can just return all entities from matching archetypes
+      for (const archetype of this.cachedArchetypes) {
+        result.push(...archetype.getEntities());
+      }
     }
+
     return result;
   }
 

src/world.ts

@@ -573,13 +573,9 @@ export class World<UpdateParams extends any[] = []> {
 
     // Filter by wildcard relations
     for (const wildcard of wildcardRelations) {
-      // Keep only archetypes that have the component
+      // Keep only archetypes that have the component (including dontFragment relations)
       matchingArchetypes = matchingArchetypes.filter((archetype) =>
-        archetype.componentTypes.some((archetypeType) => {
-          if (!isRelationId(archetypeType)) return false;
-          const decoded = decodeRelationId(archetypeType);
-          return decoded.componentId === wildcard.componentId;
-        }),
+        archetype.hasRelationWithComponentId(wildcard.componentId),
       );
     }