Use VisitMethod instead of VisitType for SerializationVisitor (Azure#51670)

Author: live1206

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

@@ -6,12 +6,16 @@
 using Azure.Generator.Management.Providers;
 using Microsoft.TypeSpec.Generator.ClientModel;
 using Microsoft.TypeSpec.Generator.ClientModel.Providers;
+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;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 
 namespace Azure.Generator.Management.Visitors;
 
@@ -169,6 +173,39 @@ private void DoPreVisitPropertyForUrlPropertyName(InputProperty property, Proper
         return base.VisitMethod(method);
     }
 
+    // TODO: we will remove this manual updated when https://github.com/microsoft/typespec/issues/8079 is resolved
+    protected override MethodBodyStatement? VisitExpressionStatement(ExpressionStatement statement, MethodProvider method)
+    {
+        if (method.EnclosingType is MrwSerializationTypeDefinition serializationTypeDefinition
+            && _deserializationRename.TryGetValue(serializationTypeDefinition, out var newName)
+            && method.Signature.Name == "JsonModelCreateCore"
+            && statement.Expression is KeywordExpression keyword && keyword.Keyword == "return"
+            && keyword.Expression is InvokeMethodExpression invokeMethod)
+        {
+            invokeMethod.Update(methodName: newName);
+        }
+        return base.VisitExpressionStatement(statement, method);
+    }
+
+    // TODO: we will remove this manual updated when https://github.com/microsoft/typespec/issues/8079 is resolved
+    protected override SwitchCaseStatement? VisitSwitchCaseStatement(SwitchCaseStatement statement, MethodProvider method)
+    {
+        if (method.EnclosingType is MrwSerializationTypeDefinition serializationTypeDefinition
+            && _deserializationRename.TryGetValue(serializationTypeDefinition, out var newName)
+            && method.Signature.Name == "PersistableModelCreateCore")
+        {
+            if (statement.Statement.AsStatement().FirstOrDefault() is UsingScopeStatement usingScopeStatement
+                && usingScopeStatement.Body.AsStatement().FirstOrDefault() is ExpressionStatement expression
+                && expression.Expression is KeywordExpression keywordExpression
+                && keywordExpression.Keyword == "return"
+                && keywordExpression.Expression is InvokeMethodExpression invokeMethod)
+            {
+                invokeMethod.Update(methodName: newName);
+            }
+        }
+        return base.VisitSwitchCaseStatement(statement, method);
+    }
+
     private bool TryTransformUrlToUri(string name, [MaybeNullWhen(false)] out string newName)
     {
         const char i = 'i';

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

@@ -14,29 +14,22 @@ internal class SerializationVisitor : ScmLibraryVisitor
     internal const string FromResponseMethodName = "FromResponse";
 
     /// <inheritdoc/>
-    protected override TypeProvider? VisitType(TypeProvider type)
+    protected override MethodProvider? VisitMethod(MethodProvider method)
     {
-        if (type is MrwSerializationTypeDefinition serializationTypeDefinition)
+        if (method.EnclosingType is MrwSerializationTypeDefinition && method.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Operator))
         {
-            foreach (var method in serializationTypeDefinition.Methods)
+            var modifiers = method.Signature.Modifiers & ~MethodSignatureModifiers.Operator & ~MethodSignatureModifiers.Public | MethodSignatureModifiers.Internal;
+            if (modifiers.HasFlag(MethodSignatureModifiers.Implicit))
             {
-                if (method.Signature.Modifiers.HasFlag(MethodSignatureModifiers.Operator))
-                {
-                    var modifiers = method.Signature.Modifiers & ~MethodSignatureModifiers.Operator & ~MethodSignatureModifiers.Public | MethodSignatureModifiers.Internal;
-                    if (modifiers.HasFlag(MethodSignatureModifiers.Implicit))
-                    {
-                        modifiers &= ~MethodSignatureModifiers.Implicit;
-                        method.Signature.Update(modifiers: modifiers, name: ToRequestContentMethodName);
-                    }
-                    else if (modifiers.HasFlag(MethodSignatureModifiers.Explicit))
-                    {
-                        modifiers &= ~MethodSignatureModifiers.Explicit;
-                        method.Signature.Update(modifiers: modifiers, name: FromResponseMethodName);
-                    }
-                }
+                modifiers &= ~MethodSignatureModifiers.Implicit;
+                method.Signature.Update(modifiers: modifiers, name: ToRequestContentMethodName);
+            }
+            else if (modifiers.HasFlag(MethodSignatureModifiers.Explicit))
+            {
+                modifiers &= ~MethodSignatureModifiers.Explicit;
+                method.Signature.Update(modifiers: modifiers, name: FromResponseMethodName);
             }
-            return type;
         }
-        return base.VisitType(type);
+        return base.VisitMethod(method);
     }
 }

