Address PR feedback

src/Compilers/CSharp/Test/Semantic/Semantics/QueryTests.cs

@@ -3958,5 +3958,99 @@ static void Main()
 
             VerifyOperationTreeAndDiagnosticsForTest<QueryExpressionSyntax>(source, expectedOperationTree, expectedDiagnostics);
         }
+
+        [CompilerTrait(CompilerFeature.IOperation)]
+        [Fact, WorkItem(17838, "https://github.com/dotnet/roslyn/issues/17838")]
+        public void IOperationForQueryClause()
+        {
+            string source = @"
+using System.Collections.Generic;
+using System.Linq;
+
+class Query
+{
+    public static void Main(string[] args)
+    {
+        List<int> c = new List<int>() { 1, 2, 3, 4, 5, 6, 7 };
+        var r = from i in c /*<bind>*/select i + 1/*</bind>*/;
+    }
+}
+";
+            string expectedOperationTree = @"
+IInvocationExpression (System.Collections.Generic.IEnumerable<System.Int32> System.Linq.Enumerable.Select<System.Int32, System.Int32>(this System.Collections.Generic.IEnumerable<System.Int32> source, System.Func<System.Int32, System.Int32> selector)) (OperationKind.InvocationExpression, Type: System.Collections.Generic.IEnumerable<System.Int32>) (Syntax: 'select i + 1')
+  Instance Receiver: null
+  Arguments(2):
+      IArgument (ArgumentKind.Explicit, Matching Parameter: source) (OperationKind.Argument) (Syntax: 'from i in c')
+        IConversionExpression (Implicit, TryCast: False, Unchecked) (OperationKind.ConversionExpression, Type: System.Collections.Generic.IEnumerable<System.Int32>) (Syntax: 'from i in c')
+          Conversion: CommonConversion (Exists: True, IsIdentity: False, IsNumeric: False, IsReference: True, IsUserDefined: False) (MethodSymbol: null)
+          Operand: ILocalReferenceExpression: c (OperationKind.LocalReferenceExpression, Type: System.Collections.Generic.List<System.Int32>) (Syntax: 'c')
+        InConversion: null
+        OutConversion: null
+      IArgument (ArgumentKind.Explicit, Matching Parameter: selector) (OperationKind.Argument) (Syntax: 'i + 1')
+        IConversionExpression (Implicit, TryCast: False, Unchecked) (OperationKind.ConversionExpression, Type: System.Func<System.Int32, System.Int32>) (Syntax: 'i + 1')
+          Conversion: CommonConversion (Exists: True, IsIdentity: False, IsNumeric: False, IsReference: False, IsUserDefined: False) (MethodSymbol: null)
+          Operand: IAnonymousFunctionExpression (Symbol: lambda expression) (OperationKind.AnonymousFunctionExpression, Type: null) (Syntax: 'i + 1')
+              IBlockStatement (1 statements) (OperationKind.BlockStatement) (Syntax: 'i + 1')
+                IReturnStatement (OperationKind.ReturnStatement) (Syntax: 'i + 1')
+                  ReturnedValue: IBinaryOperatorExpression (BinaryOperatorKind.Add) (OperationKind.BinaryOperatorExpression, Type: System.Int32) (Syntax: 'i + 1')
+                      Left: IOperation:  (OperationKind.None) (Syntax: 'i')
+                      Right: ILiteralExpression (OperationKind.LiteralExpression, Type: System.Int32, Constant: 1) (Syntax: '1')
+        InConversion: null
+        OutConversion: null
+";
+            var expectedDiagnostics = DiagnosticDescription.None;
+
+            VerifyOperationTreeAndDiagnosticsForTest<SelectClauseSyntax>(source, expectedOperationTree, expectedDiagnostics);
+        }
+
+        [CompilerTrait(CompilerFeature.IOperation)]
+        [Fact, WorkItem(17838, "https://github.com/dotnet/roslyn/issues/17838")]
+        public void IOperationForRangeVariableDefinition()
+        {
+            string source = @"
+using System.Collections.Generic;
+using System.Linq;
+
+class Query
+{
+    public static void Main(string[] args)
+    {
+        List<int> c = new List<int>() { 1, 2, 3, 4, 5, 6, 7 };
+        var r = /*<bind>*/from i in c/*</bind>*/ select i + 1;
+    }
+}
+";
+            string expectedOperationTree = @"
+ILocalReferenceExpression: c (OperationKind.LocalReferenceExpression, Type: System.Collections.Generic.List<System.Int32>) (Syntax: 'c')
+";
+            var expectedDiagnostics = DiagnosticDescription.None;
+
+            VerifyOperationTreeAndDiagnosticsForTest<FromClauseSyntax>(source, expectedOperationTree, expectedDiagnostics);
+        }
+
+        [CompilerTrait(CompilerFeature.IOperation)]
+        [Fact, WorkItem(17838, "https://github.com/dotnet/roslyn/issues/17838")]
+        public void IOperationForRangeVariableReference()
+        {
+            string source = @"
+using System.Collections.Generic;
+using System.Linq;
+
+class Query
+{
+    public static void Main(string[] args)
+    {
+        List<int> c = new List<int>() { 1, 2, 3, 4, 5, 6, 7 };
+        var r = from i in c select /*<bind>*/i/*</bind>*/ + 1;
+    }
+}
+";
+            string expectedOperationTree = @"
+IOperation:  (OperationKind.None) (Syntax: 'i')
+";
+            var expectedDiagnostics = DiagnosticDescription.None;
+
+            VerifyOperationTreeAndDiagnosticsForTest<IdentifierNameSyntax>(source, expectedOperationTree, expectedDiagnostics);
+        }
     }
 }

