Optimize indexed OR/Contains predicates to prevent cross-collection scan leakage in single-file mode

src/BLite.Core/Collections/DocumentCollection.cs

@@ -2457,9 +2457,9 @@ private async Task<TResult> BsonAggregateFieldAsync<TResult>(BsonAggregator agg,
     /// </summary>
     [RequiresDynamicCode("Count-by-predicate uses index optimization and Expression.Compile() which require dynamic code generation.")]
     [RequiresUnreferencedCode("Count-by-predicate uses reflection to resolve members at runtime. Ensure all entity types are preserved.")]
-    internal async Task<int> CountByPredicateAsync(
-        System.Linq.Expressions.LambdaExpression whereClause,
-        CancellationToken ct = default)
+    internal async Task<int> CountByPredicateAsync(
+        System.Linq.Expressions.LambdaExpression whereClause,
+        CancellationToken ct = default)
     {
         // Strategy 1: Index key-only scan — no data-page reads at all.
         // Applicable whenever the predicate targets an indexed field AND the index fully
@@ -2469,19 +2469,31 @@ internal async Task<int> CountByPredicateAsync(
         // When HasResiduePredicate=true (compound AND with a non-indexed clause), we must
         // fall through to FetchAsync so the residue predicate is applied per document.
         var indexOpt = Query.IndexOptimizer.TryOptimize<T>(whereClause, GetIndexes(), ConverterRegistry);
-        if (indexOpt != null
-            && !indexOpt.IsVectorSearch
-            && !indexOpt.IsSpatialSearch
-            && !indexOpt.HasResiduePredicate)
-        {
-            var index = _indexManager.GetIndex(indexOpt.IndexName);
-            if (index != null)
-            {
-                // Use the per-bound inclusivity flags from OptimizationResult.
-                // These are set correctly for every operator (==, >=, >, <=, <) and
-                // propagated through AND-merges, so compound predicates like
-                // x.Price > 50 && x.Price < 90 get both boundaries exclusive.
-                return index.CountRange(indexOpt.MinValue, indexOpt.MaxValue,
+        if (indexOpt != null
+            && !indexOpt.IsVectorSearch
+            && !indexOpt.IsSpatialSearch
+            && !indexOpt.HasResiduePredicate)
+        {
+            var index = _indexManager.GetIndex(indexOpt.IndexName);
+            if (index != null)
+            {
+                if (indexOpt.InValues != null)
+                {
+                    int inCount = 0;
+                    var seen = new HashSet<object?>(s_inProbeKeyComparer);
+                    foreach (var key in indexOpt.InValues)
+                    {
+                        if (!seen.Add(key)) continue;
+                        inCount += index.CountRange(key, key, true, true, null);
+                    }
+                    return inCount;
+                }
+
+                // Use the per-bound inclusivity flags from OptimizationResult.
+                // These are set correctly for every operator (==, >=, >, <=, <) and
+                // propagated through AND-merges, so compound predicates like
+                // x.Price > 50 && x.Price < 90 get both boundaries exclusive.
+                return index.CountRange(indexOpt.MinValue, indexOpt.MaxValue,
                     indexOpt.StartInclusive, indexOpt.EndInclusive, null);
             }
         }
@@ -2656,13 +2668,26 @@ Func<T, bool> GetCompiled() =>
                     await foreach (var item in spatialSeq)
                         if (indexOpt.FilterCompleteness == Query.IndexOptimizer.FilterCompleteness.Exact || GetCompiled()(item)) { yield return item; if (++yielded >= fetchLimit) yield break; }
                 }
-                else
-                {
-                    await foreach (var item in QueryIndexAsync(indexOpt.IndexName, indexOpt.MinValue, indexOpt.MaxValue, true, 0, int.MaxValue, transaction, ct))
-                        if (indexOpt.FilterCompleteness == Query.IndexOptimizer.FilterCompleteness.Exact || GetCompiled()(item)) { yield return item; if (++yielded >= fetchLimit) yield break; }
-                }
-                yield break;
-            }
+                else
+                {
+                    if (indexOpt.InValues != null)
+                    {
+                        var seen = new HashSet<object?>(s_inProbeKeyComparer);
+                        foreach (var key in indexOpt.InValues)
+                        {
+                            if (!seen.Add(key)) continue;
+                            await foreach (var item in QueryIndexAsync(indexOpt.IndexName, key, key, true, 0, int.MaxValue, transaction, ct))
+                                if (indexOpt.FilterCompleteness == Query.IndexOptimizer.FilterCompleteness.Exact || GetCompiled()(item)) { yield return item; if (++yielded >= fetchLimit) yield break; }
+                        }
+                    }
+                    else
+                    {
+                        await foreach (var item in QueryIndexAsync(indexOpt.IndexName, indexOpt.MinValue, indexOpt.MaxValue, true, 0, int.MaxValue, transaction, ct))
+                            if (indexOpt.FilterCompleteness == Query.IndexOptimizer.FilterCompleteness.Exact || GetCompiled()(item)) { yield return item; if (++yielded >= fetchLimit) yield break; }
+                    }
+                }
+                yield break;
+            }
 
             // ── Strategy 2: BSON-level predicate scan ─────────────────────────
             // Filters at raw-BSON level before deserializing — no compiled Func<T,bool> needed.
@@ -2690,9 +2715,34 @@ Func<T, bool> GetCompiled() =>
                 if (++yielded >= fetchLimit) yield break;
             }
         }
-    }
-
-    #endregion
+    }
+
+    private static readonly IEqualityComparer<object?> s_inProbeKeyComparer = new InProbeKeyComparer();
+
+    private sealed class InProbeKeyComparer : IEqualityComparer<object?>
+    {
+        bool IEqualityComparer<object?>.Equals(object? x, object? y)
+        {
+            if (ReferenceEquals(x, y)) return true;
+            if (x is null || y is null) return false;
+            if (x is byte[] xb && y is byte[] yb) return xb.AsSpan().SequenceEqual(yb);
+            return x.Equals(y);
+        }
+
+        int IEqualityComparer<object?>.GetHashCode(object? obj)
+        {
+            if (obj is null) return 0;
+            if (obj is byte[] bytes)
+            {
+                var hc = new HashCode();
+                foreach (var b in bytes) hc.Add(b);
+                return hc.ToHashCode();
+            }
+            return obj.GetHashCode();
+        }
+    }
+
+    #endregion
 
     /// <summary>
     /// Serializes an entity with adaptive buffer sizing (Stepped Retry).