Remove internal setter from collection-type property

eng/packages/http-client-csharp-mgmt/emitter/src/emitter.ts

@@ -9,8 +9,9 @@ import {
   $onEmit as $onAzureEmit,
   AzureEmitterOptions
 } from "@azure-typespec/http-client-csharp";
-import { azureSDKContextOptions } from "./sdk-context-options.js";
+import { azureSDKContextOptions, flattenPropertyDecorator } from "./sdk-context-options.js";
 import { updateClients } from "./resource-detection.js";
+import { DecoratorInfo } from "@azure-tools/typespec-client-generator-core";
 
 export async function $onEmit(context: EmitContext<AzureEmitterOptions>) {
   context.options["generator-name"] ??= "ManagementClientGenerator";
@@ -25,6 +26,22 @@ export async function $onEmit(context: EmitContext<AzureEmitterOptions>) {
     sdkContext: CSharpEmitterContext
   ): CodeModel {
     updateClients(codeModel, sdkContext);
+    setFlattenProperty(codeModel, sdkContext);
     return codeModel;
   }
 }
+
+function setFlattenProperty(codeModel: CodeModel, sdkContext: CSharpEmitterContext): void {
+  for (const model of sdkContext.sdkPackage.models) {
+    for (const property of model.properties) {
+      if (property.flatten ) {
+
+          const flattenPropertyMetadataDecorator: DecoratorInfo = {
+            name: flattenPropertyDecorator,
+            arguments: {}
+          };
+        property.decorators.push(flattenPropertyMetadataDecorator);
+      }
+    }
+  }
+}

eng/packages/http-client-csharp-mgmt/emitter/src/sdk-context-options.ts

@@ -76,9 +76,6 @@ export const resourceGroupResource =
 const resourceGroupResourceRegex =
   "Azure\\.ResourceManager\\.@resourceGroupResource";
 
-const flattenPropertyRegex =
-  "Azure\\.ClientGenerator\\.Core\\.@flattenProperty";
-
 // TODO: add this decorator to TCGC
 export const resourceMetadata = "Azure.ClientGenerator.Core.@resourceSchema";
 const resourceMetadataRegex =
@@ -88,6 +85,8 @@ export const nonResourceMethodMetadata =
 const nonResourceMethodMetadataRegex =
   "Azure\\.ClientGenerator\\.Core\\.@nonResourceMethodSchema";
 
