Remove unexpected methods from ModelFactory for management (#51205)

* Add ModelFacotryVisitor to omit methods for mgmt

* fix for safe flatten

* Update modelFactory for safe-flatten

* Address comments

Author: live1206

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/ManagementClientGenerator.cs

@@ -58,6 +58,7 @@ protected override void Configure()
             AddVisitor(new SerializationVisitor());
             AddVisitor(new PaginationVisitor());
             AddVisitor(new SafeFlattenVisitor());
+            AddVisitor(new ModelFactoryVisitor());
         }
     }
 }

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/ManagementOutputLibrary.cs

@@ -18,7 +18,8 @@ namespace Azure.Generator.Management
 {
     /// <inheritdoc/>
     public class ManagementOutputLibrary : AzureOutputLibrary
-    {   private ManagementLongRunningOperationProvider? _armOperation;
+    {
+        private ManagementLongRunningOperationProvider? _armOperation;
         internal ManagementLongRunningOperationProvider ArmOperation => _armOperation ??= new ManagementLongRunningOperationProvider(false);
 
         private ManagementLongRunningOperationProvider? _genericArmOperation;
@@ -28,6 +29,55 @@ public class ManagementOutputLibrary : AzureOutputLibrary
         private HashSet<CSharpType>? _resourceTypes;
         private HashSet<CSharpType> ResourceTypes => _resourceTypes ??= BuildResourceModels();
 
+        // TODO: replace this with CSharpType to TypeProvider mapping and move this logic to ModelFactoryVisitor
+        private HashSet<CSharpType>? _modelFactoryModels;
+        private HashSet<CSharpType> ModelFactoryModels => _modelFactoryModels ??= BuildModelFactoryModels();
+        internal HashSet<CSharpType> BuildModelFactoryModels()
+        {
+            var result = new HashSet<CSharpType>();
+            foreach (var inputModel in ManagementClientGenerator.Instance.InputLibrary.InputNamespace.Models)
+            {
+                var model = ManagementClientGenerator.Instance.TypeFactory.CreateModel(inputModel);
+                if (model is not null && IsModelFactoryModel(model))
+                {
+                    result.Add(model.Type);
+                }
+            }
+            return result;
+        }
+
+        private static bool IsModelFactoryModel(ModelProvider model)
+        {
+            // A model is a model factory model if it is public and it has at least one public property without a setter.
+            return model.DeclarationModifiers.HasFlag(TypeSignatureModifiers.Public) && EnumerateAllPublicProperties(model).Any(prop => !prop.Body.HasSetter);
+
+            IEnumerable<PropertyProvider> EnumerateAllPublicProperties(ModelProvider current)
+            {
+                var currentModel = current;
+                foreach (var property in currentModel.Properties)
+                {
+                    if (property.Modifiers.HasFlag(MethodSignatureModifiers.Public))
+                    {
+                        yield return property;
+                    }
+                }
+
+                while (currentModel.BaseModelProvider is not null)
+                {
+                    currentModel = currentModel.BaseModelProvider;
+                    foreach (var property in currentModel.Properties)
+                    {
+                        if (property.Modifiers.HasFlag(MethodSignatureModifiers.Public))
+                        {
+                            yield return property;
+                        }
+                    }
+                }
+            }
+        }
+
+        internal bool IsModelFactoryModelType(CSharpType type) => ModelFactoryModels.Contains(type);
+
         private HashSet<CSharpType> BuildResourceModels()
         {
             var resourceTypes = new HashSet<CSharpType>();

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Visitors/ModelFactoryVisitor.cs

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Microsoft.TypeSpec.Generator.ClientModel;
+using Microsoft.TypeSpec.Generator.Providers;
+using System.Collections.Generic;
+
+namespace Azure.Generator.Management.Visitors
+{
+    internal class ModelFactoryVisitor : ScmLibraryVisitor
+    {
+        protected override TypeProvider? VisitType(TypeProvider type)
+        {
+            if (type is ModelFactoryProvider modelFactory)
+            {
+                var updatedMethods = new List<MethodProvider>();
+                foreach (var method in modelFactory.Methods)
+                {
+                    var returnType = method.Signature.ReturnType;
+                    if (returnType is not null && ManagementClientGenerator.Instance.OutputLibrary.IsModelFactoryModelType(returnType))
+                    {
+                        updatedMethods.Add(method);
+                    }
+                }
+                modelFactory.Update(methods: updatedMethods);
+                return modelFactory;
+            }
+            return base.VisitType(type);
+        }
+    }
+}

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Visitors/SafeFlattenVisitor.cs

@@ -7,6 +7,7 @@
 using Microsoft.TypeSpec.Generator.Primitives;
 using Microsoft.TypeSpec.Generator.Providers;
 using Microsoft.TypeSpec.Generator.Statements;
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
@@ -15,8 +16,8 @@ namespace Azure.Generator.Management.Visitors
 {
     internal class SafeFlattenVisitor : ScmLibraryVisitor
     {
-        private Dictionary<ModelProvider, (HashSet<PropertyProvider> FlattenedProperties, HashSet<PropertyProvider> InternalizedProperties)> _flattenedModels = new();
-        private HashSet<CSharpType> _flattenedTypes = new();
+        private readonly Dictionary<ModelProvider, (HashSet<PropertyProvider> FlattenedProperties, HashSet<PropertyProvider> InternalizedProperties)> _flattenedModels = new();
+        private readonly Dictionary<CSharpType, Dictionary<CSharpType, PropertyProvider>> _flattenedModelTypes = new();
 
         // TODO: we can't check property count of property types in VisitType since we don't have TypeProvider from CSharpType.
         // Once we have CSharpType to TypeProvider mapping, we can remove this and have this logic in VisitType instead
@@ -27,6 +28,7 @@ internal class SafeFlattenVisitor : ScmLibraryVisitor
                 return null;
             }
 
+            var flattenedPropertyMap = new Dictionary<CSharpType, PropertyProvider>();
             var flattenedProperties = new HashSet<PropertyProvider>();
             var internalizedProperties = new HashSet<PropertyProvider>();
             foreach (var property in model.Properties)
@@ -43,7 +45,6 @@ internal class SafeFlattenVisitor : ScmLibraryVisitor
                         // make the current property internal
                         var internalSingleProperty = type!.Properties.Single(p => p.Type.AreNamesEqual(propertyTypeProvider.Type)); // type equal not working here, so we use AreNamesEqual
                         internalizedProperties.Add(internalSingleProperty);
-                        _flattenedTypes.Add(propertyTypeProvider.Type!);
 
                         // flatten the single property to public and associate it with the internal property
                         var flattenPropertyName = $"{singleProperty.Name}"; // TODO: handle name conflicts
@@ -60,19 +61,99 @@ internal class SafeFlattenVisitor : ScmLibraryVisitor
                             });
                         var flattenedProperty = new PropertyProvider(singleProperty.Description, singleProperty.Modifiers, singleProperty.Type, flattenPropertyName, flattenPropertyBody, type, singleProperty.ExplicitInterface, singleProperty.WireInfo, singleProperty.Attributes);
                         flattenedProperties.Add(flattenedProperty);
+                        flattenedPropertyMap.Add(internalSingleProperty.Type, flattenedProperty);
                     }
                 }
             }
             if (flattenedProperties.Count > 0)
             {
                 _flattenedModels[type!] = (flattenedProperties, internalizedProperties);
+                _flattenedModelTypes.Add(type.Type, flattenedPropertyMap);
             }
             return base.PreVisitModel(model, type);
         }
 
         protected override TypeProvider? VisitType(TypeProvider type)
         {
-            if (type is ModelProvider model && _flattenedModels.TryGetValue(model, out var value))
+            if (type is ModelProvider model)
+            {
+                UpdateModel(type, model);
+            }
+
+            if (type is ModelFactoryProvider modelFactory)
+            {
+                // If this is a model factory, we need to ensure that the flattened properties are also included in the factory methods.
+                UpdateModelFactory(modelFactory);
+            }
+
+            return base.VisitType(type);
+        }
+
+        private void UpdateModelFactory(ModelFactoryProvider modelFactory)
+        {
+            foreach (var method in modelFactory.Methods)
+            {
+                var returnType = method.Signature.ReturnType;
+                if (returnType is not null && _flattenedModelTypes.TryGetValue(returnType, out var value))
+                {
+                    UpdateModelFactoryMethod(method, returnType, value);
+                }
+            }
+        }
+
+        private void UpdateModelFactoryMethod(MethodProvider method, CSharpType returnType, Dictionary<CSharpType, PropertyProvider> propertyMap)
+        {
+            var updatedParameters = new List<ParameterProvider>(method.Signature.Parameters.Count);
+            var updated = false;
+            foreach (var parameter in method.Signature.Parameters)
+            {
+                if (propertyMap.TryGetValue(parameter.Type, out var flattenedProperty))
+                {
+                    updated = true;
+                    var updatedParameter = flattenedProperty.AsParameter;
+                    if (flattenedProperty.Type.IsNullable)
+                    {
+                        updatedParameter.DefaultValue = Default; // Ensure that the default value is set to null for nullable types
+                    }
+                    updatedParameters.Add(updatedParameter);
+                }
+                else
+                {
+                    updatedParameters.Add(parameter);
+                }
+            }
+            if (updated)
+            {
+                method.Signature.Update(parameters: updatedParameters);
+
+                // The model factory method return a new instance of the model type, update the constructor arguments with the flattened properties of internalized properties.
+                var instanceExpression = (method.BodyStatements?.OfType<ExpressionStatement>().Single()?.Expression as KeywordExpression)?.Expression as NewInstanceExpression;
+                var updatedInstanceParamters = new List<ValueExpression>(instanceExpression!.Parameters.Count);
+                foreach (var instanceParemter in instanceExpression!.Parameters)
+                {
+                    // If the instance parameter is a variable expression, we can check if it is a flattened property and update it accordingly.
+                    if (instanceParemter is VariableExpression variable && propertyMap.TryGetValue(variable.Type, out var flattenedProperty))
+                    {
+                        updatedInstanceParamters.Add(
+                            new TernaryConditionalExpression(
+                                flattenedProperty.AsParameter.Is(Null),
+                                Default,
+                                New.Instance(
+                                    variable.Type,
+                                    [flattenedProperty.AsParameter, New.Instance(new CSharpType(typeof(Dictionary<,>), typeof(string), typeof(BinaryData)))]))); // TODO: handle additional parameters properly or should it be nullable?
+                    }
+                    else
+                    {
+                        updatedInstanceParamters.Add(instanceParemter);
+                    }
+                }
+                method.Update(signature: method.Signature, bodyStatements: Return(New.Instance(instanceExpression.Type!, updatedInstanceParamters)));
+            }
+        }
+
+        private void UpdateModel(TypeProvider type, ModelProvider model)
+        {
+            if (_flattenedModels.TryGetValue(model, out var value))
             {
                 var (flattenedProperties, internalizedProperties) = value;
                 type.Update(properties: [.. type.Properties, .. flattenedProperties]);
@@ -82,13 +163,6 @@ internal class SafeFlattenVisitor : ScmLibraryVisitor
                     internalProperty.Update(modifiers: internalProperty.Modifiers & ~MethodSignatureModifiers.Public | MethodSignatureModifiers.Internal);
                 }
             }
-
-            // Since we make the properties with flattend types internal, the flattened types will be internalized during post processing, We need to keep the flattened types public
-            if (_flattenedTypes.Any(x => x.AreNamesEqual(type.Type)))
-            {
-                ManagementClientGenerator.Instance.AddTypeToKeep(type);
-            }
-            return base.VisitType(type);
         }
     }
 }