A game about forced loneliness, made by TACStudios
1using System;
2using System.Collections.Generic;
3using System.Linq;
4using System.Text;
5using Mono.Cecil;
6using Mono.Cecil.Cil;
7
8namespace zzzUnity.Burst.CodeGen
9{
10 internal delegate void LogDelegate(string message);
11 internal delegate void ErrorDiagnosticDelegate(MethodDefinition method, Instruction instruction, string message);
12
13 /// <summary>
14 /// Main class for post processing assemblies. The post processing is currently performing:
15 /// - Replace C# call from C# to Burst functions with attributes [BurstCompile] to a call to the compiled Burst function
16 /// In both editor and standalone scenarios. For DOTS Runtime, this is done differently at BclApp level by patching
17 /// DllImport.
18 /// - Replace calls to `SharedStatic.GetOrCreate` with `SharedStatic.GetOrCreateUnsafe`, and calculate the hashes during ILPP time
19 /// rather than in static constructors at runtime.
20 /// </summary>
21 internal class ILPostProcessingLegacy
22 {
23 private AssemblyDefinition _burstAssembly;
24 private TypeDefinition _burstCompilerTypeDefinition;
25 private MethodReference _burstCompilerIsEnabledMethodDefinition;
26 private MethodReference _burstCompilerCompileILPPMethod;
27 private MethodReference _burstCompilerGetILPPMethodFunctionPointer;
28 private MethodReference _burstDiscardAttributeConstructor;
29 private MethodReference _burstCompilerCompileUnsafeStaticMethodReinitialiseAttributeCtor;
30 private TypeSystem _typeSystem;
31 private TypeReference _systemType;
32 private TypeReference _systemDelegateType;
33 private TypeReference _systemASyncCallbackType;
34 private TypeReference _systemIASyncResultType;
35 private AssemblyDefinition _assemblyDefinition;
36 private bool _modified;
37 private readonly StringBuilder _builder = new StringBuilder(1024);
38 private readonly List<Instruction> _instructionsToReplace = new List<Instruction>(4);
39 private readonly List<MethodDefinition> _directCallInitializeMethods = new List<MethodDefinition>();
40
41 private const string PostfixBurstDirectCall = "$BurstDirectCall";
42 private const string PostfixBurstDelegate = "$PostfixBurstDelegate";
43 private const string PostfixManaged = "$BurstManaged";
44 private const string GetFunctionPointerName = "GetFunctionPointer";
45 private const string GetFunctionPointerDiscardName = "GetFunctionPointerDiscard";
46 private const string InvokeName = "Invoke";
47
48 public ILPostProcessingLegacy(AssemblyResolver loader, bool isForEditor, ErrorDiagnosticDelegate error, LogDelegate log = null, int logLevel = 0, bool skipInitializeOnLoad = false)
49 {
50 _skipInitializeOnLoad = skipInitializeOnLoad;
51 Loader = loader;
52 IsForEditor = isForEditor;
53 }
54
55 public bool _skipInitializeOnLoad;
56
57 public bool IsForEditor { get; private set; }
58
59 private AssemblyResolver Loader { get; }
60
61 public bool Run(AssemblyDefinition assemblyDefinition)
62 {
63 _assemblyDefinition = assemblyDefinition;
64 _typeSystem = assemblyDefinition.MainModule.TypeSystem;
65
66 _modified = false;
67 var types = assemblyDefinition.MainModule.GetTypes().ToArray();
68 foreach (var type in types)
69 {
70 ProcessType(type);
71 }
72
73 // If we processed any direct-calls, then generate a single [RuntimeInitializeOnLoadMethod] method
74 // for the whole assembly, which will initialize each individual direct-call class.
75 if (_directCallInitializeMethods.Count > 0)
76 {
77 GenerateInitializeOnLoadMethod();
78 }
79
80 return _modified;
81 }
82
83 private void GenerateInitializeOnLoadMethod()
84 {
85 // [UnityEngine.RuntimeInitializeOnLoadMethod(UnityEngine.RuntimeInitializeLoadType.AfterAssembliesLoaded)]
86 // [UnityEditor.InitializeOnLoadMethod] // When its an editor assembly
87 // private static void Initialize()
88 // {
89 // DirectCallA.Initialize();
90 // DirectCallB.Initialize();
91 // }
92 const string initializeOnLoadClassName = "$BurstDirectCallInitializer";
93 var initializeOnLoadClass = _assemblyDefinition.MainModule.Types.FirstOrDefault(x => x.Name == initializeOnLoadClassName);
94 if (initializeOnLoadClass != null)
95 {
96 // If there's already a class with this name, remove it,
97 // This would mean that we're postprocessing an already-postprocessed assembly;
98 // I don't think that ever happens, but no sense in breaking if it does.
99 _assemblyDefinition.MainModule.Types.Remove(initializeOnLoadClass);
100 }
101 initializeOnLoadClass = new TypeDefinition(
102 "",
103 initializeOnLoadClassName,
104 TypeAttributes.NotPublic |
105 TypeAttributes.AutoLayout |
106 TypeAttributes.AnsiClass |
107 TypeAttributes.Abstract |
108 TypeAttributes.Sealed |
109 TypeAttributes.BeforeFieldInit)
110 {
111 BaseType = _typeSystem.Object
112 };
113 _assemblyDefinition.MainModule.Types.Add(initializeOnLoadClass);
114 var initializeOnLoadMethod = new MethodDefinition("Initialize", MethodAttributes.Private | MethodAttributes.HideBySig | MethodAttributes.Static, _typeSystem.Void)
115 {
116 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
117 DeclaringType = initializeOnLoadClass
118 };
119
120 var processor = initializeOnLoadMethod.Body.GetILProcessor();
121 foreach (var initializeMethod in _directCallInitializeMethods)
122 {
123 processor.Emit(OpCodes.Call, initializeMethod);
124 }
125 processor.Emit(OpCodes.Ret);
126 initializeOnLoadClass.Methods.Add(FixDebugInformation(initializeOnLoadMethod));
127
128 var attribute = new CustomAttribute(_unityEngineInitializeOnLoadAttributeCtor);
129 attribute.ConstructorArguments.Add(new CustomAttributeArgument(_unityEngineRuntimeInitializeLoadType, _unityEngineRuntimeInitializeLoadAfterAssemblies.Constant));
130 initializeOnLoadMethod.CustomAttributes.Add(attribute);
131
132 if (IsForEditor && !_skipInitializeOnLoad)
133 {
134 // Need to ensure the editor tag for initialize on load is present, otherwise edit mode tests will not call Initialize
135 attribute = new CustomAttribute(_unityEditorInitilizeOnLoadAttributeCtor);
136 initializeOnLoadMethod.CustomAttributes.Add(attribute);
137 }
138 }
139
140 private static bool CanComputeCompileTimeHash(TypeReference typeRef)
141 {
142 if (typeRef.ContainsGenericParameter)
143 {
144 return false;
145 }
146
147 var assemblyNameReference = typeRef.Scope as AssemblyNameReference ?? typeRef.Module.Assembly?.Name;
148
149 if (assemblyNameReference == null)
150 {
151 return false;
152 }
153
154 switch (assemblyNameReference.Name)
155 {
156 case "netstandard":
157 case "mscorlib":
158 return false;
159 }
160
161 return true;
162 }
163
164 private void ProcessType(TypeDefinition type)
165 {
166 if (!type.HasGenericParameters && TryGetBurstCompilerAttribute(type, out _))
167 {
168 // Make a copy because we are going to modify it
169 var methodCount = type.Methods.Count;
170 for (var j = 0; j < methodCount; j++)
171 {
172 var method = type.Methods[j];
173 if (!method.IsStatic || method.HasGenericParameters || !TryGetBurstCompilerAttribute(method, out var methodBurstCompileAttribute)) continue;
174
175 bool isDirectCallDisabled = false;
176 bool foundProperty = false;
177 if (methodBurstCompileAttribute.HasProperties)
178 {
179 foreach (var property in methodBurstCompileAttribute.Properties)
180 {
181 if (property.Name == "DisableDirectCall")
182 {
183 isDirectCallDisabled = (bool)property.Argument.Value;
184 foundProperty = true;
185 break;
186 }
187 }
188 }
189
190 // If the method doesn't have a direct call specified, try the assembly level, do one last check for any assembly level [BurstCompile] instead.
191 if (foundProperty == false && TryGetBurstCompilerAttribute(method.Module.Assembly, out var assemblyBurstCompileAttribute))
192 {
193 if (assemblyBurstCompileAttribute.HasProperties)
194 {
195 foreach (var property in assemblyBurstCompileAttribute.Properties)
196 {
197 if (property.Name == "DisableDirectCall")
198 {
199 isDirectCallDisabled = (bool)property.Argument.Value;
200 break;
201 }
202 }
203 }
204 }
205
206 foreach (var customAttribute in method.CustomAttributes)
207 {
208 if (customAttribute.AttributeType.FullName == "System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute")
209 {
210 // Can't / shouldn't enable direct call for [UnmanagedCallersOnly] methods -
211 // these can't be called from managed code.
212 isDirectCallDisabled = true;
213 break;
214 }
215 }
216
217#if !UNITY_DOTSPLAYER // Direct call is not Supported for dots runtime via this pre-processor, its handled elsewhere, this code assumes a Unity Editor based burst
218 if (!isDirectCallDisabled)
219 {
220 if (_burstAssembly == null)
221 {
222 var resolved = methodBurstCompileAttribute.Constructor.DeclaringType.Resolve();
223 InitializeBurstAssembly(resolved.Module.Assembly);
224 }
225
226 ProcessMethodForDirectCall(method);
227 _modified = true;
228 }
229#endif
230 }
231 }
232
233 if (TypeHasSharedStaticInIt(type))
234 {
235 foreach (var method in type.Methods)
236 {
237 // Skip anything that isn't the static constructor.
238 if (method.Name != ".cctor")
239 {
240 continue;
241 }
242
243 try
244 {
245#if DEBUG
246 if (_instructionsToReplace.Count != 0)
247 {
248 throw new InvalidOperationException("Instructions to replace wasn't cleared properly!");
249 }
250#endif
251
252 foreach (var instruction in method.Body.Instructions)
253 {
254 // Skip anything that isn't a call.
255 if (instruction.OpCode != OpCodes.Call)
256 {
257 continue;
258 }
259
260 var calledMethod = (MethodReference)instruction.Operand;
261
262 if (calledMethod.Name != "GetOrCreate")
263 {
264 continue;
265 }
266
267 // Skip anything that isn't member of the `SharedStatic` class.
268 if (!TypeIsSharedStatic(calledMethod.DeclaringType))
269 {
270 continue;
271 }
272
273 // We only handle the `GetOrCreate` calls with a single parameter (the alignment).
274 if (calledMethod.Parameters.Count != 1)
275 {
276 continue;
277 }
278
279 // We only post-process the generic versions of `GetOrCreate`.
280 if (!(calledMethod is GenericInstanceMethod genericInstanceMethod))
281 {
282 continue;
283 }
284
285 var atLeastOneArgumentCanBeComputed = false;
286
287 foreach (var genericArgument in genericInstanceMethod.GenericArguments)
288 {
289 if (CanComputeCompileTimeHash(genericArgument))
290 {
291 atLeastOneArgumentCanBeComputed = true;
292 }
293 }
294
295 // We cannot post-process a shared static with all arguments being open generic.
296 // We cannot post-process a shared static where all of its types are in core libraries.
297 if (!atLeastOneArgumentCanBeComputed)
298 {
299 continue;
300 }
301
302 _instructionsToReplace.Add(instruction);
303 }
304
305 if (_instructionsToReplace.Count > 0)
306 {
307 _modified = true;
308 }
309
310 foreach (var instruction in _instructionsToReplace)
311 {
312 var calledMethod = (GenericInstanceMethod)instruction.Operand;
313
314 var hashCode64 = CalculateHashCode64(calledMethod.GenericArguments[0]);
315
316 long subHashCode64 = 0;
317
318 var useCalculatedHashCode = true;
319 var useCalculatedSubHashCode = true;
320
321 if (calledMethod.GenericArguments.Count == 2)
322 {
323 subHashCode64 = CalculateHashCode64(calledMethod.GenericArguments[1]);
324
325 useCalculatedHashCode = CanComputeCompileTimeHash(calledMethod.GenericArguments[0]);
326 useCalculatedSubHashCode = CanComputeCompileTimeHash(calledMethod.GenericArguments[1]);
327 }
328
329#if DEBUG
330 if (!useCalculatedHashCode && !useCalculatedSubHashCode)
331 {
332 throw new InvalidOperationException("Cannot replace when both hashes are invalid!");
333 }
334#endif
335
336 var methodToCall = "GetOrCreateUnsafe";
337 TypeReference genericArgument = null;
338
339 if (!useCalculatedHashCode)
340 {
341 methodToCall = "GetOrCreatePartiallyUnsafeWithSubHashCode";
342 genericArgument = calledMethod.GenericArguments[0];
343 }
344 else if (!useCalculatedSubHashCode)
345 {
346 methodToCall = "GetOrCreatePartiallyUnsafeWithHashCode";
347 genericArgument = calledMethod.GenericArguments[1];
348 }
349
350 var getOrCreateUnsafe = _assemblyDefinition.MainModule.ImportReference(
351 calledMethod.DeclaringType.Resolve().Methods.First(m => m.Name == methodToCall));
352
353 getOrCreateUnsafe.DeclaringType = calledMethod.DeclaringType;
354
355 if (genericArgument != null)
356 {
357 var genericInstanceMethod = new GenericInstanceMethod(getOrCreateUnsafe);
358
359 genericInstanceMethod.GenericArguments.Add(genericArgument);
360
361 getOrCreateUnsafe = genericInstanceMethod;
362 }
363
364 var processor = method.Body.GetILProcessor();
365
366 if (useCalculatedHashCode)
367 {
368 processor.InsertBefore(instruction, processor.Create(OpCodes.Ldc_I8, hashCode64));
369 }
370
371 if (useCalculatedSubHashCode)
372 {
373 processor.InsertBefore(instruction, processor.Create(OpCodes.Ldc_I8, subHashCode64));
374 }
375
376 processor.Replace(instruction, processor.Create(OpCodes.Call, getOrCreateUnsafe));
377 }
378 }
379 finally
380 {
381 _instructionsToReplace.Clear();
382 }
383 }
384 }
385 }
386
387 // WARNING: This **must** be kept in sync with the definition in BurstRuntime.cs!
388 private static long HashStringWithFNV1A64(string text)
389 {
390 // Using http://www.isthe.com/chongo/tech/comp/fnv/index.html#FNV-1a
391 // with basis and prime:
392 const ulong offsetBasis = 14695981039346656037;
393 const ulong prime = 1099511628211;
394
395 ulong result = offsetBasis;
396
397 foreach(var c in text)
398 {
399 result = prime * (result ^ (byte)(c & 255));
400 result = prime * (result ^ (byte)(c >> 8));
401 }
402
403 return (long)result;
404 }
405
406 private long CalculateHashCode64(TypeReference type)
407 {
408 try
409 {
410#if DEBUG
411 if (_builder.Length != 0)
412 {
413 throw new InvalidOperationException("StringBuilder wasn't cleared properly!");
414 }
415#endif
416
417 type.BuildAssemblyQualifiedName(_builder);
418 return HashStringWithFNV1A64(_builder.ToString());
419 }
420 finally
421 {
422 _builder.Clear();
423 }
424 }
425
426 private static bool TypeIsSharedStatic(TypeReference typeRef)
427 {
428 if (typeRef.Namespace != "Unity.Burst")
429 {
430 return false;
431 }
432
433 if (typeRef.Name != "SharedStatic`1")
434 {
435 return false;
436 }
437
438 return true;
439 }
440
441 private static bool TypeHasSharedStaticInIt(TypeDefinition typeDef)
442 {
443 foreach (var field in typeDef.Fields)
444 {
445 if (TypeIsSharedStatic(field.FieldType))
446 {
447 return true;
448 }
449 }
450
451 return false;
452 }
453
454 private TypeDefinition InjectDelegate(TypeDefinition declaringType, string originalName, MethodDefinition managed, string uniqueSuffix)
455 {
456 var injectedDelegateType = new TypeDefinition(declaringType.Namespace, $"{originalName}{uniqueSuffix}{PostfixBurstDelegate}",
457 TypeAttributes.NestedPublic |
458 TypeAttributes.AutoLayout |
459 TypeAttributes.AnsiClass |
460 TypeAttributes.Sealed
461 )
462 {
463 DeclaringType = declaringType,
464 BaseType = _systemDelegateType
465 };
466
467 declaringType.NestedTypes.Add(injectedDelegateType);
468
469 {
470 var constructor = new MethodDefinition(".ctor",
471 MethodAttributes.Public |
472 MethodAttributes.HideBySig |
473 MethodAttributes.SpecialName |
474 MethodAttributes.RTSpecialName,
475 _typeSystem.Void)
476 {
477 HasThis = true,
478 IsManaged = true,
479 IsRuntime = true,
480 DeclaringType = injectedDelegateType
481 };
482
483 constructor.Parameters.Add(new ParameterDefinition(_typeSystem.Object));
484 constructor.Parameters.Add(new ParameterDefinition(_typeSystem.IntPtr));
485 injectedDelegateType.Methods.Add(constructor);
486 }
487
488 {
489 var invoke = new MethodDefinition("Invoke",
490 MethodAttributes.Public |
491 MethodAttributes.HideBySig |
492 MethodAttributes.NewSlot |
493 MethodAttributes.Virtual,
494 managed.ReturnType)
495 {
496 HasThis = true,
497 IsManaged = true,
498 IsRuntime = true,
499 DeclaringType = injectedDelegateType
500 };
501
502 foreach (var parameter in managed.Parameters)
503 {
504 invoke.Parameters.Add(parameter);
505 }
506
507 injectedDelegateType.Methods.Add(invoke);
508 }
509
510 {
511 var beginInvoke = new MethodDefinition("BeginInvoke",
512 MethodAttributes.Public |
513 MethodAttributes.HideBySig |
514 MethodAttributes.NewSlot |
515 MethodAttributes.Virtual,
516 _systemIASyncResultType)
517 {
518 HasThis = true,
519 IsManaged = true,
520 IsRuntime = true,
521 DeclaringType = injectedDelegateType
522 };
523
524 foreach (var parameter in managed.Parameters)
525 {
526 beginInvoke.Parameters.Add(parameter);
527 }
528
529 beginInvoke.Parameters.Add(new ParameterDefinition(_systemASyncCallbackType));
530 beginInvoke.Parameters.Add(new ParameterDefinition(_typeSystem.Object));
531
532 injectedDelegateType.Methods.Add(beginInvoke);
533 }
534
535 {
536 var endInvoke = new MethodDefinition("EndInvoke",
537 MethodAttributes.Public |
538 MethodAttributes.HideBySig |
539 MethodAttributes.NewSlot |
540 MethodAttributes.Virtual,
541 managed.ReturnType)
542 {
543 HasThis = true,
544 IsManaged = true,
545 IsRuntime = true,
546 DeclaringType = injectedDelegateType
547 };
548
549 endInvoke.Parameters.Add(new ParameterDefinition(_systemIASyncResultType));
550
551 injectedDelegateType.Methods.Add(endInvoke);
552 }
553
554 return injectedDelegateType;
555 }
556
557 private MethodDefinition CreateGetFunctionPointerDiscardMethod(TypeDefinition cls, FieldDefinition pointerField, FieldDefinition deferredCompilationField, MethodDefinition managedFallbackMethod, TypeDefinition injectedDelegate)
558 {
559 // Create GetFunctionPointer method:
560 //
561 // [BurstDiscard]
562 // public static void GetFunctionPointerDiscard(ref IntPtr ptr) {
563 // if (Pointer == null) {
564 // Pointer = BurstCompiler.GetILPPMethodFunctionPointer2(DeferredCompilation, managedFallbackMethod, DelegateType);
565 // }
566 //
567 // ptr = Pointer
568 // }
569 var getFunctionPointerDiscardMethod = new MethodDefinition(GetFunctionPointerDiscardName, MethodAttributes.Private | MethodAttributes.HideBySig | MethodAttributes.Static, _typeSystem.Void)
570 {
571 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
572 DeclaringType = cls
573 };
574
575 getFunctionPointerDiscardMethod.Parameters.Add(new ParameterDefinition(new ByReferenceType(_typeSystem.IntPtr)));
576
577 var processor = getFunctionPointerDiscardMethod.Body.GetILProcessor();
578 processor.Emit(OpCodes.Ldsfld, pointerField);
579 var branchPosition = processor.Body.Instructions[processor.Body.Instructions.Count - 1];
580
581 processor.Emit(OpCodes.Ldsfld, deferredCompilationField);
582 processor.Emit(OpCodes.Ldtoken, managedFallbackMethod);
583 processor.Emit(OpCodes.Ldtoken, injectedDelegate);
584 processor.Emit(OpCodes.Call, _burstCompilerGetILPPMethodFunctionPointer);
585 processor.Emit(OpCodes.Stsfld, pointerField);
586
587 processor.Emit(OpCodes.Ldarg_0);
588 processor.InsertAfter(branchPosition, Instruction.Create(OpCodes.Brtrue, processor.Body.Instructions[processor.Body.Instructions.Count - 1]));
589 processor.Emit(OpCodes.Ldsfld, pointerField);
590 processor.Emit(OpCodes.Stind_I);
591 processor.Emit(OpCodes.Ret);
592
593 cls.Methods.Add(FixDebugInformation(getFunctionPointerDiscardMethod));
594
595 getFunctionPointerDiscardMethod.CustomAttributes.Add(new CustomAttribute(_burstDiscardAttributeConstructor));
596
597 return getFunctionPointerDiscardMethod;
598 }
599
600 private MethodDefinition CreateGetFunctionPointerMethod(TypeDefinition cls, MethodDefinition getFunctionPointerDiscardMethod)
601 {
602 // Create GetFunctionPointer method:
603 //
604 // public static IntPtr GetFunctionPointer() {
605 // var ptr;
606 // GetFunctionPointerDiscard(ref ptr);
607 // return ptr;
608 // }
609 var getFunctionPointerMethod = new MethodDefinition(GetFunctionPointerName, MethodAttributes.Private | MethodAttributes.HideBySig | MethodAttributes.Static, _typeSystem.IntPtr)
610 {
611 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
612 DeclaringType = cls
613 };
614
615 getFunctionPointerMethod.Body.Variables.Add(new VariableDefinition(_typeSystem.IntPtr));
616 getFunctionPointerMethod.Body.InitLocals = true;
617
618 var processor = getFunctionPointerMethod.Body.GetILProcessor();
619
620 processor.Emit(OpCodes.Ldc_I4_0);
621 processor.Emit(OpCodes.Conv_I);
622 processor.Emit(OpCodes.Stloc_0);
623 processor.Emit(OpCodes.Ldloca_S, (byte)0);
624 processor.Emit(OpCodes.Call, getFunctionPointerDiscardMethod);
625 processor.Emit(OpCodes.Ldloc_0);
626
627 processor.Emit(OpCodes.Ret);
628
629 cls.Methods.Add(FixDebugInformation(getFunctionPointerMethod));
630
631 return getFunctionPointerMethod;
632 }
633
634 private void ProcessMethodForDirectCall(MethodDefinition burstCompileMethod)
635 {
636 var declaringType = burstCompileMethod.DeclaringType;
637
638 var uniqueSuffix = $"_{burstCompileMethod.MetadataToken.RID:X8}";
639
640 var injectedDelegate = InjectDelegate(declaringType, burstCompileMethod.Name, burstCompileMethod, uniqueSuffix);
641
642 // Create a copy of the original method that will be the actual managed method
643 // The original method is patched at the end of this method to call
644 // the dispatcher that will go to the Burst implementation or the managed method (if in the editor and Burst is disabled)
645 var managedFallbackMethod = new MethodDefinition($"{burstCompileMethod.Name}{PostfixManaged}", burstCompileMethod.Attributes, burstCompileMethod.ReturnType)
646 {
647 DeclaringType = declaringType,
648 ImplAttributes = burstCompileMethod.ImplAttributes,
649 MetadataToken = burstCompileMethod.MetadataToken,
650 };
651
652 // Ensure the CustomAttributes are the same
653 managedFallbackMethod.CustomAttributes.Clear();
654 foreach (var attr in burstCompileMethod.CustomAttributes)
655 {
656 managedFallbackMethod.CustomAttributes.Add(attr);
657 }
658
659 declaringType.Methods.Add(managedFallbackMethod);
660
661 foreach (var parameter in burstCompileMethod.Parameters)
662 {
663 managedFallbackMethod.Parameters.Add(parameter);
664 }
665
666 // Copy the body from the original burst method to the managed fallback, we'll replace the burstCompileMethod body later.
667 managedFallbackMethod.Body.InitLocals = burstCompileMethod.Body.InitLocals;
668 managedFallbackMethod.Body.LocalVarToken = burstCompileMethod.Body.LocalVarToken;
669 managedFallbackMethod.Body.MaxStackSize = burstCompileMethod.Body.MaxStackSize;
670
671 foreach (var variable in burstCompileMethod.Body.Variables)
672 {
673 managedFallbackMethod.Body.Variables.Add(variable);
674 }
675
676 foreach (var instruction in burstCompileMethod.Body.Instructions)
677 {
678 managedFallbackMethod.Body.Instructions.Add(instruction);
679 }
680
681 foreach (var exceptionHandler in burstCompileMethod.Body.ExceptionHandlers)
682 {
683 managedFallbackMethod.Body.ExceptionHandlers.Add(exceptionHandler);
684 }
685
686 managedFallbackMethod.ImplAttributes &= MethodImplAttributes.NoInlining;
687 // 0x0100 is AggressiveInlining
688 managedFallbackMethod.ImplAttributes |= (MethodImplAttributes)0x0100;
689
690 // The method needs to be public because we query for it in the ILPP code.
691 managedFallbackMethod.Attributes &= ~MethodAttributes.Private;
692 managedFallbackMethod.Attributes |= MethodAttributes.Public;
693
694 // private static class (Name_RID.$Postfix)
695 var cls = new TypeDefinition(declaringType.Namespace, $"{burstCompileMethod.Name}{uniqueSuffix}{PostfixBurstDirectCall}",
696 TypeAttributes.NestedAssembly |
697 TypeAttributes.AutoLayout |
698 TypeAttributes.AnsiClass |
699 TypeAttributes.Abstract |
700 TypeAttributes.Sealed |
701 TypeAttributes.BeforeFieldInit
702 )
703 {
704 DeclaringType = declaringType,
705 BaseType = _typeSystem.Object
706 };
707
708 declaringType.NestedTypes.Add(cls);
709
710 // Create Field:
711 //
712 // private static IntPtr Pointer;
713 var pointerField = new FieldDefinition("Pointer", FieldAttributes.Static | FieldAttributes.Private, _typeSystem.IntPtr)
714 {
715 DeclaringType = cls
716 };
717 cls.Fields.Add(pointerField);
718
719 // Create Field:
720 //
721 // private static IntPtr DeferredCompilation;
722 var deferredCompilationField = new FieldDefinition("DeferredCompilation", FieldAttributes.Static | FieldAttributes.Private, _typeSystem.IntPtr)
723 {
724 DeclaringType = cls
725 };
726 cls.Fields.Add(deferredCompilationField);
727
728 var getFunctionPointerDiscardMethod = CreateGetFunctionPointerDiscardMethod(cls, pointerField, deferredCompilationField, managedFallbackMethod, injectedDelegate);
729 var getFunctionPointerMethod = CreateGetFunctionPointerMethod(cls, getFunctionPointerDiscardMethod);
730
731 var asmAttribute = new CustomAttribute(_burstCompilerCompileUnsafeStaticMethodReinitialiseAttributeCtor);
732 asmAttribute.ConstructorArguments.Add(new CustomAttributeArgument(_systemType, cls));
733 _assemblyDefinition.CustomAttributes.Add(asmAttribute);
734
735 // Create the static Constructor Method (called via .cctor and via reflection on burst compilation enable)
736 // private static void Constructor() {
737 // deferredCompilation = CompileILPPMethod2(burstCompileMethod);
738 // }
739
740 var constructor = new MethodDefinition("Constructor", MethodAttributes.Public | MethodAttributes.HideBySig | MethodAttributes.Static, _typeSystem.Void)
741 {
742 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
743 DeclaringType = cls
744 };
745
746 var processor = constructor.Body.GetILProcessor();
747 // In the editor we'll ask for the fallback method, it will be effectively redirected to the burstCompileMethod
748 // While in the player managedFallbackMethod won't be in the compiled delegate cache, but burstCompileMethod
749 // will, and it's safe to use the burstCompileMethod because it will always be the Burst compiled one
750 processor.Emit(OpCodes.Ldtoken, IsForEditor ? managedFallbackMethod : burstCompileMethod);
751 processor.Emit(OpCodes.Call, _burstCompilerCompileILPPMethod);
752 processor.Emit(OpCodes.Stsfld, deferredCompilationField);
753 processor.Emit(OpCodes.Ret);
754
755 cls.Methods.Add(FixDebugInformation(constructor));
756
757 // Create an Initialize method
758 // This will be called from the single [RuntimeInitializeOnLoadMethod]
759 // method that we'll generate for this assembly.
760 // Its only job is to cause the .cctor to run.
761 //
762 // public static void Initialize() { }
763 var initializeMethod = new MethodDefinition("Initialize", MethodAttributes.Public | MethodAttributes.HideBySig | MethodAttributes.Static, _typeSystem.Void)
764 {
765 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
766 DeclaringType = cls
767 };
768
769 processor = initializeMethod.Body.GetILProcessor();
770 processor.Emit(OpCodes.Ret);
771 cls.Methods.Add(FixDebugInformation(initializeMethod));
772
773 var currentInitializer = initializeMethod;
774 var currentDeclaringType = declaringType;
775
776 // If our target method is hidden behind one or more nested private classes, then
777 // create a method on the parent type that calls said method (for each private nested class)
778 while (currentDeclaringType.DeclaringType != null)
779 {
780 var parentType = currentDeclaringType.DeclaringType;
781 if (((currentDeclaringType.Attributes & TypeAttributes.NestedPrivate) == TypeAttributes.NestedPrivate) ||
782 ((currentDeclaringType.Attributes & TypeAttributes.NestedFamily) == TypeAttributes.NestedFamily))
783 {
784 var redirectingInitializer = new MethodDefinition($"Initialize${declaringType.Name}_{cls.Name}",
785 MethodAttributes.Public | MethodAttributes.HideBySig | MethodAttributes.Static,
786 _typeSystem.Void)
787 {
788 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
789 DeclaringType = parentType
790 };
791 processor = redirectingInitializer.Body.GetILProcessor();
792 processor.Emit(OpCodes.Call, currentInitializer);
793 processor.Emit(OpCodes.Ret);
794 parentType.Methods.Add(redirectingInitializer);
795 currentInitializer = redirectingInitializer;
796 }
797
798 currentDeclaringType = parentType;
799 }
800
801 _directCallInitializeMethods.Add(currentInitializer);
802
803 // Create the static constructor
804 //
805 // public static .cctor() {
806 // Constructor();
807 // }
808 var cctor = new MethodDefinition(".cctor", MethodAttributes.Private | MethodAttributes.HideBySig | MethodAttributes.SpecialName | MethodAttributes.RTSpecialName | MethodAttributes.Static, _typeSystem.Void)
809 {
810 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
811 DeclaringType = cls,
812 };
813
814 processor = cctor.Body.GetILProcessor();
815 processor.Emit(OpCodes.Call, constructor);
816 processor.Emit(OpCodes.Ret);
817
818 cls.Methods.Add(FixDebugInformation(cctor));
819
820 // Create the Invoke method based on the original method (same signature)
821 //
822 // public static XXX Invoke(...args) {
823 // if (BurstCompiler.IsEnabled)
824 // {
825 // var funcPtr = GetFunctionPointer();
826 // if (funcPtr != null) return funcPtr(...args);
827 // }
828 // return OriginalMethod(...args);
829 // }
830 var invokeAttributes = managedFallbackMethod.Attributes;
831 invokeAttributes &= ~MethodAttributes.Private;
832 invokeAttributes |= MethodAttributes.Public;
833 var invoke = new MethodDefinition(InvokeName, invokeAttributes, burstCompileMethod.ReturnType)
834 {
835 ImplAttributes = MethodImplAttributes.IL | MethodImplAttributes.Managed,
836 DeclaringType = cls
837 };
838
839 var signature = new CallSite(burstCompileMethod.ReturnType)
840 {
841 CallingConvention = MethodCallingConvention.C
842 };
843
844 foreach (var parameter in burstCompileMethod.Parameters)
845 {
846 invoke.Parameters.Add(parameter);
847 signature.Parameters.Add(parameter);
848 }
849
850 invoke.Body.Variables.Add(new VariableDefinition(_typeSystem.IntPtr));
851 invoke.Body.InitLocals = true;
852
853 processor = invoke.Body.GetILProcessor();
854 processor.Emit(OpCodes.Call, _burstCompilerIsEnabledMethodDefinition);
855 var branchPosition0 = processor.Body.Instructions[processor.Body.Instructions.Count - 1];
856
857 processor.Emit(OpCodes.Call, getFunctionPointerMethod);
858 processor.Emit(OpCodes.Stloc_0);
859 processor.Emit(OpCodes.Ldloc_0);
860 var branchPosition1 = processor.Body.Instructions[processor.Body.Instructions.Count - 1];
861
862 EmitArguments(processor, invoke);
863 processor.Emit(OpCodes.Ldloc_0);
864 processor.Emit(OpCodes.Calli, signature);
865 processor.Emit(OpCodes.Ret);
866 var previousRet = processor.Body.Instructions[processor.Body.Instructions.Count - 1];
867
868 EmitArguments(processor, invoke);
869 processor.Emit(OpCodes.Call, managedFallbackMethod);
870 processor.Emit(OpCodes.Ret);
871
872 // Insert the branch once we have emitted the instructions
873 processor.InsertAfter(branchPosition0, Instruction.Create(OpCodes.Brfalse, previousRet.Next));
874 processor.InsertAfter(branchPosition1, Instruction.Create(OpCodes.Brfalse, previousRet.Next));
875 cls.Methods.Add(FixDebugInformation(invoke));
876
877 // Final patching of the original method
878 // public static XXX OriginalMethod(...args) {
879 // Name_RID.$Postfix.Invoke(...args);
880 // ret;
881 // }
882 burstCompileMethod.Body = new MethodBody(burstCompileMethod);
883 processor = burstCompileMethod.Body.GetILProcessor();
884 EmitArguments(processor, burstCompileMethod);
885 processor.Emit(OpCodes.Call, invoke);
886 processor.Emit(OpCodes.Ret);
887 FixDebugInformation(burstCompileMethod);
888 }
889
890 private static MethodDefinition FixDebugInformation(MethodDefinition method)
891 {
892 method.DebugInformation.Scope = new ScopeDebugInformation(method.Body.Instructions.First(), method.Body.Instructions.Last());
893 return method;
894 }
895
896 private AssemblyDefinition GetAsmDefinitionFromFile(AssemblyResolver loader, string assemblyName)
897 {
898 if (loader.TryResolve(AssemblyNameReference.Parse(assemblyName), out var result))
899 {
900 return result;
901 }
902 return null;
903 }
904
905 private MethodReference _unityEngineInitializeOnLoadAttributeCtor;
906 private TypeReference _unityEngineRuntimeInitializeLoadType;
907 private FieldDefinition _unityEngineRuntimeInitializeLoadAfterAssemblies;
908 private MethodReference _unityEditorInitilizeOnLoadAttributeCtor;
909
910 private void InitializeBurstAssembly(AssemblyDefinition burstAssembly)
911 {
912 _burstAssembly = burstAssembly;
913 _burstCompilerTypeDefinition = burstAssembly.MainModule.GetType("Unity.Burst", "BurstCompiler");
914
915 _burstCompilerIsEnabledMethodDefinition = _assemblyDefinition.MainModule.ImportReference(_burstCompilerTypeDefinition.Methods.FirstOrDefault(x => x.Name == "get_IsEnabled"));
916 _burstCompilerCompileILPPMethod = _assemblyDefinition.MainModule.ImportReference(_burstCompilerTypeDefinition.Methods.FirstOrDefault(x => x.Name == "CompileILPPMethod2"));
917 _burstCompilerGetILPPMethodFunctionPointer = _assemblyDefinition.MainModule.ImportReference(_burstCompilerTypeDefinition.Methods.FirstOrDefault(x => x.Name == "GetILPPMethodFunctionPointer2"));
918
919 var reinitializeAttribute = _burstCompilerTypeDefinition.NestedTypes.FirstOrDefault(x => x.Name == "StaticTypeReinitAttribute");
920 _burstCompilerCompileUnsafeStaticMethodReinitialiseAttributeCtor = _assemblyDefinition.MainModule.ImportReference(reinitializeAttribute.Methods.FirstOrDefault(x=>x.Name == ".ctor" && x.HasParameters));
921
922 var corLibrary = Loader.Resolve((AssemblyNameReference)_typeSystem.CoreLibrary);
923 _systemType = _assemblyDefinition.MainModule.ImportReference(corLibrary.MainModule.GetType("System.Type"));
924 _systemDelegateType = _assemblyDefinition.MainModule.ImportReference(corLibrary.MainModule.GetType("System.MulticastDelegate"));
925 _systemASyncCallbackType = _assemblyDefinition.MainModule.ImportReference(corLibrary.MainModule.GetType("System.AsyncCallback"));
926 _systemIASyncResultType = _assemblyDefinition.MainModule.ImportReference(corLibrary.MainModule.GetType("System.IAsyncResult"));
927
928 var asmDef = GetAsmDefinitionFromFile(Loader, "UnityEngine.CoreModule");
929 var runtimeInitializeOnLoadMethodAttribute = asmDef.MainModule.GetType("UnityEngine", "RuntimeInitializeOnLoadMethodAttribute");
930 var runtimeInitializeLoadType = asmDef.MainModule.GetType("UnityEngine", "RuntimeInitializeLoadType");
931
932 var burstDiscardType = asmDef.MainModule.GetType("Unity.Burst", "BurstDiscardAttribute");
933 _burstDiscardAttributeConstructor = _assemblyDefinition.MainModule.ImportReference(burstDiscardType.Methods.First(method => method.Name == ".ctor"));
934
935 _unityEngineInitializeOnLoadAttributeCtor = _assemblyDefinition.MainModule.ImportReference(runtimeInitializeOnLoadMethodAttribute.Methods.FirstOrDefault(x => x.Name == ".ctor" && x.HasParameters));
936 _unityEngineRuntimeInitializeLoadType = _assemblyDefinition.MainModule.ImportReference(runtimeInitializeLoadType);
937 _unityEngineRuntimeInitializeLoadAfterAssemblies = runtimeInitializeLoadType.Fields.FirstOrDefault(x => x.Name=="AfterAssembliesLoaded");
938
939 if (IsForEditor && !_skipInitializeOnLoad)
940 {
941 asmDef = GetAsmDefinitionFromFile(Loader, "UnityEditor.CoreModule");
942 if (asmDef == null)
943 asmDef = GetAsmDefinitionFromFile(Loader, "UnityEditor");
944 var initializeOnLoadMethodAttribute = asmDef.MainModule.GetType("UnityEditor", "InitializeOnLoadMethodAttribute");
945
946 _unityEditorInitilizeOnLoadAttributeCtor = _assemblyDefinition.MainModule.ImportReference(initializeOnLoadMethodAttribute.Methods.FirstOrDefault(x => x.Name == ".ctor" && !x.HasParameters));
947 }
948 }
949
950 private static void EmitArguments(ILProcessor processor, MethodDefinition method)
951 {
952 for (var i = 0; i < method.Parameters.Count; i++)
953 {
954 switch (i)
955 {
956 case 0:
957 processor.Emit(OpCodes.Ldarg_0);
958 break;
959 case 1:
960 processor.Emit(OpCodes.Ldarg_1);
961 break;
962 case 2:
963 processor.Emit(OpCodes.Ldarg_2);
964 break;
965 case 3:
966 processor.Emit(OpCodes.Ldarg_3);
967 break;
968 default:
969 if (i <= 255)
970 {
971 processor.Emit(OpCodes.Ldarg_S, (byte)i);
972 }
973 else
974 {
975 processor.Emit(OpCodes.Ldarg, i);
976 }
977 break;
978 }
979 }
980 }
981
982 private static bool TryGetBurstCompilerAttribute(ICustomAttributeProvider provider, out CustomAttribute customAttribute)
983 {
984 if (provider.HasCustomAttributes)
985 {
986 foreach (var customAttr in provider.CustomAttributes)
987 {
988 if (customAttr.Constructor.DeclaringType.Name == "BurstCompileAttribute")
989 {
990 customAttribute = customAttr;
991 return true;
992 }
993 }
994 }
995 customAttribute = null;
996 return false;
997 }
998 }
999}