Support extension operators in Linq Expression Trees (#78942)

Author: AlekseyTs

src/Compilers/CSharp/Portable/CSharpResources.resx

@@ -8206,4 +8206,7 @@ To remove the warning, you can use /reference instead (set the Embed Interop Typ
   <data name="ERR_InstanceOperatorExtensionWrongReceiverType" xml:space="preserve">
     <value>Cannot declare instance extension operator for a type that is not known to be a struct and is not known to be a class</value>
   </data>
+  <data name="ERR_ExpressionTreeContainsExtensionBasedConditionalLogicalOperator" xml:space="preserve">
+    <value>An expression tree may not contain '&amp;&amp;' or '||' operators that use extension user defined operators.</value>
+  </data>
 </root>

src/Compilers/CSharp/Portable/Errors/ErrorCode.cs

@@ -2417,6 +2417,7 @@ internal enum ErrorCode
         ERR_OperatorInExtensionOfStaticClass = 9555,
         ERR_InstanceOperatorStructExtensionWrongReceiverRefKind = 9556,
         ERR_InstanceOperatorExtensionWrongReceiverType = 9557,
+        ERR_ExpressionTreeContainsExtensionBasedConditionalLogicalOperator = 9558,
 
         // Note: you will need to do the following after adding errors:
         //  1) Update ErrorFacts.IsBuildOnlyDiagnostic (src/Compilers/CSharp/Portable/Errors/ErrorFacts.cs)

src/Compilers/CSharp/Portable/Errors/ErrorFacts.cs

@@ -2528,6 +2528,7 @@ or ErrorCode.ERR_BadExtensionShiftOperatorSignature
                 or ErrorCode.ERR_OperatorInExtensionOfStaticClass
                 or ErrorCode.ERR_InstanceOperatorStructExtensionWrongReceiverRefKind
                 or ErrorCode.ERR_InstanceOperatorExtensionWrongReceiverType
+                or ErrorCode.ERR_ExpressionTreeContainsExtensionBasedConditionalLogicalOperator
                     => false,
             };
 #pragma warning restore CS8524 // The switch expression does not handle some values of its input type (it is not exhaustive) involving an unnamed enum value.

src/Compilers/CSharp/Portable/Lowering/DiagnosticsPass_ExpressionTrees.cs

@@ -755,6 +755,23 @@ public override BoundNode VisitUserDefinedConditionalLogicalOperator(BoundUserDe
                 {
                     Error(ErrorCode.ERR_ExpressionTreeContainsAbstractStaticMemberAccess, node);
                 }
+
+                if (binary.GetIsNewExtensionMember())
+                {
+                    // PROTOTYPE: Communicate to LDM
+                    // An expression tree factory isn't happy in this case. It throws
+                    //            System.ArgumentException : The user-defined operator method 'op_BitwiseOr' for operator 'OrElse' must have associated boolean True and False operators.
+                    // or
+                    //            System.ArgumentException : The user-defined operator method 'op_BitwiseAnd' for operator 'AndAlso' must have associated boolean True and False operators.
+                    //
+                    // from Expression.ValidateUserDefinedConditionalLogicOperator(ExpressionType nodeType, Type left, Type right, MethodInfo method)
+                    Error(ErrorCode.ERR_ExpressionTreeContainsExtensionBasedConditionalLogicalOperator, node);
+                }
+                else
+                {
+                    Debug.Assert(!node.TrueOperator.GetIsNewExtensionMember());
+                    Debug.Assert(!node.FalseOperator.GetIsNewExtensionMember());
+                }
             }
 
             return base.VisitUserDefinedConditionalLogicalOperator(node);

src/Compilers/CSharp/Portable/Lowering/ExtensionMethodBodyRewriter.cs

@@ -198,5 +198,15 @@ public override BoundNode VisitFunctionPointerLoad(BoundFunctionPointerLoad node
             Debug.Assert(symbol?.GetIsNewExtensionMember() != true);
             return base.VisitPropertySymbol(symbol);
         }
+
+        public override BoundNode VisitUnaryOperator(BoundUnaryOperator node)
+        {
+            return ExtensionMethodReferenceRewriter.VisitUnaryOperator(this, node);
+        }
+
+        protected override BoundBinaryOperator.UncommonData? VisitBinaryOperatorData(BoundBinaryOperator node)
+        {
+            return ExtensionMethodReferenceRewriter.VisitBinaryOperatorData(this, node);
+        }
     }
 }

