fix: update visitors to account for jsonpatch changes

Author: jorgerangel-msft

codegen/generator/src/OpenAILibraryVisitor.cs

@@ -6,6 +6,7 @@
 using Microsoft.TypeSpec.Generator.Snippets;
 using Microsoft.TypeSpec.Generator.Statements;
 using System;
+using System.ClientModel.Primitives;
 using System.Collections.Generic;
 using System.Linq;
 using static Microsoft.TypeSpec.Generator.Snippets.Snippet;
@@ -46,6 +47,8 @@ public class OpenAILibraryVisitor : ScmLibraryVisitor
         ["ReasoningResponseItem"] = [_readonlyStatusReplacementInfo],
         ["WebSearchCallResponseItem"] = [_readonlyStatusReplacementInfo],
     };
+    private static readonly SingleLineCommentStatement OptionalDefinedCheckComment =
+        new("Plugin customization: apply Optional.Is*Defined() check based on type name dictionary lookup");
 
     protected override TypeProvider VisitType(TypeProvider type)
     {
@@ -118,101 +121,161 @@ protected override FieldProvider VisitField(FieldProvider field)
 
     protected override MethodProvider VisitMethod(MethodProvider method)
     {
-        if (method.Signature.Name != JsonModelWriteCoreMethodName)
+        // If there are no body statements, or the body statements are not MethodBodyStatements,
+        // return the method as is return the method as is
+        if (method.Signature.Name != JsonModelWriteCoreMethodName ||
+            method.BodyStatements is not MethodBodyStatements statements)
         {
             return method;
         }
 
-        // If there are no body statements, return the method as is
-        if (method.BodyStatements == null)
-        {
-            return method;
-        }
+        var updatedStatements = new List<MethodBodyStatement>();
+        var flattenedStatements = new List<MethodBodyStatement>();
 
-        // If the body statements are not MethodBodyStatements, return the method as is
-        if (method.BodyStatements is not MethodBodyStatements statements)
+        foreach (var stmt in statements)
         {
-            return method;
+            if (stmt is SuppressionStatement { Inner: not null } suppressionStatement)
+            {
+                // TO-DO: remove once enumerable logic is updated to handle nested suppression statements
+                flattenedStatements.Add(suppressionStatement.DisableStatement);
+                flattenedStatements.AddRange(suppressionStatement.Inner);
+                flattenedStatements.Add(suppressionStatement.RestoreStatement);
+            }
+            else
+            {
+                flattenedStatements.Add(stmt);
+            }
         }
 
-        var updatedStatements = new List<MethodBodyStatement>();
-        var flattenedStatements = statements.ToArray();
-
         List<WritePropertyNameAdditionalReplacementInfo> additionalConditionsForWritingType
             = TypeNameToWritePropertyNameAdditionalConditionMap.GetValueOrDefault(method.EnclosingType.Name) ?? [];
 
-        for (int line = 0; line < flattenedStatements.Length; line++)
+        for (int line = 0; line < flattenedStatements.Count; line++)
         {
             var statement = flattenedStatements[line];
 
             // Much of the customization centers around treatment of WritePropertyName
             string? writePropertyNameTarget = GetWritePropertyNameTargetFromStatement(statement);
 
-            if (statement is IfStatement ifStatement)
+            switch (statement)
             {
-                // If we already have an if statement that contains property writing, we need to add the condition to the existing if statement
-                if (writePropertyNameTarget is not null)
-                {
-                    ifStatement.Update(condition: ifStatement.Condition.As<bool>().And(GetContainsKeyCondition(writePropertyNameTarget)));
-                }
+                // If we already have an if statement that contains property writing, we need to add the condition to the existing if statement.
+                // For dynamic models, we can skip adding the SARD condition.
+                case IfStatement ifStatement:
+                    ProcessIfStatement(ifStatement, writePropertyNameTarget, additionalConditionsForWritingType, updatedStatements);
+                    break;
+                case IfElseStatement ifElseStatement when GetPatchContainsExpression(ifElseStatement.If.Condition) != null:
+                    ProcessIfElseStatement(ifElseStatement, writePropertyNameTarget, additionalConditionsForWritingType, updatedStatements);
+                    break;
+                case var _ when writePropertyNameTarget is not null:
+                    line = ProcessWritePropertyNameStatement(statement, writePropertyNameTarget, additionalConditionsForWritingType, flattenedStatements, line, updatedStatements);
+                    break;
+                default:
+                    updatedStatements.Add(statement);
+                    break;
+            }
+        }
 
-                // Handle writing AdditionalProperties
-                else if (ifStatement.Body.First() is ForEachStatement foreachStatement)
-                {
-                    foreachStatement.Body.Insert(
-                        0,
-                        new IfStatement(
-                            Static(new ModelSerializationExtensionsDefinition().Type).Invoke(
-                                IsSentinelValueMethodName,
-                                foreachStatement.ItemVariable.Property("Value")))
-                        {
-                            Continue
-                        });
-                }
+        method.Update(bodyStatements: updatedStatements);
+        return method;
+    }
+
+    private static void ProcessIfStatement(
+        IfStatement ifStatement,
+        string? writePropertyNameTarget,
+        List<WritePropertyNameAdditionalReplacementInfo> additionalConditionsForWritingType,
+        List<MethodBodyStatement> updatedStatements)
+    {
+        if (writePropertyNameTarget is not null)
+        {
+            ValueExpression? patchContainsCondition = GetPatchContainsExpression(ifStatement.Condition);
 
-                updatedStatements.Add(ifStatement);
+            if (patchContainsCondition is null)
+            {
+                ifStatement.Update(condition: ifStatement.Condition.As<bool>().And(GetContainsKeyCondition(writePropertyNameTarget)));
             }
-            else if (writePropertyNameTarget is not null)
+            else if (additionalConditionsForWritingType.FirstOrDefault(additionalCondition => additionalCondition.JsonName == writePropertyNameTarget) is var matchingReplacementInfo && matchingReplacementInfo != null)
             {
-                ScopedApi<bool> enclosingIfCondition = GetContainsKeyCondition(writePropertyNameTarget);
-
-                if (additionalConditionsForWritingType
-                    .FirstOrDefault(additionalCondition => additionalCondition.JsonName == writePropertyNameTarget)
-                    is WritePropertyNameAdditionalReplacementInfo matchingReplacementInfo)
+                updatedStatements.Add(OptionalDefinedCheckComment);
+                ifStatement.Update(condition: GetOptionalIsCollectionDefinedCondition(matchingReplacementInfo).And(ifStatement.Condition));
+            }
+        }
+        // Handle writing AdditionalProperties
+        else if (ifStatement.Body.First() is ForEachStatement foreachStatement)
+        {
+            foreachStatement.Body.Insert(
+                0,
+                new IfStatement(
+                    Static(new ModelSerializationExtensionsDefinition().Type).Invoke(
+                        IsSentinelValueMethodName,
+                        foreachStatement.ItemVariable.Property("Value")))
                 {
-                    MethodBodyStatement commentStatement
-                        = new SingleLineCommentStatement("Plugin customization: apply Optional.Is*Defined() check based on type name dictionary lookup");
-                    updatedStatements.Add(commentStatement);
-                    enclosingIfCondition = GetOptionalIsCollectionDefinedCondition(matchingReplacementInfo)
-                        .And(enclosingIfCondition);
-                }
+                    Continue
+                });
+        }
 
-                var ifSt = new IfStatement(enclosingIfCondition) { statement };
+        updatedStatements.Add(ifStatement);
+    }
 
