Cleanup of TPC handling (#32891)

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

@@ -1,7 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Diagnostics.CodeAnalysis;
 using Microsoft.EntityFrameworkCore.Query.SqlExpressions;
 
 namespace Microsoft.EntityFrameworkCore.Query.Internal;
@@ -23,22 +22,30 @@ public sealed class TpcTablesExpression : TableExpressionBase
     public TpcTablesExpression(
         string? alias,
         IEntityType entityType,
-        IReadOnlyList<SelectExpression> subSelectExpressions)
+        IReadOnlyList<SelectExpression> subSelectExpressions,
+        ColumnExpression discriminatorColumn,
+        List<string> discriminatorValues)
         : base(alias)
     {
         EntityType = entityType;
         SelectExpressions = subSelectExpressions;
+        DiscriminatorColumn = discriminatorColumn;
+        DiscriminatorValues = discriminatorValues;
     }
 
     private TpcTablesExpression(
         string? alias,
         IEntityType entityType,
         IReadOnlyList<SelectExpression> subSelectExpressions,
+        ColumnExpression discriminatorColumn,
+        List<string> discriminatorValues,
         IReadOnlyDictionary<string, IAnnotation>? annotations)
         : base(alias, annotations)
     {
         EntityType = entityType;
         SelectExpressions = subSelectExpressions;
+        DiscriminatorColumn = discriminatorColumn;
+        DiscriminatorValues = discriminatorValues;
     }
 
     /// <summary>
@@ -66,6 +73,24 @@ public override string Alias
     /// </summary>
     public IReadOnlyList<SelectExpression> SelectExpressions { get; }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public ColumnExpression DiscriminatorColumn { get; }
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    // Note: this gets mutated from SelectExpression.ApplyPredicate, during which the SelectExpression is still in mutable state;
+    // so that's "OK".
+    public List<string> DiscriminatorValues { get; internal set; }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -82,7 +107,7 @@ public TpcTablesExpression Prune(IReadOnlyList<string> discriminatorValues)
 
         Check.DebugAssert(subSelectExpressions.Count > 0, "TPC must have at least 1 table selected.");
 
-        return new TpcTablesExpression(Alias, EntityType, subSelectExpressions, Annotations);
+        return new TpcTablesExpression(Alias, EntityType, subSelectExpressions, DiscriminatorColumn, DiscriminatorValues, Annotations);
     }
 
     /// <summary>
@@ -91,9 +116,14 @@ public TpcTablesExpression Prune(IReadOnlyList<string> discriminatorValues)
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    // This is implementation detail hence visitors are not supposed to see inside unless they really need to.
     protected override Expression VisitChildren(ExpressionVisitor visitor)
-        => this;
+    {
+        // This is implementation detail hence visitors are not supposed to see inside the sub-selects unless they really need to.
+        var visitedColumn = (ColumnExpression)visitor.Visit(DiscriminatorColumn);
+        return visitedColumn == DiscriminatorColumn
+            ? this
+            : new(Alias, EntityType, SelectExpressions, visitedColumn, DiscriminatorValues, Annotations);
+    }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -102,24 +132,18 @@ protected override Expression VisitChildren(ExpressionVisitor visitor)
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     protected override TpcTablesExpression WithAnnotations(IReadOnlyDictionary<string, IAnnotation> annotations)
-        => new(Alias, EntityType, SelectExpressions, annotations);
+        => new(Alias, EntityType, SelectExpressions, DiscriminatorColumn, DiscriminatorValues, annotations);
 
     /// <inheritdoc />
     public override TpcTablesExpression WithAlias(string newAlias)
-        => new(newAlias, EntityType, SelectExpressions, Annotations);
+        => new(newAlias, EntityType, SelectExpressions, DiscriminatorColumn, DiscriminatorValues, Annotations);
 
     /// <inheritdoc />
     public override TableExpressionBase Clone(string? alias, ExpressionVisitor cloningExpressionVisitor)
     {
-        // Deep clone
         var subSelectExpressions = SelectExpressions.Select(cloningExpressionVisitor.Visit).ToList<SelectExpression>();
-        var newTpcTable = new TpcTablesExpression(alias, EntityType, subSelectExpressions);
-        foreach (var annotation in GetAnnotations())
-        {
-            newTpcTable.AddAnnotation(annotation.Name, annotation.Value);
-        }
-
-        return newTpcTable;
+        var discriminatorColumn = (ColumnExpression)cloningExpressionVisitor.Visit(DiscriminatorColumn);
+        return new TpcTablesExpression(alias, EntityType, subSelectExpressions, discriminatorColumn, DiscriminatorValues, Annotations);
     }
 
     /// <summary>

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.CreateSelect.cs

@@ -218,8 +218,7 @@ [new TableExpression(alias, table)],
                             subSelectExpressions.Add(SelectExpression.CreateImmutable(alias: null!, [tableExpression], projections));
                         }
 
-                        var tpcTableAlias = _sqlAliasManager.GenerateTableAlias("t");
-                        var tpcTables = new TpcTablesExpression(tpcTableAlias, entityType, subSelectExpressions);
+                        var tpcTableAlias = _sqlAliasManager.GenerateTableAlias("union");
 
                         var firstSelectExpression = subSelectExpressions[0];
                         var columns = new Dictionary<IProperty, ColumnExpression>();
@@ -236,18 +235,15 @@ [new TableExpression(alias, table)],
                         }
 
                         var discriminatorColumn = CreateColumnExpression(firstSelectExpression.Projection[^1], tpcTableAlias);
-                        var tpcDiscriminatorValues = new Dictionary<TpcTablesExpression, (ColumnExpression, List<string>)>
-                        {
-                            [tpcTables] = (discriminatorColumn, discriminatorValues)
-                        };
+                        var tpcTables = new TpcTablesExpression(
+                            tpcTableAlias, entityType, subSelectExpressions, discriminatorColumn, discriminatorValues);
                         var tableMap = tables.ToDictionary(t => t, _ => tpcTableAlias);
 
                         return new SelectExpression(
                             [tpcTables],
                             new StructuralTypeProjectionExpression(entityType, columns, tableMap, nullable: false, discriminatorColumn),
                             identifier,
-                            _sqlAliasManager,
-                            tpcDiscriminatorValues);
+                            _sqlAliasManager);
                     }
                 }