src/Compilers/CSharp/Portable/Lowering/ExtensionMethodReferenceRewriter.cs

@@ -176,7 +176,6 @@ method.OriginalDefinition is ErrorMethodSymbol ||
                              { Name: nameof(VisitReadOnlySpanFromArray) } => method is { Name: "op_Implicit", IsExtensionMethod: false }, // Conversion operator from array to span cannot be an extension method
                              { Name: nameof(VisitLoweredConditionalAccess) } => // Nullable.HasValue cannot be an extension method
                                             method.ContainingAssembly.GetSpecialTypeMember(SpecialMember.System_Nullable_T_get_HasValue) == (object)method.OriginalDefinition,
-                             { Name: nameof(VisitUnaryOperator) } => !method.IsExtensionMethod, // Expression tree context. At the moment an operator cannot be an extension method
                              { Name: nameof(VisitUserDefinedConditionalLogicalOperator) } => !method.IsExtensionMethod, // Expression tree context. At the moment an operator cannot be an extension method
                              { Name: nameof(VisitCollectionElementInitializer) } => !method.IsExtensionMethod, // Expression tree context. At the moment an extension method cannot be used in expression tree here.
                              { Name: nameof(VisitAwaitableInfo) } => method is { Name: "GetResult", IsExtensionMethod: false }, // Cannot be an extension method
@@ -231,10 +230,19 @@ public static BoundNode VisitFunctionPointerLoad(BoundTreeRewriter rewriter, Bou
 
         protected override BoundBinaryOperator.UncommonData? VisitBinaryOperatorData(BoundBinaryOperator node)
         {
-            Debug.Assert(node.Method is null ||
-                         (!node.Method.IsExtensionMethod && !node.Method.GetIsNewExtensionMember())); // Expression tree context. At the moment an operator cannot be an extension method
+            return VisitBinaryOperatorData(this, node);
+        }
 
-            return base.VisitBinaryOperatorData(node);
+        public static BoundBinaryOperator.UncommonData? VisitBinaryOperatorData(BoundTreeRewriter rewriter, BoundBinaryOperator node)
+        {
+            // Local rewriter should have already rewritten interpolated strings into their final form of calls and gotos
+            Debug.Assert(node.InterpolatedStringHandlerData is null);
+
+            return BoundBinaryOperator.UncommonData.CreateIfNeeded(
+                node.ConstantValueOpt,
+                VisitMethodSymbolWithExtensionRewrite(rewriter, node.Method),
+                rewriter.VisitType(node.ConstrainedToType),
+                node.OriginalUserDefinedOperatorsOpt);
         }
 
         [return: NotNullIfNotNull(nameof(symbol))]
@@ -243,5 +251,20 @@ public static BoundNode VisitFunctionPointerLoad(BoundTreeRewriter rewriter, Bou
             Debug.Assert(symbol?.GetIsNewExtensionMember() != true);
             return base.VisitPropertySymbol(symbol);
         }
+
+        public override BoundNode VisitUnaryOperator(BoundUnaryOperator node)
+        {
+            return VisitUnaryOperator(this, node);
+        }
+
+        public static BoundNode VisitUnaryOperator(BoundTreeRewriter rewriter, BoundUnaryOperator node)
+        {
+            MethodSymbol? methodOpt = VisitMethodSymbolWithExtensionRewrite(rewriter, node.MethodOpt);
+            ImmutableArray<MethodSymbol> originalUserDefinedOperatorsOpt = rewriter.VisitSymbols<MethodSymbol>(node.OriginalUserDefinedOperatorsOpt);
+            BoundExpression operand = (BoundExpression)rewriter.Visit(node.Operand);
+            TypeSymbol? constrainedToTypeOpt = rewriter.VisitType(node.ConstrainedToTypeOpt);
+            TypeSymbol? type = rewriter.VisitType(node.Type);
+            return node.Update(node.OperatorKind, operand, node.ConstantValueOpt, methodOpt, constrainedToTypeOpt, node.ResultKind, originalUserDefinedOperatorsOpt, type);
+        }
     }
 }