-                // If this is a plain expression statement, we need to add the next statement as well which
-                // will either write the property value or start writing an array
-                if (statement is ExpressionStatement)
-                {
-                    ifSt.Add(flattenedStatements[++line]);
-                    // Include array writing in the if statement
-                    if (flattenedStatements[line + 1] is ForEachStatement)
-                    {
-                        // Foreach
-                        ifSt.Add(flattenedStatements[++line]);
-                        // End array
-                        ifSt.Add(flattenedStatements[++line]);
-                    }
-                }
-                updatedStatements.Add(ifSt);
-            }
-            else
+    private static void ProcessIfElseStatement(
+        IfElseStatement ifElseStatement,
+        string? writePropertyNameTarget,
+        List<WritePropertyNameAdditionalReplacementInfo> additionalConditionsForWritingType,
+        List<MethodBodyStatement> updatedStatements)
+    {
+        if (ifElseStatement.Else is null)
+        {
+            updatedStatements.Add(ifElseStatement);
+            return;
+        }
+
+        if (additionalConditionsForWritingType.FirstOrDefault(additionalCondition => additionalCondition.JsonName == writePropertyNameTarget) is var matchingReplacementInfo && matchingReplacementInfo != null)
+        {
+            var enclosingCondition = GetOptionalIsCollectionDefinedCondition(matchingReplacementInfo);
+            var updatedCondition = new IfStatement(enclosingCondition) { ifElseStatement.Else };
+
+            ifElseStatement.Update(elseStatement: new MethodBodyStatements([OptionalDefinedCheckComment, updatedCondition]));
+        }
+
+        updatedStatements.Add(ifElseStatement);
+    }
+
+    private static int ProcessWritePropertyNameStatement(
+        MethodBodyStatement statement,
+        string writePropertyNameTarget,
+        List<WritePropertyNameAdditionalReplacementInfo> additionalConditionsForWritingType,
+        List<MethodBodyStatement> flattenedStatements,
+        int currentLine,
+        List<MethodBodyStatement> updatedStatements)
+    {
+        var line = currentLine;
+        ScopedApi<bool> enclosingIfCondition = GetContainsKeyCondition(writePropertyNameTarget);
+
+        if (additionalConditionsForWritingType.FirstOrDefault(additionalCondition => additionalCondition.JsonName == writePropertyNameTarget) is var matchingReplacementInfo && matchingReplacementInfo != null)
+        {
+            updatedStatements.Add(OptionalDefinedCheckComment);
+            enclosingIfCondition = GetOptionalIsCollectionDefinedCondition(matchingReplacementInfo).And(enclosingIfCondition);
+        }
+
+        var ifSt = new IfStatement(enclosingIfCondition) { statement };
+
+        // If this is a plain expression statement, we need to add the next statement as well which
+        // will either write the property value or start writing an array
+        if (statement is ExpressionStatement)
+        {
+            ifSt.Add(flattenedStatements[++line]);
+            // Include array writing in the if statement
+            if (flattenedStatements[line + 1] is ForEachStatement)
             {
-                updatedStatements.Add(statement);
+                // Foreach
+                ifSt.Add(flattenedStatements[++line]);
+                // End array
+                ifSt.Add(flattenedStatements[++line]);
             }
         }
-        
-        method.Update(bodyStatements: updatedStatements);
-        return method;
+
+        updatedStatements.Add(ifSt);
+        return line;
     }
 
     private static ScopedApi<bool> GetContainsKeyCondition(string propertyName)
