From a1a471cff1cbeb50fd20fd76856a64f84600f044 Mon Sep 17 00:00:00 2001
From: Eirik Tsarpalis <eirik.tsarpalis@gmail.com>
Date: Sat, 26 Oct 2024 08:48:03 +0100
Subject: [PATCH] Fix source generation for standalone enums (#30)

---
 .../TypeDataModelGenerator.Enum.cs              | 14 +++++++++-----
 .../ModelGenerator/TypeDataModelGenerator.cs    |  2 +-
 .../CompilationTests.cs                         | 17 +++++++++++++++++
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.Enum.cs b/src/TypeShape.Roslyn/ModelGenerator/TypeDataModelGenerator.Enum.cs
index 730e04e..822665a 100644
--- 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
index f3960dc..8e137a8 100644
--- 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
index 8337896..ab1c197 100644
--- 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);
+    }
 }