Update generator for LRO rehydration

Packages.Data.props

@@ -33,8 +33,10 @@
   <!-- All csproj dependencies. -->
   <ItemGroup>
     <PackageReference Update="NUnit" Version="3.13.2" />
-    <PackageReference Update="System.ClientModel" Version="1.0.0-beta.3" />
-    <PackageReference Update="Azure.Core" Version="1.36.0" />
+    <PackageReference Update="System.ClientModel" Version="1.0.0-beta.3" >
+      <NoWarn>NU1605</NoWarn>
+    </PackageReference>
+    <PackageReference Update="Azure.Core" Version="1.38.0-alpha.20240130.3" />
     <PackageReference Update="Azure.Core.Experimental" Version="0.1.0-preview.18" />
     <PackageReference Update="Azure.ResourceManager" Version="1.9.0" />
     <PackageReference Update="Azure.Core.Expressions.DataFactory" Version="1.0.0-beta.4" />

src/AutoRest.CSharp/Common/Generation/Writers/RequestWriterHelpers.cs

@@ -25,8 +25,9 @@ namespace AutoRest.CSharp.Generation.Writers
 {
     internal static class RequestWriterHelpers
     {
-        public static void WriteRequestCreation(CodeWriter writer, RestClientMethod clientMethod, string methodAccessibility, ClientFields? fields, string? responseClassifierType, bool writeSDKUserAgent, IReadOnlyList<Parameter>? clientParameters = null)
+        public static void WriteRequestAndUriCreation(CodeWriter writer, RestClientMethod clientMethod, string methodAccessibility, ClientFields? fields, string? responseClassifierType, bool writeSDKUserAgent, IReadOnlyList<Parameter>? clientParameters = null)
         {
+            WriteUriCreation(writer, clientMethod, methodAccessibility, fields, clientParameters);
             using var methodScope = writer.AmbientScope();
             var parameters = clientMethod.Parameters;
 
@@ -107,7 +108,7 @@ public static void WriteRequestCreation(CodeWriter writer, RestClientMethod clie
                         using (WriteValueNullCheck(writer, body.Value))
                         {
                             WriteHeaders(writer, clientMethod, request, content: true, fields);
-                            WriteSerializeContent(writer, request, body.Serialization, GetConstantOrParameter(body.Value, ignoreNullability: true, convertBinaryDataToArray: false));
+                            WriteSerializeContent(writer, request, body.Serialization, GetConstantOrParameter(body.Value, ignoreNullability: true, convertBinaryDataToArray: false), body.Value.Type);
                         }
 
                         break;
@@ -183,7 +184,7 @@ public static void WriteRequestCreation(CodeWriter writer, RestClientMethod clie
                                 flattenedSchemaRequestBody.ObjectType.InitializationConstructor,
                                 initializers);
 
-                        WriteSerializeContent(writer, request, flattenedSchemaRequestBody.Serialization, $"{modelVariable:I}");
+                        WriteSerializeContent(writer, request, flattenedSchemaRequestBody.Serialization, $"{modelVariable:I}", flattenedSchemaRequestBody.ObjectType.Type);
                         break;
                     case UrlEncodedBody urlEncodedRequestBody:
                         var urlContent = new CodeWriterDeclaration("content");
@@ -218,6 +219,50 @@ public static void WriteRequestCreation(CodeWriter writer, RestClientMethod clie
             writer.Line();
         }
 
+        private static void WriteUriCreation(CodeWriter writer, RestClientMethod clientMethod, string methodAccessibility, ClientFields? fields, IReadOnlyList<Parameter>? clientParameters = null)
+        {
+            using var methodScope = writer.AmbientScope();
+            var methodName = CreateRequestUriMethodName($"{clientMethod.Name}");
+            writer.Append($"{methodAccessibility} {typeof(RequestUriBuilder)} {methodName}(");
+            foreach (Parameter clientParameter in clientMethod.Parameters)
+            {
+                writer.Append($"{clientParameter.Type} {clientParameter.Name:D},");
+            }
+            writer.RemoveTrailingComma();
+            writer.Line($")");
+            using (writer.Scope())
+            {
+                var uri = new CodeWriterDeclaration("uri");
+                writer.Line($"var {uri:D} = new RawRequestUriBuilder();");
+                foreach (var segment in clientMethod.Request.PathSegments)
+                {
+                    var value = GetFieldReference(fields, segment.Value);
+                    if (value.Type.IsFrameworkType && value.Type.FrameworkType == typeof(Uri))
+                    {
+                        writer.Append($"{uri}.Reset(");
+                        WriteConstantOrParameter(writer, value, enumAsString: !segment.IsRaw);
+                        writer.Line($");");
+                    }
+                    else if (!value.IsConstant && value.Reference.Name == "nextLink")
+                    {
+                        WritePathSegment(writer, uri, segment, value, "AppendRawNextLink");
+                    }
+                    else
+                    {
+                        WritePathSegment(writer, uri, segment, value);
+                    }
+                }
+
+                //TODO: Duplicate code between query and header parameter processing logic
+                foreach (var queryParameter in clientMethod.Request.Query)
+                {
+                    WriteQueryParameter(writer, uri, queryParameter, fields, clientParameters);
+                }
+                writer.Line($"return {uri};");
+            }
+            writer.Line();
+        }
+
         private static ReferenceOrConstant GetFieldReference(ClientFields? fields, ReferenceOrConstant value) =>
             fields != null && !value.IsConstant ? fields.GetFieldByParameter(value.Reference.Name, value.Reference.Type) ?? value : value;
 
@@ -247,8 +292,9 @@ public static void WriteHeaders(CodeWriter writer, RestClientMethod clientMethod
 
         public static string CreateRequestMethodName(RestClientMethod method) => CreateRequestMethodName(method.Name);
         public static string CreateRequestMethodName(string name) => $"Create{name}Request";
+        public static string CreateRequestUriMethodName(string name) => $"Create{name}RequestUri";
 
-        private static void WriteSerializeContent(CodeWriter writer, CodeWriterDeclaration request, ObjectSerialization bodySerialization, FormattableString value)
+        private static void WriteSerializeContent(CodeWriter writer, CodeWriterDeclaration request, ObjectSerialization bodySerialization, FormattableString value, CSharpType type)
         {
             writer.WriteMethodBodyStatement(GetRequestContentForSerialization(request, bodySerialization, value));
         }

src/AutoRest.CSharp/Common/Generation/Writers/RestClientWriter.cs

@@ -31,7 +31,7 @@ public void WriteClient(CodeWriter writer, DataPlaneRestClient restClient)
                     WriteClientCtor(writer, restClient);
                     foreach (var method in restClient.Methods)
                     {
-                        RequestWriterHelpers.WriteRequestCreation(writer, method, "internal", restClient.Fields, null, false, restClient.Parameters);
+                        RequestWriterHelpers.WriteRequestAndUriCreation(writer, method, "internal", restClient.Fields, null, false, restClient.Parameters);
                         WriteOperation(writer, method, restClient.Fields.PipelineField.Name, true, WriteFuncBodyWithSend, null, MethodSignatureModifiers.Public, WriteStatusCodeSwitch, restClient.Fields.GetFieldByParameter(KnownParameters.ClientDiagnostics));
                         WriteOperation(writer, method, restClient.Fields.PipelineField.Name, false, WriteFuncBodyWithSend, null, MethodSignatureModifiers.Public, WriteStatusCodeSwitch, restClient.Fields.GetFieldByParameter(KnownParameters.ClientDiagnostics));
                         var protocolMethod = restClient.ProtocolMethods.FirstOrDefault(m => m.RequestMethod.Operation.Equals(method.Operation));

src/AutoRest.CSharp/Common/Generation/Writers/SerializationWriter.cs

@@ -10,6 +10,7 @@
 using AutoRest.CSharp.Common.Output.Builders;
 using AutoRest.CSharp.Common.Output.Models.Types;
 using AutoRest.CSharp.Generation.Types;
+using AutoRest.CSharp.Mgmt.Output;
 using AutoRest.CSharp.Output.Models.Serialization.Json;
 using AutoRest.CSharp.Output.Models.Serialization.Xml;
 using AutoRest.CSharp.Output.Models.Types;
@@ -23,6 +24,12 @@ public void WriteSerialization(CodeWriter writer, TypeProvider schema)
         {
             switch (schema)
             {
+                // The corresponding ResoruceData does not exist for PartialResource, so we do not need to generate serialization methods for it.
+                case PartialResource:
+                    break;
+                case Resource resource:
+                    WriteResourceJsonSerialization(writer, resource);
+                    break;
                 case SerializableObjectType obj:
                     if (obj.IncludeSerializer || obj.IncludeDeserializer)
                     {
@@ -35,6 +42,25 @@ public void WriteSerialization(CodeWriter writer, TypeProvider schema)
             }
         }
 
+        private void WriteResourceJsonSerialization(CodeWriter writer, Resource resource)
+        {
+            var declaration = resource.Declaration;
+
+            using (writer.Namespace(declaration.Namespace))
+            {
+                var resourceDataType = resource.ResourceData.Type;
+                writer.Append($"{declaration.Accessibility} partial class {declaration.Name} : IJsonModel<{resourceDataType}>");
+
+                using (writer.Scope())
+                {
+                    foreach (var method in JsonSerializationMethodsBuilder.BuildResourceJsonSerializationMethods(resource))
+                    {
+                        writer.WriteMethod(method);
+                    }
+                }
+            }
+        }
+
         private void WriteObjectSerialization(CodeWriter writer, SerializableObjectType model)
         {
             var declaration = model.Declaration;

src/AutoRest.CSharp/Common/Output/Builders/JsonSerializationMethodsBuilder.cs

@@ -43,6 +43,61 @@ namespace AutoRest.CSharp.Common.Output.Builders
 {
     internal static class JsonSerializationMethodsBuilder
     {
+        public static IEnumerable<Method> BuildResourceJsonSerializationMethods(Resource resource)
+        {
+            var resourceDataType = resource.ResourceData.Type;
+            var jsonModelInterface = new CSharpType(typeof(IJsonModel<>), resourceDataType);
+            var options = new ModelReaderWriterOptionsExpression(KnownParameters.Serializations.Options);
+            var modelReaderWriter = new CSharpType(typeof(ModelReaderWriter));
+            var iModelTInterface = new CSharpType(typeof(IPersistableModel<>), resourceDataType);
+            var data = new BinaryDataExpression(KnownParameters.Serializations.Data);
+
+
+            // void IJsonModel<T>.Write(Utf8JsonWriter writer, ModelReaderWriterOptions options)
+            var writer = new Utf8JsonWriterExpression(KnownParameters.Serializations.Utf8JsonWriter);
+            yield return new Method(
+                new MethodSignature(nameof(IJsonModel<object>.Write), null, null, MethodSignatureModifiers.None, null, null, new[] { KnownParameters.Serializations.Utf8JsonWriter, KnownParameters.Serializations.Options }, ExplicitInterface: jsonModelInterface),
+                new MethodBodyStatement[]
+                {
+                    // writer.WriteStringValue(ModelReaderWriter.Write(Data, options));
+                    new InvokeInstanceMethodStatement(writer, "WriteStringValue", new ValueExpression[] { new InvokeStaticMethodExpression(modelReaderWriter, "Write", new ValueExpression[] { new MemberExpression(This, "Data"), options }, new List<CSharpType> { resourceDataType }) }, false)
+                });
+
+            // T IJsonModel<T>.Create(ref Utf8JsonReader reader, ModelReaderWriterOptions options)
+            var reader = new Parameter("reader", null, typeof(Utf8JsonReader), null, CSharp.Output.Models.Shared.ValidationType.None, null);
+            yield return new Method(
+                new MethodSignature(nameof(IJsonModel<object>.Create), null, null, MethodSignatureModifiers.None, resourceDataType, null, new[] { KnownParameters.Serializations.Utf8JsonReader, KnownParameters.Serializations.Options }, ExplicitInterface: jsonModelInterface),
+                new MethodBodyStatement[]
+                {
+                    // return ModelReaderWriter.Read<ResourceData>(new BinaryData(reader.ValueSequence));
+                    Return(new InvokeStaticMethodExpression(modelReaderWriter, "Read", new List<ValueExpression> { New.Instance(new CSharpType(typeof(BinaryData)), new MemberExpression(reader, "ValueSequence")), options }, new List<CSharpType> { resourceDataType }))
+                });
+
+            // BinaryData IPersistableModel<T>.Write(ModelReaderWriterOptions options)
+            yield return new Method(
+                new MethodSignature(nameof(IPersistableModel<object>.Write), null, null, MethodSignatureModifiers.None, typeof(BinaryData), null, new[] { KnownParameters.Serializations.Options }, ExplicitInterface: iModelTInterface),
+                new MethodBodyStatement[]
+                {
+                    // return ModelReaderWriter.Write<ResourceData>(Data, options);
+                    Return(new InvokeStaticMethodExpression(modelReaderWriter, "Write", new List<ValueExpression> { new MemberExpression(This, "Data"), options}, new List<CSharpType>{ resourceDataType }))
+                });
+
+            // T IPersistableModel<T>.Create(BinaryData data, ModelReaderWriterOptions options)
+            yield return new Method(
+                new MethodSignature(nameof(IPersistableModel<object>.Create), null, null, MethodSignatureModifiers.None, resourceDataType, null, new[] { KnownParameters.Serializations.Data, KnownParameters.Serializations.Options }, ExplicitInterface: iModelTInterface),
+                new MethodBodyStatement[]
+                {
+                    // return ModelReaderWriter.Read<ResourceData>(new BinaryData(reader.ValueSequence));
+                    Return(new InvokeStaticMethodExpression(modelReaderWriter, "Read", new List<ValueExpression> { data, options }, new List<CSharpType> { resourceDataType }))
+                });
+
+            // ModelReaderWriterFormat IPersistableModel<T>.GetFormatFromOptions(ModelReaderWriterOptions options)
+            yield return new Method(
+                new MethodSignature(nameof(IPersistableModel<object>.GetFormatFromOptions), null, null, MethodSignatureModifiers.None, typeof(string), null, new[] { KnownParameters.Serializations.Options }, ExplicitInterface: iModelTInterface),
+                Serializations.JsonFormat
+                );
+        }
+
         public static IEnumerable<Method> BuildJsonSerializationMethods(SerializableObjectType model, JsonObjectSerialization json)
         {
             var jsonModelInterface = json.IJsonModelInterface;

src/AutoRest.CSharp/LowLevel/Generation/DpgClientWriter.cs

@@ -515,7 +515,7 @@ private void WriteSubClientFactoryMethod()
 
         public static void WriteRequestCreationMethod(CodeWriter writer, RestClientMethod restMethod, ClientFields fields)
         {
-            RequestWriterHelpers.WriteRequestCreation(writer, restMethod, "internal", fields, restMethod.ResponseClassifierType.Name, false);
+            RequestWriterHelpers.WriteRequestAndUriCreation(writer, restMethod, "internal", fields, restMethod.ResponseClassifierType.Name, false);
         }
 
         public static void WriteResponseClassifierMethod(CodeWriter writer, IEnumerable<ResponseClassifierType> responseClassifierTypes)

src/AutoRest.CSharp/Mgmt/AutoRest/MgmtTarget.cs

@@ -178,10 +178,18 @@ public static async Task ExecuteAsync(GeneratedCodeWorkspace project, CodeModel
                 var writer = ResourceWriter.GetWriter(resource);
                 writer.Write();
 
+                var name = resource.Type.Name;
+
                 var ri = new ResourceItem(resource, MgmtReport.Instance.TransformSection);
                 MgmtReport.Instance.ResourceSection.Add(ri.Name, ri);
 
-                AddGeneratedFile(project, $"{resource.Type.Name}.cs", writer.ToString());
+                AddGeneratedFile(project, $"{name}.cs", writer.ToString());
+
+                // we do not need this if model reader writer feature is not enabled
+                if (Configuration.UseModelReaderWriter)
+                {
+                    WriteSerialization(project, resource, serializeWriter, $"Models/{name}.Serialization.cs");
+                }
             }
 
             // write extension class
@@ -289,6 +297,11 @@ private static void WriteArmModel(GeneratedCodeWorkspace project, TypeProvider m
                     return;
             }
 
+            WriteSerialization(project, model, serializeWriter, serializationFileName);
+        }
+
+        private static void WriteSerialization(GeneratedCodeWorkspace project, TypeProvider model, SerializationWriter serializeWriter, string serializationFileName)
+        {
             var serializerCodeWriter = new CodeWriter();
             serializeWriter.WriteSerialization(serializerCodeWriter, model);
             AddGeneratedFile(project, serializationFileName, serializerCodeWriter.ToString());

src/AutoRest.CSharp/Mgmt/Generation/MgmtClientBaseWriter.cs

@@ -627,6 +627,24 @@ protected virtual void WriteLROMethodBranch(MgmtRestOperation operation, IEnumer
             WriteArguments(_writer, parameterMapping);
             _writer.Line($"cancellationToken){GetConfigureAwait(async)};");
 
+            if (operation.IsFakeLongRunningOperation)
+            {
+                _writer.Append($"var uri = ");
+                _writer.Append($"{GetRestClientName(operation)}.{RequestWriterHelpers.CreateRequestUriMethodName(operation.Method.Name)}(");
+                WriteArguments(_writer, parameterMapping);
+                _writer.RemoveTrailingComma();
+                _writer.Line($");");
+
+                _writer.Append($"var rehydrationToken = {typeof(NextLinkOperationImplementation)}.GetRehydrationToken(");
+
+                _writer.Append($"{typeof(RequestMethod)}.{new CultureInfo("en-US", false).TextInfo.ToTitleCase(operation.Method.Request.HttpMethod.ToString().ToLower())}, ");
+                _writer.Append($"uri.ToUri(), ");
+                _writer.Append($"uri.ToString(), ");
+                _writer.Append($"{typeof(NextLinkOperationImplementation)}.HeaderSource.None.ToString(), ");
+                _writer.Append($"null, ");
+                _writer.Line($"{typeof(OperationFinalStateVia)}.OriginalUri.ToString());");
+            }
+
             WriteLROResponse(GetDiagnosticReference(operation).Name, PipelineProperty, operation, parameterMapping, async);
         }
 
@@ -656,6 +674,7 @@ protected virtual void WriteLROResponse(string diagnosticsVariableName, string p
                 {
                     _writer.Append($"{Configuration.ApiTypes.ResponseParameterName}");
                 }
+                _writer.Append($", rehydrationToken");
             }
             else
             {