EFCore 3 full scan workaround

EFCoreHacks.cs

// Workaround for https://github.com/dotnet/efcore/issues/17936
// Requires https://www.nuget.org/packages/LinqKit/

// EFCoreHacks.Rewrite is a custom "expression optimizer" for LinqKit, which should be set to LinqKitExtension.QueryOptimizer
// I do it in my DbContext static constructor. It allows us to use SelectEF and FirstOrDefaultEF extensions methods,
// which are to be rewritten before the expression is passed to EF query translator. Those force EF to use OUTER APPLY with proper predicate.
//
// SelectEF(exp) -> .Select(exp).Take(Int32.MaxValue).ToArray();
// FirstOrDefaultEF -> .Take(1).ToArray().FirstOrDefault();
//
// We need to call AsExpandable on our IQueryable to be able to use those:
// var items = dbContext.Items
//    .AsExpandable()
//    .Select(i => new 
//    {
//        Id = i.Id,
//        Subitems = i.SubItems.SelectEF(s => new {s.Id, s.Name}).ToArray(),
//        TopSubitem = i.SubItems.Where(s => s.IsTop).Select(s => {s. Id, s.Name}).FirstOrDefaultEF()
//    })
//    .Where(i => i.Id = 1)
//    .First();
//
public static class EFCoreHacks
{
  private static readonly EFCoreVisitor Visitor = new EFCoreVisitor();

  public static T FirstOrDefaultEF<T>(this IEnumerable<T> source) =>
    throw new NotSupportedException(
      $"It should have been rewritten by {nameof(Rewrite)}. " +
      "Make sure you have AsExpandable() called for your source");

  public static IEnumerable<TOut> SelectEF<T, TOut>(this IEnumerable<T> source, Func<T, TOut> selector) =>
    throw new NotSupportedException(
      $"It should have been rewritten by {nameof(Rewrite)}. " +
      "Make sure you have AsExpandable() called for your source");

  public static readonly Func<Expression, Expression> Rewrite = expression =>
    Visitor.Visit(expression);

  private sealed class EFCoreVisitor : ExpressionVisitor
  {
    private static readonly MethodInfo Take =
      typeof(Enumerable).GetMethod(nameof(Enumerable.Take));

    private static readonly MethodInfo ToArray =
      typeof(Enumerable).GetMethod(nameof(Enumerable.ToArray));

    private static readonly MethodInfo FirstOrDefault =
      typeof(Enumerable)
        .GetMethods()
        .First(m => m.Name == nameof(Enumerable.FirstOrDefault) &&
                    m.GetParameters().Length == 1);

    private static readonly MethodInfo Select = GetSelectMethod();

    private static MethodInfo GetSelectMethod()
    {
      Expression<Func<IEnumerable<int>, IEnumerable<int>>> exp =
        s => s.Select(_ => _);

      return ((MethodCallExpression)exp.Body)
        .Method
        .GetGenericMethodDefinition();
    }

    protected override Expression VisitMethodCall(MethodCallExpression node) =>
      TryRewriteFirstOrDefault(node, out MethodCallExpression e1) ? base.VisitMethodCall(e1)
      : TryRewriteSelect(node, out MethodCallExpression e2) ? base.VisitMethodCall(e2)
      : base.VisitMethodCall(node);

    // Replace .FirstOrDefaultEF by .Take(1).ToArray().FirstOrDefault()
    private static bool TryRewriteFirstOrDefault(MethodCallExpression node, out MethodCallExpression rewritten)
    {
      if (!(node.Method.Name == nameof(FirstOrDefaultEF) &&
            node.Method.DeclaringType == typeof(EFCoreHacks)))
      {
        rewritten = null;
        return false;
      }

      Expression queryExp = node.Arguments.Single();
      MethodCallExpression takeExp =
        Expression.Call(null,
          Take.MakeGenericMethod(node.Type),
          queryExp,
          Expression.Constant(1));

      MethodCallExpression toArrayExp =
        Expression.Call(null,
          ToArray.MakeGenericMethod(node.Type),
          takeExp);

      MethodCallExpression firstOrDefaultExp =
        Expression.Call(null,
          FirstOrDefault.MakeGenericMethod(node.Type),
          toArrayExp);

      rewritten = firstOrDefaultExp;
      return true;
    }

    // Replace .SelectEF() by .Select().Take(Int32.MaxValue).ToArray()
    private static bool TryRewriteSelect(MethodCallExpression node, out MethodCallExpression rewritten)
    {
      if (!(node.Method.Name == nameof(SelectEF) &&
            node.Method.DeclaringType == typeof(EFCoreHacks)))
      {
        rewritten = null;
        return false;
      }

      Expression queryExp = node.Arguments.First();
      Expression selectorExp = node.Arguments.Last();
      Type[] selectorTypes = node.Method.GetGenericArguments();
      if (selectorTypes.Length != 2)
      {
        throw new InvalidOperationException("SelectEF method should have 2 generic arguments");
      }

      Type selectedType = selectorTypes[1];

      MethodCallExpression selectExp =
        Expression.Call(null,
          Select.MakeGenericMethod(selectorTypes),
          queryExp,
          selectorExp);

      MethodCallExpression takeExp =
        Expression.Call(null,
          Take.MakeGenericMethod(selectedType),
          selectExp,
          Expression.Constant(int.MaxValue));

      MethodCallExpression toArrayExp =
        Expression.Call(null,
          ToArray.MakeGenericMethod(selectedType),
          takeExp);

      rewritten = toArrayExp;
      return true;
    }
  }
}