Rewrite(SemanticModel model, SyntaxNode originalLambdaExpression)
{
var explicitThisRewriter = new LambdaBodySyntaxReplacer(model);
explicitThisRewriter.Visit(originalLambdaExpression);
var allTrackedNodes = explicitThisRewriter._lineStatementNodesNeedingReplacement.Concat(explicitThisRewriter._thisNodesNeedingReplacement);
allTrackedNodes = allTrackedNodes.Concat(explicitThisRewriter._constantNodesNeedingReplacement);
if (!allTrackedNodes.Any())
{
return(originalLambdaExpression, new List <SyntaxNode>());
}
// Track nodes for later rewriting
var rewrittenBody = originalLambdaExpression.TrackNodes(allTrackedNodes);
// Replace use of constants in the lambda with the actual constant value
var currentConstantNodesNeedingReplacement = rewrittenBody.GetCurrentNodes(explicitThisRewriter._constantNodesNeedingReplacement);
var oldToNewConstantNodesNeedingReplacement =
currentConstantNodesNeedingReplacement.Zip(explicitThisRewriter._constantNodesNeedingReplacement,
(k, v) => new { key = k, value = v }).ToDictionary(x => x.key, x => x.value);
rewrittenBody = rewrittenBody.ReplaceNodes(currentConstantNodesNeedingReplacement, (originalNode, rewrittenNode) =>
ReplaceConstantNodesWithConstantValue(model, originalNode, oldToNewConstantNodesNeedingReplacement));
// Replace those locations to access through "__this" instead (since they need to access through a stored field on job struct).
// Also replace data access methods on SystemBase that need to be patched (GetComponent/SetComponent/etc)
var currentThisNodesNeedingReplacement = rewrittenBody.GetCurrentNodes(explicitThisRewriter._thisNodesNeedingReplacement);
rewrittenBody = rewrittenBody.ReplaceNodes(currentThisNodesNeedingReplacement, (originalNode, rewrittenNode) =>
ReplaceNodeWithAccessThroughLocalThis(originalNode));
// Also rewrite any remaining references to `this` to our local this
rewrittenBody = rewrittenBody.ReplaceNodes(rewrittenBody.DescendantNodes().OfType <ThisExpressionSyntax>(),
(originalNode, rewrittenNode) => LocalThisFieldSyntax);
// Replace line statements with one's with line directive trivia
foreach (var originalNode in explicitThisRewriter._lineStatementNodesNeedingReplacement)
{
var currentNode = rewrittenBody.GetCurrentNode(originalNode);
rewrittenBody = rewrittenBody.ReplaceNode(currentNode, currentNode.WithLineTrivia(originalNode));
}
return(rewrittenBody, explicitThisRewriter._thisNodesNeedingReplacement);
}
Rewrite(LambdaJobDescription description)
{
SemanticModel model = description.Model;
SyntaxNode originalLambdaExpression = description.OriginalLambdaExpression;
List <LambdaCapturedVariableDescription> variablesCapturedOnlyByLocals = description.VariablesCapturedOnlyByLocals;
var lambdaBodyRewriter = new LambdaBodyRewriter();
// Find all locations where we are accessing a member on the declaring SystemBase
// and change them to access through "__this" instead.
// This also annotates the changed nodes so that we can find them later for patching (and get their original symbols).
var rewrittenLambdaBodyData = LambdaBodySyntaxReplacer.Rewrite(model, originalLambdaExpression);
var rewrittenLambdaExpression = rewrittenLambdaBodyData.rewrittenLambdaExpression;
// Go through all changed nodes and check to see if they are a component access method that we need to patch (GetComponent/SetComponent/etc)
// Only need to do this if we are not doing structural changes (in which case we can't as structural changes will invalidate)
if (!description.WithStructuralChanges)
{
foreach (var originalNode in rewrittenLambdaBodyData.thisAccessNodesNeedingReplacement)
{
var originalInvocation = originalNode.Ancestors().OfType <InvocationExpressionSyntax>().First();
var currentNode = rewrittenLambdaExpression.GetCurrentNode(originalNode);
var currentMemberAccessExpression = currentNode.Ancestors().OfType <MemberAccessExpressionSyntax>().First();
var currentNodeInvocationExpression = currentNode.Ancestors().OfType <InvocationExpressionSyntax>().FirstOrDefault();
if (currentNodeInvocationExpression == null)
{
continue;
}
var replacementMemberAccessExpression = lambdaBodyRewriter.GenerateReplacementMemberAccessExpression(originalInvocation,
currentMemberAccessExpression, model);
if (replacementMemberAccessExpression != null)
{
rewrittenLambdaExpression = rewrittenLambdaExpression.ReplaceNode(currentNodeInvocationExpression, replacementMemberAccessExpression);
}
}
}
// Go through all local declaration nodes and replace them with assignment nodes (or remove) if they are now captured variables that live in job struct
// This is needed for variables captured for local methods
foreach (var localDeclarationSyntax in rewrittenLambdaExpression.DescendantNodes().OfType <LocalDeclarationStatementSyntax>())
{
var variableDeclaration = localDeclarationSyntax.DescendantNodes().OfType <VariableDeclarationSyntax>().FirstOrDefault();
if (variableDeclaration != null &&
variablesCapturedOnlyByLocals.Any(variable => variable.OriginalVariableName == variableDeclaration.Variables.First().Identifier.Text))
{
if (variableDeclaration.DescendantTokens().Any(token => token.Kind() == SyntaxKind.EqualsToken))
{
var variableIdentifier = variableDeclaration.Variables.First().Identifier;
var nodeAfter = variableDeclaration.NodeAfter(node => node.Kind() == SyntaxKind.EqualsToken);
rewrittenLambdaExpression = rewrittenLambdaExpression.ReplaceNode(localDeclarationSyntax,
SyntaxFactory.ExpressionStatement(SyntaxFactory.AssignmentExpression(
SyntaxKind.SimpleAssignmentExpression,
SyntaxFactory.IdentifierName(variableIdentifier.Text),
(ExpressionSyntax)nodeAfter)));
}
else
{
rewrittenLambdaExpression = rewrittenLambdaExpression.RemoveNode(localDeclarationSyntax, SyntaxRemoveOptions.KeepExteriorTrivia);
}
}
}
// Go through all local function statements and omit them as method declarations on the job struct
// (local methods accessing fields on this are not allowed in C#)
// https://sharplab.io/#v2:EYLgtghgzgLgpgJwDQxNGAfAAgJgIwCwAUFgMwAEu5AwuQN7HlOUVYAs5AsgBQCU9jZgF9BTUeVgIArgGMY5AKIA7GAEsYAT3oiizCTGlyaAezAAHY0rgqAKhrNwAEhCUATADZwAPDYB828UlZeWpTCysVABEIGAgAMQRTZTVNH386HT0gowUZKBs4WGjY5PUtcQZdPWYyRRUy8gA3CHcpODwAbnFM5mz5XPzCmGKIASrqlkU8gqKYiG5VFSaWtv4M7vFN8fJF+AQAMwgZOHIASQApY2BqAAspJQBrAO2+8gAZCDBgVwhL4AB9AAM/z+5BAZz+t3uDwq4j0tVC5ks1hgdgczjcni8AxmwzmpU0/n+MFccDRThcHjgXW28IoiPCKJGCSS9VSOKGIyJ/w5s1i/xZYAJGhpEzEtJqHAA8ghVABzRYtD5fH4AIWMrg03GF5BRZSQOyUU0GfIgOpJcF4cOYlTF1V25AA4nAYMKAGorODatlaS3Wia2u1irAAdnI/x5005cwFiSFPoA2nAALoAOmarWp/uqOmz8I4AGUXe7Pd6Ur6DQ6M6s8zbaxMI7y8fzBcKk8nyABechWADuxtxI241ctoqDTB64/r5CLrp9Hsz3D1mgNzrn5YXbW4FvTnrwvF4Y4muYlTAA9AAqbO1e5QCD7E7sRQADzgMik8G4AEEEDIbi77DgKFHnIP9oUrJYwMeU43DgZ8IPkfZVAQWBhRg0lnytU8xnHJhmgQXUfQABQMLtyDLMoL14KCHlTNcADkYlURo4B/BAIC1f5lw0ckMSpXg6JdABVJQ7wfAAlOAIFcKUlHcDQSIQPgjyDfDyAtL8ZGOKAoGMAju21KNTWFKiaMEmBGLUFi2I47hiVJXjKU8AS1xEsS4Ek6TZPkxTlOnB1uNCe55G7MygpUFS7X2PTyIdVQyMBDodnIABCbtAuMYKktUABqHKsNwnDCuYGV5UVdxlW+CB1U1bgLyXYjSJynZeANeqNK0wpdII5rVAPSKxUnHNswvM8NiIIQgA===
var localFunctions = rewrittenLambdaExpression.DescendantNodes().OfType <LocalFunctionStatementSyntax>();
rewrittenLambdaExpression = rewrittenLambdaExpression.RemoveNodes(localFunctions, SyntaxRemoveOptions.KeepNoTrivia);
var methodsForLocalFunctions = new List <MethodDeclarationSyntax>();
foreach (var localFunction in localFunctions)
{
methodsForLocalFunctions.Add((MethodDeclarationSyntax)SyntaxFactory.ParseMemberDeclaration(localFunction.ToString()));
}
return(rewrittenLambdaExpression, lambdaBodyRewriter.DataFromEntityFields.Values.ToList(), methodsForLocalFunctions);
}