src/Compilers/Core/Portable/Generated/Operations.xml.Generated.cs

@@ -5617,8 +5617,11 @@ internal sealed partial class LazyCollectionElementInitializerExpression : BaseC
 
     /// <summary>
     /// Represents an unrolled/lowered query expression in C# and VB.
-    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the select clause is the last clause of the unrolled query expression,
-    /// with the where clause as one of its descendant, and the from clause as the descendant of the where clause.
+    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the Operation tree has the following shape:
+    ///   ITranslatedQueryExpression
+    ///     IInvocationExpression ('Select' invocation for "select x.Name")
+    ///       IInvocationExpression ('Where' invocation for "where x.Name != null")
+    ///         IInvocationExpression ('From' invocation for "from x in set")
     /// </summary>
     internal abstract partial class BaseTranslatedQueryExpression : Operation, ITranslatedQueryExpression
     {
@@ -5650,8 +5653,11 @@ public override void Accept(OperationVisitor visitor)
 
     /// <summary>
     /// Represents an unrolled/lowered query expression in C# and VB.
-    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the select clause is the last clause of the unrolled query expression,
-    /// with the where clause as one of its descendant, and the from clause as the descendant of the where clause.
+    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the Operation tree has the following shape:
+    ///   ITranslatedQueryExpression
+    ///     IInvocationExpression ('Select' invocation for "select x.Name")
+    ///       IInvocationExpression ('Where' invocation for "where x.Name != null")
+    ///         IInvocationExpression ('From' invocation for "from x in set")
     /// </summary>
     internal sealed partial class TranslatedQueryExpression : BaseTranslatedQueryExpression, ITranslatedQueryExpression
     {
@@ -5665,8 +5671,11 @@ internal sealed partial class TranslatedQueryExpression : BaseTranslatedQueryExp
 
     /// <summary>
     /// Represents an unrolled/lowered query expression in C# and VB.
-    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the select clause is the last clause of the unrolled query expression,
-    /// with the where clause as one of its descendant, and the from clause as the descendant of the where clause.
+    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the Operation tree has the following shape:
+    ///   ITranslatedQueryExpression
+    ///     IInvocationExpression ('Select' invocation for "select x.Name")
+    ///       IInvocationExpression ('Where' invocation for "where x.Name != null")
+    ///         IInvocationExpression ('From' invocation for "from x in set")
     /// </summary>
     internal sealed partial class LazyTranslatedQueryExpression : BaseTranslatedQueryExpression, ITranslatedQueryExpression
     {

src/Compilers/Core/Portable/Operations/ITranslatedQueryExpression.cs

@@ -4,8 +4,11 @@ namespace Microsoft.CodeAnalysis.Semantics
 {
     /// <summary>
     /// Represents an unrolled/lowered query expression in C# and VB.
-    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the select clause is the last clause of the unrolled query expression,
-    /// with the where clause as one of its descendant, and the from clause as the descendant of the where clause.
+    /// For example, for the query expression "from x in set where x.Name != null select x.Name", the Operation tree has the following shape:
+    ///   ITranslatedQueryExpression
+    ///     IInvocationExpression ('Select' invocation for "select x.Name")
+    ///       IInvocationExpression ('Where' invocation for "where x.Name != null")
+    ///         IInvocationExpression ('From' invocation for "from x in set")
     /// </summary>
     /// <remarks>
     /// This interface is reserved for implementation by its associated APIs. We reserve the right to

src/Compilers/VisualBasic/Portable/Lowering/LocalRewriter/LocalRewriter_Query.vb

@@ -49,7 +49,14 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
 
             _instrumentTopLevelNonCompilerGeneratedExpressionsInQuery = Not instrumentQueryLambdaBody
 
-            Dim returnstmt = RewriteQueryLambdaBody(AddressOf VisitExpressionNode, node, _rangeVariableMap, _instrumenterOpt, instrumentQueryLambdaBody:=instrumentQueryLambdaBody AndAlso Instrument)
+            Dim rewrittenBody As BoundExpression = VisitExpressionNode(node.Expression)
+            Dim returnstmt = CreateReturnStatementForQueryLambdaBody(rewrittenBody, node)
+
+            If instrumentQueryLambdaBody AndAlso Instrument Then
+                returnstmt = _instrumenterOpt.InstrumentQueryLambdaBody(node, returnstmt)
+            End If
+
+            RemoveRangeVariables(node, _rangeVariableMap)
 
             _instrumentTopLevelNonCompilerGeneratedExpressionsInQuery = save_createSequencePointsForTopLevelNonCompilerGeneratedExpressions
 
@@ -90,14 +97,18 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
                     Dim paramRef As New BoundParameter(node.Syntax,
                                                        parameter,
                                                        False,
-                                                       parameter.Type,
-                                                       hasErrors:=parameter.Type.IsErrorType())
+                                                       parameter.Type)
 
-                    If Not paramRef.HasErrors AndAlso isReservedName AndAlso Not String.Equals(parameterName, StringConstants.Group, StringComparison.Ordinal) Then
+                    If isReservedName AndAlso Not String.Equals(parameterName, StringConstants.Group, StringComparison.Ordinal) Then
+                        If parameter.Type.IsErrorType() Then
+                            ' Skip adding to the range variable map for error case.
+                            firstUnmappedRangeVariable += 1
+                        Else
                             ' Compound variable.
                             ' Each range variable is an Anonymous Type property.
                             Debug.Assert(parameterName.Equals(StringConstants.It) OrElse parameterName.Equals(StringConstants.It1) OrElse parameterName.Equals(StringConstants.It2))
                             PopulateRangeVariableMapForAnonymousType(node.Syntax, paramRef, nodeRangeVariables, firstUnmappedRangeVariable, rangeVariableMap, inExpressionLambda)
+                        End If
                     Else
                         ' Simple case, range variable is a lambda parameter.
                         Debug.Assert(IdentifierComparison.Equals(parameterName, nodeRangeVariables(firstUnmappedRangeVariable).Name))
@@ -158,28 +169,21 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
             Next
         End Sub
 
-        Friend Shared Function RewriteQueryLambdaBody(
-            expressionRewriter As Func(Of BoundExpression, BoundExpression),
-            originalNode As BoundQueryLambda,
-            rangeVariableMap As Dictionary(Of RangeVariableSymbol, BoundExpression),
-            Optional instrumenterOpt As Instrumenter = Nothing,
-            Optional instrumentQueryLambdaBody As Boolean = False) As BoundStatement
+        Friend Shared Function CreateReturnStatementForQueryLambdaBody(
+            rewrittenBody As BoundExpression,
+            originalNode As BoundQueryLambda) As BoundStatement
 
-            Dim returnstmt As BoundStatement = New BoundReturnStatement(originalNode.Syntax,
-                                                                        expressionRewriter(originalNode.Expression),
+            Return New BoundReturnStatement(originalNode.Syntax,
+                                            rewrittenBody,
                                             Nothing,
                                             Nothing)
+        End Function
 
-            If instrumentQueryLambdaBody Then
-                returnstmt = instrumenterOpt.InstrumentQueryLambdaBody(originalNode, returnstmt)
-            End If
-
+        Friend Shared Sub RemoveRangeVariables(originalNode As BoundQueryLambda, rangeVariableMap As Dictionary(Of RangeVariableSymbol, BoundExpression))
             For Each rangeVar As RangeVariableSymbol In originalNode.RangeVariables
                 rangeVariableMap.Remove(rangeVar)
             Next
-
-            Return returnstmt
-        End Function
+        End Sub
 
         Friend Shared Function RewriteQueryLambda(rewrittenBody As BoundStatement, originalNode As BoundQueryLambda) As BoundLambda
             Dim lambdaBody = New BoundBlock(originalNode.Syntax,

src/Compilers/VisualBasic/Portable/Operations/VisualBasicOperationFactory_QueryLambdaRewriter.vb

 ' Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 Imports Microsoft.CodeAnalysis.VisualBasic
+Imports Microsoft.CodeAnalysis.VisualBasic.BoundTreeVisitor
 Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
 
 Namespace Microsoft.CodeAnalysis.Semantics
@@ -9,11 +10,17 @@ Namespace Microsoft.CodeAnalysis.Semantics
             ' We rewrite query lambda into regular lambda with 2 passes.
             ' Pass 1 uses helper methods from LocalRewriter to do the lowering. This introduces large number of DAGs.
             ' Pass 2 walks over the lowered tree and replaces duplicate bound nodes in the tree with their clones - this is a requirement for the Operation tree.
+            ' Note that the rewriter also rewrites all the query lambdas inside the body of this query lambda.
+
+            Try
                 Dim pass1Rewriter As New QueryLambdaRewriterPass1
                 Dim rewrittenLambda As BoundLambda = DirectCast(pass1Rewriter.VisitQueryLambda(node), BoundLambda)
 
                 Dim pass2Rewriter As New QueryLambdaRewriterPass2
                 Return pass2Rewriter.VisitLambda(rewrittenLambda)
+            Catch ex As CancelledByStackGuardException
+                Return node
+            End Try
         End Function
 
         Private NotInheritable Class QueryLambdaRewriterPass1
@@ -25,9 +32,11 @@ Namespace Microsoft.CodeAnalysis.Semantics
             End Sub
 
             Public Overrides Function VisitQueryLambda(node As BoundQueryLambda) As BoundNode
-                LocalRewriter.PopulateRangeVariableMapForQueryLambdaRewrite(node, _rangeVariableMap, inExpressionLambda:=False)
-                Dim rewrittenBody As BoundStatement = LocalRewriter.RewriteQueryLambdaBody(AddressOf VisitExpressionWithStackGuard, node, _rangeVariableMap)
-                Return LocalRewriter.RewriteQueryLambda(rewrittenBody, node)
+                LocalRewriter.PopulateRangeVariableMapForQueryLambdaRewrite(node, _rangeVariableMap, inExpressionLambda:=True)
+                Dim rewrittenBody As BoundExpression = VisitExpressionWithStackGuard(node.Expression)
+                Dim rewrittenStatement As BoundStatement = LocalRewriter.CreateReturnStatementForQueryLambdaBody(rewrittenBody, node)
+                LocalRewriter.RemoveRangeVariables(node, _rangeVariableMap)
+                Return LocalRewriter.RewriteQueryLambda(rewrittenStatement, node)
             End Function
 
             Public Overrides Function VisitRangeVariable(node As BoundRangeVariable) As BoundNode
@@ -39,13 +48,12 @@ Namespace Microsoft.CodeAnalysis.Semantics
                 End If
 
 #If DEBUG Then
-                ' Range variable reference should be rewritten to a parameter reference, or a call or a property access.
+                ' Range variable reference should be rewritten to a parameter reference or a property access.
                 ' We clone these bound nodes in QueryLambdaRewriterPass2 to avoid dag in the generated bound tree.
                 ' If the LocalRewriter is changed to generate more kind of bound nodes for range variables, we should handle these in QueryLambdaRewriterPass2.
                 ' Below assert helps us to stay in sync with the LocalRewriter.
                 Select Case expression.Kind
                     Case BoundKind.Parameter
-                    Case BoundKind.Call
                     Case BoundKind.PropertyAccess
                         Exit Select
                     Case Else
@@ -59,34 +67,16 @@ Namespace Microsoft.CodeAnalysis.Semantics
 
         Private NotInheritable Class QueryLambdaRewriterPass2
             Inherits BoundTreeRewriterWithStackGuard
-            Private ReadOnly _uniqueNodes As HashSet(Of BoundExpression)
-
-            Public Sub New()
-                _uniqueNodes = New HashSet(Of BoundExpression)
-            End Sub
-
-            Private Function HandleNode(Of T As BoundExpression)(node As T, cloneNode As Func(Of T, T)) As T
-                If Not _uniqueNodes.Add(node) Then
-                    node = cloneNode(node)
-                    _uniqueNodes.Add(node)
-                End If
-
-                Return node
-            End Function
+            Private ReadOnly _uniqueNodes As New HashSet(Of BoundParameter)
 
             Public Overrides Function VisitParameter(node As BoundParameter) As BoundNode
                 node = DirectCast(MyBase.VisitParameter(node), BoundParameter)
-                Return HandleNode(node, cloneNode:=Function(n) New BoundParameter(node.Syntax, node.ParameterSymbol, node.IsLValue, node.SuppressVirtualCalls, node.Type, node.HasErrors))
-            End Function
 
-            Public Overrides Function VisitCall(node As BoundCall) As BoundNode
-                node = DirectCast(MyBase.VisitCall(node), BoundCall)
-                Return HandleNode(node, cloneNode:=Function(n) New BoundCall(node.Syntax, node.Method, node.MethodGroupOpt, node.ReceiverOpt, node.Arguments, node.ConstantValueOpt, node.IsLValue, node.SuppressObjectClone, node.Type, node.HasErrors))
-            End Function
+                If node.ParameterSymbol?.ContainingSymbol.IsQueryLambdaMethod AndAlso Not _uniqueNodes.Add(node) Then
+                    node = New BoundParameter(node.Syntax, node.ParameterSymbol, node.IsLValue, node.SuppressVirtualCalls, node.Type, node.HasErrors)
+                End If
 
-            Public Overrides Function VisitPropertyAccess(node As BoundPropertyAccess) As BoundNode
-                node = DirectCast(MyBase.VisitPropertyAccess(node), BoundPropertyAccess)
-                Return HandleNode(node, cloneNode:=Function(n) New BoundPropertyAccess(node.Syntax, node.PropertySymbol, node.PropertyGroupOpt, node.AccessKind, node.IsWriteable, node.IsLValue, node.ReceiverOpt, node.Arguments, node.Type, node.HasErrors))
+                Return node
             End Function
         End Class
     End Class

src/Compilers/VisualBasic/Test/Semantic/IOperation/IOperationTests_IParameterReferenceExpression.vb

@@ -238,12 +238,10 @@ ITranslatedQueryExpression (OperationKind.TranslatedQueryExpression, Type: Syste
                     IReturnStatement (OperationKind.ReturnStatement) (Syntax: 'Group By w  ... nto Count()')
                       ReturnedValue: IAnonymousObjectCreationExpression (OperationKind.AnonymousObjectCreationExpression, Type: <anonymous type: Key w As System.String(), Key z As System.String, Key Count As System.Int32>) (Syntax: 'Group By w  ... nto Count()')
                           Initializers(3):
-                              IInvocationExpression ( Function <anonymous type: Key w As System.String(), Key z As System.String>.get_w() As System.String()) (OperationKind.InvocationExpression, Type: System.String()) (Syntax: 'Group By w  ... nto Count()')
+                              IPropertyReferenceExpression: ReadOnly Property <anonymous type: Key w As System.String(), Key z As System.String>.w As System.String() (OperationKind.PropertyReferenceExpression, Type: System.String()) (Syntax: 'Group By w  ... nto Count()')
                                 Instance Receiver: IParameterReferenceExpression: $VB$It (OperationKind.ParameterReferenceExpression, Type: <anonymous type: Key w As System.String(), Key z As System.String>) (Syntax: 'Group By w  ... nto Count()')
-                                Arguments(0)
-                              IInvocationExpression ( Function <anonymous type: Key w As System.String(), Key z As System.String>.get_z() As System.String) (OperationKind.InvocationExpression, Type: System.String) (Syntax: 'Group By w  ... nto Count()')
+                              IPropertyReferenceExpression: ReadOnly Property <anonymous type: Key w As System.String(), Key z As System.String>.z As System.String (OperationKind.PropertyReferenceExpression, Type: System.String) (Syntax: 'Group By w  ... nto Count()')
                                 Instance Receiver: IParameterReferenceExpression: $VB$It (OperationKind.ParameterReferenceExpression, Type: <anonymous type: Key w As System.String(), Key z As System.String>) (Syntax: 'Group By w  ... nto Count()')
-                                Arguments(0)
                               IInvocationExpression ( Function System.Collections.Generic.IEnumerable(Of System.String).Count() As System.Int32) (OperationKind.InvocationExpression, Type: System.Int32) (Syntax: 'Count()')
                                 Instance Receiver: IParameterReferenceExpression: $VB$ItAnonymous (OperationKind.ParameterReferenceExpression, Type: System.Collections.Generic.IEnumerable(Of System.String)) (Syntax: 'Group By w  ... nto Count()')
                                 Arguments(0)