Fix source generation for standalone enums (#30)

eiriktsarpalis · Oct 26, 2024 · a1a471c · a1a471c
1 parent ce639f0
commit a1a471c
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.Enum.cs b/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.Enum.cs
@@ -1,10 +1,11 @@
-using Microsoft.CodeAnalysis;
+using System.Diagnostics;
+using Microsoft.CodeAnalysis;
 
 namespace TypeShape.Roslyn;
 
 public partial class TypeDataModelGenerator
 {
-    private static bool TryMapEnum(ITypeSymbol type, out TypeDataModel? model, out TypeDataModelGenerationStatus status)
+    private bool TryMapEnum(ITypeSymbol type, ref TypeDataModelGenerationContext ctx, out TypeDataModel? model, out TypeDataModelGenerationStatus status)
     {
         if (type.TypeKind is not TypeKind.Enum)
         {
@@ -13,13 +14,16 @@ private static bool TryMapEnum(ITypeSymbol type, out TypeDataModel? model, out T
             return false;
         }
 
+        INamedTypeSymbol underlyingType = ((INamedTypeSymbol)type).EnumUnderlyingType!;
+        status = IncludeNestedType(underlyingType, ref ctx);
+        Debug.Assert(status is TypeDataModelGenerationStatus.Success);
+
         model = new EnumDataModel
         {
             Type = type,
-            UnderlyingType = ((INamedTypeSymbol)type).EnumUnderlyingType!,
+            UnderlyingType = underlyingType,
         };
-
-        status = TypeDataModelGenerationStatus.Success;
+
         return true;
     }
 }
diff --git a/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.cs b/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.cs
@@ -160,7 +160,7 @@ protected TypeDataModelGenerationStatus IncludeNestedType(ITypeSymbol type, ref
     /// </remarks>
     protected virtual TypeDataModelGenerationStatus MapType(ITypeSymbol type, ref TypeDataModelGenerationContext ctx, out TypeDataModel? model)
     {
-        if (TryMapEnum(type, out model, out TypeDataModelGenerationStatus status))
+        if (TryMapEnum(type, ref ctx, out model, out TypeDataModelGenerationStatus status))
         {
             return status;
         }

diff --git a/tests/TypeShape.SourceGenerator.UnitTests/CompilationTests.cs b/tests/TypeShape.SourceGenerator.UnitTests/CompilationTests.cs
@@ -208,4 +208,21 @@ public partial class MyWitness;
         TypeShapeSourceGeneratorResult result = CompilationHelpers.RunTypeShapeSourceGenerator(compilation);
         Assert.Empty(result.Diagnostics);
     }
+
+    [Fact]
+    public static void EnumGeneration_NoErrors()
+    {
+        // Regression test for https://github.com/eiriktsarpalis/typeshape-csharp/issues/29
+        Compilation compilation = CompilationHelpers.CreateCompilation("""
+           using TypeShape;
+           
+           enum MyEnum { A, B, C }
+
+           [GenerateShape<MyEnum>]
+           partial class Witness { }
+           """);
+
+        TypeShapeSourceGeneratorResult result = CompilationHelpers.RunTypeShapeSourceGenerator(compilation);
+        Assert.Empty(result.Diagnostics);
+    }
 }