@@ -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 ! ;
0 commit comments