Extend associated type support to all type shape kinds. (#251)

eiriktsarpalis · web-flow · commit ecc6c98fe1dd · 2025-09-04T17:06:01.000+03:00
diff --git a/src/PolyType.Roslyn/ModelGenerator/TypeDataModelGenerator.Object.cs b/src/PolyType.Roslyn/ModelGenerator/TypeDataModelGenerator.Object.cs
@@ -121,7 +121,13 @@ type.TypeKind is not (TypeKind.Struct or TypeKind.Class or TypeKind.Interface) |
         return true;
     }
 
-    private void IncludeAssociatedShapes(ITypeSymbol type, ImmutableArray<AssociatedTypeModel> associatedTypes, ref TypeDataModelGenerationContext ctx)
+    /// <summary>
+    /// Includes the specified associated types to the type graph traversal.
+    /// </summary>
+    /// <param name="type">The type from which the associated types originate.</param>
+    /// <param name="associatedTypes">The associated types to include.</param>
+    /// <param name="ctx">The type graph traversal context.</param>
+    protected void IncludeAssociatedShapes(ITypeSymbol type, ImmutableArray<AssociatedTypeModel> associatedTypes, ref TypeDataModelGenerationContext ctx)
     {
         _associatedTypes = _associatedTypes.SetItem(
             type,
diff --git a/src/PolyType.SourceGenerator/Parser/Parser.Diagnostics.cs b/src/PolyType.SourceGenerator/Parser/Parser.Diagnostics.cs
@@ -119,14 +119,6 @@ public sealed partial class Parser
         defaultSeverity: DiagnosticSeverity.Error,
         isEnabledByDefault: true);
 
-    private static DiagnosticDescriptor AssociatedTypeInExternalAssembly { get; } = new DiagnosticDescriptor(
-        id: "PT0017",
-        title: "Associated type must be declared in same assembly.",
-        messageFormat: "The associated type '{0}' is not declared in the same assembly as the attribute annotation specifying it.",
-        category: "PolyType.SourceGenerator",
-        defaultSeverity: DiagnosticSeverity.Error,
-        isEnabledByDefault: true);
-
     internal static DiagnosticDescriptor ConflictingMarshalers { get; } = new DiagnosticDescriptor(
         id: "PT0018",
         title: "Multiple marshalers specified.",
diff --git a/src/PolyType.SourceGenerator/Parser/Parser.ModelMapper.cs b/src/PolyType.SourceGenerator/Parser/Parser.ModelMapper.cs
@@ -278,12 +278,6 @@ private ImmutableEquatableSet<AssociatedTypeId> CollectAssociatedTypes(TypeDataM
         ITypeSymbol[] typeArgs = namedType?.GetRecursiveTypeArguments() ?? [];
         foreach ((INamedTypeSymbol openType, IAssemblySymbol associatedAssembly, Location? location, TypeShapeRequirements requirements) in associatedTypeSymbols)
         {
-            if (!SymbolEqualityComparer.Default.Equals(openType.ContainingAssembly, associatedAssembly))
-            {
-                ReportDiagnostic(AssociatedTypeInExternalAssembly, location, openType.GetFullyQualifiedName());
-                continue;
-            }
-
             if (!IsAccessibleSymbol(openType))
             {
                 // Skip types that are not accessible in the current scope
diff --git a/src/PolyType.SourceGenerator/Parser/Parser.cs b/src/PolyType.SourceGenerator/Parser/Parser.cs
@@ -685,9 +685,12 @@ private TypeDataModelGenerationStatus MapSurrogateType(ITypeSymbol type, ITypeSy
                 MarshalerType = namedMarshaler,
                 Methods = MapMethods(type, ref ctx, methodFlags),
                 Events = MapEvents(type, ref ctx, methodFlags),
+                AssociatedTypes = associatedTypes,
             };
         }
 
+        // Include any associated shapes.
+        IncludeAssociatedShapes(type, associatedTypes, ref ctx);
         return status;
 
         TypeDataModelGenerationStatus ReportInvalidMarshalerAndExit()
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Enum.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Enum.cs
@@ -13,13 +13,14 @@ private void FormatEnumTypeShapeFactory(SourceWriter writer, string methodName,
                 return new global::PolyType.SourceGenModel.SourceGenEnumTypeShape<{{enumShapeType.Type.FullyQualifiedName}}, {{enumShapeType.UnderlyingType.FullyQualifiedName}}>
                 {
                     UnderlyingType = {{GetShapeModel(enumShapeType.UnderlyingType).SourceIdentifier}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(enumShapeType)}},
                     Provider = this,
                     Members = new global::System.Collections.Generic.Dictionary<string, {{enumShapeType.UnderlyingType.FullyQualifiedName}}>({{enumShapeType.Members.Count}}, global::System.StringComparer.Ordinal)
                     {
                         {{string.Join("\n            ", enumShapeType.Members.Select(member => $"[\"{member.Key}\"] = {member.Value},"))}}
                     }
                 };
             }
-            """);
+            """, trimNullAssignmentLines: true);
     }
 }
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.FSharpUnion.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.FSharpUnion.cs
@@ -21,6 +21,7 @@ private void FormatFSharpUnionTypeShapeFactory(SourceWriter writer, string metho
                     GetUnionCaseIndexFunc = {{FormatFSharpUnionTagReader(unionShapeModel)}},
                     CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},
                     CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(unionShapeModel)}},
                     Provider = this,
                 };
             }
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Function.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Function.cs
@@ -26,6 +26,7 @@ private void FormatFunctionTypeShapeFactory(SourceWriter writer, string methodNa
                     FunctionInvoker = {{FormatFunctionInvoker(functionShapeModel, functionArgumentStateFQN)}},
                     FromDelegateFunc = {{FormatNull(FormatFromDelegateFunc(functionShapeModel, functionArgumentStateFQN, requiredParametersMaskFieldName, requireAsync: false))}},
                     FromAsyncDelegateFunc = {{FormatNull(FormatFromDelegateFunc(functionShapeModel, functionArgumentStateFQN, requiredParametersMaskFieldName, requireAsync: true))}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(functionShapeModel)}},
                     Provider = this,
                 };
             }
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Optional.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Optional.cs
@@ -23,10 +23,11 @@ private void FormatOptionalTypeShapeFactory(SourceWriter writer, string methodNa
                     Deconstructor = {{FormatDeconstructor()}},
                     CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},
                     CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(optionalShapeModel)}},
                     Provider = this,
                 };
             }