+export const flattenPropertyDecorator = "Azure.ResourceManager.@flattenProperty";
+
 export const azureSDKContextOptions: CreateSdkContextOptions = {
   versioning: {
     previewStringRegex: /-preview$/
@@ -104,7 +103,6 @@ export const azureSDKContextOptions: CreateSdkContextOptions = {
     armResourceOperationsRegex,
     armResourceUpdateRegex,
     armResourceReadRegex,
-    flattenPropertyRegex,
     parentResourceRegex,
     resourceGroupResourceRegex,
     singletonRegex,

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

@@ -15,7 +15,7 @@ public class ManagementInputLibrary : InputLibrary
     {
         private const string ResourceMetadataDecoratorName = "Azure.ClientGenerator.Core.@resourceSchema";
         private const string NonResourceMethodMetadata = "Azure.ClientGenerator.Core.@nonResourceMethodSchema";
-        private const string FlattenPropertyDecoratorName = "Azure.ClientGenerator.Core.@flattenProperty";
+        private const string FlattenPropertyDecoratorName = "Azure.ResourceManager.@flattenProperty";
 
         private IReadOnlyDictionary<string, InputServiceMethod>? _inputServiceMethodsByCrossLanguageDefinitionId;
         private IReadOnlyDictionary<InputServiceMethod, InputClient>? _intMethodClientMap;

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

@@ -80,30 +80,36 @@ private static bool HasDefaultPublicCtor(ModelProvider? innerModel)
             return false;
         }
 
-        public static MethodBodyStatement BuildGetter(bool? includeGetterNullCheck, PropertyProvider internalProperty, TypeProvider innerModel, PropertyProvider singleProperty)
+        public static MethodBodyStatement BuildGetter(bool? includeGetterNullCheck, PropertyProvider internalProperty, TypeProvider innerModel, PropertyProvider innerProperty)
         {
             var checkNullExpression = This.Property(internalProperty.Name).Is(Null);
+            // For collection types, we do not do null check and initialization in getter, they have been initialized in constructor.
+            if (innerProperty.Type.IsCollection)
+            {
+                return new List<MethodBodyStatement>() { Return(new MemberExpression(internalProperty, innerProperty.Name)) };
+            }
+
             if (includeGetterNullCheck == true)
             {
                 return new List<MethodBodyStatement> {
                     new IfStatement(checkNullExpression)
                     {
                         internalProperty.Assign(New.Instance(innerModel.Type)).Terminate()
                     },
-                    Return(new MemberExpression(internalProperty, singleProperty.Name))
+                    Return(new MemberExpression(internalProperty, innerProperty.Name))
                 };
             }
             else if (includeGetterNullCheck == false)
             {
-                return Return(new TernaryConditionalExpression(checkNullExpression, Default, new MemberExpression(internalProperty, singleProperty.Name)));
+                return Return(new TernaryConditionalExpression(checkNullExpression, Default, new MemberExpression(internalProperty, innerProperty.Name)));
             }
             else
             {
                 if (innerModel.Type.IsNullable)
                 {
-                    return Return(new MemberExpression(internalProperty.AsVariableExpression.NullConditional(), singleProperty.Name));
+                    return Return(new MemberExpression(internalProperty.AsVariableExpression.NullConditional(), innerProperty.Name));
                 }
-                return Return(new MemberExpression(internalProperty, singleProperty.Name));
+                return Return(new MemberExpression(internalProperty, innerProperty.Name));
             }
         }
 
@@ -122,24 +128,6 @@ public static MethodBodyStatement BuildSetterForPropertyFlatten(ModelProvider in
             return setter;
         }
 
-        public static Dictionary<ValueExpression, ValueExpression> PopulateCollectionProperties(IEnumerable<PropertyProvider> collectionTypeProperties)
-        {
-            var result = new Dictionary<ValueExpression, ValueExpression>();
-            foreach (var property in collectionTypeProperties)
-            {
-                var propertyValue = Value.Property(property.Name);
-                if (property.Type.IsList)
-                {
-                    result.Add(Identifier(property.Name), New.Instance(ManagementClientGenerator.Instance.TypeFactory.ListInitializationType.MakeGenericType(property.Type.Arguments)));
-                }
-                if (property.Type.IsDictionary)
-                {
-                    result.Add(Identifier(property.Name), New.Instance(ManagementClientGenerator.Instance.TypeFactory.DictionaryInitializationType.MakeGenericType(property.Type.Arguments)));
-                }
-            }
-            return result;
-        }
-
         public static MethodBodyStatement BuildSetterForSafeFlatten(bool includeSetterCheck, ModelProvider innerModel, PropertyProvider internalProperty, PropertyProvider innerProperty)
         {
             var isOverriddenValueType = IsOverriddenValueType(innerProperty);

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

@@ -9,6 +9,7 @@
 using Microsoft.TypeSpec.Generator.Providers;
 using Microsoft.TypeSpec.Generator.Snippets;
 using Microsoft.TypeSpec.Generator.Statements;
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
@@ -36,41 +37,55 @@ internal class FlattenPropertyVisitor : ScmLibraryVisitor
 
             if (type is ModelProvider model && _collectionTypeProperties.TryGetValue(model, out var value))
             {
-                foreach (var (internalProperty, collectionProperties) in value)
+                foreach (var internalProperty in value)
                 {
-                    var innerCollectionProperties = collectionProperties.Select(x => x.InnerProperty);
-                    var initializationMethod = BuildInitializationMethod(innerCollectionProperties, internalProperty, model);
-                    var publicConstructor = model.Constructors.Single(m => m.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Public));
-                    var invokeInitialization = This.Invoke(initializationMethod.Signature.Name).Terminate();
+                    var publicConstructor = model.Constructors.SingleOrDefault(m => m.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Public));
+                    if (publicConstructor is null)
+                    {
+                        continue;
+                    }
 
+                    var internalPropertyTypeConstructor = ManagementClientGenerator.Instance.TypeFactory.CSharpTypeMap[internalProperty.Type]!.Constructors.Single(c => c.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Public));
+                    var initializationParameters = PopulateInitializationParameters(publicConstructor, internalPropertyTypeConstructor);
+                    var initialization = internalProperty.Assign(New.Instance(internalProperty.Type, initializationParameters)).Terminate();
                     // If the property is a collection type, we need to ensure that it is initialized
                     if (publicConstructor.BodyStatements is null)
                     {
-                        publicConstructor.Update(bodyStatements: new List<MethodBodyStatement> { invokeInitialization });
+                        publicConstructor.Update(bodyStatements: new List<MethodBodyStatement> { initialization });
                     }
                     else
                     {
                         var body = publicConstructor.BodyStatements.ToList();
-                        body.Add(invokeInitialization);
+                        body.Add(initialization);
                         publicConstructor.Update(bodyStatements: body);
                     }
-                    model.Update(methods: [.. model.Methods, initializationMethod]);
                 }
             }
 
             return base.PostVisitType(type);
         }
 
