Redo SQL tree pruning (dotnet#32672)

Closes dotnet#31083
Fixes dotnet#31407

src/EFCore.Relational/Query/QuerySqlGenerator.cs

@@ -182,9 +182,8 @@ private static bool IsNonComposedSetOperation(SelectExpression selectExpression)
             && selectExpression.Projection.Count == setOperation.Source1.Projection.Count
             && selectExpression.Projection.Select(
                     (pe, index) => pe.Expression is ColumnExpression column
-                        && string.Equals(column.TableAlias, setOperation.Alias, StringComparison.Ordinal)
-                        && string.Equals(
-                            column.Name, setOperation.Source1.Projection[index].Alias, StringComparison.Ordinal))
+                        && column.TableAlias == setOperation.Alias
+                        && column.Name == setOperation.Source1.Projection[index].Alias)
                 .All(e => e);
 
     /// <inheritdoc />
@@ -237,26 +236,8 @@ protected override Expression VisitSelect(SelectExpression selectExpression)
             }
 
             GenerateTop(selectExpression);
-
-            if (selectExpression.Projection.Any())
-            {
-                GenerateList(selectExpression.Projection, e => Visit(e));
-            }
-            else
-            {
-                GenerateEmptyProjection(selectExpression);
-            }
-
-            if (selectExpression.Tables.Any())
-            {
-                _relationalCommandBuilder.AppendLine().Append("FROM ");
-
-                GenerateList(selectExpression.Tables, e => Visit(e), sql => sql.AppendLine());
-            }
-            else
-            {
-                GeneratePseudoFromClause();
-            }
+            GenerateProjection(selectExpression);
+            GenerateTables(selectExpression);
 
             if (selectExpression.Predicate != null)
             {
@@ -1060,6 +1041,40 @@ protected virtual void GenerateTop(SelectExpression selectExpression)
     {
     }
 
+    /// <summary>
+    ///     Generates the projection in the relational command
+    /// </summary>
+    /// <param name="selectExpression">A select expression to use.</param>
+    protected virtual void GenerateProjection(SelectExpression selectExpression)
+    {
+        if (selectExpression.Projection.Any())
+        {
+            GenerateList(selectExpression.Projection, e => Visit(e));
+        }
+        else
+        {
+            GenerateEmptyProjection(selectExpression);
+        }
+    }
+
+    /// <summary>
+    ///     Generates the tables in the relational command
+    /// </summary>
+    /// <param name="selectExpression">A select expression to use.</param>
+    protected virtual void GenerateTables(SelectExpression selectExpression)
+    {
+        if (selectExpression.Tables.Any())
+        {
+            _relationalCommandBuilder.AppendLine().Append("FROM ");
+
+            GenerateList(selectExpression.Tables, e => Visit(e), sql => sql.AppendLine());
+        }
+        else
+        {
+            GeneratePseudoFromClause();
+        }
+    }
+
     /// <summary>
     ///     Generates an ORDER BY clause in the relational command
     /// </summary>

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 = Prune(query);
 
 #if DEBUG
         // Verifies that all SelectExpression are marked as immutable after this point.
@@ -56,6 +57,13 @@ public override Expression Process(Expression query)
         return query;
     }
 
+    /// <summary>
+    /// Prunes unnecessarily objects from the SQL tree, e.g. tables which aren't referenced by any column.
+    /// Can be overridden by providers for provider-specific pruning.
+    /// </summary>
+    protected virtual Expression Prune(Expression query)
+        => _pruner.Prune(query);
+
 #if DEBUG
     private sealed class SelectExpressionMutableVerifyingExpressionVisitor : ExpressionVisitor
     {
@@ -64,7 +72,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 +131,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;