[Fusion] Fixed possible lookup selection (#8625)

Co-authored-by: tobias-tengler <[email protected]>

Author: michaelstaib

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

@@ -1,5 +1,6 @@
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
+using HotChocolate.Fusion.Types.Metadata;
 
 namespace HotChocolate.Fusion.Types.Collections;
 

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

@@ -2,6 +2,7 @@
 using System.Collections.Frozen;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
+using HotChocolate.Fusion.Types.Metadata;
 
 namespace HotChocolate.Fusion.Types.Collections;
 

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

@@ -1,3 +1,5 @@
+using HotChocolate.Fusion.Types.Metadata;
+
 namespace HotChocolate.Fusion.Types.Collections;
 
 public class SourceObjectFieldCollection(IEnumerable<SourceOutputField> members)

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

@@ -1,5 +1,6 @@
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
+using HotChocolate.Fusion.Types.Metadata;
 
 namespace HotChocolate.Fusion.Types.Collections;
 

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

@@ -1,3 +1,5 @@
+using HotChocolate.Fusion.Types.Metadata;
+
 namespace HotChocolate.Fusion.Types.Collections;
 
 public class SourceUnionTypeCollection(IEnumerable<SourceUnionType> members)

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

@@ -3,6 +3,7 @@
 using HotChocolate.Fusion.Language;
 using HotChocolate.Fusion.Types.Collections;
 using HotChocolate.Fusion.Types.Directives;
+using HotChocolate.Fusion.Types.Metadata;
 using HotChocolate.Language;
 using HotChocolate.Types;
 
@@ -97,22 +98,28 @@ public static SourceObjectTypeCollection CreateSourceObjectTypeCollection(
     {
         var types = TypeDirectiveParser.Parse(typeDef.Directives);
         var lookupDirectives = LookupDirectiveParser.Parse(typeDef.Directives);
-        var temp = new SourceObjectType[types.Length];
+        var sourceObjectType = new SourceObjectType[types.Length];
+        var sourceImplements = ImplementsDirectiveParser.Parse(typeDef.Directives);
+        var sourceUnionMembersOf = context.GetSourceUnionMembers(typeDef.Name.Value);
 
         for (var i = 0; i < types.Length; i++)
         {
             var type = types[i];
             var schemaName = context.GetSchemaName(type.SchemaKey);
             var lookups = GetLookupBySchema(lookupDirectives, schemaName, typeDef.Name.Value, context);
             context.RegisterForCompletionRange(lookups);
+            sourceImplements.TryGetValue(type.SchemaKey, out var implements);
+            sourceUnionMembersOf.TryGetValue(type.SchemaKey, out var unionTypes);
 
-            temp[i] = new SourceObjectType(
+            sourceObjectType[i] = new SourceObjectType(
                 typeDef.Name.Value,
                 schemaName,
-                lookups);
+                lookups,
+                implements ?? ImmutableHashSet<string>.Empty,
+                unionTypes ?? ImmutableHashSet<string>.Empty);
         }
 
-        return new SourceObjectTypeCollection(temp);
+        return new SourceObjectTypeCollection(sourceObjectType);
     }
 
     public static SourceInterfaceTypeCollection CreateSourceInterfaceTypeCollection(
@@ -121,22 +128,26 @@ public static SourceInterfaceTypeCollection CreateSourceInterfaceTypeCollection(
     {
         var types = TypeDirectiveParser.Parse(typeDef.Directives);
         var lookupDirectives = LookupDirectiveParser.Parse(typeDef.Directives);
-        var temp = new SourceInterfaceType[types.Length];
+        var sourceInterfaceType = new SourceInterfaceType[types.Length];
+        var sourceImplements = ImplementsDirectiveParser.Parse(
+            typeDef.Directives);
 
         for (var i = 0; i < types.Length; i++)
         {
             var type = types[i];
             var schemaName = context.GetSchemaName(type.SchemaKey);
             var lookups = GetLookupBySchema(lookupDirectives, schemaName, typeDef.Name.Value, context);
             context.RegisterForCompletionRange(lookups);
+            sourceImplements.TryGetValue(type.SchemaKey, out var implements);
 
-            temp[i] = new SourceInterfaceType(
+            sourceInterfaceType[i] = new SourceInterfaceType(
                 typeDef.Name.Value,
                 schemaName,
-                lookups);
+                lookups,
+                implements ?? ImmutableHashSet<string>.Empty);
         }
 
-        return new SourceInterfaceTypeCollection(temp);
+        return new SourceInterfaceTypeCollection(sourceInterfaceType);
     }
 
     public static SourceUnionTypeCollection CreateSourceUnionTypeCollection(

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

@@ -6,6 +6,7 @@
 using HotChocolate.Fusion.Rewriters;
 using HotChocolate.Fusion.Types.Collections;
 using HotChocolate.Fusion.Types.Directives;
+using HotChocolate.Fusion.Types.Metadata;
 using HotChocolate.Language;
 using HotChocolate.Types;
 using Microsoft.Extensions.DependencyInjection;
@@ -123,6 +124,7 @@ private static CompositeSchemaBuilderContext CreateTypes(
         features ??= new FeatureCollection();
 
         return new CompositeSchemaBuilderContext(
+            schemaDocument,
             name,
             description,
             services,

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

@@ -2,6 +2,7 @@
 using HotChocolate.Features;
 using HotChocolate.Fusion.Types.Collections;
 using HotChocolate.Fusion.Types.Directives;
+using HotChocolate.Fusion.Types.Metadata;
 using HotChocolate.Language;
 using HotChocolate.Types;
 using DirectiveLocation = HotChocolate.Types.DirectiveLocation;
@@ -10,14 +11,17 @@ namespace HotChocolate.Fusion.Types.Completion;
 
 internal sealed class CompositeSchemaBuilderContext : ICompositeSchemaBuilderContext
 {
+    private readonly DocumentNode _document;
     private readonly Dictionary<ITypeNode, IType> _compositeTypes = new(SyntaxComparer.BySyntax);
     private readonly Dictionary<string, ITypeDefinition> _typeDefinitionLookup;
     private ImmutableDictionary<string, ITypeDefinitionNode> _typeDefinitionNodeLookup;
     private readonly Dictionary<string, FusionDirectiveDefinition> _directiveDefinitionLookup;
     private ImmutableDictionary<string, DirectiveDefinitionNode> _directiveDefinitionNodeLookup;
     private readonly List<INeedsCompletion> _completions = [];
+    private readonly ImmutableDictionary<string, ImmutableDictionary<SchemaKey, ImmutableHashSet<string>>> _sourceUnions;
 
     public CompositeSchemaBuilderContext(
+        DocumentNode document,
         string name,
         string? description,
         IServiceProvider services,
@@ -33,6 +37,9 @@ public CompositeSchemaBuilderContext(
         IFeatureCollection features,
         CompositeTypeInterceptor interceptor)
     {
+        _document =  document;
+        _sourceUnions = UnionMemberDirectiveParser.Parse(document.Definitions.OfType<UnionTypeDefinitionNode>());
+
         _typeDefinitionLookup = typeDefinitions.ToDictionary(t => t.Name);
         _directiveDefinitionLookup = directiveDefinitions.ToDictionary(t => t.Name);
         _typeDefinitionNodeLookup = typeDefinitionNodeLookup;
@@ -183,18 +190,21 @@ private static IType CreateType(ITypeNode typeNode, ITypeDefinition compositeNam
     }
 
     public FusionDirectiveDefinition GetDirectiveType(string name)
-    {
-        if (_directiveDefinitionLookup.TryGetValue(name, out var type))
-        {
-            return type;
-        }
-
-        throw new InvalidOperationException();
-    }
+        => _directiveDefinitionLookup.TryGetValue(name, out var type)
+            ? type
+            : throw new InvalidOperationException();
 
     public string GetSchemaName(SchemaKey schemaKey)
         => SourceSchemaLookup[schemaKey.Value].Name;
 
+    public ImmutableDictionary<SchemaKey, ImmutableHashSet<string>> GetSourceUnionMembers(
+        string objectTypeName)
+    {
+        return _sourceUnions.TryGetValue(objectTypeName, out var sourceUnions)
+            ? sourceUnions
+            : ImmutableDictionary<SchemaKey, ImmutableHashSet<string>>.Empty;
+    }
+
     private void AddSpecDirectives()
     {
         var directive = CreateSkipDirective();

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

@@ -0,0 +1,57 @@
+using System.Collections.Immutable;
+using HotChocolate.Fusion.Types.Completion;
+using HotChocolate.Language;
+
+namespace HotChocolate.Fusion.Types.Directives;
+
+/*
+   directive @fusion__implements(
+     schema: fusion__Schema!
+     interface: String!
+   ) repeatable on OBJECT | INTERFACE
+*/
+internal static class ImplementsDirectiveParser
+{
+    public static ImmutableDictionary<SchemaKey, ImmutableHashSet<string>> Parse(
+        IReadOnlyList<DirectiveNode> directiveNodes)
+    {
+        Dictionary<SchemaKey, ImmutableHashSet<string>.Builder>? temp = null;
+
+        foreach (var directive in directiveNodes)
+        {
+            if (directive.Name.Value.Equals(FusionBuiltIns.Implements, StringComparison.Ordinal))
+            {
+                temp ??= new Dictionary<SchemaKey, ImmutableHashSet<string>.Builder>();
+
+                var schemaValue = directive.Arguments.FirstOrDefault(t => t.Name.Value == "schema")?.Value;
+                var interfaceValue = directive.Arguments.FirstOrDefault(t => t.Name.Value == "interface")?.Value;
+
+                if (schemaValue is not EnumValueNode { Value: { Length: > 0 } schemaName })
+                {
+                    throw new InvalidOperationException(
+                        $"The directive `@fusion__implements` has an invalid value for `schema`.\r\n{directive}");
+                }
+
+                if (interfaceValue is not StringValueNode { Value.Length: > 0 } interfaceName)
+                {
+                    throw new InvalidOperationException(
+                        $"The directive `@fusion__implements` has an invalid value for `interface`.\r\n{directive}");
+                }
+
+                var schemaKey = new SchemaKey(schemaName);
+
+                if(!temp.TryGetValue(schemaKey, out var implements))
+                {
+                    implements = ImmutableHashSet.CreateBuilder<string>();
+                    temp.Add(schemaKey, implements);
+                }
+
+                implements.Add(interfaceName.Value);
+            }
+        }
+
+        return temp is null
+            ? ImmutableDictionary<SchemaKey, ImmutableHashSet<string>>.Empty
+            : temp.ToImmutableDictionary(t => t.Key, t => t.Value.ToImmutable());
+    }
+}