Using a SymbolVisitor implementation to locate function methods from referenced assemblies

sdk/Sdk.Generators/FunctionMetadataProviderGenerator/FunctionMetadataProviderGenerator.cs

@@ -93,55 +93,10 @@ private IEnumerable<IMethodSymbol> GetEntryAssemblyFunctions(List<MethodDeclarat
         /// </summary>
         private IEnumerable<IMethodSymbol> GetDependentAssemblyFunctions(GeneratorExecutionContext context)
         {
-            foreach (var assembly in context.Compilation.SourceModule.ReferencedAssemblySymbols)
-            {
-                foreach (var methodSymbol in GetMethodsFromNamespace(context.Compilation, assembly.GlobalNamespace))
-                {
-                    yield return methodSymbol;
-                }
-            }
-        }
+            var visitor = new ReferencedAssemblyMethodVisitor(context.Compilation);
+            visitor.Visit(context.Compilation.SourceModule);
 
-        private IEnumerable<IMethodSymbol> GetMethodsFromNamespace(Compilation compilation, INamespaceSymbol namespaceSymbol)
-        {
-            foreach (var member in namespaceSymbol.GetMembers())
-            {
-                if (member is INamespaceSymbol nestedNamespace)
-                {
-                    // Recursive call for nested namespaces
-                    foreach (var methodSymbol in GetMethodsFromNamespace(compilation, nestedNamespace))
-                    {
-                        yield return methodSymbol;
-                    }
-                }
-                else if (member is INamedTypeSymbol typeSymbol)
-                {
-                    // Recursive call for nested types
-                    foreach (var methodSymbol in GetMethodsFromType(compilation, typeSymbol))
-                    {
-                        yield return methodSymbol;
-                    }
-                }
-            }
-        }
-
-        private IEnumerable<IMethodSymbol> GetMethodsFromType(Compilation compilation, INamedTypeSymbol typeSymbol)
-        {
-            foreach (var member in typeSymbol.GetMembers())
-            {
-                if (member is IMethodSymbol methodSymbol && FunctionsUtil.IsFunctionSymbol(methodSymbol, compilation))
-                {
-                    yield return methodSymbol;
-                }
-                else if (member is INamedTypeSymbol nestedType)
-                {
-                    // Recursive call for nested types
-                    foreach (var nestedMethodSymbol in GetMethodsFromType(compilation, nestedType))
-                    {
-                        yield return nestedMethodSymbol;
-                    }
-                }
-            }
+            return visitor.FunctionMethods;
         }
     }
 }

sdk/Sdk.Generators/FunctionMetadataProviderGenerator/ReferencedAssemblyMethodVisitor.cs

@@ -0,0 +1,75 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.CodeAnalysis;
+
+namespace Microsoft.Azure.Functions.Worker.Sdk.Generators
+{
+    /// <summary>
+    /// Visits all symbols from referenced assemblies and returns all methods which are valid Azure Functions.
+    /// </summary>
+    internal sealed class ReferencedAssemblyMethodVisitor : SymbolVisitor
+    {
+        private readonly Compilation _compilation;
+
+        /// <summary>
+        /// Gets all methods which are valid Azure Functions.
+        /// </summary>
+        internal readonly List<IMethodSymbol> FunctionMethods = new();
+
+        internal ReferencedAssemblyMethodVisitor(Compilation compilation)
+        {
+            _compilation = compilation ?? throw new ArgumentNullException(nameof(compilation));
+        }
+
+        public override void VisitModule(IModuleSymbol moduleSymbol)
+        {
+            foreach (var assemblySymbol in moduleSymbol.ReferencedAssemblySymbols)
+            {
+                assemblySymbol.Accept(this);
+            }
+        }
+
+        public override void VisitAssembly(IAssemblySymbol symbol)
+        {
+            var namespaceSymbol = symbol.GlobalNamespace;
+            namespaceSymbol.Accept(this);
+        }
+
+        public override void VisitNamespace(INamespaceSymbol symbol)
+        {
+            // Get classes in this namespace or child namespaces
+            var classesOrNamespaces = symbol.GetMembers()
+                .Where(a => a.Kind is SymbolKind.Namespace or SymbolKind.NamedType);
+
+            foreach (var childSymbol in classesOrNamespaces)
+            {
+                childSymbol.Accept(this);
+            }
+        }
+
+        public override void VisitNamedType(INamedTypeSymbol symbol)
+        {
+            // Get methods in this class or nested child classes
+            var methodsOrClasses = symbol.GetMembers()
+                .Where(a => a.Kind is SymbolKind.NamedType or SymbolKind.Method);
+
+            foreach (var childSymbol in methodsOrClasses)
+            {
+                childSymbol.Accept(this);
+            }
+        }
+
+        public override void VisitMethod(IMethodSymbol methodSymbol)
+        {
+            if (methodSymbol.MethodKind == MethodKind.Ordinary &&
+                FunctionsUtil.IsFunctionSymbol(methodSymbol, _compilation))
+            {
+                FunctionMethods.Add(methodSymbol);
+            }
+        }
+    }
+}