-            """);
+            """, trimNullAssignmentLines: true);
         writer.WriteLine("#pragma warning restore CS8767", disableIndentation: true);
 
         if (methodFactoryMethodName is not null)
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Surrogate.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Surrogate.cs
@@ -15,14 +15,15 @@ private void FormatSurrogateTypeShapeFactory(SourceWriter writer, string methodN
             {
                 return new global::PolyType.SourceGenModel.SourceGenSurrogateTypeShape<{{surrogateShapeModel.Type.FullyQualifiedName}}, {{surrogateShapeModel.SurrogateType.FullyQualifiedName}}>
                 {
-                    Marshaler = new {{surrogateShapeModel.MarshalerType.FullyQualifiedName}}(),
+                    Marshaler = new {{surrogateShapeModel.MarshalerType.FullyQualifiedName}}()!,
                     SurrogateType = {{GetShapeModel(surrogateShapeModel.SurrogateType).SourceIdentifier}},
                     CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},
                     CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(surrogateShapeModel)}},
                     Provider = this,
                 };
             }
-            """);
+            """, trimNullAssignmentLines: true);
 
         if (methodFactoryMethodName is not null)
         {
diff --git a/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Union.cs b/src/PolyType.SourceGenerator/SourceFormatter/SourceFormatter.Union.cs
@@ -23,10 +23,11 @@ private void FormatUnionTypeShapeFactory(SourceWriter writer, string methodName,
                     GetUnionCaseIndexFunc = {{getUnionCaseIndexMethod}},
                     CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},
                     CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},
+                    AssociatedTypeShapes = {{FormatAssociatedTypeShapes(unionShapeModel)}},
                     Provider = this,
                 };
             }
-            """);
+            """, trimNullAssignmentLines: true);
 
         writer.WriteLine();
         FormatUnionCasesFactory(writer, unionShapeModel, createUnionCasesMethodName);
diff --git a/tests/PolyType.Tests/AssociatedTypesTests.cs b/tests/PolyType.Tests/AssociatedTypesTests.cs
@@ -1,7 +1,5 @@
 ﻿using PolyType;
 using PolyType.Tests;
-using System.ComponentModel;
-using Xunit.Internal;
 
 [assembly: TypeShapeExtension(typeof(AssociatedTypesTests.GenericDataType<,>), Requirements = TypeShapeRequirements.Constructor, AssociatedTypes = [typeof(AssociatedTypesTests.GenericDataTypeVerifier<,>)])]
 [assembly: TypeShapeExtension(typeof(AssociatedTypesTests.GenericDataType<,>), Requirements = TypeShapeRequirements.Full, AssociatedTypes = [typeof(AssociatedTypesTests.ExtraShape<,>)])]
@@ -15,6 +13,8 @@
 [assembly: TypeShapeExtension(typeof(IEnumerable<>), AssociatedTypes = [typeof(AssociatedTypesTests.GenericWrapper<>)])]
 [assembly: TypeShapeExtension(typeof(Dictionary<,>), AssociatedTypes = [typeof(AssociatedTypesTests.GenericWrapper<>.GenericNested<>)])]
 
