Query: Introduce UpdateQueryExpression on ShapedQueryExpression (#27824)

This makes sure that shaper expression has right query expression

src/EFCore.Cosmos/Query/Internal/CosmosValueConverterCompensatingExpressionVisitor.cs

@@ -42,16 +42,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  };
 
  private Expression VisitShapedQueryExpression(ShapedQueryExpression shapedQueryExpression)
- {
- var selectExpression = shapedQueryExpression.QueryExpression;
- var updatedSelectExpression = Visit(selectExpression);
- return updatedSelectExpression != selectExpression
- ? shapedQueryExpression.Update(
- updatedSelectExpression,
- ReplacingExpressionVisitor.Replace(
- selectExpression, updatedSelectExpression, shapedQueryExpression.ShaperExpression))
- : shapedQueryExpression;
- }
+ => shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
 
  private Expression VisitSelect(SelectExpression selectExpression)
  {

src/EFCore.Cosmos/Query/Internal/SqlExpressionVisitor.cs

@@ -22,8 +22,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  switch (extensionExpression)
  {
  case ShapedQueryExpression shapedQueryExpression:
- return shapedQueryExpression.Update(
- Visit(shapedQueryExpression.QueryExpression), shapedQueryExpression.ShaperExpression);
+ return shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
 
  case ReadItemExpression readItemExpression:
  return readItemExpression;

src/EFCore.Relational/Query/Internal/RelationalValueConverterCompensatingExpressionVisitor.cs

@@ -46,16 +46,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  };
 
  private Expression VisitShapedQueryExpression(ShapedQueryExpression shapedQueryExpression)
- {
- var selectExpression = shapedQueryExpression.QueryExpression;
- var updatedSelectExpression = Visit(selectExpression);
- return updatedSelectExpression != selectExpression
- ? shapedQueryExpression.Update(
- updatedSelectExpression,
- ReplacingExpressionVisitor.Replace(
- selectExpression, updatedSelectExpression, shapedQueryExpression.ShaperExpression))
- : shapedQueryExpression;
- }
+ => shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
 
  private Expression VisitCase(CaseExpression caseExpression)
  {

src/EFCore.Relational/Query/Internal/SelectExpressionProjectionApplyingExpressionVisitor.cs

@@ -35,8 +35,7 @@ public SelectExpressionProjectionApplyingExpressionVisitor(QuerySplittingBehavio
  protected override Expression VisitExtension(Expression extensionExpression)
  => extensionExpression is ShapedQueryExpression shapedQueryExpression
  && shapedQueryExpression.QueryExpression is SelectExpression selectExpression
- ? shapedQueryExpression.Update(
- selectExpression,
+ ? shapedQueryExpression.UpdateShaperExpression(
  selectExpression.ApplyProjection(
  shapedQueryExpression.ShaperExpression, shapedQueryExpression.ResultCardinality, _querySplittingBehavior))
  : base.VisitExtension(extensionExpression);

src/EFCore.Relational/Query/Internal/SqlExpressionSimplifyingExpressionVisitor.cs

@@ -40,7 +40,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  {
  if (extensionExpression is ShapedQueryExpression shapedQueryExpression)
  {
- return shapedQueryExpression.Update(Visit(shapedQueryExpression.QueryExpression), shapedQueryExpression.ShaperExpression);
+ return shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
  }
 
  // Only applies to 'CASE WHEN condition...' not 'CASE operand WHEN...'

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -193,12 +193,10 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
  }
 
  translation = _sqlExpressionFactory.Exists(selectExpression, true);
+ selectExpression = _sqlExpressionFactory.Select(translation);
 
- return source.Update(
- _sqlExpressionFactory.Select(translation),
- Expression.Convert(
- new ProjectionBindingExpression(source.QueryExpression, new ProjectionMember(), typeof(bool?)),
- typeof(bool)));
+ return source.Update(selectExpression,
+ Expression.Convert(new ProjectionBindingExpression(selectExpression, new ProjectionMember(), typeof(bool?)), typeof(bool)));
  }
 
  /// <inheritdoc />
@@ -225,12 +223,10 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
  }
 
  var translation = _sqlExpressionFactory.Exists(selectExpression, false);
+ selectExpression = _sqlExpressionFactory.Select(translation);
 
- return source.Update(
- _sqlExpressionFactory.Select(translation),
- Expression.Convert(
- new ProjectionBindingExpression(source.QueryExpression, new ProjectionMember(), typeof(bool?)),
- typeof(bool)));
+ return source.Update(selectExpression,
+ Expression.Convert(new ProjectionBindingExpression(selectExpression, new ProjectionMember(), typeof(bool?)), typeof(bool)));
  }
 
  /// <inheritdoc />
@@ -290,12 +286,11 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
  selectExpression.ApplyProjection();
 
  translation = _sqlExpressionFactory.In(translation, selectExpression, false);
+ selectExpression = _sqlExpressionFactory.Select(translation);
 
- return source.Update(
- _sqlExpressionFactory.Select(translation),
+ return source.Update(selectExpression,
  Expression.Convert(
- new ProjectionBindingExpression(source.QueryExpression, new ProjectionMember(), typeof(bool?)),
- typeof(bool)));
+ new ProjectionBindingExpression(selectExpression, new ProjectionMember(), typeof(bool?)), typeof(bool)));
  }
  }
 
