perf: index-aware LINQ terminals with predicates, single-copy page writes, single-pass B-Tree insert (#179)

src/BLite.Core/Collections/DocumentCollection.cs

@@ -1920,7 +1920,7 @@ private uint AllocateOverflowPage(ReadOnlySpan<byte> data, uint nextOverflowPage
             // NEW: Buffer write in transaction or write immediately
             var writeOp = new WriteOperation(
                 documentId: ObjectId.Empty,
-                newValue: buffer.AsSpan(0, _storage.PageSize).ToArray(),
+                newValue: buffer.AsMemory(0, _storage.PageSize),
                 pageId: pageId,
                 type: OperationType.Insert
             );

src/BLite.Core/Indexing/BTreeIndex.cs

@@ -101,35 +101,35 @@ public void Insert(IndexKey key, DocumentLocation location, ulong? transactionId
     {
         var txnId = transactionId ?? 0;
 
-        if (_options.Unique && TryFind(key, out var existingLocation, txnId))
-        {
-            if (!existingLocation.Equals(location))
-                throw new InvalidOperationException("Duplicate key violation for unique index");
-
-            return;
-        }
-
         var entry = new IndexEntry(key, location);
         var path = new List<uint>();
 
-        // Find the leaf node for insertion
-        var leafPageId = FindLeafNodeWithPath(key, path, txnId);
-
         var pageBuffer = System.Buffers.ArrayPool<byte>.Shared.Rent(_storage.PageSize);
         try
         {
-            ReadPage(leafPageId, txnId, pageBuffer);
+            // Find the leaf node for insertion; the leaf is left in pageBuffer.
+            var leafPageId = FindLeafNodeWithPath(key, path, txnId, pageBuffer);
             var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32));
 
+            if (_options.Unique)
+            {
+                var existing = BinarySearchLeaf(pageBuffer, header.EntryCount, key);
+                if (existing.Found)
+                {
+                    if (!existing.Location.Equals(location))
+                        throw new InvalidOperationException("Duplicate key violation for unique index");
+                    return;
+                }
+            }
+
             // Check if we need to split
             if (header.EntryCount >= MaxEntriesPerNode)
             {
                 SplitNode(leafPageId, path, txnId);
 
                 // Re-find leaf after split to ensure we have correct node
                 path.Clear();
-                leafPageId = FindLeafNodeWithPath(key, path, txnId);
-                ReadPage(leafPageId, txnId, pageBuffer);
+                leafPageId = FindLeafNodeWithPath(key, path, txnId, pageBuffer);
             }
 
             // Insert entry into leaf
@@ -151,17 +151,24 @@ public bool TryFind(IndexKey key, out DocumentLocation location, ulong? transact
 
         var leafPageId = FindLeafNode(key, txnId);
 
-        Span<byte> pageBuffer = stackalloc byte[_storage.PageSize];
-        ReadPage(leafPageId, txnId, pageBuffer);
+        var pageBuffer = System.Buffers.ArrayPool<byte>.Shared.Rent(_storage.PageSize);
+        try
+        {
+            ReadPage(leafPageId, txnId, pageBuffer);
 
-        var header = BTreeNodeHeader.ReadFrom(pageBuffer[32..]);
-        var result = BinarySearchLeaf(pageBuffer, header.EntryCount, key);
-        if (result.Found)
+            var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32));
+            var result = BinarySearchLeaf(pageBuffer, header.EntryCount, key);
+            if (result.Found)
+            {
+                location = result.Location;
+                return true;
+            }
+            return false;
+        }
+        finally
         {
-            location = result.Location;
-            return true;
+            System.Buffers.ArrayPool<byte>.Shared.Return(pageBuffer);
         }
-        return false;
     }
 
     /// <summary>
@@ -395,31 +402,34 @@ internal uint FindLeafNode(IndexKey key, ulong transactionId)
 
     private uint FindLeafNodeWithPath(IndexKey key, List<uint> path, ulong transactionId)
     {
-        var currentPageId = _rootPageId;
         var pageBuffer = System.Buffers.ArrayPool<byte>.Shared.Rent(_storage.PageSize);
-
         try
         {
-            while (true)
-            {
-                ReadPage(currentPageId, transactionId, pageBuffer);
-                var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32));
-
-                if (header.IsLeaf)
-                {
-                    return currentPageId;
-                }
-
-                path.Add(currentPageId);
-                currentPageId = FindChildNode(pageBuffer, header, key);
-            }
+            return FindLeafNodeWithPath(key, path, transactionId, pageBuffer);
         }
         finally
         {
             System.Buffers.ArrayPool<byte>.Shared.Return(pageBuffer);
         }
     }
 
