AloneGame / Library / PackageCache / com.unity.collections / Unity.Collections.CodeGen / JobsILPostProcessor.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Threading;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Unity.CompilationPipeline.Common.Diagnostics;
using Unity.CompilationPipeline.Common.ILPostProcessing;

namespace Unity.Jobs.CodeGen
{   
    // Jobs ILPP entry point
    internal partial class JobsILPostProcessor : ILPostProcessor
    {
        AssemblyDefinition AssemblyDefinition;
        List<DiagnosticMessage> DiagnosticMessages = new List<DiagnosticMessage>();
        public HashSet<string> Defines { get; private set; }

        public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
        {
            bool madeAnyChange = false;
            Defines = new HashSet<string>(compiledAssembly.Defines);

            try 
            {
                AssemblyDefinition = AssemblyDefinitionFor(compiledAssembly);
            }
            catch (BadImageFormatException)
            {
                return new ILPostProcessResult(null, DiagnosticMessages);
            }

            try
            {
                // This only works because the PostProcessorAssemblyResolver is explicitly loading 
                // transitive dependencies (and then some) and so if we can't find a references to
                // Unity.Jobs (via EarlyInitHelpers) in there than we are confident the assembly doesn't need processing
                var earlyInitHelpers = AssemblyDefinition.MainModule.ImportReference(typeof(EarlyInitHelpers)).CheckedResolve();
            }
            catch (ResolutionException)
            {
                return new ILPostProcessResult(null, DiagnosticMessages);
            }

            madeAnyChange = PostProcessImpl();

            // Hack to remove circular references
            var selfName = AssemblyDefinition.Name.FullName;
            foreach (var referenceName in AssemblyDefinition.MainModule.AssemblyReferences)
            {
                if (referenceName.FullName == selfName)
                {
                    AssemblyDefinition.MainModule.AssemblyReferences.Remove(referenceName);
                    break;
                }
            }

            if (!madeAnyChange)
                return new ILPostProcessResult(null, DiagnosticMessages);

            bool hasError = false;
            foreach (var d in DiagnosticMessages)
            {
                if (d.DiagnosticType == DiagnosticType.Error)
                {
                    hasError = true;
                    break;
                }
            }

            if (hasError)
                return new ILPostProcessResult(null, DiagnosticMessages);

            var pe = new MemoryStream();
            var pdb = new MemoryStream();
            var writerParameters = new WriterParameters
            {
                SymbolWriterProvider = new PortablePdbWriterProvider(), SymbolStream = pdb, WriteSymbols = true
            };

            AssemblyDefinition.Write(pe, writerParameters);
            return new ILPostProcessResult(new InMemoryAssembly(pe.ToArray(), pdb.ToArray()), DiagnosticMessages);
        }

        public override ILPostProcessor GetInstance()
        {
            return this;
        }

        public override bool WillProcess(ICompiledAssembly compiledAssembly)
        {
            if (compiledAssembly.Name.EndsWith("CodeGen.Tests", StringComparison.Ordinal))
                return false;

            if (compiledAssembly.InMemoryAssembly.PdbData == null || compiledAssembly.InMemoryAssembly.PeData == null)
                return false;

            var referencesCollections = false;
            if (compiledAssembly.Name == "Unity.Collections")
            {
                return true;
            }
            for (int i = 0; i < compiledAssembly.References.Length; ++i)
            {
                var fileName = Path.GetFileNameWithoutExtension(compiledAssembly.References[i]);
                if (fileName == "Unity.Collections")
                {
                    referencesCollections = true;
                    break;
                }
            }

            if (!referencesCollections) return false;

            return true;
        }

        // *******************************************************************************
        // ** NOTE
        // ** Everything below this is a copy of the same process used in EntitiesILPostProcessor and 
        // ** should stay synced with it.
        // *******************************************************************************

        class PostProcessorAssemblyResolver : IAssemblyResolver
        {
            private readonly HashSet<string> _referenceDirectories;
            private Dictionary<string, HashSet<string>> _referenceToPathMap;
            Dictionary<string, AssemblyDefinition> _cache = new Dictionary<string, AssemblyDefinition>();
            private ICompiledAssembly _compiledAssembly;
            private AssemblyDefinition _selfAssembly;

            public PostProcessorAssemblyResolver(ICompiledAssembly compiledAssembly)
            {
                _compiledAssembly = compiledAssembly;
                _referenceToPathMap = new Dictionary<string, HashSet<string>>();
                _referenceDirectories = new HashSet<string>();
                foreach (var reference in compiledAssembly.References)
                {
                    var assemblyName = Path.GetFileNameWithoutExtension(reference);
                    if (!_referenceToPathMap.TryGetValue(assemblyName, out var fileList))
                    {
                        fileList = new HashSet<string>();
                        _referenceToPathMap.Add(assemblyName, fileList);
                    }
                    fileList.Add(reference);
                    _referenceDirectories.Add(Path.GetDirectoryName(reference));
                }
            }

            public void Dispose()
            {
            }

            public AssemblyDefinition Resolve(AssemblyNameReference name)
            {
                return Resolve(name, new ReaderParameters(ReadingMode.Deferred));
            }

            public AssemblyDefinition Resolve(AssemblyNameReference name, ReaderParameters parameters)
            {
                {
                    if (name.Name == _compiledAssembly.Name)
                        return _selfAssembly;

                    var fileName = FindFile(name);
                    if (fileName == null)
                        return null;

                    var cacheKey = fileName;

                    if (_cache.TryGetValue(cacheKey, out var result))
                        return result;

                    parameters.AssemblyResolver = this;

                    var ms = MemoryStreamFor(fileName);

                    var pdb = fileName + ".pdb";
                    if (File.Exists(pdb))
                        parameters.SymbolStream = MemoryStreamFor(pdb);

                    var assemblyDefinition = AssemblyDefinition.ReadAssembly(ms, parameters);
                    _cache.Add(cacheKey, assemblyDefinition);
                    return assemblyDefinition;
                }
            }