-        internal const string s_initializationMethodName = "Initialize";
-        private MethodProvider BuildInitializationMethod(IEnumerable<PropertyProvider> collectionTypeProperties, PropertyProvider internalProperty, ModelProvider model)
+        private ValueExpression[] PopulateInitializationParameters(ConstructorProvider publicConstructor, ConstructorProvider internalPropertyTypeConstructor)
         {
-            var signature = new MethodSignature($"{s_initializationMethodName}{internalProperty.Type.Name}", null, MethodSignatureModifiers.Private, null, null, []);
-            MethodBodyStatement[] body = [
-                    new IfStatement(This.Property(internalProperty.Name).Is(Null))
-                    {
-                        internalProperty.Assign(New.Instance(internalProperty.Type, PropertyHelpers.PopulateCollectionProperties(collectionTypeProperties))).Terminate()
-                    },];
-            return new MethodProvider(signature, body, model);
+            var parameters = new List<ValueExpression>();
+            foreach (var parameter in internalPropertyTypeConstructor.Signature.Parameters)
+            {
+                if (parameter.Type.IsList)
+                {
+                    parameters.Add(New.Instance(ManagementClientGenerator.Instance.TypeFactory.ListInitializationType.MakeGenericType(parameter.Type.Arguments)));
+                }
+                else if (parameter.Type.IsDictionary)
+                {
+                    parameters.Add(New.Instance(ManagementClientGenerator.Instance.TypeFactory.DictionaryInitializationType.MakeGenericType(parameter.Type.Arguments)));
+                }
+                else
+                {
+                    var constructorParameter = publicConstructor.Signature.Parameters.Single(p => p.Name.Equals(parameter.Name, System.StringComparison.OrdinalIgnoreCase));
+
+                    parameters.Add(constructorParameter);
+                }
+            }
+            return parameters.ToArray();
         }
 
         private void UpdateModelFactory(ModelFactoryProvider modelFactory)
@@ -256,7 +271,7 @@ int GetAdditionalPropertyIndex()
         // So that, we can use this to update the model factory methods later.
         private readonly Dictionary<CSharpType, Dictionary<string, List<(bool IsOverriddenValueType, PropertyProvider FlattenedProperty)>>> _flattenedModelTypes = new();
         // TODO: Workadound to initialize all collection-type properties in all collection-type setters, remove this once we have lazy initializtion for collection-type properties
-        private readonly Dictionary<ModelProvider, Dictionary<PropertyProvider, List<(PropertyProvider FlattenedProperty, PropertyProvider InnerProperty)>>> _collectionTypeProperties = new();
+        private readonly Dictionary<ModelProvider, HashSet<PropertyProvider>> _collectionTypeProperties = new();
         private void FlattenProperties(ModelProvider model)
         {
             var isFlattened = false;
@@ -281,6 +296,7 @@ private void FlattenProperties(ModelProvider model)
 
                         foreach (var innerProperty in innerProperties)
                         {
+                            CollectFlattenTypeCollectionProperty(property, innerProperty, model);
                             // flatten the property to public and associate it with the internal property
                             var (_, includeGetterNullCheck, _) = PropertyHelpers.GetFlags(property, innerProperty);
                             var flattenPropertyName = innerProperty.Name; // TODO: handle name conflicts
@@ -304,7 +320,6 @@ private void FlattenProperties(ModelProvider model)
                                     innerProperty.Attributes);
 
                             flattenedProperties.Add((isOverriddenValueType, flattenedProperty));
-                            AddInternalSetterForFlattenTypeCollectionProperty(property, innerProperty, flattenedProperty, model);
                         }
                         // make the internalized properties internal
                         property.Update(modifiers: property.Modifiers & ~MethodSignatureModifiers.Public | MethodSignatureModifiers.Internal);
@@ -322,28 +337,20 @@ private void FlattenProperties(ModelProvider model)
         }
 
         // TODO: workaround to add internal setter, we should remove this once we add lazy initialization for collection type properties
-        private void AddInternalSetterForFlattenTypeCollectionProperty(PropertyProvider internalProperty, PropertyProvider innerProperty, PropertyProvider flattenedProperty, ModelProvider modelProvider)
+        private void CollectFlattenTypeCollectionProperty(PropertyProvider internalProperty, PropertyProvider innerProperty, ModelProvider modelProvider)
         {
             if (innerProperty.Type.IsCollection)
             {
                 if (_collectionTypeProperties.TryGetValue(modelProvider, out var value))
                 {
-                    if (value.TryGetValue(internalProperty, out var properties))
-                    {
-                        properties.Add((flattenedProperty, innerProperty));
-                    }
-                    else
-                    {
-                        value.Add(internalProperty, [(flattenedProperty, innerProperty)]);
-                    }
+                    value.Add(internalProperty);
                 }
                 else
                 {
-                    var dict = new Dictionary<PropertyProvider, List<(PropertyProvider FlattenedProperty, PropertyProvider InnerProperty)>>();
-                    dict.Add(internalProperty, [(flattenedProperty, innerProperty)]);
-                    _collectionTypeProperties.Add(modelProvider, dict);
+                    var set = new HashSet<PropertyProvider>();
+                    set.Add(internalProperty);
+                    _collectionTypeProperties.Add(modelProvider, set);
                 }
-                innerProperty.Update(body: new AutoPropertyBody(true, MethodSignatureModifiers.Internal));
             }
         }
 

eng/packages/http-client-csharp-mgmt/generator/TestProjects/Local/Mgmt-TypeSpec/tspCodeModel.json

@@ -5109,7 +5109,7 @@
           "flatten": true,
           "decorators": [
             {
-              "name": "Azure.ClientGenerator.Core.@flattenProperty",
+              "name": "Azure.ResourceManager.@flattenProperty",
               "arguments": {}
             }
           ],