Filter custom members after all visitors have been visited

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/CSharpGen.cs

@@ -67,6 +67,17 @@ await customCodeWorkspace.GetCompilationAsync(),
                 visitor.VisitLibrary(output);
             }
 
+            foreach (var typeProvider in output.TypeProviders)
+            {
+                // Update the type with the potentially modified members, filtering out customized members
+                // after the visitors have been applied so that the filtering is done against the final version.
+                typeProvider.Update(
+                    typeProvider.FilterCustomizedMethods(typeProvider.Methods),
+                    typeProvider.FilterCustomizedConstructors(typeProvider.Constructors),
+                    typeProvider.FilterCustomizedProperties(typeProvider.Properties),
+                    typeProvider.FilterCustomizedFields(typeProvider.Fields));
+            }
+
             LoggingHelpers.LogElapsedTime("All visitors have been applied");
 
             foreach (var outputType in output.TypeProviders)

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/LibraryVisitor.cs

@@ -94,13 +94,11 @@ protected internal virtual void VisitLibrary(OutputLibrary library)
                     }
                 }
 
-                // Update the type with the potentially modified members, filtering out customized members
-                // after the visitors have been applied so that the filtering is done against the final version.
                 type.Update(
-                    type.FilterCustomizedMethods(methods),
-                    type.FilterCustomizedConstructors(constructors),
-                    type.FilterCustomizedProperties(properties),
-                    type.FilterCustomizedFields(fields),
+                    methods,
+                    constructors,
+                    properties,
+                    fields,
                     serializations,
                     nestedTypes);
                 type = PostVisitType(type);

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/OutputLibraryVisitorTests.cs

@@ -220,6 +220,8 @@ public async Task MatchingMethodSignatureIsFilteredAfterVisitorMutation()
             var visitor = new TestFilterVisitor();
             visitor.VisitLibrary(generator.Object.OutputLibrary);
 
+            typeProvider.Update(methods: typeProvider.FilterCustomizedMethods(typeProvider.Methods));
+
             Assert.AreEqual(0, typeProvider.Methods.Count);
         }
 
@@ -239,6 +241,8 @@ public async Task MatchingConstructorSignatureIsFilteredAfterVisitorMutation()
             var visitor = new TestFilterVisitor();
             visitor.VisitLibrary(generator.Object.OutputLibrary);
 
+            typeProvider.Update(constructors: typeProvider.FilterCustomizedConstructors(typeProvider.Constructors));
+
             Assert.AreEqual(0, typeProvider.Constructors.Count);
         }
 
@@ -257,6 +261,8 @@ public async Task MatchingPropertyIsFilteredAfterVisitorMutation()
             var visitor = new TestFilterVisitor();
             visitor.VisitLibrary(generator.Object.OutputLibrary);
 
+            typeProvider.Update(properties: typeProvider.FilterCustomizedProperties(typeProvider.Properties));
+
             Assert.AreEqual(0, typeProvider.Properties.Count);
         }
 
@@ -274,9 +280,77 @@ public async Task MatchingFieldIsFilteredAfterVisitorMutation()
             var visitor = new TestFilterVisitor();
             visitor.VisitLibrary(generator.Object.OutputLibrary);
 
+            typeProvider.Update(fields: typeProvider.FilterCustomizedFields(typeProvider.Fields));
+
             Assert.AreEqual(0, typeProvider.Fields.Count);
         }
 
+        [Test]
+        public async Task MultipleVisitorsMutateMember()
+        {
+            var typeProvider = new TestTypeProvider();
+            var methodProvider = new MethodProvider(
+                new MethodSignature("OriginalMethod", $"", MethodSignatureModifiers.Public, null, $"", [new ParameterProvider("param1", $"", typeof(float))]),
+                Snippet.Throw(Snippet.Null), typeProvider);
+            typeProvider.Update(methods: [methodProvider], reset: true);
+
+            var generator = await MockHelpers.LoadMockGeneratorAsync(
+                createOutputLibrary: () => new TestOutputLibrary(typeProvider),
+                compilation: async () => await Helpers.GetCompilationFromDirectoryAsync());
+
+            var visitor1 = new RenameMethodVisitor("OriginalMethod", "TestMethod");
+            var visitor2 = new ChangeParameterTypeVisitor("TestMethod", typeof(int));
+
+            visitor1.VisitLibrary(generator.Object.OutputLibrary);
+            visitor2.VisitLibrary(generator.Object.OutputLibrary);
+
+            typeProvider.Update(methods: typeProvider.FilterCustomizedMethods(typeProvider.Methods));
+
+            Assert.AreEqual(0, typeProvider.Methods.Count);
+        }
+
+        private class RenameMethodVisitor : LibraryVisitor
+        {
+            private readonly string _originalName;
+            private readonly string _newName;
+
+            public RenameMethodVisitor(string originalName, string newName)
+            {
+                _originalName = originalName;
+                _newName = newName;
+            }
+
+            protected internal override MethodProvider? VisitMethod(MethodProvider method)
+            {
+                if (method.Signature.Name == _originalName)
+                {
+                    method.Signature.Update(name: _newName);
+                }
+                return method; // Return method even if not renamed, to allow further visiting (though logic here seems to not matter much as LibraryVisitor base is used or visitor logic)
+            }
+        }
+
+        private class ChangeParameterTypeVisitor : LibraryVisitor
+        {
+            private readonly string _methodName;
+            private readonly System.Type _newType;
+
+            public ChangeParameterTypeVisitor(string methodName, System.Type newType)
+            {
+                _methodName = methodName;
+                _newType = newType;
+            }
+
+            protected internal override MethodProvider? VisitMethod(MethodProvider method)
+            {
+                if (method.Signature.Name == _methodName)
+                {
+                    method.Signature.Parameters[0].Update(type: _newType);
+                }
+                return method;
+            }
+        }
+
 
         private class TestFilterVisitor : LibraryVisitor
         {

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/TestData/OutputLibraryVisitorTests/MultipleVisitorsMutateMember/TestName.cs

@@ -0,0 +1,16 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Test
+{
+    /// <summary>
+    /// This is used to verify method replacement when a method signature is changed in a visitor to match the below custom definition.
+    /// </summary>
+    public partial class TestName
+    {
+        public void TestMethod(int param1) { }
+    }
+}