@@ -1089,9 +1084,9 @@ protected override Expression VisitExtension(Expression extensionExpression)
  var innerSelectExpression = BuildInnerSelectExpressionForOwnedTypeMappedToDifferentTable(
  entityProjectionExpression,
  navigation);
- 
- var innerShapedQuery = CreateShapedQueryExpression(
- targetEntityType, innerSelectExpression);
+
+  var innerShapedQuery = CreateShapedQueryExpression(
+  targetEntityType, innerSelectExpression);
 
  var makeNullable = foreignKey.PrincipalKey.Properties
  .Concat(foreignKey.Properties)
@@ -1181,7 +1176,7 @@ outerKey is NewArrayExpression newArrayExpression
  // Owned types don't support inheritance See https://github.com/dotnet/efcore/issues/9630
  // So there is no handling for dependent having TPT
  table = targetEntityType.GetViewOrTableMappings().Single().Table;
- var innerSelectExpression =  BuildInnerSelectExpressionForOwnedTypeMappedToDifferentTable(
+ var innerSelectExpression = BuildInnerSelectExpressionForOwnedTypeMappedToDifferentTable(
  entityProjectionExpression,
  navigation);
 

src/EFCore.Relational/Query/SqlExpressionVisitor.cs

@@ -22,8 +22,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
  switch (extensionExpression)
  {
  case ShapedQueryExpression shapedQueryExpression:
- return shapedQueryExpression.Update(
- Visit(shapedQueryExpression.QueryExpression), shapedQueryExpression.ShaperExpression);
+ return shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
 
  case CaseExpression caseExpression:
  return VisitCase(caseExpression);

src/EFCore.Relational/Query/SqlExpressions/SelectExpression.cs

@@ -2893,8 +2893,8 @@ private SqlRemappingVisitor PushdownIntoSubqueryInternal()
  {
  if (_clientProjections[i] is ShapedQueryExpression shapedQueryExpression)
  {
- _clientProjections[i] = shapedQueryExpression.Update(
- sqlRemappingVisitor.Visit(shapedQueryExpression.QueryExpression), shapedQueryExpression.ShaperExpression);
+ _clientProjections[i] = shapedQueryExpression.UpdateQueryExpression(
+ sqlRemappingVisitor.Visit(shapedQueryExpression.QueryExpression));
  }
  }
  }
@@ -3398,8 +3398,7 @@ List<T> VisitList<T>(List<T> list, bool inPlace, out bool changed)
  {
  var item = list[i];
  var newItem = item is ShapedQueryExpression shapedQueryExpression
- ? shapedQueryExpression.Update(
- visitor.Visit(shapedQueryExpression.QueryExpression), shapedQueryExpression.ShaperExpression)
+ ? shapedQueryExpression.UpdateQueryExpression(visitor.Visit(shapedQueryExpression.QueryExpression))
  : visitor.Visit(item);
  if (newItem != item
  && newList == list)

src/EFCore/Query/ShapedQueryExpression.cs

@@ -83,6 +83,18 @@ public virtual ShapedQueryExpression Update(Expression queryExpression, Expressi
  ? new ShapedQueryExpression(queryExpression, shaperExpression, ResultCardinality)
  : this;
 
+ /// <summary>
+ /// Creates a new expression that is like this one, but using the supplied query expression. If query expression is the same, it will
+ /// return this expression.
+ /// </summary>
+ /// <param name="queryExpression">The <see cref="QueryExpression" /> property of the result.</param>
+ /// <returns>This expression if shaper expression did not change, or an expression with the updated shaper expression.</returns>
+ public virtual ShapedQueryExpression UpdateQueryExpression(Expression queryExpression)
+ => !ReferenceEquals(queryExpression, QueryExpression)
+ ? new ShapedQueryExpression(queryExpression,
+ ReplacingExpressionVisitor.Replace(QueryExpression, queryExpression, ShaperExpression), ResultCardinality)
+ : this;
+
  /// <summary>
  /// Creates a new expression that is like this one, but using the supplied shaper expression. If shaper expression is the same, it will
  /// return this expression.