Add a resource data vs resource client provider map (Azure#51697)

* introduce a method to get the corresponding resource for a resource data

* wip

* finish the refactor

* clean up

* introduce strong-typed response

* fix test case failures

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

Co-authored-by: Copilot <[email protected]>

* Update eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/OperationMethodProviders/ResourceOperationMethodProvider.cs

Co-authored-by: Copilot <[email protected]>

* resolve comments

* add another static

* fix compilation error

---------

Co-authored-by: Copilot <[email protected]>

Author: ArcturusZhang

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

@@ -12,6 +12,7 @@
 using Microsoft.TypeSpec.Generator.Providers;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 
 namespace Azure.Generator.Management
@@ -23,17 +24,16 @@ public class ManagementOutputLibrary : AzureOutputLibrary
         internal ManagementLongRunningOperationProvider ArmOperation => _armOperation ??= new ManagementLongRunningOperationProvider(false);
 
         private ManagementLongRunningOperationProvider? _genericArmOperation;
-        internal ManagementLongRunningOperationProvider GenericArmOperation => _genericArmOperation ??= new ManagementLongRunningOperationProvider(true);
+        internal ManagementLongRunningOperationProvider ArmOperationOfT => _genericArmOperation ??= new ManagementLongRunningOperationProvider(true);
 
         private PageableWrapperProvider? _pageableWrapper;
         internal PageableWrapperProvider PageableWrapper => _pageableWrapper ??= new PageableWrapperProvider(false);
 
         private PageableWrapperProvider? _asyncPageableWrapper;
         internal PageableWrapperProvider AsyncPageableWrapper => _asyncPageableWrapper ??= new PageableWrapperProvider(true);
 
-        // TODO: replace this with CSharpType to TypeProvider mapping
-        private HashSet<CSharpType>? _resourceTypes;
-        private HashSet<CSharpType> ResourceTypes => _resourceTypes ??= BuildResourceModels();
+        private IReadOnlyList<ResourceClientProvider>? _resourceClients;
+        internal IReadOnlyList<ResourceClientProvider> ResourceClients => _resourceClients ??= BuildResources();
 
         // TODO: replace this with CSharpType to TypeProvider mapping and move this logic to ModelFactoryVisitor
         private HashSet<CSharpType>? _modelFactoryModels;
@@ -84,24 +84,6 @@ IEnumerable<PropertyProvider> EnumerateAllPublicProperties(ModelProvider current
 
         internal bool IsModelFactoryModelType(CSharpType type) => ModelFactoryModels.Contains(type);
 
-        private HashSet<CSharpType> BuildResourceModels()
-        {
-            var resourceTypes = new HashSet<CSharpType>();
-
-            foreach (var model in ManagementClientGenerator.Instance.InputLibrary.InputNamespace.Models)
-            {
-                if (ManagementClientGenerator.Instance.InputLibrary.IsResourceModel(model))
-                {
-                    var modelProvider = ManagementClientGenerator.Instance.TypeFactory.CreateModel(model);
-                    if (modelProvider is not null)
-                    {
-                        resourceTypes.Add(modelProvider.Type);
-                    }
-                }
-            }
-            return resourceTypes;
-        }
-
         private IReadOnlyList<ResourceClientProvider> BuildResources()
         {
             var resources = new List<ResourceClientProvider>();
@@ -179,14 +161,14 @@ private IReadOnlyList<TypeProvider> BuildExtensions(IReadOnlyList<ResourceClient
         /// <inheritdoc/>
         protected override TypeProvider[] BuildTypeProviders()
         {
-            var resources = BuildResources();
+            var resources = ResourceClients;
             var collections = resources.Select(r => r.ResourceCollection).WhereNotNull();
             var extensions = BuildExtensions(resources);
 
             return [
                 .. base.BuildTypeProviders().Where(t => t is not SystemObjectModelProvider),
                 ArmOperation,
-                GenericArmOperation,
+                ArmOperationOfT,
                 .. resources,
                 .. collections,
                 .. extensions,
@@ -196,6 +178,13 @@ .. resources.Select(r => r.Source),
                 .. resources.SelectMany(r => r.SerializationProviders)];
         }
 
-        internal bool IsResourceModelType(CSharpType type) => ResourceTypes.Contains(type);
+        internal bool IsResourceModelType(CSharpType type) => TryGetResourceClientProvider(type, out _);
+
+        private IReadOnlyDictionary<CSharpType, ResourceClientProvider>? _resourceDataTypes;
+        internal bool TryGetResourceClientProvider(CSharpType resourceDataType, [MaybeNullWhen(false)] out ResourceClientProvider resourceClientProvider)
+        {
+            _resourceDataTypes ??= ResourceClients.ToDictionary(r => r.ResourceData.Type, r => r);
+            return _resourceDataTypes.TryGetValue(resourceDataType, out resourceClientProvider);
+        }
     }
 }

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

@@ -5,12 +5,10 @@
 using Azure.Core.Pipeline;
 using Azure.Generator.Management.Utilities;
 using Azure.ResourceManager;
-using Microsoft.TypeSpec.Generator.Expressions;
 using Microsoft.TypeSpec.Generator.Primitives;
 using Microsoft.TypeSpec.Generator.Providers;
 using Microsoft.TypeSpec.Generator.Statements;
 using System;
-using System.ClientModel.Primitives;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -48,7 +46,7 @@ protected override string BuildRelativeFilePath()
         {
             return Path.Combine("src", "Generated", "LongRunningOperation", GetFileName());
 
-            string GetFileName() => _isGeneric ? $"{_serviceName}ArmOperationOfT.cs" : $"{Name}.cs";
+            string GetFileName() => _isGeneric ? $"{Name}OfT.cs" : $"{Name}.cs";
         }
 
         protected override CSharpType[] BuildImplements() => [_isGeneric ? new CSharpType(typeof(ArmOperation<>), _t) : typeof(ArmOperation)];

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/OperationMethodProviders/ExistsOperationMethodProvider.cs

@@ -2,11 +2,12 @@
 // Licensed under the MIT License.
 
 using Azure.Generator.Management.Models;
+using Azure.Generator.Management.Snippets;
 using Azure.Generator.Management.Utilities;
-using Microsoft.TypeSpec.Generator.Expressions;
 using Microsoft.TypeSpec.Generator.Input;
 using Microsoft.TypeSpec.Generator.Primitives;
 using Microsoft.TypeSpec.Generator.Providers;
+using Microsoft.TypeSpec.Generator.Snippets;
 using Microsoft.TypeSpec.Generator.Statements;
 using System.Collections.Generic;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
@@ -39,17 +40,16 @@ protected override MethodSignature CreateSignature()
                 _convenienceMethod.Signature.NonDocumentComment);
         }
 
