[Fusion] Rewrite conditionals in document (#8734)

Author: tobias-tengler

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Nodes/OperationCompiler.cs

@@ -10,7 +10,7 @@ namespace HotChocolate.Fusion.Execution.Nodes;
 public sealed class OperationCompiler
 {
     private readonly ISchemaDefinition _schema;
-    private readonly InlineFragmentOperationRewriter _inlineRewriter;
+    private readonly DocumentRewriter _documentRewriter;
     private readonly ObjectPool<OrderedDictionary<string, List<FieldSelectionNode>>> _fieldsPool;
     private readonly TypeNameField _typeNameField;
     private static readonly ArrayPool<object> s_objectArrayPool = ArrayPool<object>.Shared;
@@ -24,7 +24,7 @@ public OperationCompiler(
 
         _schema = schema;
         _fieldsPool = fieldsPool;
-        _inlineRewriter = new InlineFragmentOperationRewriter(schema, removeStaticallyExcludedSelections: true);
+        _documentRewriter = new(schema, removeStaticallyExcludedSelections: true);
         var nonNullStringType = new NonNullType(_schema.Types.GetType<IScalarTypeDefinition>(SpecScalarNames.String));
         _typeNameField = new TypeNameField(nonNullStringType);
     }
@@ -35,7 +35,7 @@ public Operation Compile(string id, string hash, OperationDefinitionNode operati
         ArgumentNullException.ThrowIfNull(operationDefinition);
 
         var document = new DocumentNode(new IDefinitionNode[] { operationDefinition });
-        document = _inlineRewriter.RewriteDocument(document);
+        document = _documentRewriter.RewriteDocument(document);
         operationDefinition = (OperationDefinitionNode)document.Definitions[0];
 
         var includeConditions = new IncludeConditionCollection();

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Execution/Pipeline/OperationPlanMiddleware.cs

@@ -11,7 +11,7 @@ namespace HotChocolate.Fusion.Execution.Pipeline;
 internal sealed class OperationPlanMiddleware
 {
     private readonly OperationPlanner _planner;
-    private readonly InlineFragmentOperationRewriter _rewriter;
+    private readonly DocumentRewriter _documentRewriter;
     private readonly IOperationPlannerInterceptor[] _interceptors;
     private readonly IFusionExecutionDiagnosticEvents _diagnosticsEvents;
 
@@ -21,7 +21,7 @@ private OperationPlanMiddleware(
         IEnumerable<IOperationPlannerInterceptor>? interceptors,
         IFusionExecutionDiagnosticEvents diagnosticsEvents)
     {
-        _rewriter = new InlineFragmentOperationRewriter(schema);
+        _documentRewriter = new(schema, removeStaticallyExcludedSelections: true);
         _planner = planner;
         _interceptors = interceptors?.ToArray() ?? [];
         _diagnosticsEvents = diagnosticsEvents;
@@ -60,8 +60,8 @@ private void PlanOperation(
 
         try
         {
-            // Before we can plan an operation, we must defragmentize it and remove statical include conditions.
-            var rewritten = _rewriter.RewriteDocument(operationDocument, context.Request.OperationName);
+            // Before we can plan an operation, we must de-fragmentize it and remove static include conditions.
+            var rewritten = _documentRewriter.RewriteDocument(operationDocument, context.Request.OperationName);
             var operation = rewritten.GetOperation(context.Request.OperationName);
 
             // After optimizing the query structure we can begin the planning process.

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationPlanner.cs

@@ -1276,6 +1276,7 @@ private OperationDefinitionNode InlineSelections(
         bool inlineInternal = false)
     {
         List<SelectionSetNode>? backlog = null;
+        var didInline = false;
 
         var rewriter = SyntaxRewriter.Create<List<ISyntaxNode>>(
             rewrite: (node, path) =>
@@ -1337,6 +1338,8 @@ private OperationDefinitionNode InlineSelections(
                         index);
                 }
 
+                didInline = true;
+
                 index.Register(originalSelectionSet, newSelectionSet);
                 return newSelectionSet;
             },
@@ -1347,7 +1350,16 @@ private OperationDefinitionNode InlineSelections(
             },
             leave: (_, path) => path.Pop());
 
-        return (OperationDefinitionNode)rewriter.Rewrite(operation, [])!;
+        var rewrittenOperation = (OperationDefinitionNode)rewriter.Rewrite(operation, [])!;
+
+        if (!didInline)
+        {
+            throw new InvalidOperationException(
+                $"Selections `{selectionsToInline}` could not be inlined into selection set of type "
+                + $"'{selectionSetType.Name}', as no selection set with the id {targetSelectionSetId} was found.");
+        }
+
+        return rewrittenOperation;
 
         static IReadOnlyList<DirectiveNode> AddInternalDirective(IHasDirectives selection)
         {

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/Partitioners/NodeFieldSelectionSetPartitioner.cs

@@ -49,40 +49,32 @@ public RootSelectionSetPartitionerResult Partition(RootSelectionSetPartitionerIn
         {
             switch (selection)
             {
-                case FieldNode fieldNode when schema.QueryType.Fields.TryGetField(fieldNode.Name.Value, out var field)
-                    && field is { Name: "node", Type: IInterfaceTypeDefinition { Name: "Node" } }:
-                    var directives = new List<DirectiveNode>(fieldNode.Directives);
-                    foreach (var fragment in context.FragmentPath)
+                case FieldNode fieldNode:
+                    var field = schema.QueryType.Fields[fieldNode.Name.Value];
+
+                    if (field.IsIntrospectionField)
                     {
-                        directives.AddRange(fragment.Directives);
+                        continue;
                     }
 
-                    context.NodeFields ??= [];
-                    context.NodeFields.Add(fieldNode.WithDirectives(directives));
-                    break;
+                    if (field is { Name: "node", Type: IInterfaceTypeDefinition { Name: "Node" } })
+                    {
+                        context.NodeFields ??= [];
+                        context.NodeFields.Add(fieldNode);
+                    }
+                    else
+                    {
+                        selections ??= [];
+                        selections.Add(selection);
+                    }
 
-                case FieldNode:
-                    selections ??= [];
-                    selections.Add(selection);
                     break;
 
                 case InlineFragmentNode inlineFragmentNode:
-                    var hasDirectives = inlineFragmentNode.Directives.Any();
-
-                    if (hasDirectives)
-                    {
-                        context.FragmentPath.Push(inlineFragmentNode);
-                    }
-
                     var fragmentSelectionSet = RewriteSelectionSet(
                         inlineFragmentNode.SelectionSet,
                         context);
 
-                    if (hasDirectives)
-                    {
-                        context.FragmentPath.Pop();
-                    }
-
                     if (fragmentSelectionSet is not null)
                     {
                         selections ??= [];
@@ -103,13 +95,6 @@ public RootSelectionSetPartitionerResult Partition(RootSelectionSetPartitionerIn
 
     private class Context
     {
-        /// <summary>
-        /// Gets the fragment path.
-        /// This is pushed to whenever we enter an inline fragment with directives,
-        /// in order to preserve those.
-        /// </summary>
-        public List<InlineFragmentNode> FragmentPath { get; } = [];
-
         public List<FieldNode>? NodeFields { get; set; }
     }
 }