Improve selection lookup on value completion

Author: michaelstaib

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/Results/ValueCompletion.cs

@@ -78,9 +78,19 @@ public bool BuildResult(
             InitializeTargetObject(source, target);
         }
 
+        var hasObjectScope = target.TryGetObjectScope(out var objectScope);
+
         foreach (var property in source.EnumerateObject())
         {
-            if (!target.TryGetProperty(property.NameSpan, out var resultField))
+            CompositeResultElement resultField;
+            if (hasObjectScope)
+            {
+                if (!objectScope.TryGetProperty(property.NameSpan, out resultField))
+                {
+                    continue;
+                }
+            }
+            else if (!target.TryGetProperty(property.NameSpan, out resultField))
             {
                 continue;
             }
@@ -812,9 +822,19 @@ private bool TryCompleteObjectValue(
             target.SetObjectValue(objectSelectionSet);
         }
 
+        var hasObjectScope = target.TryGetObjectScope(out var objectScope);
+
         foreach (var property in source.EnumerateObject())
         {
-            if (!target.TryGetProperty(property.NameSpan, out var targetProperty))
+            CompositeResultElement targetProperty;
+            if (hasObjectScope)
+            {
+                if (!objectScope.TryGetProperty(property.NameSpan, out targetProperty))
+                {
+                    continue;
+                }
+            }
+            else if (!target.TryGetProperty(property.NameSpan, out targetProperty))
             {
                 continue;
             }

src/HotChocolate/Fusion/src/Fusion.Execution/Text/Json/CompositeResultDocument.TryGetProperty.cs

@@ -1,5 +1,7 @@
 using System.Buffers;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using HotChocolate.Fusion.Execution.Nodes;
 using HotChocolate.Text.Json;
 
 namespace HotChocolate.Fusion.Text.Json;
@@ -162,6 +164,28 @@ internal bool TryGetNamedPropertyValue(
             out value);
     }
 
+    internal bool TryGetObjectScope(
+        Cursor startCursor,
+        out CompositeResultObjectScope scope)
+    {
+        ObjectDisposedException.ThrowIf(_disposed != 0, this);
+
+        var row = _metaDb.GetValue(ref startCursor);
+        CheckExpectedType(ElementTokenType.StartObject, row.TokenType);
+
+        if (row.OperationReferenceType is OperationReferenceType.SelectionSet)
+        {
+            var selectionSetId = row.OperationReferenceId;
+            var selectionSet = _operation.GetSelectionSetById(selectionSetId);
+
+            scope = new CompositeResultObjectScope(this, startCursor, selectionSet, selectionSetId);
+            return true;
+        }
+
+        scope = default;
+        return false;
+    }
+
     private bool TryGetNamedPropertyValue(
         Cursor startCursor,
         Cursor endCursor,
@@ -274,3 +298,40 @@ internal Cursor GetEndCursor(Cursor cursor)
         return cursor + _metaDb.GetNumberOfRows(cursor);
     }
 }
+
+internal readonly ref struct CompositeResultObjectScope
+{
+    private readonly CompositeResultDocument _document;
+    private readonly CompositeResultDocument.Cursor _startObject;
+    private readonly SelectionSet _selectionSet;
+    private readonly int _selectionSetId;
+
+    internal CompositeResultObjectScope(
+        CompositeResultDocument document,
+        CompositeResultDocument.Cursor startObject,
+        SelectionSet selectionSet,
+        int selectionSetId)
+    {
+        _document = document;
+        _startObject = startObject;
+        _selectionSet = selectionSet;
+        _selectionSetId = selectionSetId;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public bool TryGetProperty(
+        ReadOnlySpan<byte> name,
+        out CompositeResultElement value)
+    {
+        if (_selectionSet.TryGetSelection(name, out var selection))
+        {
+            var propertyIndex = selection.Id - _selectionSetId - 1;
+            var cursor = _startObject + (propertyIndex * 2) + 2;
+            value = new CompositeResultElement(_document, cursor);
+            return true;
+        }
+
+        value = default;
+        return false;
+    }
+}

src/HotChocolate/Fusion/src/Fusion.Execution/Text/Json/CompositeResultElement.cs

@@ -432,6 +432,13 @@ public bool TryGetProperty(ReadOnlySpan<byte> utf8PropertyName, out CompositeRes
         return _parent.TryGetNamedPropertyValue(_cursor, utf8PropertyName, out value);
     }
 
+    internal bool TryGetObjectScope(out CompositeResultObjectScope scope)
+    {
+        CheckValidInstance();
+
+        return _parent.TryGetObjectScope(_cursor, out scope);
+    }
+
     internal CompositeResultElement GetPropertyBySelectionId(int selectionId)
     {
         CheckValidInstance();

src/HotChocolate/Fusion/test/Fusion.Execution.Tests/Text/Json/CompositeResultDocumentTests.cs

@@ -80,6 +80,55 @@ fragment Product on Product {
         Assert.Equal(JsonValueKind.Object, productBySlug.ValueKind);
     }
 
+    [Fact]
+    public void TryGetObjectScope_Should_ReturnFalseAndFallback_When_ObjectHasNoSelectionSet()
+    {
+        // arrange
+        var schema = CreateCompositeSchema();
+
+        var plan = PlanOperation(
+            schema,
+            """
+            {
+                productBySlug(slug: "1") {
+                    ... Product
+                }
+            }
+
+            fragment Product on Product {
+                id
+                name
+            }
+            """);
+
+        var compositeResult = new CompositeResultDocument(CommonTestExtensions.CreateArena(), plan.Operation, 0);
+        var operation = compositeResult.Data.Operation;
+
+        var productBySlug = compositeResult.Data.GetProperty("productBySlug");
+        var productBySlugSelection = productBySlug.AssertSelection();
+        var selectionSet = operation.GetSelectionSet(productBySlugSelection);
+        productBySlug.SetObjectValue(selectionSet);
+
+        var objectStart = compositeResult.GetStartCursor(productBySlug.Cursor);
+        var row = compositeResult._metaDb.Get(objectStart);
+        compositeResult._metaDb.Replace(
+            objectStart,
+            ElementTokenType.StartObject,
+            sizeOrLength: row.SizeOrLength,
+            parentRow: row.Parent,
+            numberOfRows: row.NumberOfRows,
+            flags: row.Flags);
+
+        // act
+        var hasScope = productBySlug.TryGetObjectScope(out _);
+        var hasProperty = productBySlug.TryGetProperty("id"u8, out var id);
+
+        // assert
+        Assert.False(hasScope);
+        Assert.True(hasProperty);
+        Assert.Equal("id", id.GetPropertyName());
+    }
+
     [Fact]
     public void Add_SourceResult_Leaf_Value()
     {