-        protected override IReadOnlyList<MethodBodyStatement> BuildReturnStatements(ValueExpression responseVariable, MethodSignature signature)
+        protected override IReadOnlyList<MethodBodyStatement> BuildReturnStatements(ScopedApi<Response> responseVariable, MethodSignature signature)
         {
             // For Exists methods, we check if Value is not null and return a boolean
-            var returnValueExpression = responseVariable.Property("Value").NotEqual(Null);
+            var returnValueExpression = responseVariable.Value().NotEqual(Null);
 
             return [
                 Return(
-                    Static(typeof(Response)).Invoke(
-                        nameof(Response.FromValue),
+                    ResponseSnippets.FromValue(
                         returnValueExpression,
-                        responseVariable.Invoke("GetRawResponse")
+                        responseVariable.GetRawResponse()
                     )
                 )
             ];

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/OperationMethodProviders/GetIfExistsOperationMethodProvider.cs

@@ -2,11 +2,14 @@
 // Licensed under the MIT License.
 
 using Azure.Generator.Management.Models;
+using Azure.Generator.Management.Snippets;
 using Azure.Generator.Management.Utilities;
+using Azure.ResourceManager;
 using Microsoft.TypeSpec.Generator.Expressions;
 using Microsoft.TypeSpec.Generator.Input;
 using Microsoft.TypeSpec.Generator.Primitives;
 using Microsoft.TypeSpec.Generator.Providers;
+using Microsoft.TypeSpec.Generator.Snippets;
 using Microsoft.TypeSpec.Generator.Statements;
 using System.Collections.Generic;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
@@ -18,49 +21,44 @@ internal class GetIfExistsOperationMethodProvider(
         RestClientInfo restClientInfo,
         InputServiceMethod method,
         MethodProvider convenienceMethod,
-        bool isAsync) : ResourceOperationMethodProvider(collection, collection.ContextualPath, restClientInfo, method, convenienceMethod, isAsync)
+        bool isAsync)
+        : ResourceOperationMethodProvider(
+            collection,
+            collection.ContextualPath,
+            restClientInfo,
+            method,
+            convenienceMethod,
+            isAsync,
+            methodName: isAsync ? "GetIfExistsAsync" : "GetIfExists",
+            description: $"Tries to get details for this resource from the service.")
     {
-        protected override MethodSignature CreateSignature()
+        protected override CSharpType BuildReturnType()
         {
-            var returnType = new CSharpType(typeof(NullableResponse<>), _resource.Type)
-                .WrapAsync(_isAsync);
-
-            return new MethodSignature(
-                _isAsync ? "GetIfExistsAsync" : "GetIfExists",
-                $"Tries to get details for this resource from the service.",
-                _convenienceMethod.Signature.Modifiers,
-                returnType,
-                _convenienceMethod.Signature.ReturnDescription,
-                GetOperationMethodParameters(),
-                _convenienceMethod.Signature.Attributes,
-                _convenienceMethod.Signature.GenericArguments,
-                _convenienceMethod.Signature.GenericParameterConstraints,
-                _convenienceMethod.Signature.ExplicitInterface,
-                _convenienceMethod.Signature.NonDocumentComment);
+            return new CSharpType(typeof(NullableResponse<>), _returnBodyType!).WrapAsync(_isAsync);
         }
 
-        protected override IReadOnlyList<MethodBodyStatement> BuildReturnStatements(ValueExpression responseVariable, MethodSignature signature)
+        protected override IReadOnlyList<MethodBodyStatement> BuildReturnStatements(ScopedApi<Response> responseVariable, MethodSignature signature)
         {
+            // we need to add some null checks before we return the response.
             List<MethodBodyStatement> statements =
             [
-                new IfStatement(responseVariable.Property("Value").Equal(Null))
+                new IfStatement(responseVariable.Value().Equal(Null))
                 {
                     Return(
                         New.Instance(
-                            new CSharpType(typeof(NoValueResponse<>), _resource.Type),
-                            responseVariable.Invoke("GetRawResponse")
+                            new CSharpType(typeof(NoValueResponse<>), _returnBodyResourceClient!.Type),
+                            responseVariable.GetRawResponse()
                         )
                     )
                 }
             ];
 
-            var returnValueExpression = New.Instance(_resource.Type, This.Property("Client"), responseVariable.Property("Value"));
+            var returnValueExpression = New.Instance(_returnBodyResourceClient.Type, This.As<ArmResource>().Client(), responseVariable.Value());
             statements.Add(
                 Return(
-                    Static(typeof(Response)).Invoke(
-                        nameof(Response.FromValue),
+                    ResponseSnippets.FromValue(
                         returnValueExpression,
-                        responseVariable.Invoke("GetRawResponse")
+                        responseVariable.GetRawResponse()
                     )
                 )
             );

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/src/Providers/OperationMethodProviders/PageableOperationMethodProvider.cs

@@ -12,7 +12,6 @@
 using Microsoft.TypeSpec.Generator.Primitives;
 using Microsoft.TypeSpec.Generator.Providers;
 using Microsoft.TypeSpec.Generator.Statements;
-using System;
 using System.Collections.Generic;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
 
@@ -27,6 +26,8 @@ internal class PageableOperationMethodProvider
         private readonly MethodProvider _convenienceMethod;
         private readonly bool _isAsync;
         private readonly CSharpType _itemType;
+        private readonly CSharpType _actualItemType;
+        private ResourceClientProvider? _itemResourceClient;
         private readonly ResourceOperationKind _methodKind;
         private readonly MethodSignature _signature;
         private readonly MethodBodyStatement[] _bodyStatements;
@@ -48,11 +49,29 @@ public PageableOperationMethodProvider(
             _convenienceMethod = convenienceMethod;
             _isAsync = isAsync;
             _itemType = itemType;
+            InitializeTypeInfo(
+                itemType,
+                ref _actualItemType!,
+                ref _itemResourceClient
+            );
             _methodKind = methodKind;
             _signature = CreateSignature();
             _bodyStatements = BuildBodyStatements();
         }
 
+        private static void InitializeTypeInfo(
+            CSharpType itemType,
+            ref CSharpType actualItemType,
+            ref ResourceClientProvider? resourceClient
+            )
+        {
+            actualItemType = itemType;
+            if (ManagementClientGenerator.Instance.OutputLibrary.TryGetResourceClientProvider(itemType, out resourceClient))
+            {
+                actualItemType = resourceClient.Type;
+            }
+        }
+
         public static implicit operator MethodProvider(PageableOperationMethodProvider pageableOperationMethodProvider)
         {
             return new MethodProvider(
@@ -63,11 +82,9 @@ public static implicit operator MethodProvider(PageableOperationMethodProvider p
 
         protected MethodSignature CreateSignature()
         {
-            var actualItemType = IsResourceDataType(_itemType) ? GetResourceClientProvider().Type : _itemType;
-
             var returnType = _isAsync
-                ? new CSharpType(typeof(AsyncPageable<>), actualItemType)
-                : new CSharpType(typeof(Pageable<>), actualItemType);
+                ? new CSharpType(typeof(AsyncPageable<>), _actualItemType)
+                : new CSharpType(typeof(Pageable<>), _actualItemType);
             var methodName = _methodKind == ResourceOperationKind.List
                 ? (_isAsync ? "GetAllAsync" : "GetAll")
                 : _convenienceMethod.Signature.Name;
@@ -102,9 +119,9 @@ protected MethodBodyStatement[] BuildBodyStatements()
             arguments.AddRange(_contextualPath.PopulateArguments(This.As<ArmResource>().Id(), requestMethod.Signature.Parameters, contextVariable, _signature.Parameters));
 
             // Handle ResourceData type conversion if needed
-            if (IsResourceDataType(_itemType))
+            if (_itemResourceClient != null)
             {
-                statements.Add(BuildResourceDataConversionStatement(collectionResultOfT, arguments));
+                statements.Add(BuildResourceDataConversionStatement(collectionResultOfT, _itemResourceClient.Type, arguments));
             }
             else
             {
@@ -114,21 +131,17 @@ protected MethodBodyStatement[] BuildBodyStatements()
             return statements.ToArray();
         }
 
-        private MethodBodyStatement BuildResourceDataConversionStatement(CSharpType sourcePageable, List<ValueExpression> arguments)
+        private MethodBodyStatement BuildResourceDataConversionStatement(CSharpType sourcePageable, CSharpType typeOfResource, List<ValueExpression> arguments)
         {
-            // Get the resource client provider to access the resource type
-            var resourceClientProvider = GetResourceClientProvider();
-            var resourceType = resourceClientProvider.Type;
-
             // Create PageableWrapper instance to convert from ResourceData to Resource
             var pageableWrapperType = _isAsync ? ManagementClientGenerator.Instance.OutputLibrary.AsyncPageableWrapper : ManagementClientGenerator.Instance.OutputLibrary.PageableWrapper;
 
             // Create the concrete wrapper type with proper generic parameters
             // Since pageableWrapperType.Type represents the constructed generic type, we need to use it directly
-            var concreteWrapperType = pageableWrapperType.Type.MakeGenericType([_itemType, resourceType]);
+            var concreteWrapperType = pageableWrapperType.Type.MakeGenericType([_itemType, typeOfResource]);
 
             // Create converter function: data => new ResourceType(Client, data)
-            var converterFunc = CreateConverterFunction(_itemType, resourceType);
+            var converterFunc = CreateConverterFunction(_itemType, typeOfResource);
 
             var wrapperArguments = new List<ValueExpression>
             {
@@ -139,30 +152,6 @@ private MethodBodyStatement BuildResourceDataConversionStatement(CSharpType sour
             return Return(New.Instance(concreteWrapperType, wrapperArguments));
         }
 
-        private bool IsResourceDataType(CSharpType itemType)
-        {
-            try
-            {
-                var resourceClientProvider = GetResourceClientProvider();
-                return itemType.Equals(resourceClientProvider.ResourceData.Type);
-            }
-            catch (InvalidOperationException)
-            {
-                // If we can't get a ResourceClientProvider, then this is not a ResourceData type
-                return false;
-            }
-        }
-
-        private ResourceClientProvider GetResourceClientProvider()
-        {
-            return _enclosingType switch
-            {
-                ResourceClientProvider rcp => rcp,
-                ResourceCollectionClientProvider rccp => rccp.Resource, // Return the Resource property
-                _ => throw new InvalidOperationException($"Expected ResourceClientProvider or ResourceCollectionClientProvider, but got: {_enclosingType.GetType()}")
-            };
-        }
-
         private ValueExpression CreateConverterFunction(CSharpType fromType, CSharpType toType)
         {
             // Create a lambda expression: data => new ResourceType(Client, data)