+    /// <summary>Like the overload above, but leaves the leaf page in <paramref name="pageBuffer"/>.</summary>
+    private uint FindLeafNodeWithPath(IndexKey key, List<uint> path, ulong transactionId, byte[] pageBuffer)
+    {
+        var currentPageId = _rootPageId;
+        while (true)
+        {
+            ReadPage(currentPageId, transactionId, pageBuffer);
+            var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32));
+
+            if (header.IsLeaf)
+                return currentPageId;
+
+            path.Add(currentPageId);
+            currentPageId = FindChildNode(pageBuffer, header, key);
+        }
+    }
+
     private uint FindChildNode(Span<byte> nodeBuffer, BTreeNodeHeader header, IndexKey key)
     {
         // Internal Node Format:
@@ -680,8 +690,7 @@ private void InsertIntoLeaf(uint leafPageId, IndexEntry entry, Span<byte> pageBu
         if (insertOffset < currentEnd)
         {
             int bytesToShift = currentEnd - insertOffset;
-            for (int i = bytesToShift - 1; i >= 0; i--)
-                pageBuffer[insertOffset + entrySize + i] = pageBuffer[insertOffset + i];
+            pageBuffer.Slice(insertOffset, bytesToShift).CopyTo(pageBuffer.Slice(insertOffset + entrySize));
         }
 
         // Write new entry at sorted insertion point

src/BLite.Core/Query/BTreeExpressionVisitor.cs

@@ -44,6 +44,15 @@ protected override Expression VisitMethodCall(MethodCallExpression node)
                 case "LongCount":
                     VisitCount(node);
                     return node;
+                case "First":
+                case "FirstOrDefault":
+                case "Single":
+                case "SingleOrDefault":
+                case "Last":
+                case "LastOrDefault":
+                case "Any":
+                    VisitElementTerminal(node);
+                    return node;
                 default:
                     // GroupBy, Join, etc. — operators the direct pipeline cannot model.
                     _model.HasComplexOperators = true;
@@ -60,8 +69,11 @@ private void VisitWhere(MethodCallExpression node)
         Visit(node.Arguments[0]);
 
         var predicate = (UnaryExpression)node.Arguments[1];
-        var lambda = (LambdaExpression)predicate.Operand;
+        AddPredicate((LambdaExpression)predicate.Operand);
+    }
 
+    private void AddPredicate(LambdaExpression lambda)
+    {
         if (_model.WhereClause == null)
         {
             _model.WhereClause = lambda;
@@ -144,6 +156,36 @@ private void VisitAggregate(MethodCallExpression node)
         }
     }
 
+    /// <summary>
+    /// Handles <c>First</c>, <c>FirstOrDefault</c>, <c>Single</c>, <c>SingleOrDefault</c>,
+    /// <c>Last</c>, <c>LastOrDefault</c> and <c>Any</c>, with or without a predicate. A predicate
+    /// on the entities themselves is folded into the WHERE clause (<c>First(p)</c> is
+    /// <c>Where(p).First()</c>), so the index planner sees it, but only where
+    /// <see cref="CanFoldPredicate"/> allows it. Overloads taking a default value are left to the
+    /// fallback pipeline. <c>All(p)</c> is not handled here: it is not <c>Where(p).All()</c>.
+    /// </summary>
+    private void VisitElementTerminal(MethodCallExpression node)
+    {
+        Visit(node.Arguments[0]); // recurse into source
+
+        if (node.Arguments.Count == 1)
+        {
+            _model.ElementTerminal = node.Method.Name;
+        }
+        else if (node.Arguments.Count == 2
+                 && node.Arguments[1] is UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression predicate }
+                 && CanFoldPredicate(node.Arguments[0], predicate))
+        {
+            AddPredicate(predicate);
+            _model.ElementTerminal = node.Method.Name;
+            _model.FoldedTerminal = node;
+        }
+        else
+        {
+            _model.HasComplexOperators = true;
+        }
+    }
+
     /// <summary>
     /// Handles <c>Count()</c> / <c>LongCount()</c> as a push-down scalar.
     /// When there is no predicate, the provider can use the primary BTree key count instead
