Support Customizing Discriminator Types as Fixed Enums (#4766)

This PR fixes an issue where a discriminator property was customized via
custom code to a custom fixed enum type and the generator was still
generating the discriminator value using the original type.

fixes: https://github.com/microsoft/typespec/issues/4636,
https://github.com/microsoft/typespec/issues/4636

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/src/Providers/CanonicalTypeProvider.cs

@@ -66,6 +66,7 @@ namespace Microsoft.Generator.CSharp.Providers
                     && specToCustomPropertiesMap.TryGetValue(candidateSpecProperty, out var mappedProperty) && mappedProperty == customProperty)
                 {
                     specProperty = candidateSpecProperty;
+                    customProperty.IsDiscriminator = specProperty.IsDiscriminator;
                     customProperty.WireInfo = new PropertyWireInformation(specProperty);
                 }
 

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/src/Providers/ModelProvider.cs

@@ -527,23 +527,32 @@ namespace Microsoft.Generator.CSharp.Providers
                     {
                         return GetUnknownDiscriminatorExpression(discriminator);
                     }
-                    else
+
+                    if (!type.IsFrameworkType && type.IsEnum)
                     {
-                        if (!type.IsFrameworkType && type.IsEnum && _inputModel.BaseModel.DiscriminatorProperty!.Type is InputEnumType inputEnumType)
+                        if (_inputModel.BaseModel.DiscriminatorProperty!.Type is InputEnumType inputEnumType)
                         {
-                            /* TODO: when customize the discriminator type to a enum, then we may not be able to get the correct TypeProvider in this way.
-                             * We will handle this when issue https://github.com/microsoft/typespec/issues/4313 is resolved.
-                             * */
                             var discriminatorProvider = CodeModelPlugin.Instance.TypeFactory.CreateEnum(enumType: inputEnumType);
-                            var enumMember = discriminatorProvider!.EnumValues.FirstOrDefault(e => e.Value.ToString() == _inputModel.DiscriminatorValue) ?? throw new InvalidOperationException($"invalid discriminator value {_inputModel.DiscriminatorValue}");
+                            var enumMember = discriminatorProvider!.EnumValues.FirstOrDefault(e => e.Value.ToString() == _inputModel.DiscriminatorValue)
+                                ?? throw new InvalidOperationException($"invalid discriminator value {_inputModel.DiscriminatorValue}");
                             /* {KindType}.{enumMember} */
-                            return TypeReferenceExpression.FromType(type).Property(enumMember.Name);
+                            return Static(type).Property(enumMember.Name);
                         }
-                        else
+
+                        // Handle custom fixed enum discriminator
+                        if (discriminator.CustomProvider?.Value?.IsEnum == true)
                         {
-                            return Literal(_inputModel.DiscriminatorValue);
+                            var enumMember = discriminator.CustomProvider.Value.Fields
+                                .FirstOrDefault(f => f.Name.Equals(_inputModel.DiscriminatorValue, StringComparison.OrdinalIgnoreCase));
+                            if (enumMember != null)
+                            {
+                                return Static(type).Property(enumMember.Name);
+                            }
                         }
                     }
+
+                    // fallback to the default value
+                    return Literal(_inputModel.DiscriminatorValue);
                 }
             }
             return null;

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/src/Providers/NamedTypeSymbolProvider.cs

@@ -117,7 +117,10 @@ namespace Microsoft.Generator.CSharp.Providers
                     new AutoPropertyBody(propertySymbol.SetMethod is not null),
                     this)
                 {
-                    OriginalName = GetOriginalName(propertySymbol)
+                    OriginalName = GetOriginalName(propertySymbol),
+                    CustomProvider = new(() => propertySymbol.Type is INamedTypeSymbol propertyNamedTypeSymbol
+                        ? new NamedTypeSymbolProvider(propertyNamedTypeSymbol)
+                        : null)
                 };
                 properties.Add(propertyProvider);
             }

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/src/Providers/PropertyProvider.cs

@@ -2,10 +2,8 @@
 // Licensed under the MIT License.
 
 using System;
-using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
-using Microsoft.CodeAnalysis;
 using Microsoft.Generator.CSharp.Expressions;
 using Microsoft.Generator.CSharp.Input;
 using Microsoft.Generator.CSharp.Primitives;
@@ -29,7 +27,7 @@ namespace Microsoft.Generator.CSharp.Providers
         public CSharpType? ExplicitInterface { get; }
         public XmlDocProvider XmlDocs { get; private set; }
         public PropertyWireInformation? WireInfo { get; internal set; }
-        public bool IsDiscriminator { get; }
+        public bool IsDiscriminator { get; internal set; }
         public bool IsAdditionalProperties { get; init; }
 
         public FieldProvider? BackingField { get; set; }
