[Fusion] Support abstract lookups

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceInterfaceTypeCollection.cs

@@ -3,16 +3,10 @@
 
 namespace HotChocolate.Fusion.Types.Collections;
 
-public class SourceInterfaceTypeCollection
-    : SourceMemberCollection<SourceInterfaceType>
-    , ISourceComplexTypeCollection<SourceInterfaceType>
-    , ISourceComplexTypeCollection<ISourceComplexType>
+public class SourceInterfaceTypeCollection(IEnumerable<SourceInterfaceType> members)
+    : SourceMemberCollection<SourceInterfaceType>(members), ISourceComplexTypeCollection<SourceInterfaceType>
+        , ISourceComplexTypeCollection<ISourceComplexType>
 {
-    public SourceInterfaceTypeCollection(IEnumerable<SourceInterfaceType> members)
-        : base(members)
-    {
-    }
-
     ISourceComplexType ISourceMemberCollection<ISourceComplexType>.this[string schemaName]
         => this[schemaName];
 
@@ -22,12 +16,9 @@ public ImmutableArray<SourceInterfaceType> Types
     ImmutableArray<ISourceComplexType> ISourceComplexTypeCollection<ISourceComplexType>.Types
         => [.. Members];
 
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out SourceInterfaceType? type)
-        => TryGetMember(schemaName, out type);
-
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
+    public bool TryGetMember(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
     {
-        if (TryGetMember(schemaName, out var member))
+        if (base.TryGetMember(schemaName, out var member))
         {
             type = member;
             return true;

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceMemberCollection.cs

@@ -11,7 +11,7 @@ public class SourceMemberCollection<TMember>
 {
     private readonly FrozenDictionary<string, TMember> _members;
 
-    protected SourceMemberCollection(IEnumerable<TMember> members)
+    public SourceMemberCollection(IEnumerable<TMember> members)
     {
         _members = members.ToFrozenDictionary(t => t.SchemaName);
     }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceObjectTypeCollection.cs

@@ -17,12 +17,9 @@ public ImmutableArray<SourceObjectType> Types
     ImmutableArray<ISourceComplexType> ISourceComplexTypeCollection<ISourceComplexType>.Types
         => [.. Members];
 
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out SourceObjectType? type)
-        => TryGetMember(schemaName, out type);
-
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
+    public bool TryGetMember(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
     {
-        if (TryGetMember(schemaName, out var member))
+        if (base.TryGetMember(schemaName, out var member))
         {
             type = member;
             return true;

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceUnionTypeCollection.cs

@@ -0,0 +1,6 @@
+namespace HotChocolate.Fusion.Types.Collections;
+
+public class SourceUnionTypeCollection(IEnumerable<SourceUnionType> members)
+    : SourceMemberCollection<SourceUnionType>(members)
+{
+}

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompletionTools.cs

@@ -137,6 +137,29 @@ public static SourceInterfaceTypeCollection CreateSourceInterfaceTypeCollection(
         return new SourceInterfaceTypeCollection(temp);
     }
 
+    public static SourceUnionTypeCollection CreateSourceUnionTypeCollection(
+        UnionTypeDefinitionNode typeDef,
+        CompositeSchemaBuilderContext context)
+    {
+        var types = TypeDirectiveParser.Parse(typeDef.Directives);
+        var lookupDirectives = LookupDirectiveParser.Parse(typeDef.Directives);
+        var temp = new SourceUnionType[types.Length];
+
+        for (var i = 0; i < types.Length; i++)
+        {
+            var type = types[i];
+            var lookups = GetLookupBySchema(lookupDirectives, type.SchemaName, typeDef.Name.Value);
+            context.RegisterForCompletionRange(lookups);
+
+            temp[i] = new SourceUnionType(
+                typeDef.Name.Value,
+                type.SchemaName,
+                lookups);
+        }
+
+        return new SourceUnionTypeCollection(temp);
+    }
+
     private static ImmutableArray<Lookup> GetLookupBySchema(
         ImmutableArray<LookupDirective> allLookups,
         string schemaName,

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeSchemaBuilder.cs

@@ -532,7 +532,8 @@ private static void CompleteUnionType(
     {
         var directives = CompletionTools.CreateDirectiveCollection(typeDef.Directives, context);
         var types = CompletionTools.CreateObjectTypeCollection(typeDef.Types, context);
-        type.Complete(new CompositeUnionTypeCompletionContext(types, directives, FeatureCollection.Empty));
+        var sources = CompletionTools.CreateSourceUnionTypeCollection(typeDef, context);
+        type.Complete(new CompositeUnionTypeCompletionContext(types, directives, sources, FeatureCollection.Empty));
     }
 
     private static void CompleteOutputField(

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeUnionTypeCompletionContext.cs

@@ -6,11 +6,14 @@ namespace HotChocolate.Fusion.Types.Completion;
 internal readonly ref struct CompositeUnionTypeCompletionContext(
     FusionObjectTypeDefinitionCollection types,
     FusionDirectiveCollection directives,
+    SourceUnionTypeCollection sources,
     IFeatureCollection features)
 {
     public FusionDirectiveCollection Directives { get; } = directives;
 
     public FusionObjectTypeDefinitionCollection Types { get; } = types;
 
+    public SourceUnionTypeCollection Sources { get; } = sources;
+
     public IFeatureCollection Features { get; } = features;
 }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceComplexTypeCollection.cs

@@ -7,7 +7,5 @@ public interface ISourceComplexTypeCollection<TType>
     : ISourceMemberCollection<TType>
     where TType : ISourceComplexType
 {
-    bool TryGetType(string schemaName, [NotNullWhen(true)] out TType? type);
-
     ImmutableArray<TType> Types { get; }
 }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceMemberCollection.cs

@@ -1,8 +1,9 @@
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 
 namespace HotChocolate.Fusion.Types;
 
-public interface ISourceMemberCollection<out TMember>
+public interface ISourceMemberCollection<TMember>
     : IEnumerable<TMember>
     where TMember : ISourceMember
 {
@@ -13,4 +14,6 @@ public interface ISourceMemberCollection<out TMember>
     bool ContainsSchema(string schemaName);
 
     ImmutableArray<string> Schemas { get; }
+
+    bool TryGetMember(string schemaName, [NotNullWhen(true)] out TMember? type);
 }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/FusionUnionTypeDefinition.cs

@@ -32,6 +32,23 @@ public FusionUnionTypeDefinition(string name, string? description)
 
     public SchemaCoordinate Coordinate => new(Name, ofDirective: false);
 
+    /// <summary>
+    /// Gets the source type definition of this type.
+    /// </summary>
+    /// <value>
+    /// The source type definition of this type.
+    /// </value>
+    public SourceUnionTypeCollection Sources
+    {
+        get;
+        private set
+        {
+            EnsureNotSealed(_completed);
+
+            field = value;
+        }
+    } = null!;
+
     public FusionObjectTypeDefinitionCollection Types
     {
         get;
@@ -73,6 +90,7 @@ internal void Complete(CompositeUnionTypeCompletionContext context)
 
         Directives = context.Directives;
         Types = context.Types;
+        Sources = context.Sources;
         Features = context.Features;
 
         _completed = true;

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Metadata/Lookup.cs

@@ -17,7 +17,7 @@ public sealed class Lookup : INeedsCompletion
     /// </summary>
     /// <param name="schemaName">The name of the source schema.</param>
     /// <param name="declaringTypeName">The name of the type that declares the field.</param>
-    /// <param name="name">The name of the lookup field.</param>
+    /// <param name="fieldName">The name of the lookup field.</param>
     /// <param name="arguments">The arguments that represent field requirements.</param>
     /// <param name="fields">The paths to the field that are required.</param>
     /// <exception cref="ArgumentException">
@@ -30,13 +30,13 @@ public sealed class Lookup : INeedsCompletion
     public Lookup(
         string schemaName,
         string declaringTypeName,
-        string name,
+        string fieldName,
         ImmutableArray<LookupArgument> arguments,
         ImmutableArray<IValueSelectionNode> fields)
     {
         ArgumentException.ThrowIfNullOrEmpty(schemaName);
         ArgumentException.ThrowIfNullOrEmpty(declaringTypeName);
-        ArgumentException.ThrowIfNullOrEmpty(name);
+        ArgumentException.ThrowIfNullOrEmpty(fieldName);
 
         if (arguments.Length == 0)
         {
@@ -50,7 +50,7 @@ public Lookup(
 
         _declaringTypeName = declaringTypeName;
         SchemaName = schemaName;
-        Name = name;
+        FieldName = fieldName;
         Arguments = arguments;
         Fields = fields;
     }
@@ -63,7 +63,12 @@ public Lookup(
     /// <summary>
     /// Gets the name of the lookup field.
     /// </summary>
-    public string Name { get; }
+    public string FieldName { get; }
+
+    /// <summary>
+    /// Get the name of lookup field type.
+    /// </summary>
+    public string FieldType => _declaringTypeName;
 
     /// <summary>
     /// Gets the arguments that represent field requirements.

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Metadata/SourceUnionType.cs

@@ -0,0 +1,15 @@
+using System.Collections.Immutable;
+
+namespace HotChocolate.Fusion.Types;
+
+public sealed class SourceUnionType(
+    string name,
+    string schemaName,
+    ImmutableArray<Lookup> lookups) : ISourceMember
+{
+    public string Name { get; } = name;
+
+    public string SchemaName { get; } = schemaName;
+
+    public ImmutableArray<Lookup> Lookups { get; } = lookups;
+}

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

@@ -76,7 +76,7 @@ public OperationDefinitionBuilder SetSelectionSet(SelectionSetNode selectionSet)
             }
 
             var lookupField = new FieldNode(
-                new NameNode(_lookup.Name),
+                new NameNode(_lookup.FieldName),
                 null,
                 [],
                 arguments,
@@ -87,7 +87,7 @@ public OperationDefinitionBuilder SetSelectionSet(SelectionSetNode selectionSet)
             var indexBuilder = index.ToBuilder();
             indexBuilder.Register(selectionSet);
             index = indexBuilder;
-            selectionPath = selectionPath.AppendField(_lookup.Name);
+            selectionPath = selectionPath.AppendField(_lookup.FieldName);
         }
 
         var definition = new OperationDefinitionNode(

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

@@ -295,6 +295,28 @@ lookup is null
                 requirements = requirements.Add(argumentRequirementKey, operationRequirement);
             }
 
+            // If the lookup returns an abstract type, we might need to insert a type refinement
+            // around our selection set.
+            if (_schema.Types.TryGetType(lookup.FieldType, out var lookupFieldType)
+                && lookupFieldType != workItem.SelectionSet.Type)
+            {
+                var typeRefinement =
+                    new InlineFragmentNode(
+                        null,
+                        new NamedTypeNode(workItem.SelectionSet.Type.Name),
+                        [],
+                        resolvable);
+                var selectionSetWithTypeRefinement = new SelectionSetNode(null, [typeRefinement]);
+
+                var indexBuilder = index.ToBuilder();
+
+                indexBuilder.Register(resolvable, selectionSetWithTypeRefinement);
+
+                index = indexBuilder;
+
+                operationBuilder.SetSelectionSet(selectionSetWithTypeRefinement);
+            }
+
             operationBuilder.SetLookup(lookup, requirementKey);
         }
 
@@ -1039,7 +1061,7 @@ public static void Enqueue(
                 break;
 
             case FieldRequirementWorkItem wi:
-                possiblePlans.EnqueueRequirePlanNodes(planNodeTemplate, wi);
+                possiblePlans.EnqueueRequirePlanNodes(planNodeTemplate, wi, compositeSchema);
                 break;
 
             default:
@@ -1076,7 +1098,7 @@ private static void EnqueueLookupPlanNodes(
 
         foreach (var (schemaName, resolutionCost) in compositeSchema.GetPossibleSchemas(workItem.SelectionSet))
         {
-            foreach (var lookup in workItem.SelectionSet.Type.GetPossibleLookups(schemaName))
+            foreach (var lookup in workItem.SelectionSet.Type.GetPossibleLookups(schemaName, compositeSchema))
             {
                 possiblePlans.Enqueue(
                     planNodeTemplate with
@@ -1093,7 +1115,8 @@ planNodeTemplate with
     private static void EnqueueRequirePlanNodes(
         this PriorityQueue<PlanNode, double> possiblePlans,
         PlanNode planNodeTemplate,
-        FieldRequirementWorkItem workItem)
+        FieldRequirementWorkItem workItem,
+        FusionSchemaDefinition compositeSchema)
     {
         var backlog = planNodeTemplate.Backlog.Pop();
 
@@ -1109,7 +1132,7 @@ planNodeTemplate with
                         Backlog = backlog.Push(workItem)
                     });
 
-                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName))
+                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName, compositeSchema))
                 {
                     possiblePlans.Enqueue(
                         planNodeTemplate with
@@ -1123,7 +1146,7 @@ planNodeTemplate with
             }
             else
             {
-                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName))
+                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName, compositeSchema))
                 {
                     possiblePlans.Enqueue(
                         planNodeTemplate with
@@ -1208,12 +1231,42 @@ static void TrackSchema(Dictionary<string, int> possibleSchemas, string schemaNa
         }
     }
 
-    private static IEnumerable<Lookup> GetPossibleLookups(this ITypeDefinition type, string schemaName)
+    private static IEnumerable<Lookup> GetPossibleLookups(
+        this ITypeDefinition type,
+        string schemaName,
+        FusionSchemaDefinition compositeSchema)
     {
-        if (type is FusionComplexTypeDefinition complexType
-            && complexType.Sources.TryGetType(schemaName, out var source))
+        if (type is FusionComplexTypeDefinition complexType)
         {
-            return source.Lookups;
+            var lookups = new List<Lookup>();
+
+            if (complexType.Sources.TryGetMember(schemaName, out var source))
+            {
+                lookups.AddRange(source.Lookups);
+            }
+
+            foreach (var interfaceType in complexType.Implements)
+            {
+                if (interfaceType.Sources.TryGetMember(schemaName, out var interfaceSource))
+                {
+                    lookups.AddRange(interfaceSource.Lookups);
+                }
+            }
+
+            var unionTypes = compositeSchema.Types
+                .OfType<FusionUnionTypeDefinition>()
+                .Where(u => u.Types.Contains(type))
+                .ToArray();
+
+            foreach (var unionType in unionTypes)
+            {
+                if (unionType.Sources.TryGetMember(schemaName, out var unionSource))
+                {
+                    lookups.AddRange(unionSource.Lookups);
+                }
+            }
+
+            return lookups;
         }
 
         return [];