eng/packages/http-client-csharp-mgmt/generator/Azure.Generator.Management/test/NameVisitorTests.cs

@@ -34,11 +34,11 @@ public void TestTransformUrlToUri()
                 decorators: []);
 
             var plugin = ManagementMockHelpers.LoadMockPlugin(inputModels: () => [model], clients: () => [client]);
-            var visitor = new TestVisitor();
+
+            // PreVisitModel is called during the model creation
             var type = plugin.Object.TypeFactory.CreateModel(model);
-            var transformedModel = visitor.InvokeVisit(model, type);
-            Assert.That(transformedModel?.Name, Is.EqualTo(TestModelName.Replace("Url", "Uri")));
-            Assert.That(transformedModel?.Properties[0].Name, Is.EqualTo(TestProtyName.Replace("Url", "Uri")));
+            Assert.That(type?.Name, Is.EqualTo(TestModelName.Replace("Url", "Uri")));
+            Assert.That(type?.Properties[0].Name, Is.EqualTo(TestProtyName.Replace("Url", "Uri")));
         }
 
         [Test]
@@ -58,26 +58,18 @@ public void TestPrependResourceProviderName()
                 decorators: []);
 
             var plugin = ManagementMockHelpers.LoadMockPlugin(inputModels: () => [model], clients: () => [client]);
-            var visitor = new TestVisitor();
+
+            // PreVisitModel is called during the model creation
             var type = plugin.Object.TypeFactory.CreateModel(model);
-            var transformedModel = visitor.InvokeVisit(model, type);
             var resourceProviderName = ManagementClientGenerator.Instance.TypeFactory.ResourceProviderName;
             var updatedSkuModelName = $"{resourceProviderName}{skuModelName}";
-            Assert.AreEqual(transformedModel?.Name, updatedSkuModelName);
-            Assert.AreEqual(transformedModel!.Constructors[0].Signature.Name, $"{resourceProviderName}{skuModelName}");
-            var serializationProvider = transformedModel?.SerializationProviders.SingleOrDefault();
+            Assert.AreEqual(type?.Name, updatedSkuModelName);
+            Assert.AreEqual(type!.Constructors[0].Signature.Name, $"{resourceProviderName}{skuModelName}");
+            var serializationProvider = type?.SerializationProviders.SingleOrDefault();
             Assert.NotNull(serializationProvider);
             Assert.AreEqual(serializationProvider!.Name, updatedSkuModelName);
-            var deserializationMethod = serializationProvider.Methods.SingleOrDefault(m => m.Signature.Name.Equals($"Deserialize{updatedSkuModelName}"));
-            Assert.NotNull(deserializationMethod);
-        }
-
-        private class TestVisitor : NameVisitor
-        {
-            public ModelProvider? InvokeVisit(InputModelType model, ModelProvider? type)
-            {
-                return base.PreVisitModel(model, type);
-            }
+            var deserializationMethod = serializationProvider.Methods.SingleOrDefault(m => m.Signature.Name.StartsWith("Deserialize"));
+            Assert.AreEqual("DeserializeSamplesSku", deserializationMethod!.Signature.Name);
         }
     }
 }