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);
}
}
}