Skip to content

Commit 6dfacf1

Browse files
author
george.lemeshko
committed
Fix RIDER-92392 - update baker generation code due to new DOTS api
1 parent 8cde29c commit 6dfacf1

File tree

6 files changed

+183
-16
lines changed

6 files changed

+183
-16
lines changed

resharper/resharper-unity/src/Unity/CSharp/ExpressionReferenceUtils.cs

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,41 @@ public static bool IsTagProperty([CanBeNull] this IReferenceExpression expressio
3535
return false;
3636
}
3737

38+
public static bool IsAnyAddComponentMethod(this IInvocationExpression invocationExpression)
39+
{
40+
if (invocationExpression.Reference.Resolve().DeclaredElement is not IMethod method)
41+
return false;
42+
43+
if (!method.ShortName.Equals("AddComponent"))
44+
return false;
45+
46+
if (method.ContainingType != null && method.ContainingType.GetClrName().Equals(KnownTypes.IBaker))
47+
return true;
48+
49+
return false;
50+
}
51+
52+
public static bool IsBakerGetPrimaryEntityMethod(this IInvocationExpression invocationExpression)
53+
{
54+
if (invocationExpression.Reference.Resolve().DeclaredElement is not IMethod method)
55+
return false;
56+
57+
if (!method.ShortName.Equals("GetEntity"))
58+
return false;
59+
60+
if (method.ContainingType == null || !method.ContainingType.GetClrName().Equals(KnownTypes.IBaker))
61+
return false;
62+
63+
var parameters = method.Parameters;
64+
if (parameters.Count != 1)
65+
return false;
66+
67+
var parameter = parameters[0];
68+
69+
return parameter.Type is IDeclaredType declaredType
70+
&& declaredType.GetClrName().Equals(KnownTypes.TransformUsageFlags);
71+
}
72+
3873
public static bool IsCompareTagMethod(this IInvocationExpression expr)
3974
{
4075
return IsSpecificMethod(expr, KnownTypes.Component, "CompareTag")

resharper/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/BakerGeneratorUtils.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ public static class BakerGeneratorUtils
1212
{
1313
private static readonly Dictionary<IClrTypeName, ConversionData> ourComponentDataToAuthoringTypesConversion = new()
1414
{
15-
{KnownTypes.Entity, new ConversionData(KnownTypes.GameObject, "GetEntity($0.$1)")},
15+
{KnownTypes.Entity, new ConversionData(KnownTypes.GameObject, "GetEntity($0.$1, TransformUsageFlags.Dynamic)")},
1616
{KnownTypes.Random, new ConversionData(PredefinedType.UINT_FQN, "Unity.Mathematics.Random.CreateFromIndex($0.$1)")}
1717
};
1818

@@ -31,7 +31,7 @@ public static class BakerGeneratorUtils
3131
{KnownTypes.Vector3, new ConversionData(KnownTypes.Float3, "$0.$1")},
3232
};
3333

34-
private static readonly ConversionData ourComponentToEntityConversions = new(KnownTypes.Entity, "GetEntity($0.$1)");
34+
private static readonly ConversionData ourComponentToEntityConversions = new(KnownTypes.Entity, "GetEntity($0.$1, TransformUsageFlags.Dynamic)");
3535

3636
public static ConversionData? ConvertAuthoringToComponentField(IClrTypeName clrTypeName, IPsiModule psiModule)
3737
{

resharper/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndAuthoringActionBuilder.cs

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,69 @@ protected override void Process(CSharpGeneratorContext context, IProgressIndicat
140140
var asNested = selectedBaker.Equals(Strings.UnityDots_GenerateBakerAndAuthoring_NewBaker_As_Nested);
141141
return (null, asNested);
142142
}
143+
144+
private static ITreeNode GetOrCreateGetEntityExpression(ICSharpFunctionDeclaration bakeMethodExpression, CSharpElementFactory factory)
145+
{
146+
var entityNode = TryGetExistingEntityCreationNode(bakeMethodExpression);
147+
if (entityNode != null)
148+
return entityNode;
149+
150+
//return any AddComponent(...)
151+
var anyAddComponentMethodExpression = bakeMethodExpression.Body.FindNextNode(node =>
152+
{
153+
if (node is IMethodDeclaration)
154+
return TreeNodeActionType.IGNORE_SUBTREE;
155+
156+
if (node is not IInvocationExpression invocationExpression)
157+
return TreeNodeActionType.CONTINUE;
158+
159+
return invocationExpression.IsAnyAddComponentMethod()
160+
? TreeNodeActionType.ACCEPT
161+
: TreeNodeActionType.CONTINUE;
162+
});
163+
164+
165+
if (anyAddComponentMethodExpression is IInvocationExpression { Arguments: { Count: > 0 } } methodExpression)
166+
{
167+
return methodExpression.Arguments[0].Value;
168+
}
169+
170+
//var entity = GetEntity(TransformUsageFlags.Dynamic);
171+
var getEntityExpression =
172+
(IDeclarationStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("var entity = GetEntity(TransformUsageFlags.Dynamic);"),
173+
null);
174+
175+
//returns "entity" node
176+
177+
return getEntityExpression.VariableDeclarations.SingleItem!.FirstChild!;
178+
}
179+
180+
private static ITreeNode? TryGetExistingEntityCreationNode(ICSharpFunctionDeclaration bakeMethodExpression)
181+
{
182+
var existingExpression = bakeMethodExpression.Body.FindNextNode(node =>
183+
{
184+
if (node is IMethodDeclaration)
185+
return TreeNodeActionType.IGNORE_SUBTREE;
186+
187+
if (node is not IInvocationExpression invocationExpression)
188+
return TreeNodeActionType.CONTINUE;
189+
190+
var localVariableDeclaration = invocationExpression.GetContainingNode<ILocalVariableDeclaration>();
191+
192+
if (localVariableDeclaration == null)
193+
return TreeNodeActionType.CONTINUE;
194+
195+
return invocationExpression.IsBakerGetPrimaryEntityMethod()
196+
? TreeNodeActionType.ACCEPT
197+
: TreeNodeActionType.CONTINUE;
198+
});
199+
200+
var variableDeclaration = existingExpression?.GetContainingNode<ILocalVariableDeclaration>();
201+
202+
var variableNameNode = variableDeclaration?.DeclaredElement.GetDeclarations().SingleItem()?.FirstChild;
203+
204+
return variableNameNode;
205+
}
143206

144207
private static void GenerateBaker(IGeneratorContext context, Dictionary<string, string> componentToAuthoringFieldNames, BakerGenerationInfo generationInfo)
145208
{
@@ -148,7 +211,8 @@ private static void GenerateBaker(IGeneratorContext context, Dictionary<string,
148211
: CreateBakerClassDeclaration(generationInfo);
149212

150213
var bakeMethodExpression = GetOrCreateBakeMethodExpression(bakerClassDeclarations, generationInfo.Factory, generationInfo, out var authoringParameterName);
151-
var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentStructDeclaration.DeclaredElement!);
214+
var entityExpression = GetOrCreateGetEntityExpression(bakeMethodExpression, generationInfo.Factory);
215+
var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentStructDeclaration.DeclaredElement!, entityExpression);
152216
if(context.InputElements.Count != 0)
153217
{
154218
var creationExpressionInitializer = GetOrCreateInitializer(componentCreationExpression, generationInfo.Factory);
@@ -261,7 +325,7 @@ private static IMethodDeclaration GetOrCreateBakeMethodExpression(IClassLikeDecl
261325
}
262326

263327
private static IObjectCreationExpression GetOrCreateComponentCreationExpression(CSharpElementFactory factory,
264-
IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType)
328+
IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType, ITreeNode entityExpression)
265329
{
266330
var existingCreationExpression = bakeMethodExpression.Body.FindNextNode( node =>
267331
{
@@ -280,11 +344,15 @@ private static IObjectCreationExpression GetOrCreateComponentCreationExpression(
280344
//AddComponent/AddComponentObject(new ComponentData{})
281345
var addComponentMethod = componentDeclaredType is IStruct ? "AddComponent();" : "AddComponentObject();";
282346
var addComponentStatement =
283-
(IExpressionStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement(addComponentMethod),
347+
(IExpressionStatement)bakeMethodExpression.Body.AddStatementBefore(factory.CreateStatement(addComponentMethod),
284348
null);
285349
var addComponentExpression = (addComponentStatement.Expression as IInvocationExpression).NotNull();
350+
351+
var entityArgument = factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("$0", entityExpression));
352+
entityArgument = addComponentExpression.AddArgumentAfter(entityArgument, null);
353+
286354
var creationArgument = addComponentExpression.AddArgumentAfter(
287-
factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), null);
355+
factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), entityArgument);
288356

289357
var componentCreationExpression = creationArgument.Value as IObjectCreationExpression;
290358
return componentCreationExpression!;

resharper/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionBuilder.cs

Lines changed: 73 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,8 @@ private static void GenerateBaker(IGeneratorContext context, Dictionary<string,
183183
: CreateBakerClassDeclaration(generationInfo);
184184

185185
var bakeMethodExpression = GetOrCreateBakeMethodExpression(bakerClassDeclarations, generationInfo.Factory, generationInfo, out var authoringParameterName);
186-
var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentDataDeclaration.DeclaredElement!);
186+
var entityExpression = GetOrCreateGetEntityExpression(bakeMethodExpression, generationInfo.Factory);
187+
var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentDataDeclaration.DeclaredElement!, entityExpression);
187188
if(context.InputElements.Count != 0)
188189
{
189190
var creationExpressionInitializer = GetOrCreateInitializer(componentCreationExpression, generationInfo.Factory);
@@ -305,10 +306,73 @@ private static IMethodDeclaration GetOrCreateBakeMethodExpression(IClassLikeDecl
305306
return bakeMethodExpression;
306307
}
307308

309+
private static ITreeNode GetOrCreateGetEntityExpression(ICSharpFunctionDeclaration bakeMethodExpression, CSharpElementFactory factory)
310+
{
311+
var entityNode = TryGetExistingEntityCreationNode(bakeMethodExpression);
312+
if (entityNode != null)
313+
return entityNode;
314+
315+
//return any AddComponent(...)
316+
var anyAddComponentMethodExpression = bakeMethodExpression.Body.FindNextNode(node =>
317+
{
318+
if (node is IMethodDeclaration)
319+
return TreeNodeActionType.IGNORE_SUBTREE;
320+
321+
if (node is not IInvocationExpression invocationExpression)
322+
return TreeNodeActionType.CONTINUE;
323+
324+
return invocationExpression.IsAnyAddComponentMethod()
325+
? TreeNodeActionType.ACCEPT
326+
: TreeNodeActionType.CONTINUE;
327+
});
328+
329+
330+
if (anyAddComponentMethodExpression is IInvocationExpression { Arguments: { Count: > 0 } } methodExpression)
331+
{
332+
return methodExpression.Arguments[0].Value;
333+
}
334+
335+
//var entity = GetEntity(TransformUsageFlags.Dynamic);
336+
var getEntityExpression =
337+
(IDeclarationStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("var entity = GetEntity(TransformUsageFlags.Dynamic);"),
338+
null);
339+
340+
//returns "entity" node
341+
342+
return getEntityExpression.VariableDeclarations.SingleItem!.FirstChild!;
343+
}
344+
345+
private static ITreeNode? TryGetExistingEntityCreationNode(ICSharpFunctionDeclaration bakeMethodExpression)
346+
{
347+
var existingExpression = bakeMethodExpression.Body.FindNextNode(node =>
348+
{
349+
if (node is IMethodDeclaration)
350+
return TreeNodeActionType.IGNORE_SUBTREE;
351+
352+
if (node is not IInvocationExpression invocationExpression)
353+
return TreeNodeActionType.CONTINUE;
354+
355+
var localVariableDeclaration = invocationExpression.GetContainingNode<ILocalVariableDeclaration>();
356+
357+
if (localVariableDeclaration == null)
358+
return TreeNodeActionType.CONTINUE;
359+
360+
return invocationExpression.IsBakerGetPrimaryEntityMethod()
361+
? TreeNodeActionType.ACCEPT
362+
: TreeNodeActionType.CONTINUE;
363+
});
364+
365+
var variableDeclaration = existingExpression?.GetContainingNode<ILocalVariableDeclaration>();
366+
367+
var variableNameNode = variableDeclaration?.DeclaredElement.GetDeclarations().SingleItem()?.FirstChild;
368+
369+
return variableNameNode;
370+
}
371+
308372
private static IObjectCreationExpression GetOrCreateComponentCreationExpression(CSharpElementFactory factory,
309-
IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType)
373+
IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType, ITreeNode entityExpression)
310374
{
311-
var existingCreationExpression = bakeMethodExpression.Body.FindNextNode( node =>
375+
var existingCreationExpression = bakeMethodExpression.Body.FindNextNode(node =>
312376
{
313377
if (node is IMethodDeclaration)
314378
return TreeNodeActionType.IGNORE_SUBTREE;
@@ -321,14 +385,17 @@ private static IObjectCreationExpression GetOrCreateComponentCreationExpression(
321385

322386
if (existingCreationExpression != null)
323387
return (IObjectCreationExpression)existingCreationExpression;
324-
388+
325389
//AddComponent(new ComponentData{})
326390
var addComponentStatement =
327-
(IExpressionStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("AddComponent();"),
391+
(IExpressionStatement)bakeMethodExpression.Body.AddStatementBefore(factory.CreateStatement("AddComponent();"),
328392
null);
329393
var addComponentExpression = (addComponentStatement.Expression as IInvocationExpression).NotNull();
394+
var entityArgument = factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("$0", entityExpression));
395+
entityArgument = addComponentExpression.AddArgumentAfter(entityArgument, null);
396+
330397
var creationArgument = addComponentExpression.AddArgumentAfter(
331-
factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), null);
398+
factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), entityArgument);
332399

333400
var componentCreationExpression = creationArgument.Value as IObjectCreationExpression;
334401
return componentCreationExpression!;

resharper/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionWorkflow.cs

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
using JetBrains.Application.DataContext;
2-
using JetBrains.ProjectModel;
3-
using JetBrains.ProjectModel.DataContext;
42
using JetBrains.ReSharper.Feature.Services.Generate;
53
using JetBrains.ReSharper.Feature.Services.Generate.Actions;
64
using JetBrains.ReSharper.Feature.Services.Generate.Workflows;
7-
using JetBrains.ReSharper.Plugins.Unity.Core.ProjectModel;
85
using JetBrains.ReSharper.Plugins.Unity.CSharp.Daemon.Stages.Dots;
96
using JetBrains.ReSharper.Plugins.Unity.Resources;
107
using JetBrains.ReSharper.Plugins.Unity.Resources.Icons;
11-
using JetBrains.ReSharper.Plugins.Unity.UnityEditorIntegration.Packages;
128

139
namespace JetBrains.ReSharper.Plugins.Unity.CSharp.Feature.Services.Generate.Dots
1410
{

resharper/resharper-unity/src/Unity/UnityEditorIntegration/Api/KnownTypes.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ public static class KnownTypes
121121
public static readonly IClrTypeName EnabledRefRO = new ClrTypeName("Unity.Entities.EnabledRefRO`1");
122122
public static readonly IClrTypeName EnabledRefRW = new ClrTypeName("Unity.Entities.EnabledRefRW`1");
123123
public static readonly IClrTypeName IEnableableComponent = new ClrTypeName("Unity.Entities.IEnableableComponent");
124+
public static readonly IClrTypeName TransformUsageFlags = new ClrTypeName("Unity.Entities.TransformUsageFlags");
124125

125126
//Unity.Mathematics
126127
public static readonly IClrTypeName Random = new ClrTypeName("Unity.Mathematics.Random");

0 commit comments

Comments
 (0)