Query: Dedupe & simplify materializer

src/EFCore.InMemory/Extensions/InMemoryServiceCollectionExtensions.cs

@@ -77,6 +77,7 @@ public static IServiceCollection AddEntityFrameworkInMemoryDatabase([NotNull] th
  .TryAdd<IShapedQueryCompilingExpressionVisitorFactory, InMemoryShapedQueryCompilingExpressionVisitorFactory>()
  .TryAdd<IQueryableMethodTranslatingExpressionVisitorFactory, InMemoryQueryableMethodTranslatingExpressionVisitorFactory>()
  .TryAdd<IEntityQueryableTranslatorFactory, InMemoryEntityQueryableTranslatorFactory>()
+ .TryAdd<IShapedQueryOptimizerFactory, InMemoryShapedQueryOptimizerFactory>()
 
 
  .TryAdd<ISingletonOptions, IInMemorySingletonOptions>(p => p.GetService<IInMemorySingletonOptions>())

src/EFCore.InMemory/Query/PipeLine/InMemoryQueryExpression.cs

@@ -117,7 +117,8 @@ public void ApplyProjection(IDictionary<ProjectionMember, Expression> projection
  {
  foreach (var property in entityValuesExpression.EntityType.GetProperties())
  {
- _valueBufferSlots.Add(CreateReadValueExpression(property.ClrType, currentIndex + property.GetIndex(), property));
+ _valueBufferSlots.Add(
+ CreateReadValueExpression(property.ClrType, entityValuesExpression.StartIndex + property.GetIndex(), property));
  }
 
  _projectionMapping[member] = new EntityProjectionExpression(entityValuesExpression.EntityType, currentIndex);

src/EFCore.InMemory/Query/PipeLine/InMemoryShapedQueryOptimizer.cs

@@ -0,0 +1,18 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Query.Pipeline;
+
+namespace Microsoft.EntityFrameworkCore.InMemory.Query.Pipeline
+{
+ public class InMemoryShapedQueryOptimizer : ShapedQueryOptimizer
+ {
+ public override Expression Visit(Expression query)
+ {
+ query = base.Visit(query);
+
+ return query;
+ }
+ }
+}

src/EFCore.InMemory/Query/PipeLine/InMemoryShapedQueryOptimizerFactory.cs

@@ -0,0 +1,15 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using Microsoft.EntityFrameworkCore.Query.Pipeline;
+
+namespace Microsoft.EntityFrameworkCore.InMemory.Query.Pipeline
+{
+ public class InMemoryShapedQueryOptimizerFactory : ShapedQueryOptimizerFactory
+ {
+ public override ShapedQueryOptimizer Create(QueryCompilationContext2 queryCompilationContext)
+ {
+ return new InMemoryShapedQueryOptimizer();
+ }
+ }
+}

src/EFCore.Relational/Query/PipeLine/RelationalShapedQueryOptimizingExpressionVisitors.cs

@@ -18,6 +18,7 @@ public RelationalShapedQueryOptimizer(QueryCompilationContext2 queryCompilationC
  public override Expression Visit(Expression query)
  {
  query = base.Visit(query);
+ query = new ShaperExpressionDedupingExpressionVisitor().Process(query);
  query = new SelectExpressionProjectionApplyingExpressionVisitor().Visit(query);
  query = new SelectExpressionTableAliasUniquifyingExpressionVisitor().Visit(query);
  query = new NullComparisonTransformingExpressionVisitor().Visit(query);

src/EFCore.Relational/Query/PipeLine/ShaperExpressionDedupingExpressionVisitor.cs

@@ -0,0 +1,86 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Query.Pipeline;
+using Microsoft.EntityFrameworkCore.Relational.Query.Pipeline.SqlExpressions;
+
+namespace Microsoft.EntityFrameworkCore.Relational.Query.Pipeline
+{
+ public class ShaperExpressionDedupingExpressionVisitor : ExpressionVisitor
+ {
+ private SelectExpression _selectExpression;
+ private IDictionary<Expression, IList<Expression>> _duplicateShapers;
+
+ public Expression Process(Expression expression)
+ {
+ if (expression is ShapedQueryExpression shapedQueryExpression)
+ {
+ _selectExpression = (SelectExpression)shapedQueryExpression.QueryExpression;
+ _duplicateShapers = new Dictionary<Expression, IList<Expression>>();
+ Visit(shapedQueryExpression.ShaperExpression);
+
+ var variables = new List<ParameterExpression>();
+ var expressions = new List<Expression>();
+ var replacements = new Dictionary<Expression, Expression>();
+ var index = 0;
+
+ foreach (var kvp in _duplicateShapers)
+ {
+ if (kvp.Value.Count > 1)
+ {
+ var firstShaper = kvp.Value[0];
+ var entityParameter = Expression.Parameter(firstShaper.Type, $"entity{index++}");
+ variables.Add(entityParameter);
+ expressions.Add(Expression.Assign(
+ entityParameter,
+ firstShaper));
+
+ foreach (var shaper in kvp.Value)
+ {
+ replacements[shaper] = entityParameter;
+ }
+ }
+ }
+
+ if (variables.Count == 0)
+ {
+ return shapedQueryExpression;
+ }
+
+ expressions.Add(new ReplacingExpressionVisitor(replacements)
+ .Visit(shapedQueryExpression.ShaperExpression));
+
+ shapedQueryExpression.ShaperExpression = Expression.Block(variables, expressions);
+
+ return shapedQueryExpression;
+ }
+
+ return expression;
+ }
+
+ protected override Expression VisitExtension(Expression extensionExpression)
+ {
+ if (extensionExpression is EntityShaperExpression entityShaperExpression)
+ {
+ var serverProjection = _selectExpression.GetProjectionExpression(
+ entityShaperExpression.ValueBufferExpression.ProjectionMember);
+
+ if (_duplicateShapers.ContainsKey(serverProjection))
+ {
+ _duplicateShapers[serverProjection].Add(entityShaperExpression);
+ }
+ else
+ {
+ _duplicateShapers[serverProjection] = new List<Expression> { entityShaperExpression };
+ }
+
+ return entityShaperExpression;
+ }
+
+ return base.VisitExtension(extensionExpression);
+ }
+ }
+}

src/EFCore/Query/PipeLine/ProjectionBindingExpression.cs

@@ -20,6 +20,11 @@ public ProjectionBindingExpression(ProjectionMember projectionMember, Type type)
  public override Type Type { get; }
  public override ExpressionType NodeType => ExpressionType.Extension;
 
+ protected override Expression VisitChildren(ExpressionVisitor visitor)
+ {
+ return this;
+ }
+
  public void Print(ExpressionPrinter expressionPrinter)
  {
  expressionPrinter.StringBuilder.Append(nameof(ProjectionBindingExpression) + ": " + ProjectionMember);

src/EFCore/Query/PipeLine/ShapedQueryExpressionVisitor.cs

@@ -146,9 +146,6 @@ protected virtual Expression InjectEntityMaterializer(Expression expression)
 
  private class EntityMaterializerInjectingExpressionVisitor : ExpressionVisitor
  {
- private readonly IDictionary<EntityShaperExpression, ParameterExpression> _entityCache
- = new Dictionary<EntityShaperExpression, ParameterExpression>();
-
  private static readonly ConstructorInfo _materializationContextConstructor
  = typeof(MaterializationContext).GetConstructors().Single(ci => ci.GetParameters().Length == 2);
 
@@ -187,23 +184,10 @@ public EntityMaterializerInjectingExpressionVisitor(
  public Expression Inject(Expression expression)
  {
  var modifiedBody = Visit(expression);
- if (_async)
- {
- var resultVariable = Expression.Variable(typeof(Task<>).MakeGenericType(expression.Type), "result");
- _variables.Add(resultVariable);
- _expressions.Add(Expression.Assign(resultVariable,
- Expression.Call(
- _taskFromResultMethodInfo.MakeGenericMethod(expression.Type),
- modifiedBody)));
- _expressions.Add(resultVariable);
- }
- else
- {
- var resultVariable = Expression.Variable(expression.Type, "result");
- _variables.Add(resultVariable);
- _expressions.Add(Expression.Assign(resultVariable, modifiedBody));
- _expressions.Add(resultVariable);
- }
+ _expressions.Add(
+ _async
+ ? Expression.Call(_taskFromResultMethodInfo.MakeGenericMethod(expression.Type), modifiedBody)
+ : modifiedBody);
 
  return Expression.Block(_variables, _expressions);
  }
@@ -212,89 +196,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  {
  if (extensionExpression is EntityShaperExpression entityShaperExpression)
  {
- if (_entityCache.TryGetValue(entityShaperExpression, out var existingInstance))
- {
- return existingInstance;
- }
-
- _currentEntityIndex++;
- var entityType = entityShaperExpression.EntityType;
- var valueBuffer = entityShaperExpression.ValueBufferExpression;
- var primaryKey = entityType.FindPrimaryKey();
-
- if (_trackQueryResults && primaryKey == null)
- {
- throw new InvalidOperationException();
- }
-
- var result = Expression.Variable(entityType.ClrType, "result" + _currentEntityIndex);
- _variables.Add(result);
-
- if (_trackQueryResults)
- {
- var entry = Expression.Variable(typeof(InternalEntityEntry), "entry" + _currentEntityIndex);
- var hasNullKey = Expression.Variable(typeof(bool), "hasNullKey" + _currentEntityIndex);
- _variables.Add(entry);
- _variables.Add(hasNullKey);
-
- _expressions.Add(
- Expression.Assign(
- entry,
- Expression.Call(
- Expression.MakeMemberAccess(
- QueryCompilationContext2.QueryContextParameter,
- _stateManagerMemberInfo),
- _tryGetEntryMethodInfo,
- Expression.Constant(primaryKey),
- Expression.NewArrayInit(
- typeof(object),
- primaryKey.Properties
- .Select(p => _entityMaterializerSource.CreateReadValueExpression(
- entityShaperExpression.ValueBufferExpression,
- typeof(object),
- p.GetIndex(),
- p))),
- Expression.Constant(!entityShaperExpression.Nullable),
- hasNullKey)));
-
- _expressions.Add(
- Expression.Assign(
- result,
- Expression.Condition(
- hasNullKey,
- Expression.Constant(null, entityType.ClrType),
- Expression.Condition(
- Expression.NotEqual(
- entry,
- Expression.Constant(default(InternalEntityEntry), typeof(InternalEntityEntry))),
- Expression.Convert(
- Expression.MakeMemberAccess(entry, _entityMemberInfo),
- entityType.ClrType),
- MaterializeEntity(entityType, valueBuffer)))));
- }
- else
- {
- _expressions.Add(
- Expression.Assign(
- result,
- Expression.Condition(
- primaryKey.Properties
- .Select(p =>
- Expression.Equal(
- _entityMaterializerSource.CreateReadValueExpression(
- entityShaperExpression.ValueBufferExpression,
- typeof(object),
- p.GetIndex(),
- p),
- Expression.Constant(null)))
- .Aggregate((a, b) => Expression.OrElse(a, b)),
- Expression.Constant(null, entityType.ClrType),
- MaterializeEntity(entityType, valueBuffer))));
- }
-
- _entityCache[entityShaperExpression] = result;
-
- return result;
+ return ProcessEntityShaper(entityShaperExpression);
  }
 
  if (extensionExpression is EntityValuesExpression entityValuesExpression)
@@ -330,6 +232,81 @@ protected override Expression VisitExtension(Expression extensionExpression)
  return base.VisitExtension(extensionExpression);
  }
 
+ private Expression ProcessEntityShaper(EntityShaperExpression entityShaperExpression)
+ {
+ _currentEntityIndex++;
+ var expressions = new List<Expression>();
+ var variables = new List<ParameterExpression>();
+
+ var entityType = entityShaperExpression.EntityType;
+ var valueBuffer = entityShaperExpression.ValueBufferExpression;
+
+ var primaryKey = entityType.FindPrimaryKey();
+
+ if (_trackQueryResults && primaryKey == null)
+ {
+ throw new InvalidOperationException();
+ }
+
+ if (_trackQueryResults)
+ {
+ var entry = Expression.Variable(typeof(InternalEntityEntry), "entry" + _currentEntityIndex);
+ var hasNullKey = Expression.Variable(typeof(bool), "hasNullKey" + _currentEntityIndex);
+ variables.Add(entry);
+ variables.Add(hasNullKey);
+
+ expressions.Add(
+ Expression.Assign(
+ entry,
+ Expression.Call(
+ Expression.MakeMemberAccess(
+ QueryCompilationContext2.QueryContextParameter,
+ _stateManagerMemberInfo),
+ _tryGetEntryMethodInfo,
+ Expression.Constant(primaryKey),
+ Expression.NewArrayInit(
+ typeof(object),
+ primaryKey.Properties
+ .Select(p => _entityMaterializerSource.CreateReadValueExpression(
+ entityShaperExpression.ValueBufferExpression,
+ typeof(object),
+ p.GetIndex(),
+ p))),
+ Expression.Constant(!entityShaperExpression.Nullable),
+ hasNullKey)));
+
+ expressions.Add(Expression.Condition(
+ hasNullKey,
+ Expression.Constant(null, entityType.ClrType),
+ Expression.Condition(
+ Expression.NotEqual(
+ entry,
+ Expression.Constant(default(InternalEntityEntry), typeof(InternalEntityEntry))),
+ Expression.Convert(
+ Expression.MakeMemberAccess(entry, _entityMemberInfo),
+ entityType.ClrType),
+ MaterializeEntity(entityType, valueBuffer))));
+ }
+ else
+ {
+ expressions.Add(Expression.Condition(
+ primaryKey.Properties
+ .Select(p =>
+ Expression.Equal(
+ _entityMaterializerSource.CreateReadValueExpression(
+ entityShaperExpression.ValueBufferExpression,
+ typeof(object),
+ p.GetIndex(),
+ p),
+ Expression.Constant(null)))
+ .Aggregate((a, b) => Expression.OrElse(a, b)),
+ Expression.Constant(null, entityType.ClrType),
+ MaterializeEntity(entityType, valueBuffer)));
+ }
+
+ return Expression.Block(variables, expressions);
+ }
+
  private Expression MaterializeEntity(IEntityType entityType, Expression valueBuffer)
  {
  var expressions = new List<Expression>();

test/EFCore.Specification.Tests/Query/SimpleQueryTestBase.JoinGroupJoin.cs

@@ -48,6 +48,24 @@ join o in os on c.CustomerID equals o.CustomerID
  entryCount: 919);
  }
 
+ [ConditionalTheory]
+ [MemberData(nameof(IsAsyncData))]
+ public virtual Task Join_customers_orders_entities_same_entity_twice(bool isAsync)
+ {
+ return AssertQuery<Customer, Order>(
+ isAsync,
+ (cs, os) =>
+ from c in cs
+ join o in os on c.CustomerID equals o.CustomerID
+ select new
+ {
+ A = c,
+ B = c
+ },
+ e => e.A.CustomerID + " " + e.B.CustomerID,
+ entryCount: 89);
+ }
+
  [ConditionalTheory(Skip = "QueryIssue")]
  [MemberData(nameof(IsAsyncData))]
  public virtual Task Join_select_many(bool isAsync)