@@ -235,6 +298,10 @@ private static ScopedApi<bool> GetContainsKeyCondition(string propertyName)
         {
             return stringLiteralExpression.Literal?.ToString();
         }
+        if (statement is SuppressionStatement suppressionStatement)
+        {
+            return GetWritePropertyNameTargetFromStatement(suppressionStatement.Inner);
+        }
         else if (statement is MethodBodyStatements compoundStatements)
         {
             foreach (MethodBodyStatement innerStatement in compoundStatements.Statements)
@@ -270,4 +337,45 @@ public class WritePropertyNameAdditionalReplacementInfo(string propertyName, str
         public string JsonName { get; set; } = jsonName;
         public bool IsCollection { get; set; } = isCollection;
     }
+
+
+    /// <summary>
+    /// Recursively checks if the given expression or any of its sub-expressions is a call to Patch.Contains().
+    /// Handles various wrapping scenarios including unary operators, binary operators, and nested expressions.
+    /// </summary>
+    private static ValueExpression? GetPatchContainsExpression(ValueExpression? expression)
+    {
+        if (expression is null)
+        {
+            return null;
+        }
+
+#pragma warning disable SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+        return expression switch
+        {
+            // Case 1: Direct Patch.Contains() call
+            ScopedApi<bool> { Original: InvokeMethodExpression { InstanceReference: ScopedApi<JsonPatch> } } => expression,
+
+            // Case 2: !Patch.Contains() call
+            ScopedApi<bool> { Original: UnaryOperatorExpression { Operator: "!", Operand: ScopedApi<bool> { Original: InvokeMethodExpression { InstanceReference: ScopedApi<JsonPatch> } } } } => expression,
+
+            // Case 3 & 4: Binary operator expression (wrapped or unwrapped)
+            ScopedApi<bool> { Original: BinaryOperatorExpression binaryExpr } =>
+                GetPatchContainsExpression(binaryExpr.Left) ?? GetPatchContainsExpression(binaryExpr.Right),
+
+            BinaryOperatorExpression binaryExpr =>
+                GetPatchContainsExpression(binaryExpr.Left) ?? GetPatchContainsExpression(binaryExpr.Right),
+
+            // Case 5: Direct UnaryOperatorExpression (not wrapped in ScopedApi)
+            UnaryOperatorExpression { Operator: "!" } unaryExpr =>
+                GetPatchContainsExpression(unaryExpr.Operand) != null ? expression : null,
+
+            // Case 6: Direct InvokeMethodExpression (not wrapped in ScopedApi)
+            InvokeMethodExpression { InstanceReference: ScopedApi<JsonPatch> } => expression,
+
+            _ => null
+        };
+
+#pragma warning restore SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+    }
 }

codegen/generator/src/Visitors/VisitorHelpers.cs

@@ -27,6 +27,17 @@ List<MethodBodyStatement> foreachBodyStatements
                 foreachStatement.Body.Clear();
                 foreachStatement.Body.Add(new MethodBodyStatements(foreachBodyStatements));
             }
