terra-llm/Assets/IntelligentProjectAnalyzer/Analyzer/RoslynTypeDependencyAnalyzer.cs

using System.IO;
using System.Linq;
using UnityEngine;
using System.Collections.Generic;
using System.Collections.Immutable;

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace IntelligentProjectAnalyzer.Analyzer
{
    public class RoslynTypeDependencyAnalyzer
    {
        private readonly CSharpCompilation _compilation;
        private readonly HashSet<string> _visitedTypes = new();
        private readonly Dictionary<string, List<string>> _dependencyTree = new();
        private readonly List<string> _systemTypePrefixes;
        
        /// <summary>
        /// A map that stores the file path for each type name.
        /// This is built once during initialization for efficient lookups.
        /// </summary>
        public IReadOnlyDictionary<string, string> TypeToPathMap { get; private set; }
        
        /// <summary>
        /// A thread-safe constructor that accepts pre-fetched system type prefixes.
        /// </summary>
        public RoslynTypeDependencyAnalyzer(string[] sourceFiles, string[] references, string[] preprocessorSymbols, List<string> systemTypePrefixes)
        {
            _systemTypePrefixes = systemTypePrefixes; // Use the provided prefixes

            // The rest of the constructor logic remains the same
            var parseOptions = new CSharpParseOptions(
                languageVersion: LanguageVersion.Default,
                preprocessorSymbols: preprocessorSymbols ?? Enumerable.Empty<string>()
            );
    
            var syntaxTrees = new List<SyntaxTree>();
            GenerateSyntaxTrees(sourceFiles, syntaxTrees, parseOptions);

            var metadataReferences = new List<MetadataReference>();
            GenerateMetaDataReferences(references, metadataReferences);

            _compilation = CSharpCompilation.Create(
                "AnalysisAssembly",
                syntaxTrees,
                metadataReferences,
                new CSharpCompilationOptions(OutputKind.DynamicallyLinkedLibrary));
        
            BuildTypeToPathMap();
        }

        private static void GenerateSyntaxTrees(string[] sourceFiles, List<SyntaxTree> syntaxTrees, CSharpParseOptions parseOptions)
        {
            foreach (var file in sourceFiles)
            {
                try
                {
                    var sourceText = File.ReadAllText(file);
                    syntaxTrees.Add(CSharpSyntaxTree.ParseText(sourceText, options: parseOptions, path: file));
                }
                catch (IOException ex)
                {
                    // Gracefully handle cases where a source file cannot be read.
                    Debug.LogError($"Error reading source file {file}: {ex.Message}");
                }
                catch (System.Exception ex)
                {
                    Debug.LogError($"An unexpected error occurred while processing file {file}: {ex.Message}");
                }
            }
        }

        private static void GenerateMetaDataReferences(string[] references, List<MetadataReference> metadataReferences)
        {
            foreach (var reference in references)
            {
                try
                {
                    metadataReferences.Add(MetadataReference.CreateFromFile(reference));
                }
                catch (IOException ex)
                {
                    // Gracefully handle cases where a reference DLL cannot be found or read.
                    Debug.LogError($"Error loading reference assembly {reference}: {ex.Message}");
                }
                catch (System.Exception ex)
                {
                    Debug.LogError($"An unexpected error occurred while loading reference {reference}: {ex.Message}");
                }
            }
        }
        
        /// <summary>
        /// Iterates through all syntax trees to build a map of every type to the file path where it is declared.
        /// </summary>
        private void BuildTypeToPathMap()
        {
            var map = new Dictionary<string, string>();
            // Parallelizing this can speed up the process on projects with many scripts.
            foreach (var syntaxTree in _compilation.SyntaxTrees)
            {
                var semanticModel = _compilation.GetSemanticModel(syntaxTree);
                // We look for any declaration that can define a type (class, struct, interface, enum).
                var typeDeclarations = syntaxTree.GetRoot().DescendantNodes().OfType<BaseTypeDeclarationSyntax>();

                foreach (var typeDecl in typeDeclarations)
                {
                    if (semanticModel.GetDeclaredSymbol(typeDecl) is { } symbol)
                    {
                        // symbol.ToDisplayString() provides the fully qualified name (e.g., MyNamespace.MyClass),
                        // which is a reliable key. We map this to the file it's in.
                        map[symbol.ToDisplayString()] = syntaxTree.FilePath;
                    }
                }
            }
            TypeToPathMap = map;
        }

        /// <summary>
        /// Analyze dependencies for multiple types, populating the internal dependency tree.
        /// This is the method that should be called from DependencyStoreBuilder.
        /// </summary>
        public void AnalyzeTypeDependenciesForFile(IEnumerable<string> typeNames)
        {
            foreach (var typeName in typeNames)
            {
                var typeSymbol = FindTypeSymbol(typeName) ?? _compilation.GetTypeByMetadataName(typeName);
                if (typeSymbol != null)
                {
                    AnalyzeTypeSymbol(typeSymbol);
                }
            }
        }

        /// <summary>
        /// Traverses the populated dependency tree to get all recursive dependencies for a type.
        /// </summary>
        /// <param name="typeName">The fully qualified name of the type.</param>
        /// <returns>A HashSet containing all unique, recursive dependency type names.</returns>
        public HashSet<string> GetRecursiveDependencies(string typeName)
        {
            var allDependencies = new HashSet<string>();
            var toProcess = new Queue<string>();
            var processed = new HashSet<string>();

            toProcess.Enqueue(typeName);

            while (toProcess.Count > 0)
            {
                var currentType = toProcess.Dequeue();
                if (!processed.Add(currentType))
                {
                    continue;
                }

                if (_dependencyTree.TryGetValue(currentType, out var directDependencies))
                {
                    foreach (var dep in directDependencies)
                    {
                        if (allDependencies.Add(dep)) // Add to the final set and queue for processing if it's a new dependency
                        {
                            toProcess.Enqueue(dep);
                        }
                    }
                }
            }
            return allDependencies;
        }


        /// <summary>
        /// The main recursive method that builds the dependency list for a given type symbol.
        /// </summary>
        private void AnalyzeTypeSymbol(INamedTypeSymbol typeSymbol)
        {
            var typeName = typeSymbol.ToDisplayString();

            if (!_visitedTypes.Add(typeName))
                return;

            var dependencies = new List<string>();

            // Analyze base type
            if (typeSymbol.BaseType != null && !IsSystemType(typeSymbol.BaseType))
            {
                var baseTypeName = typeSymbol.BaseType.ToDisplayString();
                dependencies.Add(baseTypeName);
                AnalyzeTypeSymbol(typeSymbol.BaseType);
            }

            // Analyze interfaces
            foreach (var interfaceSymbol in typeSymbol.Interfaces)
            {
                if (IsSystemType(interfaceSymbol)) continue;
                var interfaceName = interfaceSymbol.ToDisplayString();
                dependencies.Add(interfaceName);
                AnalyzeTypeSymbol(interfaceSymbol);
            }

            // Analyze members
            foreach (var member in typeSymbol.GetMembers())
            {
                switch (member)
                {
                    case IFieldSymbol field:
                        AnalyzeMemberType(field.Type, dependencies);
                        AnalyzeAttributes(field.GetAttributes(), dependencies);
                        break;
                    case IPropertySymbol property:
                        AnalyzeMemberType(property.Type, dependencies);
                        AnalyzeAttributes(property.GetAttributes(), dependencies);
                        break;
                    case IMethodSymbol method:
                        AnalyzeMemberType(method.ReturnType, dependencies);
                        foreach (var parameter in method.Parameters)
                        {
                            AnalyzeMemberType(parameter.Type, dependencies);
                        }
                        AnalyzeAttributes(method.GetAttributes(), dependencies);
                        AnalyzeMethodBody(method, dependencies);
                        break;
                }
            }

            _dependencyTree[typeName] = dependencies.Distinct().ToList();
        }
        
        /// <summary>
        /// Helper method to analyze the type of a member (field, property, parameter, etc.).
        /// </summary>
        private void AnalyzeMemberType(ITypeSymbol typeSymbol, List<string> dependencies)
        {
            while (true)
            {
                if (typeSymbol == null || IsSystemType(typeSymbol)) return;

                var typeName = typeSymbol.ToDisplayString();
                dependencies.Add(typeName);

                switch (typeSymbol)
                {
                    case INamedTypeSymbol namedType:
                        if (namedType.IsGenericType)
                        {
                            foreach (var typeArg in namedType.TypeArguments)
                            {
                                AnalyzeMemberType(typeArg, dependencies);
                            }
                        }
                        // FIX: Recursively analyze the type itself to add its dependencies to the main tree
                        AnalyzeTypeSymbol(namedType); 
                        break;
                    case IArrayTypeSymbol arrayType:
                        typeSymbol = arrayType.ElementType;
                        continue; // Continue loop with the element type
                }
                break; // Exit loop
            }
        }

        private INamedTypeSymbol FindTypeSymbol(string typeName)
        {
            // First, try to get the type by its fully qualified metadata name.
            var symbol = _compilation.GetTypeByMetadataName(typeName);
            if (symbol != null) return symbol;

            // If that fails, search through all types in the global namespace.
            // This is less efficient and might be ambiguous, but serves as a fallback.
            return _compilation.GlobalNamespace.GetNamespaceMembers()
                .SelectMany(n => n.GetTypeMembers())
                .FirstOrDefault(t => t.Name == typeName);
        }

        private void AnalyzeAttributes(ImmutableArray<AttributeData> attributes, List<string> dependencies)
        {
            foreach (var attribute in attributes)
            {
                if (attribute.AttributeClass != null && !IsSystemType(attribute.AttributeClass))
                {
                    var attributeName = attribute.AttributeClass.ToDisplayString();
                    dependencies.Add(attributeName);
                    AnalyzeTypeSymbol(attribute.AttributeClass);
                }

                // Analyze attribute constructor arguments
                foreach (var arg in attribute.ConstructorArguments)
                {
                    AnalyzeAttributeArgument(arg, dependencies);
                }

                // Analyze named arguments
                foreach (var namedArg in attribute.NamedArguments)
                {
                    AnalyzeAttributeArgument(namedArg.Value, dependencies);
                }
            }
        }

        private void AnalyzeAttributeArgument(TypedConstant argument, List<string> dependencies)
        {
            if (argument.Kind == TypedConstantKind.Error) return;
            
            if (argument.Type != null && !IsSystemType(argument.Type))
            {
                var typeName = argument.Type.ToDisplayString();
                dependencies.Add(typeName);

                if (argument.Type is INamedTypeSymbol namedType)
                {
                    AnalyzeTypeSymbol(namedType);
                }
            }

            // Handle array arguments
            if (argument.Kind != TypedConstantKind.Array) return;
            foreach (var element in argument.Values)
            {
                AnalyzeAttributeArgument(element, dependencies);
            }
        }

        private void AnalyzeMethodBody(IMethodSymbol method, List<string> dependencies)
        {
            // Get syntax references for the method
            foreach (var syntaxNode in method.DeclaringSyntaxReferences.Select(syntaxRef => syntaxRef.GetSyntax()))
            {
                AnalyzeMethodSyntax(syntaxNode, dependencies);
            }
        }

        private void AnalyzeMethodSyntax(SyntaxNode syntaxNode, List<string> dependencies)
        {
            if (syntaxNode == null) return;

            var walker = new ReflectionDependencyWalker(_compilation.GetSemanticModel(syntaxNode.SyntaxTree), dependencies, this);
            walker.Visit(syntaxNode);
        }

        private bool IsSystemType(ITypeSymbol typeSymbol)
        {
            if (typeSymbol == null) return true;
            var namespaceName = typeSymbol.ContainingNamespace?.ToDisplayString() ?? "";
            return _systemTypePrefixes.Any(prefix => namespaceName.StartsWith(prefix));
        }

        // Inner class to walk a syntax tree and find reflection dependencies
        private class ReflectionDependencyWalker : CSharpSyntaxWalker
        {
            private readonly SemanticModel _semanticModel;
            private readonly List<string> _dependencies;
            private readonly RoslynTypeDependencyAnalyzer _analyzer;

            public ReflectionDependencyWalker(SemanticModel semanticModel, List<string> dependencies,
                RoslynTypeDependencyAnalyzer analyzer)
            {
                _semanticModel = semanticModel;
                _dependencies = dependencies;
                _analyzer = analyzer;
            }

            public override void VisitInvocationExpression(InvocationExpressionSyntax node)
            {
                var symbolInfo = _semanticModel.GetSymbolInfo(node.Expression);
                if (symbolInfo.Symbol is IMethodSymbol methodSymbol)
                {
                    // Check if the method is a known reflection-based method that accepts string literals as type names.
                    if (IsReflectionMethod(methodSymbol.Name))
                    {
                        AnalyzeReflectionCall(node);
                    }
                }

                base.VisitInvocationExpression(node);
            }

            public override void VisitObjectCreationExpression(ObjectCreationExpressionSyntax node)
            {
                var typeInfo = _semanticModel.GetTypeInfo(node.Type);
                if (typeInfo.Type != null)
                {
                    _analyzer.AnalyzeMemberType(typeInfo.Type, _dependencies);
                }
                base.VisitObjectCreationExpression(node);
            }


            public override void VisitMemberAccessExpression(MemberAccessExpressionSyntax node)
            {
                // Check for typeof() expressions
                if (node.Expression is TypeOfExpressionSyntax typeOfExpr)
                {
                    var typeInfo = _semanticModel.GetTypeInfo(typeOfExpr.Type);
                    if (typeInfo.Type != null && !_analyzer.IsSystemType(typeInfo.Type))
                    {
                        var typeName = typeInfo.Type.ToDisplayString();
                        _dependencies.Add(typeName);

                        if (typeInfo.Type is INamedTypeSymbol namedType)
                        {
                            _analyzer.AnalyzeTypeSymbol(namedType);
                        }
                    }
                }

                base.VisitMemberAccessExpression(node);
            }

            private static bool IsReflectionMethod(string methodName)
            {
                return methodName switch
                {
                    // System.Type methods
                    "GetType" => true,
                    // System.Activator methods
                    "CreateInstance" => true,
                    "CreateInstanceFrom" => true,
                    // Unity-specific methods
                    "AddComponent" => true, // Covers the obsolete AddComponent(string)
                    "GetComponent" => true,
                    "GetComponentInChildren" => true,
                    "GetComponentInParent" => true,
                    // Asset loading can sometimes imply type dependency
                    "Load" => true, // e.g., Resources.Load("path/to/asset", typeof(MyType))
                    "LoadAsset" => true, // e.g., AssetDatabase.LoadAssetAtPath("path", typeof(MyType))
                    _ => false
                };
            }

            private void AnalyzeReflectionCall(InvocationExpressionSyntax node)
            {
                foreach (var arg in node.ArgumentList.Arguments)
                {
                    switch (arg.Expression)
                    {
                        case TypeOfExpressionSyntax typeOfExpr:
                        {
                            var typeInfo = _semanticModel.GetTypeInfo(typeOfExpr.Type);
                            if (typeInfo.Type != null && !_analyzer.IsSystemType(typeInfo.Type))
                            {
                                var typeName = typeInfo.Type.ToDisplayString();
                                _dependencies.Add(typeName);

                                if (typeInfo.Type is INamedTypeSymbol namedType)
                                {
                                    _analyzer.AnalyzeTypeSymbol(namedType);
                                }
                            }

                            break;
                        }
                        case LiteralExpressionSyntax literalExpr when literalExpr.IsKind(SyntaxKind.StringLiteralExpression):
                        {
                            var typeNameFromString = literalExpr.Token.ValueText;
                            // Attempt to find the type in the compilation
                            var potentialType = _analyzer._compilation.GetTypeByMetadataName(typeNameFromString);
                            if (potentialType != null && !_analyzer.IsSystemType(potentialType))
                            {
                                var typeName = potentialType.ToDisplayString();
                                _dependencies.Add(typeName);
                                _analyzer.AnalyzeTypeSymbol(potentialType);
                            }

                            break;
                        }
                    }
                }
                if (node.Expression is MemberAccessExpressionSyntax memberAccess && memberAccess.Name is GenericNameSyntax genericName)
                {
                    foreach (var typeArgSyntax in genericName.TypeArgumentList.Arguments)
                    {
                        var typeInfo = _semanticModel.GetTypeInfo(typeArgSyntax);
                        if (typeInfo.Type != null && !_analyzer.IsSystemType(typeInfo.Type))
                        {
                            var typeName = typeInfo.Type.ToDisplayString();
                            _dependencies.Add(typeName);
                
                            if (typeInfo.Type is INamedTypeSymbol namedType)
                            {
                                _analyzer.AnalyzeTypeSymbol(namedType);
                            }
                        }
                    }
                }
            }
        }
    }
}