private void ApplyPostProcessingOnJobCode(MethodDefinition[] methodUsedByLambdaJobs, LambdaParameterValueInformations lambdaParameterValueInformations) { var forEachInvocations = new List <(MethodDefinition, Instruction)>(); var methodDefinition = methodUsedByLambdaJobs.First(); forEachInvocations.AddRange(methodDefinition.Body.Instructions.Where(IsChunkEntitiesForEachInvocation).Select(i => (methodDefinition, i))); foreach (var methodUsedByLambdaJob in methodUsedByLambdaJobs) { var methodBody = methodUsedByLambdaJob.Body; var displayClassVariable = methodBody.Variables.SingleOrDefault(v => v.VariableType.Name.Contains("DisplayClass")); if (displayClassVariable != null) { TypeDefinition displayClass = displayClassVariable.VariableType.Resolve(); bool allDelegatesAreGuaranteedNotToOutliveMethod = displayClass.IsValueType || CecilHelpers.AllDelegatesAreGuaranteedNotToOutliveMethodFor(methodUsedByLambdaJob); if (!displayClass.IsValueType && allDelegatesAreGuaranteedNotToOutliveMethod) { CecilHelpers.PatchMethodThatUsedDisplayClassToTreatItAsAStruct(methodBody, displayClassVariable, displayClass); CecilHelpers.PatchDisplayClassToBeAStruct(displayClass); } } } int counter = 1; foreach (var(methodUsedByLambdaJob, instruction) in forEachInvocations) { var methodBody = methodUsedByLambdaJob.Body; var(ldFtn, newObj) = FindClosureCreatingInstructions(methodBody, instruction); var newType = new TypeDefinition("", "InlineEntitiesForEachInvocation" + counter++, TypeAttributes.NestedPublic | TypeAttributes.SequentialLayout, methodUsedByLambdaJob.Module.ImportReference(typeof(ValueType))) { DeclaringType = methodUsedByLambdaJob.DeclaringType }; methodUsedByLambdaJob.DeclaringType.NestedTypes.Add(newType); CloneLambdaMethodAndItsLocalMethods(); var iterateEntitiesMethod = CreateIterateEntitiesMethod(lambdaParameterValueInformations); var variable = new VariableDefinition(newType); methodBody.Variables.Add(variable); InstructionExtensions.MakeNOP(ldFtn.Previous); InstructionExtensions.MakeNOP(ldFtn); newObj.OpCode = OpCodes.Ldnull; newObj.Operand = null; var displayClassVariable = methodBody.Variables.SingleOrDefault(v => v.VariableType.Name.Contains("DisplayClass")); if (displayClassVariable == null) { continue; } var ilProcessor = methodBody.GetILProcessor(); ilProcessor.InsertAfter(instruction, new List <Instruction> { //no need to drop the delegate from the stack, because we just removed the function that placed it on the stack in the first place. //do not drop the description from the stack, as the original method returns it, and we want to maintain stack behaviour. //call our new method Instruction.Create(OpCodes.Ldloca, variable), Instruction.Create(OpCodes.Initobj, newType), Instruction.Create(OpCodes.Ldloca, variable), Instruction.Create(OpCodes.Ldloca, displayClassVariable), Instruction.Create(OpCodes.Call, ReadFromDisplayClassMethod), Instruction.Create(OpCodes.Ldloca, variable), Instruction.Create(OpCodes.Ldarga, methodBody.Method.Parameters.First(p => p.ParameterType.Name == nameof(ArchetypeChunk))), Instruction.Create(OpCodes.Ldarg_0), Instruction.Create(OpCodes.Ldfld, lambdaParameterValueInformations._runtimesField), Instruction.Create(OpCodes.Call, (MethodReference)iterateEntitiesMethod), }); #if ENABLE_DOTS_COMPILER_CHUNKS var chunkEntitiesInvocation = LambdaJobDescriptionConstruction.FindInstructionThatPushedArg(methodBody.Method, 0, instruction); if (chunkEntitiesInvocation.Operand is MethodReference mr && mr.Name == "get_" + nameof(ArchetypeChunk.Entities) && mr.DeclaringType.Name == nameof(ArchetypeChunk)) { CecilHelpers.EraseMethodInvocationFromInstructions(ilProcessor, chunkEntitiesInvocation); } #endif CecilHelpers.EraseMethodInvocationFromInstructions(ilProcessor, instruction); } }
public static (JobStructForLambdaJob, List <DiagnosticMessage>) Rewrite(MethodDefinition methodContainingLambdaJob, LambdaJobDescriptionConstruction lambdaJobDescriptionConstruction) { var diagnosticMessages = new List <DiagnosticMessage>(); if (methodContainingLambdaJob.DeclaringType.CustomAttributes.Any(c => (c.AttributeType.Name == "ExecuteAlways" && c.AttributeType.Namespace == "UnityEngine"))) { diagnosticMessages.Add(UserError.DC0032(methodContainingLambdaJob.DeclaringType, methodContainingLambdaJob, lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction)); } if (keepUnmodifiedVersionAroundForDebugging) { CecilHelpers.CloneMethodForDiagnosingProblems(methodContainingLambdaJob); } //in order to make it easier and safer to operate on IL, we're going to "Simplify" it. This means transforming opcodes that have short variants, into the full one //so that when you analyze, you can assume (ldarg, 0) opcode + operand pair, and don't have to go check for the shorthand version of ldarg_0, without an operand. //more importantly, it also rewrites branch.s (branch to a instruction that is closeby enough that the offset fits within a byte), to a regular branch instructions. //this is the only sane thing to do, because when we rewrite IL code, we add instructions, so it can happen that by adding instructions, what was possible to be a short //branch instruction, now is no longer a valid short branch target, and cecil doesn't warn against that, it will just generate broken IL, and you'll spend a long time //figuring out what is going on. methodContainingLambdaJob.Body.SimplifyMacros(); var methodLambdaWasEmittedAs = lambdaJobDescriptionConstruction.MethodLambdaWasEmittedAs; if (methodLambdaWasEmittedAs.DeclaringType.TypeReferenceEquals(methodContainingLambdaJob.DeclaringType) && !lambdaJobDescriptionConstruction.AllowReferenceTypes) { // Sometimes roslyn emits the lambda as an instance method in the same type of the method that contains the lambda expression. // it does this only in the situation where the lambda captures a field _and_ does not capture any locals. In this case // there's no displayclass being created. We should figure out exactly what instruction caused this behaviour, and tell the user // she can't read a field like that. // Here is an example: https://sharplab.io/#v2:D4AQTAjAsAUCDMACciDCiDetE8QMwBsB7AQwBd8BLAUwIBNEBeReAOgAY8BubXXnBMgAsiACoALSgGcA4tQB21AE7lqUgLLUAtgCNlIAKxlKReVIAUASiwwAkLYD0DsZKmICJXXRKIA7pQICRABzBWVVRB8tbT0lfABXeQBjY1NENLJxamQwNFYXbKVqEik06QAuWHsnHABaAvdPHW90+QIAT0QkkgAHMniitx88Gnp8gEkzMmKGIjwu3v6lSnlgxEzpRGoADx6CSiTKMg6AGirHZ1xfbO75ELCVacjEaMyiWbv0MiJEPS3t6hJeLTBgrKTTEh0fi4HAAESIIAgYHMViYAD4bDCsThUKZSgRqKwAOrLabmEa0OiWLiIaEwgC+1LpfBg2JwNQkmw8Xh8/kC90Uj2yURiygSyVSdwyWRyeQaRRKZSklVZbJqiHqEmy3OaPlMHUQRQAjvFKINEAByDZSC2ReQMHolKRqRBHdY/YaJFImeSsZlwhFIlGWdGYtm4eGc1YWa1M1XYxk8eOIelVaGCEAiTmyB6qKQAQVQHikFhDNmqzi1zvWvh+Ou8biyRQF4SePnA+S1huKpTumxK+CIgSIvmV53V9Uy2WdSVMDHrPm6fQGLp8xG6QRI9vW4nIg6USRdU66RC0PQCYu+Wy0R3Hlxw7dyV5IKXiJECnXEQ4Yx/BETmO7akQG53nUgFUEo4KNDyrQGkuSyrlQlJ2j+MqzmeF5xLW8T0Ig8RSG+H6IAAVvhZCgbg2huiKuhingXqSpEbgrOBIyQRQ3TOicvzAogUgrIegHNu+Cp0J00gUQ+sp4EQcQLkM8jtL4JDtNxbp8kE/FngaRT4dkmR7qYhLnPCiLIqijAYucti4mYQ6EiSRzUOSoxUjS/opnG4YeSsFDbEwiDsEm4amUGFlWXY9i2fiDmks5FL0NStKRTZeL2cScXmNsXlsom0Kpsm6YiKFyJmZEKRlgVMLphAABswiIJGkjRuY6BJJVsD0kAA= var illegalFieldRead = methodLambdaWasEmittedAs.Body.Instructions.FirstOrDefault(IsIllegalFieldRead); if (illegalFieldRead != null) { UserError.DC0001(methodContainingLambdaJob, illegalFieldRead, (FieldReference)illegalFieldRead.Operand).Throw(); } // Similarly, we could end up here because the lambda captured neither a field nor a local but simply // uses the this-reference. This could be the case if the user calls a function that takes this as a // parameter, either explicitly (static) or implicitly (extension or member method). var illegalInvocations = methodLambdaWasEmittedAs.Body.Instructions.Where(IsIllegalInvocation).ToArray(); if (illegalInvocations.Any()) { foreach (var illegalInvocation in illegalInvocations) { if (!IsPermittedIllegalInvocation(illegalInvocation)) { UserError.DC0002(methodContainingLambdaJob, illegalInvocation, (MethodReference)illegalInvocation.Operand, methodLambdaWasEmittedAs.DeclaringType).Throw(); } } // We only had illegal invocations where we were invoking an allowed method on our declaring type. // This is allowed as we will rewrite these invocations. // When we clone this method we need to rewrite it correctly to reflect that this is the case. lambdaJobDescriptionConstruction.HasAllowedMethodInvokedWithThis = true; } // This should never hit, but is here to make sure that in case we have a bug in detecting // why roslyn emitted it like this, we can at least report an error, instead of silently generating invalid code. // We do need to allow this case when we have an allowed method invoked with this that is later replaced with codegen. if (!lambdaJobDescriptionConstruction.HasAllowedMethodInvokedWithThis) { InternalCompilerError.DCICE001(methodContainingLambdaJob).Throw(); } bool IsIllegalFieldRead(Instruction i) { if (i.Previous == null) { return(false); } if (i.OpCode != OpCodes.Ldfld && i.OpCode != OpCodes.Ldflda) { return(false); } return(i.Previous.OpCode == OpCodes.Ldarg_0); } bool IsIllegalInvocation(Instruction i) { if (!i.IsInvocation(out _)) { return(false); } var declaringType = methodContainingLambdaJob.DeclaringType; var method = (MethodReference)i.Operand; // is it an instance method? var resolvedMethod = method.Resolve(); if (declaringType.TypeReferenceEqualsOrInheritsFrom(method.DeclaringType) && !resolvedMethod.IsStatic) { return(true); } // is it a method that potentially takes this as a parameter? foreach (var param in method.Parameters) { if (declaringType.TypeReferenceEqualsOrInheritsFrom(param.ParameterType) || declaringType.TypeImplements(param.ParameterType)) { return(true); } } return(false); } // Check for permitted illegal invocations // These are due to calling a method that we later stub out with codegen or also can be due to calling // a local method that contains a method that we stub out with codegen. bool IsPermittedIllegalInvocation(Instruction instruction) { // Check to see if this method is permitted if (instruction.OpCode == OpCodes.Call || instruction.OpCode == OpCodes.Callvirt) { var methodRef = instruction.Operand as MethodReference; if (JobStructForLambdaJob.IsPermittedMethodToInvokeWithThis(methodRef)) { return(true); } else { // Recurse into methods if they are compiler generated local methods var methodDef = methodRef.Resolve(); if (methodDef != null && methodDef.CustomAttributes.Any(c => c.AttributeType.Name == nameof(CompilerGeneratedAttribute) && c.AttributeType.Namespace == typeof(CompilerGeneratedAttribute).Namespace)) { foreach (var methodInstruction in methodDef.Body.Instructions) { if (IsIllegalInvocation(methodInstruction) && !IsPermittedIllegalInvocation(methodInstruction)) { return(false); } } return(true); } } } return(false); } } var moduleDefinition = methodContainingLambdaJob.Module; var body = methodContainingLambdaJob.Body; var ilProcessor = body.GetILProcessor(); VariableDefinition displayClassVariable = null; if (lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { bool allDelegatesAreGuaranteedNotToOutliveMethod = lambdaJobDescriptionConstruction.DisplayClass.IsValueType() || CecilHelpers.AllDelegatesAreGuaranteedNotToOutliveMethodFor(methodContainingLambdaJob); displayClassVariable = body.Variables.Single(v => v.VariableType.TypeReferenceEquals(lambdaJobDescriptionConstruction.DisplayClass)); //in this step we want to get rid of the heap allocation for the delegate. In order to make the rest of the code easier to reason about and write, //we'll make sure that while we do this, we don't change the total stackbehaviour. Because this used to push a delegate onto the evaluation stack, //we also have to write something to the evaluation stack. Later in this method it will be popped, so it doesn't matter what it is really. I use Ldc_I4_0, //as I found it introduced the most reasonable artifacts when the code is decompiled back into C#. lambdaJobDescriptionConstruction.DelegateProducingSequence.RewriteToKeepDisplayClassOnEvaluationStack(); if (allDelegatesAreGuaranteedNotToOutliveMethod) { ChangeAllDisplayClassesToStructs(methodContainingLambdaJob); } } else { //if the lambda is not capturing, roslyn will recycle the delegate in a static field. not so great for us. let's nop out all that code. var instructionThatPushedDelegate = CecilHelpers.FindInstructionThatPushedArg(methodContainingLambdaJob, 1, lambdaJobDescriptionConstruction.WithCodeInvocationInstruction); var result = CecilHelpers.MatchesDelegateProducingPattern(methodContainingLambdaJob, instructionThatPushedDelegate, CecilHelpers.DelegateProducingPattern.MatchSide.Start); result?.RewriteToProduceSingleNullValue(); } FieldDefinition entityQueryField = null; if (lambdaJobDescriptionConstruction.Kind != LambdaJobDescriptionKind.Job) { entityQueryField = InjectAndInitializeEntityQueryField.InjectAndInitialize(methodContainingLambdaJob, lambdaJobDescriptionConstruction, methodLambdaWasEmittedAs.Parameters); } var generatedJobStruct = JobStructForLambdaJob.CreateNewJobStruct(lambdaJobDescriptionConstruction); if (generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst != null) { var constructorInfo = generatedJobStruct.ExecuteDelegateType.GetConstructors().First(c => c.GetParameters().Length == 2); var instructions = new List <Instruction>() { Instruction.Create(OpCodes.Ldnull), Instruction.Create(OpCodes.Ldftn, generatedJobStruct.RunWithoutJobSystemMethod), Instruction.Create(OpCodes.Newobj, moduleDefinition.ImportReference(constructorInfo)), Instruction.Create(OpCodes.Stsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst) }; if (generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst != null) { instructions.Add(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); var methodInfo = typeof(InternalCompilerInterface) .GetMethods(BindingFlags.Static | BindingFlags.Public) .Where(m => m.Name == nameof(InternalCompilerInterface.BurstCompile)) .Single(m => m.GetParameters().First().ParameterType == generatedJobStruct.ExecuteDelegateType); instructions.Add(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(methodInfo))); instructions.Add(Instruction.Create(OpCodes.Stsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst)); } InjectAndInitializeEntityQueryField.InsertIntoOnCreateForCompilerMethod(methodContainingLambdaJob.DeclaringType, instructions.ToArray()); } IEnumerable <Instruction> InstructionsToReplaceScheduleInvocationWith() { var newJobStructVariable = new VariableDefinition(generatedJobStruct.TypeDefinition); body.Variables.Add(newJobStructVariable); bool storeJobHandleInVariable = (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule || lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.ScheduleParallel); VariableDefinition tempStorageForJobHandle = null; if (storeJobHandleInVariable) { tempStorageForJobHandle = new VariableDefinition(moduleDefinition.ImportReference(typeof(JobHandle))); body.Variables.Add(tempStorageForJobHandle); // If we aren't using an implicit system dependency and we're replacing the .Schedule() function on the description, // the lambdajobdescription and the jobhandle argument to that function will be on the stack. // we're going to need the jobhandle later when we call JobChunkExtensions.Schedule(), so lets stuff it in a variable. // If we are using implicit system dependency, lets put that in our temp instead. if (lambdaJobDescriptionConstruction.UseImplicitSystemDependency) { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(SystemBase).GetMethod("get_Dependency", BindingFlags.Instance | BindingFlags.NonPublic)))); } yield return(Instruction.Create(OpCodes.Stloc, tempStorageForJobHandle)); } //pop the Description struct off the stack, its services are no longer required yield return(Instruction.Create(OpCodes.Pop)); yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); yield return(Instruction.Create(OpCodes.Initobj, generatedJobStruct.TypeDefinition)); // Call ScheduleTimeInitializeMethod yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); yield return(Instruction.Create(OpCodes.Ldarg_0)); if (lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { //only when the lambda is capturing, did we emit the ScheduleTimeInitialize method to take a displayclass argument var opcode = methodLambdaWasEmittedAs.DeclaringType.IsValueType() ? OpCodes.Ldloca : OpCodes.Ldloc; yield return(Instruction.Create(opcode, displayClassVariable)); } yield return(Instruction.Create(OpCodes.Call, generatedJobStruct.ScheduleTimeInitializeMethod)); MethodInfo FindRunOrScheduleMethod() { switch (lambdaJobDescriptionConstruction.Kind) { case LambdaJobDescriptionKind.Entities: case LambdaJobDescriptionKind.Chunk: if (lambdaJobDescriptionConstruction.IsInSystemBase) { switch (lambdaJobDescriptionConstruction.ExecutionMode) { case ExecutionMode.Run: return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunJobChunk))); case ExecutionMode.Schedule: return(typeof(JobChunkExtensions).GetMethod(nameof(JobChunkExtensions.ScheduleSingle))); case ExecutionMode.ScheduleParallel: return(typeof(JobChunkExtensions).GetMethod(nameof(JobChunkExtensions.ScheduleParallel))); default: throw new ArgumentOutOfRangeException(); } } else { // Keep legacy behaviour in JobComponentSystems intact (aka "Schedule" equals "ScheduleParallel") if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { return(typeof(JobChunkExtensions).GetMethod(nameof(JobChunkExtensions.ScheduleParallel))); } return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunJobChunk))); } case LambdaJobDescriptionKind.Job: if (lambdaJobDescriptionConstruction.IsInSystemBase) { switch (lambdaJobDescriptionConstruction.ExecutionMode) { case ExecutionMode.Run: return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunIJob))); case ExecutionMode.Schedule: return(typeof(IJobExtensions).GetMethod(nameof(IJobExtensions.Schedule))); default: throw new ArgumentOutOfRangeException(); } } else { if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { return(typeof(IJobExtensions).GetMethod(nameof(IJobExtensions.Schedule))); } return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunIJob))); } default: throw new ArgumentOutOfRangeException(); } } // Call CompleteDependency method to complete previous dependencies if we are running in SystemBase if (lambdaJobDescriptionConstruction.IsInSystemBase && lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Run) { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(SystemBase).GetMethod("CompleteDependency", BindingFlags.Instance | BindingFlags.NonPublic)))); } MethodReference runOrScheduleMethod; if (lambdaJobDescriptionConstruction.WithStructuralChanges) { runOrScheduleMethod = generatedJobStruct.TypeDefinition.Methods.First(definition => definition.Name == "Execute"); } else { runOrScheduleMethod = moduleDefinition.ImportReference(FindRunOrScheduleMethod()) .MakeGenericInstanceMethod(generatedJobStruct.TypeDefinition); } if (lambdaJobDescriptionConstruction.WithStructuralChanges) { yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); } else { yield return(Instruction.Create(runOrScheduleMethod.Parameters.First().ParameterType.IsByReference ? OpCodes.Ldloca : OpCodes.Ldloc, newJobStructVariable)); } switch (lambdaJobDescriptionConstruction.Kind) { case LambdaJobDescriptionKind.Entities: case LambdaJobDescriptionKind.Chunk: if (lambdaJobDescriptionConstruction.WithStructuralChanges) { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldfld, entityQueryField)); } else { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldfld, entityQueryField)); } break; case LambdaJobDescriptionKind.Job: //job.Schedule() takes no entityQuery... break; } // Store returned JobHandle in temp varaible or back in SystemBase.depenedency if (storeJobHandleInVariable) { yield return(Instruction.Create(OpCodes.Ldloc, tempStorageForJobHandle)); } else if (lambdaJobDescriptionConstruction.UseImplicitSystemDependency) { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(SystemBase).GetMethod("get_Dependency", BindingFlags.Instance | BindingFlags.NonPublic)))); } if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Run && !lambdaJobDescriptionConstruction.WithStructuralChanges) { if (!lambdaJobDescriptionConstruction.UsesBurst) { yield return(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); } else { yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(JobsUtility).GetMethod("get_" + nameof(JobsUtility.JobCompilerEnabled))))); var targetInstruction = Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst); yield return(Instruction.Create(OpCodes.Brtrue, targetInstruction)); yield return(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); var finalBranchDestination = Instruction.Create(OpCodes.Nop); yield return(Instruction.Create(OpCodes.Br, finalBranchDestination)); yield return(targetInstruction); yield return(finalBranchDestination); } } yield return(Instruction.Create(OpCodes.Call, runOrScheduleMethod)); if (lambdaJobDescriptionConstruction.UseImplicitSystemDependency) { yield return(Instruction.Create(OpCodes.Stloc, tempStorageForJobHandle)); yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldloc, tempStorageForJobHandle)); yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(SystemBase).GetMethod("set_Dependency", BindingFlags.Instance | BindingFlags.NonPublic)))); } if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Run && generatedJobStruct.WriteToDisplayClassMethod != null && lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); var opcode = methodLambdaWasEmittedAs.DeclaringType.IsValueType() ? OpCodes.Ldloca : OpCodes.Ldloc; yield return(Instruction.Create(opcode, displayClassVariable)); yield return(Instruction.Create(OpCodes.Call, generatedJobStruct.WriteToDisplayClassMethod)); } } foreach (var invokedMethod in lambdaJobDescriptionConstruction.InvokedConstructionMethods) { bool invokedMethodServesNoPurposeAtRuntime = invokedMethod.MethodName != nameof(LambdaJobQueryConstructionMethods.WithSharedComponentFilter); if (invokedMethodServesNoPurposeAtRuntime) { CecilHelpers.EraseMethodInvocationFromInstructions(ilProcessor, invokedMethod.InstructionInvokingMethod); } else { // Rewrite WithSharedComponentFilter calls as they need to modify EntityQuery dynamically if (invokedMethod.MethodName == nameof(LambdaJobQueryConstructionMethods.WithSharedComponentFilter)) { var setSharedComponentFilterOnQueryMethod = moduleDefinition.ImportReference( (lambdaJobDescriptionConstruction.Kind == LambdaJobDescriptionKind.Entities ? typeof(ForEachLambdaJobDescription_SetSharedComponent) : typeof(LambdaJobChunkDescription_SetSharedComponent)).GetMethod( nameof(LambdaJobChunkDescription_SetSharedComponent.SetSharedComponentFilterOnQuery))); var callingTypeReference = lambdaJobDescriptionConstruction.IsInSystemBase ? moduleDefinition.ImportReference(typeof(ForEachLambdaJobDescription)) : moduleDefinition.ImportReference(typeof(ForEachLambdaJobDescriptionJCS)); MethodReference genericSetSharedComponentFilterOnQueryMethod = setSharedComponentFilterOnQueryMethod.MakeGenericInstanceMethod( new TypeReference[] { callingTypeReference }.Concat(invokedMethod.TypeArguments).ToArray()); // Change invocation to invocation of helper method and add EntityQuery parameter to be modified var setSharedComponentFilterOnQueryInstructions = new List <Instruction> { Instruction.Create(OpCodes.Ldarg_0), Instruction.Create(OpCodes.Ldfld, entityQueryField), Instruction.Create(OpCodes.Call, genericSetSharedComponentFilterOnQueryMethod) }; ilProcessor.Replace(invokedMethod.InstructionInvokingMethod, setSharedComponentFilterOnQueryInstructions); } } } var scheduleInstructions = InstructionsToReplaceScheduleInvocationWith().ToList(); ilProcessor.InsertAfter(lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction, scheduleInstructions); lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction.MakeNOP(); var codegenInitializeMethod = GetOrMakeOnCreateForCompilerMethod(lambdaJobDescriptionConstruction.ContainingMethod.DeclaringType); return(generatedJobStruct, diagnosticMessages); }
public static JobStructForLambdaJob Rewrite(MethodDefinition methodContainingLambdaJob, LambdaJobDescriptionConstruction lambdaJobDescriptionConstruction, List <DiagnosticMessage> warnings) { if (methodContainingLambdaJob.DeclaringType.CustomAttributes.Any(c => (c.AttributeType.Name == "ExecuteAlways" && c.AttributeType.Namespace == "UnityEngine"))) { warnings.Add(UserError.DC0032(methodContainingLambdaJob.DeclaringType, methodContainingLambdaJob, lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction)); } if (keepUnmodifiedVersionAroundForDebugging) { CecilHelpers.CloneMethodForDiagnosingProblems(methodContainingLambdaJob); } //in order to make it easier and safer to operate on IL, we're going to "Simplify" it. This means transforming opcodes that have short variants, into the full one //so that when you analyze, you can assume (ldarg, 0) opcode + operand pair, and don't have to go check for the shorthand version of ldarg_0, without an operand. //more importantly, it also rewrites branch.s (branch to a instruction that is closeby enough that the offset fits within a byte), to a regular branch instructions. //this is the only sane thing to do, because when we rewrite IL code, we add instructions, so it can happen that by adding instructions, what was possible to be a short //branch instruction, now is no longer a valid short branch target, and cecil doesn't warn against that, it will just generate broken IL, and you'll spend a long time //figuring out what is going on. methodContainingLambdaJob.Body.SimplifyMacros(); var methodLambdaWasEmittedAs = lambdaJobDescriptionConstruction.MethodLambdaWasEmittedAs; if (methodLambdaWasEmittedAs.DeclaringType.TypeReferenceEquals(methodContainingLambdaJob.DeclaringType) && !lambdaJobDescriptionConstruction.AllowReferenceTypes) { //sometimes roslyn emits the lambda as an instance method in the same type of the method that contains the lambda expression. //it does this only in the situation where the lambda captures a field _and_ does not capture any locals. in this case //there's no displayclass being created. We should figure out exactly what instruction caused this behaviour, and tell the user //she can't read a field like that. var illegalFieldRead = methodLambdaWasEmittedAs.Body.Instructions.FirstOrDefault(i => i.OpCode == OpCodes.Ldfld && i.Previous?.OpCode == OpCodes.Ldarg_0); if (illegalFieldRead != null) { UserError.DC0001(methodContainingLambdaJob, illegalFieldRead, (FieldReference)illegalFieldRead.Operand).Throw(); } var illegalInvocation = methodLambdaWasEmittedAs.Body.Instructions.FirstOrDefault(i => i.IsInvocation() && methodContainingLambdaJob.DeclaringType.TypeReferenceEqualsOrInheritsFrom(((MethodReference)i.Operand).DeclaringType)); if (illegalInvocation != null) { UserError.DC0002(methodContainingLambdaJob, illegalInvocation, (MethodReference)illegalInvocation.Operand).Throw(); } //this should never hit, but is here to make sure that in case we have a bug in detecting why roslyn emitted it like this, we can at least report an error, instead of silently generating invalid code. InternalCompilerError.DCICE001(methodContainingLambdaJob).Throw(); } var moduleDefinition = methodContainingLambdaJob.Module; var body = methodContainingLambdaJob.Body; var ilProcessor = body.GetILProcessor(); VariableDefinition displayClassVariable = null; if (lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { bool allDelegatesAreGuaranteedNotToOutliveMethod = lambdaJobDescriptionConstruction.DisplayClass.IsValueType || CecilHelpers.AllDelegatesAreGuaranteedNotToOutliveMethodFor(methodContainingLambdaJob); displayClassVariable = body.Variables.Single(v => v.VariableType.TypeReferenceEquals(lambdaJobDescriptionConstruction.DisplayClass)); //in this step we want to get rid of the heap allocation for the delegate. In order to make the rest of the code easier to reason about and write, //we'll make sure that while we do this, we don't change the total stackbehaviour. Because this used to push a delegate onto the evaluation stack, //we also have to write something to the evaluation stack. Later in this method it will be popped, so it doesn't matter what it is really. I use Ldc_I4_0, //as I found it introduced the most reasonable artifacts when the code is decompiled back into C#. lambdaJobDescriptionConstruction.DelegateProducingSequence.RewriteToKeepDisplayClassOnEvaluationStack(); if (allDelegatesAreGuaranteedNotToOutliveMethod) { ChangeAllDisplayClassesToStructs(methodContainingLambdaJob); } } else { //if the lambda is not capturing, roslyn will recycle the delegate in a static field. not so great for us. let's nop out all that code. var instructionThatPushedDelegate = CecilHelpers.FindInstructionThatPushedArg(methodContainingLambdaJob, 1, lambdaJobDescriptionConstruction.WithCodeInvocationInstruction); var result = CecilHelpers.MatchesDelegateProducingPattern(methodContainingLambdaJob, instructionThatPushedDelegate, CecilHelpers.DelegateProducingPattern.MatchSide.Start); result?.RewriteToProduceSingleNullValue(); } FieldDefinition entityQueryField = null; if (lambdaJobDescriptionConstruction.Kind != LambdaJobDescriptionKind.Job) { entityQueryField = InjectAndInitializeEntityQueryField.InjectAndInitialize(methodContainingLambdaJob, lambdaJobDescriptionConstruction, methodLambdaWasEmittedAs.Parameters); } var generatedJobStruct = JobStructForLambdaJob.CreateNewJobStruct(lambdaJobDescriptionConstruction); if (generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst != null) { var constructorInfo = generatedJobStruct.ExecuteDelegateType.GetConstructors().First(c => c.GetParameters().Length == 2); var instructions = new List <Instruction>() { Instruction.Create(OpCodes.Ldnull), Instruction.Create(OpCodes.Ldftn, generatedJobStruct.RunWithoutJobSystemMethod), Instruction.Create(OpCodes.Newobj, moduleDefinition.ImportReference(constructorInfo)), Instruction.Create(OpCodes.Stsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst) }; if (generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst != null) { instructions.Add(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); var methodInfo = typeof(InternalCompilerInterface) .GetMethods(BindingFlags.Static | BindingFlags.Public) .Where(m => m.Name == nameof(InternalCompilerInterface.BurstCompile)) .Single(m => m.GetParameters().First().ParameterType == generatedJobStruct.ExecuteDelegateType); instructions.Add(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(methodInfo))); instructions.Add(Instruction.Create(OpCodes.Stsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst)); } InjectAndInitializeEntityQueryField.InsertIntoOnCreateForCompilerMethod(methodContainingLambdaJob.DeclaringType, instructions.ToArray()); } IEnumerable <Instruction> InstructionsToReplaceScheduleInvocationWith() { var newJobStructVariable = new VariableDefinition(generatedJobStruct.TypeDefinition); body.Variables.Add(newJobStructVariable); VariableDefinition tempStorageForJobHandle = null; if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { tempStorageForJobHandle = new VariableDefinition(moduleDefinition.ImportReference(typeof(JobHandle))); body.Variables.Add(tempStorageForJobHandle); //since we're replacing the .Schedule() function on the description, the lambdajobdescription and the jobhandle argument to that function will be on the stack. //we're going to need the jobhandle later when we call JobChunkExtensions.Schedule(), so lets stuff it in a variable. yield return(Instruction.Create(OpCodes.Stloc, tempStorageForJobHandle)); } //pop the Description struct off the stack, its services are no longer required yield return(Instruction.Create(OpCodes.Pop)); yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); yield return(Instruction.Create(OpCodes.Initobj, generatedJobStruct.TypeDefinition)); // Call ScheduleTimeInitializeMethod yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); yield return(Instruction.Create(OpCodes.Ldarg_0)); if (lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { //only when the lambda is capturing, did we emit the ScheduleTimeInitialize method to take a displayclass argument var opcode = methodLambdaWasEmittedAs.DeclaringType.IsValueType ? OpCodes.Ldloca : OpCodes.Ldloc; yield return(Instruction.Create(opcode, displayClassVariable)); } yield return(Instruction.Create(OpCodes.Call, generatedJobStruct.ScheduleTimeInitializeMethod)); MethodInfo FindRunOrScheduleMethod() { switch (lambdaJobDescriptionConstruction.Kind) { case LambdaJobDescriptionKind.Entities: case LambdaJobDescriptionKind.Chunk: if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { return(typeof(JobChunkExtensions).GetMethod(nameof(JobChunkExtensions.Schedule))); } return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunJobChunk))); case LambdaJobDescriptionKind.Job: if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { return(typeof(IJobExtensions).GetMethod(nameof(IJobExtensions.Schedule))); } return(typeof(InternalCompilerInterface).GetMethod(nameof(InternalCompilerInterface.RunIJob))); default: throw new ArgumentOutOfRangeException(); } } MethodReference runOrScheduleMethod; if (lambdaJobDescriptionConstruction.WithStructuralChanges) { runOrScheduleMethod = generatedJobStruct.TypeDefinition.Methods.First(definition => definition.Name == "Execute"); } else { runOrScheduleMethod = moduleDefinition.ImportReference(FindRunOrScheduleMethod()) .MakeGenericInstanceMethod(generatedJobStruct.TypeDefinition); } if (lambdaJobDescriptionConstruction.WithStructuralChanges) { yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); } else { yield return(Instruction.Create(runOrScheduleMethod.Parameters.First().ParameterType.IsByReference ? OpCodes.Ldloca : OpCodes.Ldloc, newJobStructVariable)); } switch (lambdaJobDescriptionConstruction.Kind) { case LambdaJobDescriptionKind.Entities: case LambdaJobDescriptionKind.Chunk: if (lambdaJobDescriptionConstruction.WithStructuralChanges) { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldfld, entityQueryField)); } else { yield return(Instruction.Create(OpCodes.Ldarg_0)); yield return(Instruction.Create(OpCodes.Ldfld, entityQueryField)); } break; case LambdaJobDescriptionKind.Job: //job.Schedule() takes no entityQuery... break; } if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Schedule) { yield return(Instruction.Create(OpCodes.Ldloc, tempStorageForJobHandle)); } if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Run && !lambdaJobDescriptionConstruction.WithStructuralChanges) { if (!lambdaJobDescriptionConstruction.UsesBurst) { yield return(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); } else { yield return(Instruction.Create(OpCodes.Call, moduleDefinition.ImportReference(typeof(JobsUtility).GetMethod("get_" + nameof(JobsUtility.JobCompilerEnabled))))); var targetInstruction = Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldBurst); yield return(Instruction.Create(OpCodes.Brtrue, targetInstruction)); yield return(Instruction.Create(OpCodes.Ldsfld, generatedJobStruct.RunWithoutJobSystemDelegateFieldNoBurst)); var finalBranchDestination = Instruction.Create(OpCodes.Nop); yield return(Instruction.Create(OpCodes.Br, finalBranchDestination)); yield return(targetInstruction); yield return(finalBranchDestination); } } yield return(Instruction.Create(OpCodes.Call, runOrScheduleMethod)); if (lambdaJobDescriptionConstruction.ExecutionMode == ExecutionMode.Run && generatedJobStruct.WriteToDisplayClassMethod != null && lambdaJobDescriptionConstruction.DelegateProducingSequence.CapturesLocals) { yield return(Instruction.Create(OpCodes.Ldloca, newJobStructVariable)); var opcode = methodLambdaWasEmittedAs.DeclaringType.IsValueType ? OpCodes.Ldloca : OpCodes.Ldloc; yield return(Instruction.Create(opcode, displayClassVariable)); yield return(Instruction.Create(OpCodes.Call, generatedJobStruct.WriteToDisplayClassMethod)); } } foreach (var invokedMethod in lambdaJobDescriptionConstruction.InvokedConstructionMethods) { bool invokedMethodServesNoPurposeAtRuntime = invokedMethod.MethodName != nameof(LambdaJobQueryConstructionMethods.WithSharedComponentFilter); if (invokedMethodServesNoPurposeAtRuntime) { CecilHelpers.EraseMethodInvocationFromInstructions(ilProcessor, invokedMethod.InstructionInvokingMethod); } else { // Rewrite WithSharedComponentFilter calls as they need to modify EntityQuery dynamically if (invokedMethod.MethodName == nameof(LambdaJobQueryConstructionMethods.WithSharedComponentFilter)) { var setSharedComponentFilterOnQueryMethod = moduleDefinition.ImportReference( (lambdaJobDescriptionConstruction.Kind == LambdaJobDescriptionKind.Entities ? typeof(ForEachLambdaJobDescription_SetSharedComponent) : typeof(LambdaJobChunkDescription_SetSharedComponent)).GetMethod( nameof(LambdaJobChunkDescription_SetSharedComponent .SetSharedComponentFilterOnQuery))); MethodReference genericSetSharedComponentFilterOnQueryMethod = setSharedComponentFilterOnQueryMethod.MakeGenericInstanceMethod(invokedMethod.TypeArguments); // Change invocation to invocation of helper method and add EntityQuery parameter to be modified var setSharedComponentFilterOnQueryInstructions = new List <Instruction> { Instruction.Create(OpCodes.Ldarg_0), Instruction.Create(OpCodes.Ldfld, entityQueryField), Instruction.Create(OpCodes.Call, genericSetSharedComponentFilterOnQueryMethod) }; ilProcessor.Replace(invokedMethod.InstructionInvokingMethod, setSharedComponentFilterOnQueryInstructions); } } } var scheduleInstructions = InstructionsToReplaceScheduleInvocationWith().ToList(); ilProcessor.InsertAfter(lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction, scheduleInstructions); lambdaJobDescriptionConstruction.ScheduleOrRunInvocationInstruction.MakeNOP(); var codegenInitializeMethod = GetOrMakeOnCreateForCompilerMethod(lambdaJobDescriptionConstruction.ContainingMethod.DeclaringType); return(generatedJobStruct); }