Fix RIDER-92392 - update baker generation code due to new DOTS api

JetBrains · Apr 20, 2023 · 6dfacf1 · 6dfacf1
1 parent 8cde29c
commit 6dfacf1
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 16 deletions.
diff --git a/resharper/resharper-unity/src/Unity/CSharp/ExpressionReferenceUtils.cs b/resharper/resharper-unity/src/Unity/CSharp/ExpressionReferenceUtils.cs
@@ -35,6 +35,41 @@ public static bool IsTagProperty([CanBeNull] this IReferenceExpression expressio
             return false;
         }
 
+        public static bool IsAnyAddComponentMethod(this IInvocationExpression invocationExpression)
+        {
+            if (invocationExpression.Reference.Resolve().DeclaredElement is not IMethod method)
+                return false;
+
+            if (!method.ShortName.Equals("AddComponent"))
+                return false;
+
+            if (method.ContainingType != null && method.ContainingType.GetClrName().Equals(KnownTypes.IBaker))
+                return true;
+
+            return false;
+        }
+
+        public static bool IsBakerGetPrimaryEntityMethod(this IInvocationExpression invocationExpression)
+        {
+            if (invocationExpression.Reference.Resolve().DeclaredElement is not IMethod method)
+                return false;
+
+            if (!method.ShortName.Equals("GetEntity"))
+                return false;
+
+            if (method.ContainingType == null || !method.ContainingType.GetClrName().Equals(KnownTypes.IBaker))
+                return false;
+
+            var parameters = method.Parameters;
+            if (parameters.Count != 1)
+                return false;
+
+            var parameter = parameters[0];
+
+            return parameter.Type is IDeclaredType declaredType
+                   && declaredType.GetClrName().Equals(KnownTypes.TransformUsageFlags);
+        }
+
         public static bool IsCompareTagMethod(this IInvocationExpression expr)
         {
             return IsSpecificMethod(expr, KnownTypes.Component, "CompareTag")

diff --git a/...er/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/BakerGeneratorUtils.cs b/...er/resharper-unity/src/Unity/CSharp/Feature/Services/Generate/Dots/BakerGeneratorUtils.cs
@@ -12,7 +12,7 @@ public static class BakerGeneratorUtils
     {
         private static readonly Dictionary<IClrTypeName, ConversionData> ourComponentDataToAuthoringTypesConversion = new() 
         {
-            {KnownTypes.Entity, new ConversionData(KnownTypes.GameObject, "GetEntity($0.$1)")},
+            {KnownTypes.Entity, new ConversionData(KnownTypes.GameObject, "GetEntity($0.$1, TransformUsageFlags.Dynamic)")},
             {KnownTypes.Random, new ConversionData(PredefinedType.UINT_FQN, "Unity.Mathematics.Random.CreateFromIndex($0.$1)")}
         };
 
@@ -31,7 +31,7 @@ public static class BakerGeneratorUtils
             {KnownTypes.Vector3, new ConversionData(KnownTypes.Float3, "$0.$1")},
         };
 
-        private static readonly ConversionData ourComponentToEntityConversions = new(KnownTypes.Entity, "GetEntity($0.$1)");
+        private static readonly ConversionData ourComponentToEntityConversions = new(KnownTypes.Entity, "GetEntity($0.$1, TransformUsageFlags.Dynamic)");
 
         public static ConversionData? ConvertAuthoringToComponentField(IClrTypeName clrTypeName, IPsiModule psiModule)
         {

diff --git a/...src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndAuthoringActionBuilder.cs b/...src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndAuthoringActionBuilder.cs
@@ -140,6 +140,69 @@ protected override void Process(CSharpGeneratorContext context, IProgressIndicat
             var asNested = selectedBaker.Equals(Strings.UnityDots_GenerateBakerAndAuthoring_NewBaker_As_Nested);
             return (null, asNested);
         }
+
+           private static ITreeNode GetOrCreateGetEntityExpression(ICSharpFunctionDeclaration bakeMethodExpression, CSharpElementFactory factory)
+        {
+            var entityNode = TryGetExistingEntityCreationNode(bakeMethodExpression);
+            if (entityNode != null) 
+                return entityNode;
+
+            //return any AddComponent(...)
+            var anyAddComponentMethodExpression = bakeMethodExpression.Body.FindNextNode(node =>
+            {
+                if (node is IMethodDeclaration)
+                    return TreeNodeActionType.IGNORE_SUBTREE;
+
+                if (node is not IInvocationExpression invocationExpression)
+                    return TreeNodeActionType.CONTINUE;
+                
+                return invocationExpression.IsAnyAddComponentMethod()
+                    ? TreeNodeActionType.ACCEPT
+                    : TreeNodeActionType.CONTINUE;
+            });
+
+
+            if (anyAddComponentMethodExpression is IInvocationExpression { Arguments: { Count: > 0 } } methodExpression)
+            {
+                return methodExpression.Arguments[0].Value;
+            }
+
+            //var entity = GetEntity(TransformUsageFlags.Dynamic);
+            var getEntityExpression =
+                (IDeclarationStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("var entity = GetEntity(TransformUsageFlags.Dynamic);"),
+                    null);
+
+            //returns "entity" node
+
+            return getEntityExpression.VariableDeclarations.SingleItem!.FirstChild!;
+        }
+
+        private static ITreeNode? TryGetExistingEntityCreationNode(ICSharpFunctionDeclaration bakeMethodExpression)
+        {
+            var existingExpression = bakeMethodExpression.Body.FindNextNode(node =>
+            {
+                if (node is IMethodDeclaration)
+                    return TreeNodeActionType.IGNORE_SUBTREE;
+
+                if (node is not IInvocationExpression invocationExpression)
+                    return TreeNodeActionType.CONTINUE;
+
+                var localVariableDeclaration = invocationExpression.GetContainingNode<ILocalVariableDeclaration>();
+
+                if (localVariableDeclaration == null)
+                    return TreeNodeActionType.CONTINUE;
+
+                return invocationExpression.IsBakerGetPrimaryEntityMethod()
+                    ? TreeNodeActionType.ACCEPT
+                    : TreeNodeActionType.CONTINUE;
+            });
+
+            var variableDeclaration = existingExpression?.GetContainingNode<ILocalVariableDeclaration>();
+
+            var variableNameNode = variableDeclaration?.DeclaredElement.GetDeclarations().SingleItem()?.FirstChild;
+
+            return variableNameNode;
+        }
 
         private static void GenerateBaker(IGeneratorContext context, Dictionary<string, string> componentToAuthoringFieldNames, BakerGenerationInfo generationInfo)
         {
@@ -148,7 +211,8 @@ private static void GenerateBaker(IGeneratorContext context, Dictionary<string,
                 : CreateBakerClassDeclaration(generationInfo);
 
             var bakeMethodExpression = GetOrCreateBakeMethodExpression(bakerClassDeclarations, generationInfo.Factory, generationInfo, out var authoringParameterName);
-            var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentStructDeclaration.DeclaredElement!);
+            var entityExpression = GetOrCreateGetEntityExpression(bakeMethodExpression, generationInfo.Factory);
+            var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentStructDeclaration.DeclaredElement!, entityExpression);
             if(context.InputElements.Count != 0)
             {
                 var creationExpressionInitializer = GetOrCreateInitializer(componentCreationExpression, generationInfo.Factory);
@@ -261,7 +325,7 @@ private static IMethodDeclaration GetOrCreateBakeMethodExpression(IClassLikeDecl
         }
 
         private static IObjectCreationExpression GetOrCreateComponentCreationExpression(CSharpElementFactory factory,
-            IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType)
+            IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType, ITreeNode entityExpression)
         {
             var existingCreationExpression = bakeMethodExpression.Body.FindNextNode( node =>
             {
@@ -280,11 +344,15 @@ private static IObjectCreationExpression GetOrCreateComponentCreationExpression(
             //AddComponent/AddComponentObject(new ComponentData{})
             var addComponentMethod = componentDeclaredType is IStruct ? "AddComponent();" :  "AddComponentObject();";
             var addComponentStatement =
-                (IExpressionStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement(addComponentMethod),
+                (IExpressionStatement)bakeMethodExpression.Body.AddStatementBefore(factory.CreateStatement(addComponentMethod),
                     null);
             var addComponentExpression = (addComponentStatement.Expression as IInvocationExpression).NotNull();
+
+            var entityArgument = factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("$0", entityExpression));
+            entityArgument = addComponentExpression.AddArgumentAfter(entityArgument, null);
+
             var creationArgument = addComponentExpression.AddArgumentAfter(
-                factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), null);
+                factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), entityArgument);
 
             var componentCreationExpression = creationArgument.Value as IObjectCreationExpression;
             return componentCreationExpression!;

diff --git a/...src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionBuilder.cs b/...src/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionBuilder.cs
@@ -183,7 +183,8 @@ private static void GenerateBaker(IGeneratorContext context, Dictionary<string,
                 : CreateBakerClassDeclaration(generationInfo);
 
             var bakeMethodExpression = GetOrCreateBakeMethodExpression(bakerClassDeclarations, generationInfo.Factory, generationInfo, out var authoringParameterName);
-            var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentDataDeclaration.DeclaredElement!);
+            var entityExpression = GetOrCreateGetEntityExpression(bakeMethodExpression, generationInfo.Factory);
+            var componentCreationExpression = GetOrCreateComponentCreationExpression(generationInfo.Factory, bakeMethodExpression, generationInfo.ComponentDataDeclaration.DeclaredElement!, entityExpression);
             if(context.InputElements.Count != 0)
             {
                 var creationExpressionInitializer = GetOrCreateInitializer(componentCreationExpression, generationInfo.Factory);
@@ -305,10 +306,73 @@ private static IMethodDeclaration GetOrCreateBakeMethodExpression(IClassLikeDecl
             return bakeMethodExpression;
         }
 
+        private static ITreeNode GetOrCreateGetEntityExpression(ICSharpFunctionDeclaration bakeMethodExpression, CSharpElementFactory factory)
+        {
+            var entityNode = TryGetExistingEntityCreationNode(bakeMethodExpression);
+            if (entityNode != null) 
+                return entityNode;
+
+            //return any AddComponent(...)
+            var anyAddComponentMethodExpression = bakeMethodExpression.Body.FindNextNode(node =>
+            {
+                if (node is IMethodDeclaration)
+                    return TreeNodeActionType.IGNORE_SUBTREE;
+
+                if (node is not IInvocationExpression invocationExpression)
+                    return TreeNodeActionType.CONTINUE;
+                
+                return invocationExpression.IsAnyAddComponentMethod()
+                    ? TreeNodeActionType.ACCEPT
+                    : TreeNodeActionType.CONTINUE;
+            });
+
+
+            if (anyAddComponentMethodExpression is IInvocationExpression { Arguments: { Count: > 0 } } methodExpression)
+            {
+                return methodExpression.Arguments[0].Value;
+            }
+
+            //var entity = GetEntity(TransformUsageFlags.Dynamic);
+            var getEntityExpression =
+                (IDeclarationStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("var entity = GetEntity(TransformUsageFlags.Dynamic);"),
+                    null);
+
+            //returns "entity" node
+
+            return getEntityExpression.VariableDeclarations.SingleItem!.FirstChild!;
+        }
+
+        private static ITreeNode? TryGetExistingEntityCreationNode(ICSharpFunctionDeclaration bakeMethodExpression)
+        {
+            var existingExpression = bakeMethodExpression.Body.FindNextNode(node =>
+            {
+                if (node is IMethodDeclaration)
+                    return TreeNodeActionType.IGNORE_SUBTREE;
+
+                if (node is not IInvocationExpression invocationExpression)
+                    return TreeNodeActionType.CONTINUE;
+
+                var localVariableDeclaration = invocationExpression.GetContainingNode<ILocalVariableDeclaration>();
+
+                if (localVariableDeclaration == null)
+                    return TreeNodeActionType.CONTINUE;
+
+                return invocationExpression.IsBakerGetPrimaryEntityMethod()
+                    ? TreeNodeActionType.ACCEPT
+                    : TreeNodeActionType.CONTINUE;
+            });
+
+            var variableDeclaration = existingExpression?.GetContainingNode<ILocalVariableDeclaration>();
+
+            var variableNameNode = variableDeclaration?.DeclaredElement.GetDeclarations().SingleItem()?.FirstChild;
+
+            return variableNameNode;
+        }
+
         private static IObjectCreationExpression GetOrCreateComponentCreationExpression(CSharpElementFactory factory,
-            IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType)
+            IMethodDeclaration bakeMethodExpression, ITypeElement componentDeclaredType, ITreeNode entityExpression)
         {
-            var existingCreationExpression = bakeMethodExpression.Body.FindNextNode( node =>
+            var existingCreationExpression = bakeMethodExpression.Body.FindNextNode(node =>
             {
                 if (node is IMethodDeclaration)
                     return TreeNodeActionType.IGNORE_SUBTREE;
@@ -321,14 +385,17 @@ private static IObjectCreationExpression GetOrCreateComponentCreationExpression(
 
             if (existingCreationExpression != null)
                 return (IObjectCreationExpression)existingCreationExpression;
-
+            
             //AddComponent(new ComponentData{})
             var addComponentStatement =
-                (IExpressionStatement)bakeMethodExpression.Body.AddStatementAfter(factory.CreateStatement("AddComponent();"),
+                (IExpressionStatement)bakeMethodExpression.Body.AddStatementBefore(factory.CreateStatement("AddComponent();"),
                     null);
             var addComponentExpression = (addComponentStatement.Expression as IInvocationExpression).NotNull();
+            var entityArgument = factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("$0", entityExpression));
+            entityArgument = addComponentExpression.AddArgumentAfter(entityArgument, null);
+
             var creationArgument = addComponentExpression.AddArgumentAfter(
-                factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), null);
+                factory.CreateArgument(ParameterKind.VALUE, factory.CreateExpression("new $0()", componentDeclaredType)), entityArgument);
 
             var componentCreationExpression = creationArgument.Value as IObjectCreationExpression;
             return componentCreationExpression!;

diff --git a/...rc/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionWorkflow.cs b/...rc/Unity/CSharp/Feature/Services/Generate/Dots/GenerateBakerAndComponentActionWorkflow.cs
@@ -1,14 +1,10 @@
 using JetBrains.Application.DataContext;
-using JetBrains.ProjectModel;
-using JetBrains.ProjectModel.DataContext;
 using JetBrains.ReSharper.Feature.Services.Generate;
 using JetBrains.ReSharper.Feature.Services.Generate.Actions;
 using JetBrains.ReSharper.Feature.Services.Generate.Workflows;
-using JetBrains.ReSharper.Plugins.Unity.Core.ProjectModel;
 using JetBrains.ReSharper.Plugins.Unity.CSharp.Daemon.Stages.Dots;
 using JetBrains.ReSharper.Plugins.Unity.Resources;
 using JetBrains.ReSharper.Plugins.Unity.Resources.Icons;
-using JetBrains.ReSharper.Plugins.Unity.UnityEditorIntegration.Packages;
 
 namespace JetBrains.ReSharper.Plugins.Unity.CSharp.Feature.Services.Generate.Dots
 {

diff --git a/resharper/resharper-unity/src/Unity/UnityEditorIntegration/Api/KnownTypes.cs b/resharper/resharper-unity/src/Unity/UnityEditorIntegration/Api/KnownTypes.cs
@@ -121,6 +121,7 @@ public static class KnownTypes
         public static readonly IClrTypeName EnabledRefRO = new ClrTypeName("Unity.Entities.EnabledRefRO`1");
         public static readonly IClrTypeName EnabledRefRW = new ClrTypeName("Unity.Entities.EnabledRefRW`1");
         public static readonly IClrTypeName IEnableableComponent = new ClrTypeName("Unity.Entities.IEnableableComponent");
+        public static readonly IClrTypeName TransformUsageFlags = new ClrTypeName("Unity.Entities.TransformUsageFlags");
 
         //Unity.Mathematics
         public static readonly IClrTypeName Random = new ClrTypeName("Unity.Mathematics.Random");