@@ -157,10 +199,57 @@ private void VisitCount(MethodCallExpression node)
             // Count() with no predicate — fast O(log n) + k BTree scan.
             _model.IsCountOnly = true;
         }
+        else if (node.Arguments[1] is UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression predicate }
+                 && CanFoldPredicate(node.Arguments[0], predicate))
+        {
+            // Count(predicate) on the entities themselves is Where(predicate).Count().
+            AddPredicate(predicate);
+            _model.IsCountOnly = true;
+            _model.FoldedTerminal = node;
+        }
         else
         {
-            // Count(predicate) — treat as a complex operator.
+            // Count(predicate) after any other operator — treat as a complex operator.
             _model.HasComplexOperators = true;
         }
     }
+
+    /// <summary>
+    /// True when a terminal's <paramref name="predicate"/> may be folded into the WHERE clause,
+    /// i.e. when <c>source.Terminal(p)</c> equals <c>source.Where(p).Terminal()</c> evaluated with
+    /// the filter applied first. That holds only when every operator between the root and the
+    /// terminal is a plain <c>Where</c> or an ordering (<c>OrderBy*</c> / <c>ThenBy*</c>): those keep
+    /// the element type and commute with a filter. Anything else, such as <c>GroupBy</c>,
+    /// <c>OfType</c>, <c>Select</c> or the order-dependent <c>Take</c>, <c>Skip</c>,
+    /// <c>SkipWhile</c> and <c>TakeWhile</c>, either changes what the predicate is applied to or
+    /// gives a different result when the predicate runs first.
+    /// </summary>
+    private static bool CanFoldPredicate(Expression source, LambdaExpression predicate)
+    {
+        var current = source;
+        while (current is MethodCallExpression call)
+        {
+            if (call.Method.DeclaringType != typeof(Queryable))
+                return false;
+
+            switch (call.Method.Name)
+            {
+                case "Where" when call.Arguments[1] is UnaryExpression { Operand: LambdaExpression { Parameters.Count: 1 } }:
+                case "OrderBy":
+                case "OrderByDescending":
+                case "ThenBy":
+                case "ThenByDescending":
+                    current = call.Arguments[0];
+                    break;
+                default:
+                    return false;
+            }
+        }
+
+        // Only element-type-preserving operators were skipped, so the predicate must take the
+        // root collection's own element type (the entity type).
+        return current is ConstantExpression { Value: IQueryable root }
+               && predicate.Parameters.Count == 1
+               && predicate.Parameters[0].Type == root.ElementType;
+    }
 }

src/BLite.Core/Query/BTreeQueryProvider.cs