@@ -44,6 +42,8 @@ namespace Microsoft.Generator.CSharp.Providers
 
         internal string? OriginalName { get; init; }
 
+        internal Lazy<NamedTypeSymbolProvider?>? CustomProvider { get; init; }
+
         // for mocking
 #pragma warning disable CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider declaring as nullable.
         protected PropertyProvider()

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/test/Providers/ModelProviders/DiscriminatorTests.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using Microsoft.Generator.CSharp.Expressions;
 using Microsoft.Generator.CSharp.Input;
 using Microsoft.Generator.CSharp.Primitives;
@@ -242,5 +243,110 @@ namespace Microsoft.Generator.CSharp.Tests.Providers.ModelProviders
             Assert.AreEqual(1, publicCtor!.Signature.Parameters.Count);
             Assert.AreEqual("other", publicCtor.Signature.Parameters[0].Name);
         }
+
+        [Test]
+        public async Task ModelWithCustomFixedEnumDiscriminator()
+        {
+            var plugin = await MockHelpers.LoadMockPluginAsync(
+                inputModelTypes: [_baseModel, _catModel],
+                compilation: async () => await Helpers.GetCompilationFromDirectoryAsync());
+            var baseModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Pet");
+            Assert.IsNotNull(baseModel);
+
+            var primaryBaseModelCtor = baseModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Protected));
+            Assert.IsNotNull(primaryBaseModelCtor);
+            Assert.AreEqual(1, primaryBaseModelCtor!.Signature.Parameters.Count);
+            Assert.AreEqual("kind", primaryBaseModelCtor.Signature.Parameters[0].Name);
+            Assert.AreEqual("CustomKind", primaryBaseModelCtor.Signature.Parameters[0].Type.Name);
+
+            var catModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Cat");
+            Assert.IsNotNull(catModel);
+
+            var catSerializationCtor = catModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Internal));
+            Assert.IsNotNull(catSerializationCtor);
+            Assert.AreEqual(3, catSerializationCtor!.Signature.Parameters.Count);
+
+            // ensure discriminator is present and is the custom type
+            var kindParam = catSerializationCtor!.Signature.Parameters.FirstOrDefault(p => p.Name == "kind");
+            Assert.IsNotNull(kindParam);
+            Assert.AreEqual("CustomKind", kindParam!.Type.Name);
+
+            // the primary ctor should call the base ctor using the custom discriminator
+            var publicCtor = catModel.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Public));
+            Assert.IsNotNull(publicCtor);
+            Assert.IsTrue(publicCtor!.Signature.Parameters.Any(p => p.Name != "kind"));
+
+            var init = publicCtor!.Signature.Initializer;
+            Assert.AreEqual(1, init!.Arguments.Count);
+            Assert.IsTrue(init.Arguments[0].ToDisplayString().Contains("CustomKind.Cat"));
+        }
+
+        [Test]
+        public async Task ModelWithCustomExtensibleEnumDiscriminator()
+        {
+            var plugin = await MockHelpers.LoadMockPluginAsync(
+                inputModelTypes: [_baseModel, _catModel],
+                compilation: async () => await Helpers.GetCompilationFromDirectoryAsync());
+            var baseModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Pet");
+            Assert.IsNotNull(baseModel);
+
+            var primaryBaseModelCtor = baseModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Protected));
+            Assert.IsNotNull(primaryBaseModelCtor);
+            Assert.AreEqual(1, primaryBaseModelCtor!.Signature.Parameters.Count);
+            Assert.AreEqual("kind", primaryBaseModelCtor.Signature.Parameters[0].Name);
+            Assert.AreEqual("CustomKind", primaryBaseModelCtor.Signature.Parameters[0].Type.Name);
+
+            var catModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Cat");
+            Assert.IsNotNull(catModel);
+
+            var catSerializationCtor = catModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Internal));
+            Assert.IsNotNull(catSerializationCtor);
+            Assert.AreEqual(3, catSerializationCtor!.Signature.Parameters.Count);
+
+            var kindParam = catSerializationCtor!.Signature.Parameters.FirstOrDefault(p => p.Name == "kind");
+            Assert.IsNotNull(kindParam);
+            Assert.AreEqual("CustomKind", kindParam!.Type.Name);
+
+            // the primary ctor should call the base ctor using the custom discriminator literal value, as there
+            // is an implicit conversion from string to CustomKind
+            var publicCtor = catModel.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Public));
+            Assert.IsNotNull(publicCtor);
+            Assert.IsTrue(publicCtor!.Signature.Parameters.Any(p => p.Name != "kind"));
+
+
+            var init = publicCtor!.Signature.Initializer;
+            Assert.AreEqual(1, init!.Arguments.Count);
+            Assert.AreEqual("\"cat\"", init.Arguments[0].ToDisplayString());
+        }
+
+        [Test]
+        public async Task CanCustomizeDiscriminator()
+        {
+            var plugin = await MockHelpers.LoadMockPluginAsync(
+                inputModelTypes: [_baseModel, _catModel],
+                compilation: async () => await Helpers.GetCompilationFromDirectoryAsync());
+
+            var baseModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Pet");
+            Assert.IsNotNull(baseModel);
+
+            var baseModelPrimaryCtor = baseModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Protected));
+            Assert.IsNotNull(baseModelPrimaryCtor);
+            Assert.AreEqual(1, baseModelPrimaryCtor!.Signature.Parameters.Count);
+
+            // the custom property should be marked as the discriminator
+            Assert.AreEqual("customName", baseModelPrimaryCtor.Signature.Parameters[0].Name);
+            Assert.IsTrue(baseModelPrimaryCtor.Signature.Parameters[0].Property?.IsDiscriminator);
+
+            var catModel = plugin.Object.OutputLibrary.TypeProviders.OfType<ModelProvider>().FirstOrDefault(t => t.Name == "Cat");
+            Assert.IsNotNull(catModel);
+
+            var catSerializationCtor = catModel!.Constructors.FirstOrDefault(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Internal));
+            Assert.IsNotNull(catSerializationCtor);
+            Assert.AreEqual(3, catSerializationCtor!.Signature.Parameters.Count);
+
+            var discriminatorParam = catSerializationCtor!.Signature.Parameters.FirstOrDefault(p => p.Name == "customName");
+            Assert.IsNotNull(discriminatorParam);
+            Assert.IsTrue(discriminatorParam!.Property?.IsDiscriminator);
+        }
     }
 }

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/test/Providers/ModelProviders/TestData/DiscriminatorTests/CanCustomizeDiscriminator/Pet.cs

