@@ -140,6 +140,69 @@ protected override void Process(CSharpGeneratorContext context, IProgressIndicat
140
140
var asNested = selectedBaker . Equals ( Strings . UnityDots_GenerateBakerAndAuthoring_NewBaker_As_Nested ) ;
141
141
return ( null , asNested ) ;
142
142
}
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
+ }
143
206
144
207
private static void GenerateBaker ( IGeneratorContext context , Dictionary < string , string > componentToAuthoringFieldNames , BakerGenerationInfo generationInfo )
145
208
{
@@ -148,7 +211,8 @@ private static void GenerateBaker(IGeneratorContext context, Dictionary<string,
148
211
: CreateBakerClassDeclaration ( generationInfo ) ;
149
212
150
213
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 ) ;
152
216
if ( context . InputElements . Count != 0 )
153
217
{
154
218
var creationExpressionInitializer = GetOrCreateInitializer ( componentCreationExpression , generationInfo . Factory ) ;
@@ -261,7 +325,7 @@ private static IMethodDeclaration GetOrCreateBakeMethodExpression(IClassLikeDecl
261
325
}
262
326
263
327
private static IObjectCreationExpression GetOrCreateComponentCreationExpression ( CSharpElementFactory factory ,
264
- IMethodDeclaration bakeMethodExpression , ITypeElement componentDeclaredType )
328
+ IMethodDeclaration bakeMethodExpression , ITypeElement componentDeclaredType , ITreeNode entityExpression )
265
329
{
266
330
var existingCreationExpression = bakeMethodExpression . Body . FindNextNode ( node =>
267
331
{
@@ -280,11 +344,15 @@ private static IObjectCreationExpression GetOrCreateComponentCreationExpression(
280
344
//AddComponent/AddComponentObject(new ComponentData{})
281
345
var addComponentMethod = componentDeclaredType is IStruct ? "AddComponent();" : "AddComponentObject();" ;
282
346
var addComponentStatement =
283
- ( IExpressionStatement ) bakeMethodExpression . Body . AddStatementAfter ( factory . CreateStatement ( addComponentMethod ) ,
347
+ ( IExpressionStatement ) bakeMethodExpression . Body . AddStatementBefore ( factory . CreateStatement ( addComponentMethod ) ,
284
348
null ) ;
285
349
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
+
286
354
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 ) ;
288
356
289
357
var componentCreationExpression = creationArgument . Value as IObjectCreationExpression ;
290
358
return componentCreationExpression ! ;
0 commit comments