Reuse previous lambda body binding if possible (#43366)

src/Compilers/CSharp/Portable/Binder/Binder.QueryUnboundLambdaState.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Immutable;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.CSharp
 {
@@ -54,6 +55,16 @@ protected override UnboundLambdaState WithCachingCore(bool includeCache)
                 return new QueryUnboundLambdaState(Binder, _rangeVariableMap, _parameters, _bodyFactory, includeCache);
             }
 
+            protected override BoundExpression GetLambdaExpressionBody(BoundBlock body)
+            {
+                return null;
+            }
+
+            protected override BoundBlock CreateBlockFromLambdaExpressionBody(Binder lambdaBodyBinder, BoundExpression expression, DiagnosticBag diagnostics)
+            {
+                throw ExceptionUtilities.Unreachable;
+            }
+
             protected override BoundBlock BindLambdaBody(LambdaSymbol lambdaSymbol, Binder lambdaBodyBinder, DiagnosticBag diagnostics)
             {
                 return _bodyFactory(lambdaSymbol, lambdaBodyBinder, diagnostics);

src/Compilers/CSharp/Portable/Binder/Binder_Statements.cs

@@ -7,8 +7,6 @@
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Linq;
-using System.Threading.Tasks;
-using System.Transactions;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.PooledObjects;
@@ -3155,7 +3153,7 @@ private static bool IsValidExpressionBody(SyntaxNode expressionSyntax, BoundExpr
         }
 
         /// <summary>
-        /// Binds an expression-bodied member with expression e as either { return e;} or { e; }.
+        /// Binds an expression-bodied member with expression e as either { return e; } or { e; }.
         /// </summary>
         internal virtual BoundBlock BindExpressionBodyAsBlock(ArrowExpressionClauseSyntax expressionBody,
                                                       DiagnosticBag diagnostics)
@@ -3179,7 +3177,7 @@ static BoundBlock bindExpressionBodyAsBlockInternal(ArrowExpressionClauseSyntax
         }
 
         /// <summary>
-        /// Binds a lambda with expression e as either { return e;} or { e; }.
+        /// Binds a lambda with expression e as either { return e; } or { e; }.
         /// </summary>
         public BoundBlock BindLambdaExpressionAsBlock(ExpressionSyntax body, DiagnosticBag diagnostics)
         {
@@ -3195,6 +3193,15 @@ public BoundBlock BindLambdaExpressionAsBlock(ExpressionSyntax body, DiagnosticB
             return bodyBinder.CreateBlockFromExpression(body, bodyBinder.GetDeclaredLocalsForScope(body), refKind, expression, expressionSyntax, diagnostics);
         }
 
+        public BoundBlock CreateBlockFromExpression(ExpressionSyntax body, BoundExpression expression, DiagnosticBag diagnostics)
+        {
+            Binder bodyBinder = this.GetBinder(body);
+            Debug.Assert(bodyBinder != null);
+
+            Debug.Assert(body.Kind() != SyntaxKind.RefExpression);
+            return bodyBinder.CreateBlockFromExpression(body, bodyBinder.GetDeclaredLocalsForScope(body), RefKind.None, expression, body, diagnostics);
+        }
+
         private BindValueKind GetRequiredReturnValueKind(RefKind refKind)
         {
             BindValueKind requiredValueKind = BindValueKind.RValue;

src/Compilers/CSharp/Portable/BoundTree/UnboundLambda.cs

@@ -33,7 +33,7 @@ internal sealed partial class BoundLocalFunctionStatement : IBoundLambdaOrFuncti
         BoundBlock IBoundLambdaOrFunction.Body { get => this.Body; }
     }
 
-    internal struct InferredLambdaReturnType
+    internal readonly struct InferredLambdaReturnType
     {
         internal readonly int NumExpressions;
         internal readonly bool HadExpressionlessReturn;
@@ -185,7 +185,9 @@ internal static InferredLambdaReturnType InferReturnType(ArrayBuilder<(BoundRetu
 
             HashSet<DiagnosticInfo> useSiteDiagnostics = null;
             var bestType = CalculateReturnType(compilation, conversions, delegateType, types, isAsync, node, ref useSiteDiagnostics);
-            return new InferredLambdaReturnType(types.Count, hasReturnWithoutArgument, refKind, bestType, useSiteDiagnostics.AsImmutableOrEmpty());
+            int numExpressions = types.Count;
+            types.Free();
+            return new InferredLambdaReturnType(numExpressions, hasReturnWithoutArgument, refKind, bestType, useSiteDiagnostics.AsImmutableOrEmpty());
         }
 
         private static TypeWithAnnotations CalculateReturnType(
@@ -467,6 +469,17 @@ internal UnboundLambdaState WithCaching(bool includeCache)
         public abstract RefKind RefKind(int index);
         protected abstract BoundBlock BindLambdaBody(LambdaSymbol lambdaSymbol, Binder lambdaBodyBinder, DiagnosticBag diagnostics);
 
+        /// <summary>
+        /// Return the bound expression if the lambda is has an expression body and can be reused easily.
+        /// This is an optimization only. Implementations can return null to skip reuse.
+        /// </summary>
+        protected abstract BoundExpression GetLambdaExpressionBody(BoundBlock body);
+
+        /// <summary>
+        /// Produce a bound block for the expression returned from GetLambdaExpressionBody.
+        /// </summary>
+        protected abstract BoundBlock CreateBlockFromLambdaExpressionBody(Binder lambdaBodyBinder, BoundExpression expression, DiagnosticBag diagnostics);
+
         public virtual void GenerateAnonymousFunctionConversionError(DiagnosticBag diagnostics, TypeSymbol targetType)
         {
             this.Binder.GenerateAnonymousFunctionConversionError(diagnostics, _unboundLambda.Syntax, _unboundLambda, targetType);
@@ -541,20 +554,39 @@ private bool DelegateNeedsReturn(MethodSymbol invokeMethod)
 
         private BoundLambda ReallyBind(NamedTypeSymbol delegateType)
         {
-            // When binding for real (not for return inference), we cannot reuse a body of a lambda
-            // previously bound for return type inference because we have not converted the returned
-            // expression(s) to the return type.
+            var invokeMethod = DelegateInvokeMethod(delegateType);
+            var returnType = DelegateReturnTypeWithAnnotations(invokeMethod, out RefKind refKind);
 
-            LambdaSymbol lambdaSymbol = CreateLambdaSymbol(
-                delegateType,
-                Binder.ContainingMemberOrLambda,
-                out MethodSymbol invokeMethod,
-                out TypeWithAnnotations returnType,
-                out DiagnosticBag diagnostics,
-                out ReturnInferenceCacheKey cacheKey);
+            LambdaSymbol lambdaSymbol;
+            Binder lambdaBodyBinder;
+            BoundBlock block;
+
+            var diagnostics = DiagnosticBag.GetInstance();
+            var compilation = Binder.Compilation;
+            var cacheKey = ReturnInferenceCacheKey.Create(delegateType, IsAsync);
 
-            CSharpCompilation compilation = Binder.Compilation;
-            Binder lambdaBodyBinder = new ExecutableCodeBinder(_unboundLambda.Syntax, lambdaSymbol, ParameterBinder(lambdaSymbol, Binder));
+            // When binding for real (not for return inference), there is still a good chance
+            // we could reuse a body of a lambda previous bound for return type inference.
+            // For simplicity, reuse is limited to expression-bodied lambdas. In those cases,
+            // we reuse the bound expression and apply any conversion to the return value
+            // since the inferred return type was not used when binding for return inference.
+            if (refKind == CodeAnalysis.RefKind.None &&
+                _returnInferenceCache.TryGetValue(cacheKey, out BoundLambda returnInferenceLambda) &&
+                GetLambdaExpressionBody(returnInferenceLambda.Body) is BoundExpression expression &&
+                (lambdaSymbol = returnInferenceLambda.Symbol).RefKind == refKind &&
+                (object)LambdaSymbol.InferenceFailureReturnType != lambdaSymbol.ReturnType &&
+                lambdaSymbol.ReturnTypeWithAnnotations.Equals(returnType, TypeCompareKind.ConsiderEverything))
+            {
+                lambdaBodyBinder = returnInferenceLambda.Binder;
+                block = CreateBlockFromLambdaExpressionBody(lambdaBodyBinder, expression, diagnostics);
+                diagnostics.AddRange(returnInferenceLambda.Diagnostics);
+            }
+            else
+            {
+                lambdaSymbol = CreateLambdaSymbol(Binder.ContainingMemberOrLambda, returnType, diagnostics, cacheKey.ParameterTypes, cacheKey.ParameterRefKinds, refKind);
+                lambdaBodyBinder = new ExecutableCodeBinder(_unboundLambda.Syntax, lambdaSymbol, ParameterBinder(lambdaSymbol, Binder));
+                block = BindLambdaBody(lambdaSymbol, lambdaBodyBinder, diagnostics);
+            }
 
             if (lambdaSymbol.RefKind == CodeAnalysis.RefKind.RefReadOnly)
             {
@@ -583,8 +615,6 @@ private BoundLambda ReallyBind(NamedTypeSymbol delegateType)
             ParameterHelpers.EnsureNullableAttributeExists(compilation, lambdaSymbol, lambdaParameters, diagnostics, modifyCompilation: false);
             // Note: we don't need to warn on annotations used in #nullable disable context for lambdas, as this is handled in binding already
 
-            BoundBlock block = BindLambdaBody(lambdaSymbol, lambdaBodyBinder, diagnostics);
-
             ((ExecutableCodeBinder)lambdaBodyBinder).ValidateIteratorMethods(diagnostics);
             ValidateUnsafeParameters(diagnostics, cacheKey.ParameterTypes);
 
@@ -626,16 +656,6 @@ private BoundLambda ReallyBind(NamedTypeSymbol delegateType)
             return result;
         }
 
-        private LambdaSymbol CreateLambdaSymbol(NamedTypeSymbol delegateType, Symbol containingSymbol, out MethodSymbol invokeMethod, out TypeWithAnnotations returnType, out DiagnosticBag diagnostics, out ReturnInferenceCacheKey cacheKey)
-        {
-            invokeMethod = DelegateInvokeMethod(delegateType);
-            returnType = DelegateReturnTypeWithAnnotations(invokeMethod, out RefKind refKind);
-            diagnostics = DiagnosticBag.GetInstance();
-            cacheKey = ReturnInferenceCacheKey.Create(delegateType, IsAsync);
-
-            return CreateLambdaSymbol(containingSymbol, returnType, diagnostics, cacheKey.ParameterTypes, cacheKey.ParameterRefKinds, refKind);
-        }
-
         internal LambdaSymbol CreateLambdaSymbol(
             Symbol containingSymbol,
             TypeWithAnnotations returnType,
@@ -655,7 +675,10 @@ internal LambdaSymbol CreateLambdaSymbol(
 
         internal LambdaSymbol CreateLambdaSymbol(NamedTypeSymbol delegateType, Symbol containingSymbol)
         {
-            return CreateLambdaSymbol(delegateType, containingSymbol, out _, out _, out _, out _);
+            var invokeMethod = DelegateInvokeMethod(delegateType);
+            var returnType = DelegateReturnTypeWithAnnotations(invokeMethod, out RefKind refKind);
+            var cacheKey = ReturnInferenceCacheKey.Create(delegateType, IsAsync);
+            return CreateLambdaSymbol(containingSymbol, returnType, new DiagnosticBag(), cacheKey.ParameterTypes, cacheKey.ParameterRefKinds, refKind);
         }
 
         private void ValidateUnsafeParameters(DiagnosticBag diagnostics, ImmutableArray<TypeWithAnnotations> targetParameterTypes)
@@ -741,6 +764,7 @@ private BoundLambda ReallyInferReturnType(
             var block = BindLambdaBody(lambdaSymbol, lambdaBodyBinder, diagnostics);
             return (lambdaSymbol, block, lambdaBodyBinder, diagnostics);
         }
+
         public BoundLambda BindForReturnTypeInference(NamedTypeSymbol delegateType)
         {
             var cacheKey = ReturnInferenceCacheKey.Create(delegateType, IsAsync);
@@ -1216,6 +1240,26 @@ protected override UnboundLambdaState WithCachingCore(bool includeCache)
             return new PlainUnboundLambdaState(unboundLambda: null, Binder, _parameterNames, _parameterIsDiscardOpt, _parameterTypesWithAnnotations, _parameterRefKinds, _isAsync, includeCache);
         }
 
+        protected override BoundExpression GetLambdaExpressionBody(BoundBlock body)
+        {
+            if (IsExpressionLambda)
+            {
+                var statements = body.Statements;
+                if (statements.Length == 1 &&
+                    // To simplify Binder.CreateBlockFromExpression (used below), we only reuse by-value return values.
+                    statements[0] is BoundReturnStatement { RefKind: Microsoft.CodeAnalysis.RefKind.None, ExpressionOpt: BoundExpression expr })
+                {
+                    return expr;
+                }
+            }
+            return null;
+        }
+
+        protected override BoundBlock CreateBlockFromLambdaExpressionBody(Binder lambdaBodyBinder, BoundExpression expression, DiagnosticBag diagnostics)
+        {
+            return lambdaBodyBinder.CreateBlockFromExpression((ExpressionSyntax)this.Body, expression, diagnostics);
+        }
+
         protected override BoundBlock BindLambdaBody(LambdaSymbol lambdaSymbol, Binder lambdaBodyBinder, DiagnosticBag diagnostics)
         {
             if (this.IsExpressionLambda)

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

@@ -285,7 +285,7 @@ static class Ext
 
         [ConditionalFactAttribute(typeof(IsRelease))]
         [WorkItem(40495, "https://github.com/dotnet/roslyn/issues/40495")]
-        public void NestedLambdas()
+        public void NestedLambdas_01()
         {
             var source =
 @"#nullable enable
@@ -302,6 +302,45 @@ static void Main()
             Enumerable.Range(0, 1).Sum(f =>
             Enumerable.Range(0, 1).Count(g => true)))))));
     }
+}";
+            var comp = CreateCompilation(source);
+            comp.VerifyDiagnostics();
+        }
+
+        // Test should complete in several seconds if UnboundLambda.ReallyBind
+        // uses results from _returnInferenceCache.
+        [ConditionalFactAttribute(typeof(IsRelease))]
+        [WorkItem(1083969, "https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1083969")]
+        public void NestedLambdas_02()
+        {
+            var source =
+@"using System.Collections.Generic;
+using System.Linq;
+class Program
+{
+    static void F(IEnumerable<int[]> x)
+    {
+        x.GroupBy(y => y[1]).SelectMany(x =>
+        x.GroupBy(y => y[2]).SelectMany(x =>
+        x.GroupBy(y => y[3]).SelectMany(x =>
+        x.GroupBy(y => y[4]).SelectMany(x =>
+        x.GroupBy(y => y[5]).SelectMany(x =>
+        x.GroupBy(y => y[6]).SelectMany(x =>
+        x.GroupBy(y => y[7]).SelectMany(x =>
+        x.GroupBy(y => y[8]).SelectMany(x =>
+        x.GroupBy(y => y[9]).SelectMany(x =>
+        x.GroupBy(y => y[0]).SelectMany(x =>
+        x.GroupBy(y => y[1]).SelectMany(x =>
+        x.GroupBy(y => y[2]).SelectMany(x =>
+        x.GroupBy(y => y[3]).SelectMany(x =>
+        x.GroupBy(y => y[4]).SelectMany(x =>
+        x.GroupBy(y => y[5]).SelectMany(x =>
+        x.GroupBy(y => y[6]).SelectMany(x =>
+        x.GroupBy(y => y[7]).SelectMany(x =>
+        x.GroupBy(y => y[8]).SelectMany(x =>
+        x.GroupBy(y => y[9]).SelectMany(x =>
+        x.GroupBy(y => y[0]).Select(x => x.Average(z => z[0])))))))))))))))))))));
+    }
 }";
             var comp = CreateCompilation(source);
             comp.VerifyDiagnostics();