+            else if (statements[i] is SuppressionStatement suppressionStatement
+                && suppressionStatement.Inner != null)
+            {
+                List<MethodBodyStatement> suppressionInnerStatement = [.. suppressionStatement.Inner.SelectMany(bodyStatement => bodyStatement)];
+                VisitExplodedMethodBodyStatements(suppressionInnerStatement!, visitorFunc);
+                var updatedSuppressionStatement = new SuppressionStatement(
+                    suppressionInnerStatement,
+                    suppressionStatement.Code,
+                    suppressionStatement.Justification);
+                statements[i] = updatedSuppressionStatement;
+            }
             else if (statements[i] is IfStatement ifStatement)
             {
                 List<MethodBodyStatement> ifBodyStatements

src/Generated/Models/Chat/ChatCompletionOptions.Serialization.cs

@@ -102,10 +102,15 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
             }
             else
             {
-                writer.WritePropertyName("messages"u8);
-                SerializeMessagesValue(writer, options);
+                // Plugin customization: apply Optional.Is*Defined() check based on type name dictionary lookup
+                if (Optional.IsCollectionDefined(Messages))
+                {
+                    writer.WritePropertyName("messages"u8);
+                    SerializeMessagesValue(writer, options);
+                }
             }
-            if (!Patch.Contains("$.model"u8))
+            // Plugin customization: apply Optional.Is*Defined() check based on type name dictionary lookup
+            if (Optional.IsDefined(Model) && !Patch.Contains("$.model"u8))
             {
                 writer.WritePropertyName("model"u8);
                 writer.WriteStringValue(Model);

src/Generated/Models/Chat/ChatMessage.Serialization.cs

@@ -31,13 +31,12 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
                 writer.WritePropertyName("role"u8);
                 writer.WriteStringValue(Role.ToSerialString());
             }
-            if (Optional.IsDefined(Content) && !Patch.Contains("$.content"u8))
+            // Plugin customization: add Content.IsInnerCollectionDefined() check
+            if (Optional.IsDefined(Content) && Content.IsInnerCollectionDefined() && !Patch.Contains("$.content"u8))
             {
                 writer.WritePropertyName("content"u8);
                 SerializeContentValue(writer, options);
             }
-
-            Patch.WriteTo(writer);
 #pragma warning restore SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates.
         }
 

src/Generated/Models/Chat/ChatMessageContentPart.Serialization.cs

@@ -21,10 +21,6 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
             {
                 throw new FormatException($"The model {nameof(ChatMessageContentPart)} does not support writing '{format}' format.");
             }
-#pragma warning disable SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates.
-
-            Patch.WriteTo(writer);
-#pragma warning restore SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates.
         }
 
         ChatMessageContentPart IJsonModel<ChatMessageContentPart>.Create(ref Utf8JsonReader reader, ModelReaderWriterOptions options) => JsonModelCreateCore(ref reader, options);

src/Generated/Models/Chat/ChatOutputPrediction.Serialization.cs

@@ -44,8 +44,6 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
                 writer.WritePropertyName("type"u8);
                 writer.WriteStringValue(Kind.ToString());
             }
-
-            Patch.WriteTo(writer);
 #pragma warning restore SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates.
         }
 

src/Generated/Models/Chat/ChatResponseFormat.Serialization.cs

@@ -31,8 +31,6 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
                 writer.WritePropertyName("type"u8);
                 writer.WriteStringValue(Kind.ToString());
             }
-
-            Patch.WriteTo(writer);
 #pragma warning restore SCME0001 // Type is for evaluation purposes only and is subject to change or removal in future updates.
         }
 

src/Generated/Models/Chat/InternalChatCompletionStreamResponseDelta.Serialization.cs

@@ -41,7 +41,8 @@ protected virtual void JsonModelWriteCore(Utf8JsonWriter writer, ModelReaderWrit
                 writer.WritePropertyName("audio"u8);
                 writer.WriteObjectValue(Audio, options);
             }
-            if (Optional.IsDefined(Content) && !Patch.Contains("$.content"u8))
+            // Plugin customization: add Content.IsInnerCollectionDefined() check
+            if (Optional.IsDefined(Content) && Content.IsInnerCollectionDefined() && !Patch.Contains("$.content"u8))
             {
                 writer.WritePropertyName("content"u8);
                 SerializeContentValue(writer, options);