+[assembly: TypeShapeExtension(typeof(AssociatedTypesTests.GenericClassWithMarshallerAndAssociatedTypes<>), Marshaler = typeof(AssociatedTypesTests.GenericMarshaler<>), AssociatedTypes = [typeof(List<>)])]
+
 namespace PolyType.Tests;
 
 [Trait("AssociatedTypes", "true")]
@@ -277,6 +277,16 @@ public void Dictionaries_HaveAssociatedTypes()
         Assert.Equal(typeof(GenericWrapper<SpecialKey>.GenericNested<string>), associatedShape.Type);
     }
 
+    [Fact]
+    public void GenericClassWithMarshallerAndAssociatedTypes_ReturnsExpectedShape()
+    {
+        // Regression test for https://github.com/eiriktsarpalis/PolyType/issues/239
+        ITypeShape? typeShape = providerUnderTest.Provider.GetShape(typeof(GenericClassWithMarshallerAndAssociatedTypes<(Guid, bool)>));
+        Assert.IsType<ISurrogateTypeShape>(typeShape, exactMatch: false);
+        ITypeShape? associatedShape = typeShape.GetAssociatedTypeShape(typeof(List<>));
+        Assert.IsType<ITypeShape<List<(Guid, bool)>>>(associatedShape, exactMatch: false);
+    }
+
     private static Func<object>? GetAssociatedTypeFactory(ITypeShape typeShape, Type associatedType)
         => (typeShape.GetAssociatedTypeShape(associatedType) as IObjectTypeShape)?.GetDefaultConstructor();
 
@@ -294,6 +304,14 @@ private void AssertPartialShape(Func<object?> partialThrows)
         }
     }
 
+    public class GenericClassWithMarshallerAndAssociatedTypes<T1>;
+
+    public class GenericMarshaler<T1> : IMarshaler<GenericClassWithMarshallerAndAssociatedTypes<T1>, object?>
+    {
+        public object? Marshal(GenericClassWithMarshallerAndAssociatedTypes<T1>? _) => null;
+        public GenericClassWithMarshallerAndAssociatedTypes<T1>? Unmarshal(object? _) => new();
+    }
+
     [GenerateShape]
     [AssociatedTypeShape(typeof(NonGenericDataTypeConverter))]
     public partial class NonGenericDataType;
@@ -353,6 +371,7 @@ public class GenericConverter<T>;
     [GenerateShapeFor<GenericDataTypeFullAndPartialPaths<int, string>>]
     [GenerateShapeFor<IEnumerable<SpecialKey>>]
     [GenerateShapeFor<Dictionary<SpecialKey, string>>]
+    [GenerateShapeFor<GenericClassWithMarshallerAndAssociatedTypes<(Guid, bool)>>]
     internal partial class Witness;
 
     [GenerateShape]

Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,13 @@ type.TypeKind is not (TypeKind.Struct or TypeKind.Class or TypeKind.Interface) \|`
`121`	`121`	`return true;`
`122`	`122`	`}`
`123`	`123`
`124`		`- private void IncludeAssociatedShapes(ITypeSymbol type, ImmutableArray<AssociatedTypeModel> associatedTypes, ref TypeDataModelGenerationContext ctx)`
	`124`	`+ /// <summary>`
	`125`	`+ /// Includes the specified associated types to the type graph traversal.`
	`126`	`+ /// </summary>`
	`127`	`+ /// <param name="type">The type from which the associated types originate.</param>`
	`128`	`+ /// <param name="associatedTypes">The associated types to include.</param>`
	`129`	`+ /// <param name="ctx">The type graph traversal context.</param>`
	`130`	`+ protected void IncludeAssociatedShapes(ITypeSymbol type, ImmutableArray<AssociatedTypeModel> associatedTypes, ref TypeDataModelGenerationContext ctx)`
