From df3c876b58bcb6e151288d3f8e18b00b178a76fb Mon Sep 17 00:00:00 2001
From: ztw <412467438@qq.com>
Date: Mon, 24 Aug 2026 16:56:28 +0800
Subject: [PATCH] - Use CollectionsMarshal.GetValueRefOrAddDefault for single
 hash lookup (.NET 6+) - Add Query.CountEntities() for allocation-free entity
 counting"

---
 src/Arch/Core/Query.cs | 17 +++++++++++++++++
 src/Arch/Core/World.cs | 21 +++++++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/Arch/Core/Query.cs b/src/Arch/Core/Query.cs
index feed4a1e..6b8359b5 100644
--- a/src/Arch/Core/Query.cs
+++ b/src/Arch/Core/Query.cs
@@ -1,3 +1,4 @@
+using System.Diagnostics.Contracts;
 using Arch.Core.Extensions;
 using Arch.Core.Extensions.Internal;
 using Arch.Core.Utils;
@@ -571,6 +572,22 @@ internal Query(Archetypes allArchetypes, QueryDescription description)
         _queryDescription = description;
     }
 
+    /// <summary>
+    ///     Counts all <see cref="Entity"/>s that match this query.
+    /// </summary>
+    [Pure]
+    public int CountEntities()
+    {
+        var counter = 0;
+        foreach (var archetype in GetArchetypeIterator())
+        {
+            var entities = archetype.EntityCount;
+            counter += entities;
+        }
+
+        return counter;
+    }
+
     /// <summary>
     ///     Checks whether the specified <see cref="BitSet"/> matches.
     /// </summary>
diff --git a/src/Arch/Core/World.cs b/src/Arch/Core/World.cs
index 82c6a5ab..38c5b062 100644
--- a/src/Arch/Core/World.cs
+++ b/src/Arch/Core/World.cs
@@ -410,17 +410,30 @@ public void Destroy(Entity entity)
     [Pure]
     public Query Query(in QueryDescription queryDescription)
     {
-        // Looping over all archetypes, their chunks and their entities.
-        var queryCache = QueryCache; // Storing locally to only access the QueryCache once
+        var queryCache = QueryCache; // Store locally to only access the field once
+
+        #if NET6_0_OR_GREATER
+        ref var queryRef = ref CollectionsMarshal.GetValueRefOrAddDefault(queryCache, queryDescription, out bool exists);
+        if (exists)
+        {
+            return queryRef!;
+        }
+        else
+        {
+            var newQuery = new Query(Archetypes, queryDescription);
+            queryRef = newQuery;
+            return newQuery;
+        }
+        #else
         if (queryCache.TryGetValue(queryDescription, out var query))
         {
             return query;
         }
 
         query = new Query(Archetypes, queryDescription);
-        queryCache[queryDescription] = query;
-
+        queryCache.Add(queryDescription, query);
         return query;
+        #endif
     }
 
     /// <summary>