A game about forced loneliness, made by TACStudios
1using System;
2using System.Collections.Generic;
3using System.Diagnostics;
4using System.IO;
5using System.Linq;
6using Mono.Cecil;
7using Mono.Cecil.Cil;
8using Unity.CompilationPipeline.Common.Diagnostics;
9using Unity.CompilationPipeline.Common.ILPostProcessing;
10
11// TODO: Once DOTS has the latest 2022.2 editor merged into their fork, this can be UNITY_2022_2_OR_NEWER!
12#if UNITY_2023_1_OR_NEWER
13using ILPostProcessorToUse = zzzUnity.Burst.CodeGen.ILPostProcessing;
14#else
15using ILPostProcessorToUse = zzzUnity.Burst.CodeGen.ILPostProcessingLegacy;
16#endif
17
18/// Deliberately named zzzUnity.Burst.CodeGen, as we need to ensure its last in the chain
19namespace zzzUnity.Burst.CodeGen
20{
21 /// <summary>
22 /// Postprocessor used to replace calls from C# to [BurstCompile] functions to direct calls to
23 /// Burst native code functions without having to go through a C# delegate.
24 /// </summary>
25 internal class BurstILPostProcessor : ILPostProcessor
26 {
27 private sealed class CachedAssemblyResolver : AssemblyResolver
28 {
29 private Dictionary<string, AssemblyDefinition> _cache = new Dictionary<string, AssemblyDefinition>();
30 private Dictionary<string, string> _knownLocations = new Dictionary<string, string>();
31
32 public void RegisterKnownLocation(string path)
33 {
34 var k = Path.GetFileNameWithoutExtension(path);
35 // If an assembly is referenced multiple times, resolve to the first one
36 if (!_knownLocations.ContainsKey(k))
37 {
38 _knownLocations.Add(k, path);
39 }
40 }
41
42 public override AssemblyDefinition Resolve(AssemblyNameReference name)
43 {
44 if (!_cache.TryGetValue(name.FullName, out var definition))
45 {
46 if (_knownLocations.TryGetValue(name.Name, out var path))
47 {
48 definition = LoadFromFile(path);
49 }
50 else
51 {
52 definition = base.Resolve(name);
53 }
54
55 _cache.Add(name.FullName, definition);
56 }
57
58 return definition;
59 }
60 }
61
62 public bool IsDebugging;
63 public int DebuggingLevel;
64
65 private void SetupDebugging()
66 {
67 // This can be setup to get more diagnostics
68 var debuggingStr = Environment.GetEnvironmentVariable("UNITY_BURST_DEBUG");
69 var debugLevel = 0;
70 IsDebugging = debuggingStr != null && int.TryParse(debuggingStr, out debugLevel) && debugLevel > 0;
71 if (IsDebugging)
72 {
73 Log("[com.unity.burst] Extra debugging is turned on.");
74 DebuggingLevel = debugLevel;
75 }
76 }
77
78 private static SequencePoint FindBestSequencePointFor(MethodDefinition method, Instruction instruction)
79 {
80 var sequencePoints = method.DebugInformation?.GetSequencePointMapping().Values.OrderBy(s => s.Offset).ToList();
81 if (sequencePoints == null || !sequencePoints.Any())
82 return null;
83
84 for (int i = 0; i != sequencePoints.Count-1; i++)
85 {
86 if (sequencePoints[i].Offset < instruction.Offset &&
87 sequencePoints[i + 1].Offset > instruction.Offset)
88 return sequencePoints[i];
89 }
90
91 return sequencePoints.FirstOrDefault();
92 }
93
94 private static DiagnosticMessage MakeDiagnosticError(MethodDefinition method, Instruction errorLocation, string message)
95 {
96 var m = new DiagnosticMessage { DiagnosticType = DiagnosticType.Error };
97 var sPoint = errorLocation != null ? FindBestSequencePointFor(method, errorLocation) : null;
98 if (sPoint!=null)
99 {
100 m.Column = sPoint.StartColumn;
101 m.Line = sPoint.StartLine;
102 m.File = sPoint.Document.Url;
103 }
104 m.MessageData = message;
105 return m;
106 }
107
108 public override unsafe ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
109 {
110 var diagnostics = new List<DiagnosticMessage>();
111 if (!WillProcess(compiledAssembly))
112 return new ILPostProcessResult(null, diagnostics);
113
114 bool wasModified = false;
115 SetupDebugging();
116 bool debugging = IsDebugging && DebuggingLevel >= 2;
117
118 var inMemoryAssembly = compiledAssembly.InMemoryAssembly;
119
120
121 var peData = inMemoryAssembly.PeData;
122 var pdbData = inMemoryAssembly.PdbData;
123
124
125 var loader = new CachedAssemblyResolver();
126 var folders = new HashSet<string>();
127 var isForEditor = compiledAssembly.Defines?.Contains("UNITY_EDITOR") ?? false;
128 foreach (var reference in compiledAssembly.References)
129 {
130 loader.RegisterKnownLocation(reference);
131 folders.Add(Path.Combine(Environment.CurrentDirectory, Path.GetDirectoryName(reference)));
132 }
133 var folderList = folders.OrderBy(x => x).ToList();
134 foreach (var folder in folderList)
135 {
136 loader.AddSearchDirectory(folder);
137 }
138
139 var clock = Stopwatch.StartNew();
140 if (debugging)
141 {
142 Log($"Start processing assembly {compiledAssembly.Name}, IsForEditor: {isForEditor}, Folders: {string.Join("\n", folderList)}");
143 }
144
145 var ilPostProcessing = new ILPostProcessorToUse(loader, isForEditor,
146 (m,i,s) => { diagnostics.Add(MakeDiagnosticError(m, i, s)); },
147 IsDebugging ? Log : (LogDelegate)null, DebuggingLevel);
148 var functionPointerProcessing = new FunctionPointerInvokeTransform(loader,
149 (m,i,s) => { diagnostics.Add(MakeDiagnosticError(m, i, s)); },
150 IsDebugging ? Log : (LogDelegate)null, DebuggingLevel);
151 try
152 {
153 // For IL Post Processing, use the builtin symbol reader provider
154 var assemblyDefinition = loader.LoadFromStream(new MemoryStream(peData), new MemoryStream(pdbData), new PortablePdbReaderProvider() );
155 wasModified |= ilPostProcessing.Run(assemblyDefinition);
156 wasModified |= functionPointerProcessing.Run(assemblyDefinition);
157 if (wasModified)
158 {
159 var peStream = new MemoryStream();
160 var pdbStream = new MemoryStream();
161 var writeParameters = new WriterParameters
162 {
163 SymbolWriterProvider = new PortablePdbWriterProvider(),
164 WriteSymbols = true,
165 SymbolStream = pdbStream
166 };
167
168 assemblyDefinition.Write(peStream, writeParameters);
169 peStream.Flush();
170 pdbStream.Flush();
171
172 peData = peStream.ToArray();
173 pdbData = pdbStream.ToArray();
174 }
175 }
176 catch (Exception ex)
177 {
178 throw new InvalidOperationException($"Internal compiler error for Burst ILPostProcessor on {compiledAssembly.Name}. Exception: {ex}");
179 }
180
181 if (debugging)
182 {
183 Log($"End processing assembly {compiledAssembly.Name} in {clock.Elapsed.TotalMilliseconds}ms.");
184 }
185
186 if (wasModified && !diagnostics.Any(d => d.DiagnosticType == DiagnosticType.Error))
187 {
188 return new ILPostProcessResult(new InMemoryAssembly(peData, pdbData), diagnostics);
189 }
190 return new ILPostProcessResult(null, diagnostics);
191 }
192
193 private static void Log(string message)
194 {
195 Console.WriteLine($"{nameof(BurstILPostProcessor)}: {message}");
196 }
197
198 public override ILPostProcessor GetInstance()
199 {
200 return this;
201 }
202
203 public override bool WillProcess(ICompiledAssembly compiledAssembly)
204 {
205 return compiledAssembly.References.Any(f => Path.GetFileName(f) == "Unity.Burst.dll");
206 }
207 }
208}