            private string FindFile(AssemblyNameReference name)
            {
                if (_referenceToPathMap.TryGetValue(name.Name, out var paths))
                {
                    if (paths.Count == 1)
                    {
                        var enumerator = paths.GetEnumerator();
                        // enumerators begin before the first element 
                        enumerator.MoveNext();
                        return enumerator.Current;
                    }

                    // If we have more than one assembly with the same name loaded we now need to figure out which one
                    // is being requested based on the AssemblyNameReference
                    foreach (var path in paths)
                    {
                        var onDiskAssemblyName = AssemblyName.GetAssemblyName(path);
                        if (onDiskAssemblyName.FullName == name.FullName)
                            return path;
                    }
                    throw new ArgumentException($"Tried to resolve a reference in assembly '{name.FullName}' however the assembly could not be found. Known references which did not match: \n{string.Join("\n",paths)}");
                }

                // Unfortunately the current ICompiledAssembly API only provides direct references.
                // It is very much possible that a postprocessor ends up investigating a type in a directly
                // referenced assembly, that contains a field that is not in a directly referenced assembly.
                // if we don't do anything special for that situation, it will fail to resolve.  We should fix this
                // in the ILPostProcessing api. As a workaround, we rely on the fact here that the indirect references
                // are always located next to direct references, so we search in all directories of direct references we
                // got passed, and if we find the file in there, we resolve to it.
                foreach (var parentDir in _referenceDirectories)
                {
                    var candidate = Path.Combine(parentDir, name.Name + ".dll");
                    if (File.Exists(candidate))
                    {
                        if (!_referenceToPathMap.TryGetValue(candidate, out var referencePaths))
                        {
                            referencePaths = new HashSet<string>();
                            _referenceToPathMap.Add(candidate, referencePaths);
                        }
                        referencePaths.Add(candidate);

                        return candidate;
                    }
                }

                return null;
            }

            static MemoryStream MemoryStreamFor(string fileName)
            {
                return Retry(10, TimeSpan.FromSeconds(1), () => {
                    byte[] byteArray;
                    using (var fs = new FileStream(fileName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
                    {
                        byteArray = new byte[fs.Length];
                        var readLength = fs.Read(byteArray, 0, (int)fs.Length);
                        if (readLength != fs.Length)
                            throw new InvalidOperationException("File read length is not full length of file.");
                    }

                    return new MemoryStream(byteArray);
                });
            }

            private static MemoryStream Retry(int retryCount, TimeSpan waitTime, Func<MemoryStream> func)
            {
                try
                {
                    return func();
                }
                catch (IOException)
                {
                    if (retryCount == 0)
                        throw;
                    Console.WriteLine($"Caught IO Exception, trying {retryCount} more times");
                    Thread.Sleep(waitTime);
                    return Retry(retryCount - 1, waitTime, func);
                }
            }

            public void AddAssemblyDefinitionBeingOperatedOn(AssemblyDefinition assemblyDefinition)
            {
                _selfAssembly = assemblyDefinition;
            }
        }

        internal static AssemblyDefinition AssemblyDefinitionFor(ICompiledAssembly compiledAssembly)
        {
            var resolver = new PostProcessorAssemblyResolver(compiledAssembly);
            var readerParameters = new ReaderParameters
            {
                SymbolStream = new MemoryStream(compiledAssembly.InMemoryAssembly.PdbData),
                SymbolReaderProvider = new PortablePdbReaderProvider(),
                AssemblyResolver = resolver,
                ReflectionImporterProvider = new PostProcessorReflectionImporterProvider(),
                ReadingMode = ReadingMode.Immediate
            };

            var peStream = new MemoryStream(compiledAssembly.InMemoryAssembly.PeData);
            var assemblyDefinition = AssemblyDefinition.ReadAssembly(peStream, readerParameters);

            //apparently, it will happen that when we ask to resolve a type that lives inside Unity.Jobs, and we
            //are also postprocessing Unity.Jobs, type resolving will fail, because we do not actually try to resolve
            //inside the assembly we are processing. Let's make sure we do that, so that we can use postprocessor features inside
            //unity.Jobs itself as well.
            resolver.AddAssemblyDefinitionBeingOperatedOn(assemblyDefinition);

            return assemblyDefinition;
        }
    }

    internal class PostProcessorReflectionImporterProvider : IReflectionImporterProvider
    {
        public IReflectionImporter GetReflectionImporter(ModuleDefinition module)
        {
            return new PostProcessorReflectionImporter(module);
        }
    }

    internal class PostProcessorReflectionImporter : DefaultReflectionImporter
    {
        private const string SystemPrivateCoreLib = "System.Private.CoreLib";
        private AssemblyNameReference _correctCorlib;

        public PostProcessorReflectionImporter(ModuleDefinition module) : base(module)
        {
            _correctCorlib = default;
            foreach (var a in module.AssemblyReferences)
            {
                if (a.Name == "mscorlib" || a.Name == "netstandard" || a.Name == SystemPrivateCoreLib)
                {
                    _correctCorlib = a;
                    break;
                }
            }
        }

        public override AssemblyNameReference ImportReference(AssemblyName reference)
        {
            if (_correctCorlib != null && reference.Name == SystemPrivateCoreLib)
                return _correctCorlib;

            return base.ImportReference(reference);
        }
    }
}