Redo SQL tree pruning

Cloess dotnet#31083

src/EFCore.Relational/Query/RelationalQueryTranslationPostprocessor.cs

@@ -10,6 +10,7 @@ namespace Microsoft.EntityFrameworkCore.Query;
 /// <inheritdoc />
 public class RelationalQueryTranslationPostprocessor : QueryTranslationPostprocessor
 {
+    private readonly SqlTreePruner _pruner = new();
     private readonly bool _useRelationalNulls;
 
     /// <summary>
@@ -39,7 +40,7 @@ public override Expression Process(Expression query)
         query = base.Process(query);
         query = new SelectExpressionProjectionApplyingExpressionVisitor(
             ((RelationalQueryCompilationContext)QueryCompilationContext).QuerySplittingBehavior).Visit(query);
-        query = new SelectExpressionPruningExpressionVisitor().Visit(query);
+        query = _pruner.Prune(query);
 
 #if DEBUG
         // Verifies that all SelectExpression are marked as immutable after this point.
@@ -64,7 +65,7 @@ private sealed class SelectExpressionMutableVerifyingExpressionVisitor : Express
         {
             switch (expression)
             {
-                case SelectExpression selectExpression when selectExpression.IsMutable():
+                case SelectExpression { IsMutable: true } selectExpression:
                     throw new InvalidDataException(selectExpression.Print());
 
                 case ShapedQueryExpression shapedQueryExpression:
@@ -123,7 +124,8 @@ public Expression EntryPoint(Expression expression)
 
                 if (expression is SelectExpression selectExpression)
                 {
-                    foreach (var alias in selectExpression.RemovedAliases())
+                    Check.DebugAssert(selectExpression.RemovedAliases is not null, "RemovedAliases not set");
+                    foreach (var alias in selectExpression.RemovedAliases)
                     {
                         _usedAliases.Add(alias);
                     }

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

@@ -56,7 +56,7 @@ protected override Expression VisitChildren(ExpressionVisitor visitor)
     /// <param name="source1">The <see cref="SetOperationBase.Source1" /> property of the result.</param>
     /// <param name="source2">The <see cref="SetOperationBase.Source2" /> property of the result.</param>
     /// <returns>This expression if no children changed, or an expression with the updated children.</returns>
-    public virtual ExceptExpression Update(SelectExpression source1, SelectExpression source2)
+    public override ExceptExpression Update(SelectExpression source1, SelectExpression source2)
         => source1 != Source1 || source2 != Source2
             ? new ExceptExpression(Alias, source1, source2, IsDistinct, GetAnnotations())
             : this;

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

@@ -24,12 +24,8 @@ public ExistsExpression(
         RelationalTypeMapping? typeMapping)
         : base(typeof(bool), typeMapping)
     {
-#if DEBUG
-        if (subquery.IsMutable())
-        {
-            throw new InvalidOperationException();
-        }
-#endif
+        Check.DebugAssert(!subquery.IsMutable, "Mutable subquery provided to ExistsExpression");
+
         Subquery = subquery;
     }
 

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

@@ -66,12 +66,8 @@ private InExpression(
         RelationalTypeMapping? typeMapping)
         : base(typeof(bool), typeMapping)
     {
-#if DEBUG
-        if (subquery?.IsMutable() == true)
-        {
-            throw new InvalidOperationException();
-        }
-#endif
+        Check.DebugAssert(subquery?.IsMutable != true, "Mutable subquery provided to ExistsExpression");
+
         Item = item;
         Subquery = subquery;
         Values = values;

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

@@ -32,15 +32,6 @@ private InnerJoinExpression(
     {
     }
 
-    /// <inheritdoc />
-    protected override Expression VisitChildren(ExpressionVisitor visitor)
-    {
-        var table = (TableExpressionBase)visitor.Visit(Table);
-        var joinPredicate = (SqlExpression)visitor.Visit(JoinPredicate);
-
-        return Update(table, joinPredicate);
-    }
-
     /// <summary>
     ///     Creates a new expression that is like this one, but using the supplied children. If all of the children are the same, it will
     ///     return this expression.

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

@@ -56,7 +56,7 @@ protected override Expression VisitChildren(ExpressionVisitor visitor)
     /// <param name="source1">The <see cref="SetOperationBase.Source1" /> property of the result.</param>
     /// <param name="source2">The <see cref="SetOperationBase.Source2" /> property of the result.</param>
     /// <returns>This expression if no children changed, or an expression with the updated children.</returns>
-    public virtual IntersectExpression Update(SelectExpression source1, SelectExpression source2)
+    public override IntersectExpression Update(SelectExpression source1, SelectExpression source2)
         => source1 != Source1 || source2 != Source2
             ? new IntersectExpression(Alias, source1, source2, IsDistinct, GetAnnotations())
             : this;

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

@@ -32,15 +32,6 @@ private LeftJoinExpression(
     {
     }
 
-    /// <inheritdoc />
-    protected override Expression VisitChildren(ExpressionVisitor visitor)
-    {
-        var table = (TableExpressionBase)visitor.Visit(Table);
-        var joinPredicate = (SqlExpression)visitor.Visit(JoinPredicate);
-
-        return Update(table, joinPredicate);
-    }
-
     /// <summary>
     ///     Creates a new expression that is like this one, but using the supplied children. If all of the children are the same, it will
     ///     return this expression.

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

@@ -44,6 +44,15 @@ protected PredicateJoinExpressionBase(
     /// </summary>
     public virtual SqlExpression JoinPredicate { get; }
 
+    /// <inheritdoc />
+    protected override Expression VisitChildren(ExpressionVisitor visitor)
+    {
+        var table = (TableExpressionBase)visitor.Visit(Table);
+        var joinPredicate = (SqlExpression)visitor.Visit(JoinPredicate);
+
+        return Update(table, joinPredicate);
+    }
+
     /// <summary>
     ///     Creates a new expression that is like this one, but using the supplied children. If all of the children are the same, it will
     ///     return this expression.

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

@@ -32,18 +32,13 @@ private ScalarSubqueryExpression(SelectExpression subquery, RelationalTypeMappin
 
     private static SelectExpression Verify(SelectExpression selectExpression)
     {
+        Check.DebugAssert(!selectExpression.IsMutable, "Mutable subquery provided to ExistsExpression");
+
         if (selectExpression.Projection.Count != 1)
         {
             throw new InvalidOperationException(CoreStrings.TranslationFailed(selectExpression.Print()));
         }
 
-#if DEBUG
-        if (selectExpression.IsMutable())
-        {
-            throw new InvalidOperationException();
-        }
-#endif
-
         return selectExpression;
     }
 

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

@@ -270,83 +270,6 @@ public EnclosingTermFindingVisitor(HashSet<SqlExpression> correlatedTerms)
         }
     }
 
-    private sealed class ColumnExpressionFindingExpressionVisitor : ExpressionVisitor
-    {
-        private Dictionary<string, HashSet<string>?>? _columnReferenced;
-        private Dictionary<string, HashSet<string>?>? _columnsUsedInJoinCondition;
-
-        public Dictionary<string, HashSet<string>?> FindColumns(SelectExpression selectExpression)
-        {
-            _columnReferenced = new Dictionary<string, HashSet<string>?>();
-            _columnsUsedInJoinCondition = new Dictionary<string, HashSet<string>?>();
-
-            foreach (var table in selectExpression.Tables)
-            {
-                var tableAlias = table is JoinExpressionBase joinExpressionBase
-                    ? joinExpressionBase.Table.Alias!
-                    : table.Alias!;
-                _columnReferenced[tableAlias] = null;
-            }
-
-            Visit(selectExpression);
-
-            foreach (var keyValuePair in _columnsUsedInJoinCondition)
-            {
-                var tableAlias = keyValuePair.Key;
-                if (_columnReferenced[tableAlias] != null
-                    && _columnsUsedInJoinCondition[tableAlias] != null)
-                {
-                    _columnReferenced[tableAlias]!.UnionWith(_columnsUsedInJoinCondition[tableAlias]!);
-                }
-            }
-
-            return _columnReferenced;
-        }
-
-        [return: NotNullIfNotNull("expression")]
-        public override Expression? Visit(Expression? expression)
-        {
-            switch (expression)
-            {
-                case ColumnExpression columnExpression:
-                    var tableAlias = columnExpression.TableAlias;
-                    if (_columnReferenced!.ContainsKey(tableAlias))
-                    {
-                        _columnReferenced[tableAlias] ??= new HashSet<string>();
-
-                        _columnReferenced[tableAlias]!.Add(columnExpression.Name);
-                    }
-
-                    // Always skip the table of ColumnExpression since it will traverse into deeper subquery
-                    return columnExpression;
-
-                case PredicateJoinExpressionBase predicateJoinExpressionBase:
-                    var predicateJoinTableAlias = predicateJoinExpressionBase.Table.Alias!;
-                    // Visiting the join predicate will add some columns for join table.
-                    // But if all the referenced columns are in join predicate only then we can remove the join table.
-                    // So if there are no referenced columns yet means there is still potential to remove this table,
-                    // In such case we moved the columns encountered in join predicate to other dictionary and later merge
-                    // if there are more references to the join table outside of join predicate.
-                    // We should also remove references to the outer if this column gets removed then that subquery can also remove projections
-                    // But currently we only remove table for TPT & entity splitting scenario
-                    // in which there are all table expressions which connects via joins.
-                    var joinOnSameLevel = _columnReferenced!.ContainsKey(predicateJoinTableAlias);
-                    var noReferences = !joinOnSameLevel || _columnReferenced[predicateJoinTableAlias] == null;
-                    base.Visit(predicateJoinExpressionBase);
-                    if (noReferences && joinOnSameLevel)
-                    {
-                        _columnsUsedInJoinCondition![predicateJoinTableAlias] = _columnReferenced[predicateJoinTableAlias];
-                        _columnReferenced[predicateJoinTableAlias] = null;
-                    }
-
-                    return predicateJoinExpressionBase;
-
-                default:
-                    return base.Visit(expression);
-            }
-        }
-    }
-
     private sealed class TableReferenceUpdatingExpressionVisitor : ExpressionVisitor
     {
         private readonly SelectExpression _oldSelect;