using System;

using System.Collections;

using System.Collections.Generic;

using System.Diagnostics.CodeAnalysis;

using System.Linq;

using System.Linq.Expressions;

using System.Reflection;

​

#nullable enable

​

namespace ExtendedLinq;

​

public class Program

{

    public static void Main()

    {

        IQueryable<Sample> source = GetSamples();

        DateTime? limit = DateTime.Now;

        var query = 

            from sample in source.Extend()

            where sample.Timestamp.AtLeast(limit)

            select new

            {

                sample.Index,

                Date = sample.Timestamp.Date,

                Time = sample.Timestamp.TimeOfDay,

            };

        var modified = query.Unextend();

        Console.WriteLine("Unmodified:");

        Console.WriteLine(query.Expression.ToString());

        Console.WriteLine();

        Console.WriteLine("Modified:");

        Console.WriteLine(modified.Expression.ToString());

        Console.WriteLine();

        var directResults = query.ToArray();

        var modifiedResults = modified.ToArray();

        bool different = directResults.Length != modifiedResults.Length;

        if (!different)

        {

            for (int i = 0; !different && i < directResults.Length; i++)

            {

                var row = directResults[i];

                Console.WriteLine($"{i}\t{row.Index}\t{row.Date}\t{row.Time}");

                different |= row.Index != modifiedResults[i].Index;

            }

        }

        Console.WriteLine($"Results {(different ? "do not " : "")}match.");

    }

​

    static readonly List<Sample> _samples = 

        Enumerable.Range(0, 72)

        .Select(i => new Sample(i, DateTime.Today.AddDays(-2).AddHours(i)))

        .ToList();

    public static IQueryable<Sample> GetSamples() => _samples.AsQueryable();

}

​

public record Sample(int Index, DateTime Timestamp);

​

public static class Extensions

{

    // Wraps the query in ExtendedQueryable<T>

    public static IQueryable<T> Extend<T>(this IQueryable<T> query) 

        => query is ExtendedQueryable<T> ? query : new ExtendedQueryable<T>(query);

    // Gets an unwrapped, transformed query from ExtendedQueryable<T>

    public static IQueryable<T> Unextend<T>(this IQueryable<T> query)

        => query is ExtendedQueryable<T> extended ? extended.GetTransformedQuery() : query;

    // Sample method, uses `_atLeast` expression when running on IQueryable<T>.

    [Expression(nameof(_atLeast))]

    public static bool AtLeast(this DateTime value, DateTime? minValue)

        => minValue is null || value >= minValue;

    // Expression substituted for AtLeast during query execution.

    private static Expression<Func<DateTime, DateTime?, bool>> _atLeast = 

        (value, minValue) => minValue == null || value >= minValue;

}

​

// The attribute that specifies an expression object to use in place of a method call

[AttributeUsage(AttributeTargets.Method)]

public sealed class ExpressionAttribute : Attribute

{

    public string MemberName { get; init; }

    public ExpressionAttribute(string memberName) 

        => MemberName = memberName;

}

​

// Wrapper class to handle modifying expressions at enumeration.

public class ExtendedQueryable<T> : IQueryable<T>, IQueryProvider

{

    private readonly IQueryable<T> _wrapped;

    #region IQueryable<T> properties

    public Type ElementType => _wrapped.ElementType;

    public Expression Expression => _wrapped.Expression;

    public IQueryProvider Provider => this;

    #endregion

    public ExtendedQueryable(IQueryable<T> wrapped)

        => _wrapped = wrapped;

    public IQueryable<T> GetTransformedQuery()

    {

        Expression transformed = TransformVisitor.Transform(_wrapped.Expression);

        return _wrapped.Provider.CreateQuery<T>(transformed);

    }

    #region IQueryable<T> methods

    public IEnumerator<T> GetEnumerator()

    {

        Expression transformed = TransformVisitor.Transform(_wrapped.Expression);

        IQueryable<T> query = _wrapped.Provider.CreateQuery<T>(transformed);

        return query.GetEnumerator();

    }

    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    #endregion

    #region IQueryProvider methods

    // This one is a hack, don't do this. Should use reflection to determine the right types.

    public IQueryable CreateQuery(Expression expression) => CreateQuery<T>(expression);

    public IQueryable<TElement> CreateQuery<TElement>(Expression expression)

    {

        var inner = _wrapped.Provider.CreateQuery<TElement>(expression);

        return new ExtendedQueryable<TElement>(inner);

    }

    // Again, should use reflection.

    public object? Execute(Expression expression) => Execute<T>(expression);

    public TResult Execute<TResult>(Expression expression)

    {

        Expression transformed = TransformVisitor.Transform(_wrapped.Expression);

        return _wrapped.Provider.Execute<TResult>(transformed);

    }

    #endregion

}

​

// This transfoms an expression by replacing calls to methods marked with `ExpressionAttribute`

internal sealed class TransformVisitor : ExpressionVisitor

{

    private TransformVisitor()

    { }

    public static T Transform<T>(T expression)

        where T : Expression

    {

        TransformVisitor visitor = new();

        return (T)visitor.Visit(expression);

    }

    // Replace calls with ExpressionAttribute decorations with their expressions.

    protected override Expression VisitMethodCall(MethodCallExpression node)

    {

        if (node.Method.GetCustomAttribute<ExpressionAttribute>() is not ExpressionAttribute attr)

            return base.VisitMethodCall(node);

        // Locate the member that supplies the expression

        MemberInfo? source = node.Method.DeclaringType

            .GetMember(attr.MemberName, BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static)

            .SingleOrDefault();

        if (source is null)

            return base.VisitMethodCall(node);

        // Get the replacement expression from the named field, property or method.

        LambdaExpression? expression = source switch

        {

            FieldInfo fi => fi.GetValue(null),

            PropertyInfo pi => pi.GetValue(null),

            MethodInfo mi => mi.Invoke(null, []),

            _ => null

        } as LambdaExpression;

        if (expression is null)

            return base.VisitMethodCall(node);

        // Generate a list of replacement expressions to exchange lambda parameters with actual arguments.

        List<Expression> args = node.Arguments.ToList();

        if (node.Object is not null)

            args.Insert(0, node.Object);

        var replacements = expression.Parameters

            .Select((p, i) => (p, a: args[i]))

            .ToDictionary(row => (Expression)row.p, row => row.a);

        // Modify the body of the lambda and return the resulting expression.

        return ReplaceVisitor.Replace(expression.Body, replacements);

    }

}

​

// Replacement visitor, used by TransformVisitor to fix parameter references during lambda substitution.

internal sealed class ReplaceVisitor : ExpressionVisitor

{

    private readonly IReadOnlyDictionary<Expression, Expression> _replacements;

    private ReplaceVisitor(IReadOnlyDictionary<Expression, Expression> replacements)

        => _replacements = replacements;

    public static T Replace<T>(T expression, IReadOnlyDictionary<Expression, Expression> replacements)

        where T : Expression

    {

        ReplaceVisitor visitor = new(replacements);

        return (T)visitor.Visit(expression);

    }

    [return: NotNullIfNotNull(nameof(node))]

    public override Expression? Visit(Expression? node)

    {

        if (node is not null && _replacements.TryGetValue(node, out var replacement))

            return replacement;

        return base.Visit(node);

    }

}

​