From af76e5b0a0af207bd070a8d587ac3109a4573857 Mon Sep 17 00:00:00 2001
From: Koen Bekkenutte <2912652+kbekkenutte@users.noreply.github.com>
Date: Fri, 26 Nov 2021 19:53:18 +0800
Subject: [PATCH] Fix for projectable instance members returning complex types

---
 .../ExpressionSyntaxRewriter.cs               | 24 ++++++--
 ...alProjectableComputedProperty.verified.txt | 16 +++++
 ...imaryConstructorAndProperties.verified.txt | 20 +++++++
 .../ProjectionExpressionGeneratorTests.cs     | 59 +++++++++++++++++++
 4 files changed, 113 insertions(+), 6 deletions(-)
 create mode 100644 tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MinimalProjectableComputedProperty.verified.txt
 create mode 100644 tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MixPrimaryConstructorAndProperties.verified.txt

diff --git a/src/EntityFrameworkCore.Projectables.Generator/ExpressionSyntaxRewriter.cs b/src/EntityFrameworkCore.Projectables.Generator/ExpressionSyntaxRewriter.cs
index 524a14d..e60f9b5 100644
--- a/src/EntityFrameworkCore.Projectables.Generator/ExpressionSyntaxRewriter.cs
+++ b/src/EntityFrameworkCore.Projectables.Generator/ExpressionSyntaxRewriter.cs
@@ -126,23 +126,35 @@ namespace EntityFrameworkCore.Projectables.Generator
             {
                 if (symbolInfo.Symbol is IMethodSymbol methodSymbol && methodSymbol.IsExtensionMethod)
                 {
-
                 }
                 else if (symbolInfo.Symbol.Kind is SymbolKind.Property or SymbolKind.Method or SymbolKind.Field && SymbolEqualityComparer.Default.Equals(symbolInfo.Symbol.ContainingType, _targetTypeSymbol))
                 {
+                    bool rewrite = true;
+
                     if (node.Parent is MemberAccessExpressionSyntax parentMemberAccessNode)
                     {
                         var targetSymbolInfo = _semanticModel.GetSymbolInfo(parentMemberAccessNode.Expression);
                         if (targetSymbolInfo.Symbol is { Kind: SymbolKind.Parameter })
                         {
-                            return base.VisitIdentifierName(node);
+                            rewrite = false;
                         }
                     }
+                    else if (node.Parent.IsKind(SyntaxKind.SimpleAssignmentExpression))
+                    {
+                        rewrite = false;
+                    }
+                    else if (node.Parent.IsKind(SyntaxKind.InvocationExpression))
+                    {
+                        rewrite = true;
+                    }
 
-                    return SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
-                        SyntaxFactory.IdentifierName("@this"),
-                        node
-                    );
+                    if (rewrite)
+                    {
+                        return SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
+                            SyntaxFactory.IdentifierName("@this"),
+                            node
+                        );
+                    }
                 }
                 else if (symbolInfo.Symbol.Kind is SymbolKind.NamedType && node.Parent?.Kind() is not SyntaxKind.QualifiedName)
                 {
diff --git a/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MinimalProjectableComputedProperty.verified.txt b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MinimalProjectableComputedProperty.verified.txt
new file mode 100644
index 0000000..3b85da5
--- /dev/null
+++ b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MinimalProjectableComputedProperty.verified.txt
@@ -0,0 +1,16 @@
+using System;
+using EntityFrameworkCore.Projectables;
+using Foo;
+
+namespace EntityFrameworkCore.Projectables.Generated
+#nullable disable
+{
+    public static class Foo_C_Foo
+    {
+        public static System.Linq.Expressions.Expression<System.Func<global::Foo.C, int>> Expression()
+        {
+            return (global::Foo.C @this) => 
+                @this.Bar;
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MixPrimaryConstructorAndProperties.verified.txt b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MixPrimaryConstructorAndProperties.verified.txt
new file mode 100644
index 0000000..9be618d
--- /dev/null
+++ b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.MixPrimaryConstructorAndProperties.verified.txt
@@ -0,0 +1,20 @@
+using System;
+using System.Linq;
+using System.Collections.Generic;
+using EntityFrameworkCore.Projectables;
+using Foo;
+
+namespace EntityFrameworkCore.Projectables.Generated
+#nullable disable
+{
+    public static class Foo_EntityExtensions_Entity_Something
+    {
+        public static System.Linq.Expressions.Expression<System.Func<global::Foo.EntityExtensions.Entity, global::Foo.EntityExtensions.Entity>> Expression()
+        {
+            return (global::Foo.EntityExtensions.Entity entity) => 
+                new Entity(entity.Id) {
+                    FullName = entity.FullName
+                };
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.cs b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.cs
index e6fa9b2..2f0dc64 100644
--- a/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.cs
+++ b/tests/EntityFrameworkCore.Projectables.Generator.Tests/ProjectionExpressionGeneratorTests.cs
@@ -76,6 +76,30 @@ namespace Foo {
             return Verifier.Verify(result.GeneratedTrees[0].ToString());
         }
 
+        [Fact]
+        public Task MinimalProjectableComputedProperty()
+        {
+            var compilation = CreateCompilation(@"
+using System;
+using EntityFrameworkCore.Projectables;
+namespace Foo {
+    class C {
+        public int Bar { get; set; }
+
+        [Projectable]
+        public int Foo => Bar;
+    }
+}
+");
+
+            var result = RunGenerator(compilation);
+
+            Assert.Empty(result.Diagnostics);
+            Assert.Single(result.GeneratedTrees);
+
+            return Verifier.Verify(result.GeneratedTrees[0].ToString());
+        }
+
         [Fact]
         public Task SimpleProjectableComputedProperty()
         {
@@ -900,6 +924,41 @@ namespace Foo {
             return Verifier.Verify(result.GeneratedTrees[0].ToString());
         }
 
+        [Fact]
+        public Task MixPrimaryConstructorAndProperties()
+        {
+            var compilation = CreateCompilation(@"
+using System;
+using System.Linq;
+using System.Collections.Generic;
+using EntityFrameworkCore.Projectables;
+
+namespace Foo {
+    public static class EntityExtensions
+    {
+        public record Entity(int Id)
+        {
+            public int Id { get; set; }
+            public string? FullName { get; set; }
+
+            [Projectable]
+            public static Entity Something(Entity entity)
+                => new Entity(entity.Id) {
+                    FullName = entity.FullName
+                };
+        }
+    }
+}
+");
+
+            var result = RunGenerator(compilation);
+
+            Assert.Empty(result.Diagnostics);
+            Assert.Single(result.GeneratedTrees);
+
+            return Verifier.Verify(result.GeneratedTrees[0].ToString());
+        }
+
         #region Helpers
 
         Compilation CreateCompilation(string source, bool expectedToCompile = true)