`125`	`131`	`{`
`126`	`132`	`_associatedTypes = _associatedTypes.SetItem(`
`127`	`133`	`type,`
Original file line number	Diff line number	Diff line change
`@@ -278,12 +278,6 @@ private ImmutableEquatableSet<AssociatedTypeId> CollectAssociatedTypes(TypeDataM`
`278`	`278`	`ITypeSymbol[] typeArgs = namedType?.GetRecursiveTypeArguments() ?? [];`
`279`	`279`	`foreach ((INamedTypeSymbol openType, IAssemblySymbol associatedAssembly, Location? location, TypeShapeRequirements requirements) in associatedTypeSymbols)`
`280`	`280`	`{`
`281`		`- if (!SymbolEqualityComparer.Default.Equals(openType.ContainingAssembly, associatedAssembly))`
`282`		`- {`
`283`		`- ReportDiagnostic(AssociatedTypeInExternalAssembly, location, openType.GetFullyQualifiedName());`
`284`		`- continue;`
`285`		`- }`
`286`		`-`
`287`	`281`	`if (!IsAccessibleSymbol(openType))`
`288`	`282`	`{`
`289`	`283`	`// Skip types that are not accessible in the current scope`
Original file line number	Diff line number	Diff line change
`@@ -13,13 +13,14 @@ private void FormatEnumTypeShapeFactory(SourceWriter writer, string methodName,`
`13`	`13`	`return new global::PolyType.SourceGenModel.SourceGenEnumTypeShape<{{enumShapeType.Type.FullyQualifiedName}}, {{enumShapeType.UnderlyingType.FullyQualifiedName}}>`
`14`	`14`	`{`
`15`	`15`	`UnderlyingType = {{GetShapeModel(enumShapeType.UnderlyingType).SourceIdentifier}},`
	`16`	`+ AssociatedTypeShapes = {{FormatAssociatedTypeShapes(enumShapeType)}},`
`16`	`17`	`Provider = this,`
`17`	`18`	`Members = new global::System.Collections.Generic.Dictionary<string, {{enumShapeType.UnderlyingType.FullyQualifiedName}}>({{enumShapeType.Members.Count}}, global::System.StringComparer.Ordinal)`
`18`	`19`	`{`
`19`	`20`	`{{string.Join("\n ", enumShapeType.Members.Select(member => $"[\"{member.Key}\"] = {member.Value},"))}}`
`20`	`21`	`}`
`21`	`22`	`};`
`22`	`23`	`}`
`23`		`- """);`
	`24`	`+ """, trimNullAssignmentLines: true);`
`24`	`25`	`}`
`25`	`26`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ private void FormatFSharpUnionTypeShapeFactory(SourceWriter writer, string metho`
`21`	`21`	`GetUnionCaseIndexFunc = {{FormatFSharpUnionTagReader(unionShapeModel)}},`
`22`	`22`	`CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},`
`23`	`23`	`CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},`
	`24`	`+ AssociatedTypeShapes = {{FormatAssociatedTypeShapes(unionShapeModel)}},`
`24`	`25`	`Provider = this,`
`25`	`26`	`};`
`26`	`27`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@ private void FormatFunctionTypeShapeFactory(SourceWriter writer, string methodNa`
`26`	`26`	`FunctionInvoker = {{FormatFunctionInvoker(functionShapeModel, functionArgumentStateFQN)}},`
`27`	`27`	`FromDelegateFunc = {{FormatNull(FormatFromDelegateFunc(functionShapeModel, functionArgumentStateFQN, requiredParametersMaskFieldName, requireAsync: false))}},`
`28`	`28`	`FromAsyncDelegateFunc = {{FormatNull(FormatFromDelegateFunc(functionShapeModel, functionArgumentStateFQN, requiredParametersMaskFieldName, requireAsync: true))}},`
	`29`	`+ AssociatedTypeShapes = {{FormatAssociatedTypeShapes(functionShapeModel)}},`
`29`	`30`	`Provider = this,`
`30`	`31`	`};`
`31`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,14 +15,15 @@ private void FormatSurrogateTypeShapeFactory(SourceWriter writer, string methodN`
`15`	`15`	`{`
`16`	`16`	`return new global::PolyType.SourceGenModel.SourceGenSurrogateTypeShape<{{surrogateShapeModel.Type.FullyQualifiedName}}, {{surrogateShapeModel.SurrogateType.FullyQualifiedName}}>`
`17`	`17`	`{`
`18`		`- Marshaler = new {{surrogateShapeModel.MarshalerType.FullyQualifiedName}}(),`
	`18`	`+ Marshaler = new {{surrogateShapeModel.MarshalerType.FullyQualifiedName}}()!,`
`19`	`19`	`SurrogateType = {{GetShapeModel(surrogateShapeModel.SurrogateType).SourceIdentifier}},`
`20`	`20`	`CreateMethodsFunc = {{FormatNull(methodFactoryMethodName)}},`
`21`	`21`	`CreateEventsFunc = {{FormatNull(eventFactoryMethodName)}},`
	`22`	`+ AssociatedTypeShapes = {{FormatAssociatedTypeShapes(surrogateShapeModel)}},`
`22`	`23`	`Provider = this,`
`23`	`24`	`};`
`24`	`25`	`}`
`25`		`- """);`
	`26`	`+ """, trimNullAssignmentLines: true);`
`26`	`27`
`27`	`28`	`if (methodFactoryMethodName is not null)`
`28`	`29`	`{`