@@ -214,6 +214,35 @@ private async Task<TResult> ExecuteAsyncCore<TResult>(Expression expression, Can
         visitor.Visit(expression);
         var model = visitor.GetModel();
 
+        // ── Element terminals: First / FirstOrDefault / Single / SingleOrDefault / Any ──
+        // With only a WHERE in front, at most one (Single*: two) matching document decides the
+        // result, so fetch just that many through the index-aware FetchAsync. Anything else in
+        // front (OrderBy, Select, Take, ...), and Last / LastOrDefault, which need every match,
+        // go through the general fallback.
+        if (model.ElementTerminal is { } terminal)
+        {
+            if (terminal is not ("Last" or "LastOrDefault")
+                && !model.HasComplexOperators && model.SelectClause == null && model.OrderByClause == null
+                && !model.Take.HasValue && !model.Skip.HasValue && model.AggregateOp == null && !model.IsCountOnly)
+            {
+                bool single = terminal is "Single" or "SingleOrDefault";
+                var found = new List<T>(single ? 2 : 1);
+                await foreach (var item in _collection.FetchAsync(model.WhereClause, single ? 2 : 1, cancellationToken))
+                    found.Add(item);
+
+                return terminal switch
+                {
+                    "First"           => (TResult)(object)found.First()!,
+                    "FirstOrDefault"  => (TResult)(object?)found.FirstOrDefault()!,
+                    "Single"          => (TResult)(object)found.Single()!,
+                    "SingleOrDefault" => (TResult)(object?)found.SingleOrDefault()!,
+                    _                 => (TResult)(object)(found.Count > 0),
+                };
+            }
+
+            model.HasComplexOperators = true;
+        }
+
         // ── Fast path: Count() / LongCount() terminal with no WHERE ─────────────
         // Uses the primary BTree key scan (O(n) key reads, zero document reads).
         // Disabled when Take/Skip are present — those shaping operators reduce the
@@ -314,8 +343,9 @@ private async Task<TResult> ExecuteAsyncCore<TResult>(Expression expression, Can
         // FetchAsync applied the WHERE, so the replay drops the root-chain Where calls folded into
         // it. Re-evaluating them would change nothing except for VectorSearch, which only the index
         // can evaluate. Every other VectorSearch in the expression is still evaluated and throws.
+        // The same holds for a terminal predicate folded into the WHERE (First(p) after OrderBy).
         if (whereAlreadyApplied)
-            expression = DropAppliedWheres(expression).Expression;
+            expression = DropAppliedWheres(DropFoldedTerminalPredicate(expression, model.FoldedTerminal)).Expression;
 
         if (model.HasComplexOperators)
             return ExecuteViaEnumerableRewriter<TResult>(expression, sourceData);
@@ -675,6 +705,34 @@ private static (Expression Expression, bool IsFetchedSubset) DropAppliedWheres(E
         return (rewritten, isFetchedSubset && s_subsetPreservingOperators.Contains(call.Method.Name));
     }
 
+    /// <summary>
+    /// Replaces the terminal whose predicate <see cref="BTreeExpressionVisitor"/> folded into
+    /// <see cref="QueryModel.WhereClause"/> (<c>Count(p)</c>, <c>First(p)</c>, ...) with its
+    /// parameterless overload. The visitor folds a predicate only when the terminal's source is the
+    /// root followed by plain Where and ordering calls, so the elements reaching the terminal are
+    /// fetched documents, which all satisfy the predicate already.
+    /// </summary>
+    [RequiresDynamicCode("BLite LINQ queries use Expression.Compile() and MakeGenericMethod which require dynamic code generation.")]
+    [RequiresUnreferencedCode("BLite LINQ queries use reflection to resolve methods and types at runtime. Ensure all entity types and their members are preserved.")]
+    private static Expression DropFoldedTerminalPredicate(Expression expression, MethodCallExpression? foldedTerminal)
+    {
+        if (foldedTerminal is null || !ReferenceEquals(expression, foldedTerminal))
+            return expression;
+
+        var name = foldedTerminal.Method.Name;
+        if (!s_parameterlessTerminals.TryGetValue(name, out var definition))
+        {
+            definition = typeof(Queryable).GetMethods().Single(m => m.Name == name && m.GetParameters().Length == 1);
+            s_parameterlessTerminals[name] = definition;
+        }
+
+        var parameterless = definition.MakeGenericMethod(foldedTerminal.Method.GetGenericArguments());
+        return Expression.Call(parameterless, foldedTerminal.Arguments[0]);
+    }
+
+    /// <summary>Generic definitions of the parameterless Queryable terminals, by name.</summary>
+    private static readonly ConcurrentDictionary<string, MethodInfo> s_parameterlessTerminals = new();
+
     /// <summary>Queryable operators whose output elements are a subset of their source's.</summary>
     private static readonly HashSet<string> s_subsetPreservingOperators =
     [

src/BLite.Core/Query/QueryModel.cs

@@ -34,4 +34,17 @@ internal class QueryModel
     /// Enables an O(n) primary-BTree key scan instead of full document deserialisation.
     /// </summary>
     public bool IsCountOnly { get; set; }
+
+    /// <summary>
+    /// Element terminal (<c>First</c>, <c>FirstOrDefault</c>, <c>Single</c>,
+    /// <c>SingleOrDefault</c>, <c>Last</c>, <c>LastOrDefault</c> or <c>Any</c>) ending the query,
+    /// or null. A predicate passed to it has already been folded into <see cref="WhereClause"/>.
+    /// </summary>
+    public string? ElementTerminal { get; set; }
+
+    /// <summary>
+    /// The terminal call (<c>Count(p)</c>, <c>First(p)</c>, ...) whose predicate was folded into
+    /// <see cref="WhereClause"/>, or null. A fallback replay must not evaluate that predicate again.
+    /// </summary>
+    public MethodCallExpression? FoldedTerminal { get; set; }
 }