@@ -0,0 +1,13 @@
+#nullable disable
+
+using System;
+using Microsoft.Generator.CSharp.Customization;
+
+namespace Sample.Models;
+
+public partial class Pet
+{
+    // CUSTOM: Changed type from string to CustomKind.
+    [CodeGenMember("Kind")]
+    internal string CustomName { get; set; }
+}

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/test/Providers/ModelProviders/TestData/DiscriminatorTests/ModelWithCustomExtensibleEnumDiscriminator/Pet.cs

@@ -0,0 +1,43 @@
+#nullable disable
+
+using System;
+using System.ComponentModel;
+using Microsoft.Generator.CSharp.Customization;
+
+namespace Sample.Models;
+
+public partial class Pet
+{
+    // CUSTOM: Changed type from string to CustomKind.
+    [CodeGenMember("Kind")]
+    internal CustomKind Kind { get; set; }
+
+    public readonly partial struct CustomKind : IEquatable<CustomKind>
+    {
+        private readonly string _value;
+        private const string DogValue = "Dog";
+        private const string CatValue = "Cat";
+
+        public CustomKind(string value)
+        {
+            Argument.AssertNotNull(value, nameof(value));
+
+            _value = value;
+        }
+
+        public static CustomKind Dog { get; } = new CustomKind(DogValue);
+        public static CustomKind Cat { get; } = new CustomKind(CatValue);
+
+        public static bool operator ==(CustomKind left, CustomKind right) => left.Equals(right);
+        public static bool operator !=(CustomKind left, CustomKind right) => !left.Equals(right);
+
+        public static implicit operator CustomKind(string value) => new CustomKind(value);
+
+        public override bool Equals(object obj) => obj is CustomKind other && Equals(other);
+
+        public bool Equals(CustomKind other) => string.Equals(_value, other._value, StringComparison.InvariantCultureIgnoreCase);
+
+        public override int GetHashCode() => _value != null ? StringComparer.InvariantCultureIgnoreCase.GetHashCode(_value) : 0;
+        public override string ToString() => _value;
+    }
+}

packages/http-client-csharp/generator/Microsoft.Generator.CSharp/test/Providers/ModelProviders/TestData/DiscriminatorTests/ModelWithCustomFixedEnumDiscriminator/Pet.cs

@@ -0,0 +1,19 @@
+#nullable disable
+
+using System;
+using Microsoft.Generator.CSharp.Customization;
+
+namespace Sample.Models;
+
+public partial class Pet
+{
+    // CUSTOM: Changed type from string to CustomKind.
+    [CodeGenMember("Kind")]
+    internal CustomKind Kind { get; set; }
+
+    public enum CustomKind
+    {
+        Cat,
+        Dog
+    }
+}