SharpGenTools.Sdk/Extensibility/ExtensibilityAssemblyLoader.cs

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Reflection;
using Microsoft.CodeAnalysis;
using SharpGenTools.Sdk.Internal;
using SharpGenTools.Sdk.Internal.Roslyn;

namespace SharpGenTools.Sdk.Extensibility
{
    internal abstract class ExtensibilityAssemblyLoader
    {
        private readonly object _guard = new();

        private static readonly ImmutableHashSet<string> WellKnownAssemblyNames =
            new[]
            {
                "SharpGen",
                "SharpGen.Platform",
                "SharpGen.Runtime",
                "SharpGenTools.Sdk"
            }.ToImmutableHashSet();

        // lock _guard to read/write
        private readonly Dictionary<string, AssemblyIdentity> _loadedAssemblyIdentitiesByPath = new();
        private readonly Dictionary<AssemblyIdentity, Assembly> _loadedAssembliesByIdentity = new();

        private readonly HashSet<string> _dependencyLocations = new(PathUtilities.Comparer); 

        protected abstract Assembly LoadFromPathImpl(string fullPath);

        #region Public API

        public void AddDependencyLocation(string fullPath)
        {
            Utilities.RequireAbsolutePath(fullPath, nameof(fullPath));

            _dependencyLocations.Add(Directory.GetParent(fullPath).FullName);
        }

        public Assembly LoadFromPath(string fullPath)
        {
            Utilities.RequireAbsolutePath(fullPath, nameof(fullPath));
            return LoadFromPathUnchecked(fullPath);
        }

        #endregion

        private Assembly LoadFromPathUnchecked(string fullPath)
        {
            return LoadFromPathUncheckedCore(fullPath);
        }

        private Assembly LoadFromPathUncheckedCore(string fullPath, AssemblyIdentity identity = null)
        {
            Debug.Assert(PathUtilities.IsAbsolute(fullPath));

            // Check if we have already loaded an assembly with the same identity or from the given path.
            Assembly loadedAssembly = null;
            lock (_guard)
            {
                identity ??= GetOrAddAssemblyIdentity(fullPath);
                if (identity != null && _loadedAssembliesByIdentity.TryGetValue(identity, out var existingAssembly))
                {
                    loadedAssembly = existingAssembly;
                }
            }

            // Otherwise, load the assembly.
            if (loadedAssembly == null)
            {
                loadedAssembly = LoadFromPathImpl(fullPath);
            }

            // Add the loaded assembly to both path and identity cache.
            return AddToCache(loadedAssembly, fullPath, identity);
        }

        private Assembly AddToCache(Assembly assembly, string fullPath, AssemblyIdentity identity)
        {
            Debug.Assert(PathUtilities.IsAbsolute(fullPath));
            Debug.Assert(assembly != null);

            identity = AddToCache(fullPath, identity ?? AssemblyIdentity.FromAssemblyDefinition(assembly));
            Debug.Assert(identity != null);

            lock (_guard)
            {
                // The same assembly may be loaded from two different full paths (e.g. when loaded from GAC, etc.),
                // or another thread might have loaded the assembly after we checked above.
                if (_loadedAssembliesByIdentity.TryGetValue(identity, out var existingAssembly))
                {
                    assembly = existingAssembly;
                }
                else
                {
                    _loadedAssembliesByIdentity.Add(identity, assembly);
                }

                return assembly;
            }
        }

        private AssemblyIdentity GetOrAddAssemblyIdentity(string fullPath)
        {
            Debug.Assert(PathUtilities.IsAbsolute(fullPath));

            lock (_guard)
            {
                if (_loadedAssemblyIdentitiesByPath.TryGetValue(fullPath, out var existingIdentity))
                {
                    return existingIdentity;
                }
            }

            var identity = AssemblyIdentityUtils.TryGetAssemblyIdentity(fullPath);
            return AddToCache(fullPath, identity);
        }

        private AssemblyIdentity AddToCache(string fullPath, AssemblyIdentity identity)
        {
            lock (_guard)
            {
                if (_loadedAssemblyIdentitiesByPath.TryGetValue(fullPath, out var existingIdentity) && existingIdentity != null)
                {
                    identity = existingIdentity;
                }
                else
                {
                    _loadedAssemblyIdentitiesByPath[fullPath] = identity;
                }
            }

            return identity;
        }

        public Assembly Load(string displayName)
        {
            if (!AssemblyIdentity.TryParseDisplayName(displayName, out var requestedIdentity))
            {
                return null;
            }

            if (WellKnownAssemblyNames.Contains(requestedIdentity.Name))
            {
                // Force our assemblies to be loaded in the default ALC
                // and unify to the current version.
                return null;
            }

            lock (_guard)
            {
                // First, check if this loader already loaded the requested assembly:
                if (_loadedAssembliesByIdentity.TryGetValue(requestedIdentity, out var existingAssembly))
                {
                    return existingAssembly;
                }
            }

            IEnumerable<string> CulturePathSelector(string cultureSubfolder)
            {
                string LocationSelector(string directory) =>
                    Path.Combine(directory, cultureSubfolder, $"{requestedIdentity.Name}.dll");

                return _dependencyLocations.Select(LocationSelector).Where(File.Exists);
            }

            var culturePathVariants = string.IsNullOrEmpty(requestedIdentity.CultureName)
                                          // If no culture is specified, attempt to load directly from
                                          // the known dependency paths.
                                          ? new[] {string.Empty}
                                          // Search for satellite assemblies in culture subdirectories
                                          // of the assembly search directories, but fall back to the
                                          // bare search directory if that fails.
                                          : new[] {requestedIdentity.CultureName, string.Empty};

            var candidatePaths = culturePathVariants.SelectMany(CulturePathSelector);

            // Multiple assemblies of the same simple name but different identities might have been registered.
            // Load the one that matches the requested identity (if any).
            foreach (var candidatePath in candidatePaths)
            {
                var candidateIdentity = GetOrAddAssemblyIdentity(candidatePath);

                if (requestedIdentity.Equals(candidateIdentity))
                {
                    return LoadFromPathUncheckedCore(candidatePath, candidateIdentity);
                }
            }

            return null;
        }
    }
}