Collection expression spread optimization

src/Compilers/CSharp/Portable/Lowering/LocalRewriter/LocalRewriter_CollectionExpression.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Immutable;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Microsoft.CodeAnalysis.CSharp.CodeGen;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
@@ -520,6 +521,43 @@ private BoundExpression CreateAndPopulateArray(BoundCollectionExpression node, A
                 throw ExceptionUtilities.UnexpectedValue(node);
             }
 
+            // Collection-expr is of the form `[..spreadExpression]`, where 'spreadExpression' has same element type as the target collection.
+            // Optimize to `spreadExpression.ToArray()` if possible.
+            if (node is { Elements: [BoundCollectionExpressionSpreadElement { Expression: { } spreadExpression } spreadElement] }
+                && spreadElement.IteratorBody is BoundExpressionStatement expressionStatement
+                && expressionStatement.Expression is not BoundConversion)
+            {
+                var spreadTypeOriginalDefinition = spreadExpression.Type!.OriginalDefinition;
+                if (tryGetToArrayMethod(spreadTypeOriginalDefinition, WellKnownType.System_Collections_Generic_List_T, WellKnownMember.System_Collections_Generic_List_T__ToArray, out MethodSymbol? listToArrayMethod))
+                {
+                    var rewrittenSpreadExpression = VisitExpression(spreadExpression);
+                    return _factory.Call(rewrittenSpreadExpression, listToArrayMethod.AsMember((NamedTypeSymbol)spreadExpression.Type!));
+                }
+
+                if (TryGetSpanConversion(spreadExpression.Type, out var asSpanMethod))
+                {
+                    var spanType = CallAsSpanMethod(spreadExpression, asSpanMethod).Type!.OriginalDefinition;
+                    if (tryGetToArrayMethod(spanType, WellKnownType.System_Span_T, WellKnownMember.System_Span_T__ToArray, out var toArrayMethod)
+                        || tryGetToArrayMethod(spanType, WellKnownType.System_ReadOnlySpan_T, WellKnownMember.System_ReadOnlySpan_T__ToArray, out toArrayMethod))
+                    {
+                        var rewrittenSpreadExpression = CallAsSpanMethod(VisitExpression(spreadExpression), asSpanMethod);
+                        return _factory.Call(rewrittenSpreadExpression, toArrayMethod.AsMember((NamedTypeSymbol)rewrittenSpreadExpression.Type!));
+                    }
+                }
+
+                bool tryGetToArrayMethod(TypeSymbol spreadTypeOriginalDefinition, WellKnownType wellKnownType, WellKnownMember wellKnownMember, [NotNullWhen(true)] out MethodSymbol? toArrayMethod)
+                {
+                    if (TypeSymbol.Equals(spreadTypeOriginalDefinition, this._compilation.GetWellKnownType(wellKnownType), TypeCompareKind.AllIgnoreOptions))
+                    {
+                        toArrayMethod = _factory.WellKnownMethod(wellKnownMember, isOptional: true);
+                        return toArrayMethod is { };
+                    }
+
+                    toArrayMethod = null;
+                    return false;
+                }
+            }
+
             if (numberIncludingLastSpread == 0)
             {
                 int knownLength = elements.Length;
@@ -599,6 +637,19 @@ private BoundExpression CreateAndPopulateArray(BoundCollectionExpression node, A
                             _factory.Binary(BinaryOperatorKind.Addition, indexTemp.Type, indexTemp, _factory.Literal(1)),
                             isRef: false,
                             indexTemp.Type));
+                },
+                (ArrayBuilder<BoundExpression> sideEffects, BoundExpression arrayTemp, BoundCollectionExpressionSpreadElement spreadElement, BoundExpression rewrittenSpreadOperand) =>
+                {
+                    if (PrepareCopyToOptimization(spreadElement, rewrittenSpreadOperand) is not var (spanSliceMethod, spreadElementAsSpan, getLengthMethod, copyToMethod))
+                        return false;
+
+                    // https://github.com/dotnet/roslyn/issues/71270
+                    // Could save the targetSpan to temp in the enclosing scope, but need to make sure we are async-safe etc.
+                    if (!TryConvertToSpanOrReadOnlySpan(arrayTemp, out var targetSpan))
+                        return false;
+
+                    PerformCopyToOptimization(sideEffects, localsBuilder, indexTemp, targetSpan, rewrittenSpreadOperand, spanSliceMethod, spreadElementAsSpan, getLengthMethod, copyToMethod);
+                    return true;
                 });
 
             var locals = localsBuilder.SelectAsArray(l => l.LocalSymbol);
@@ -612,6 +663,166 @@ private BoundExpression CreateAndPopulateArray(BoundCollectionExpression node, A
                 arrayType);
         }
 
+        /// <summary>
+        /// Returns true if type is convertible to Span or ReadOnlySpan.
+        /// If non-identity conversion, also returns a non-null asSpanMethod.
+        /// </summary>
+        private bool TryGetSpanConversion(TypeSymbol type, out MethodSymbol? asSpanMethod)
+        {
+            if (type is ArrayTypeSymbol { IsSZArray: true } arrayType
+                && _factory.WellKnownMethod(WellKnownMember.System_Span_T__ctor_Array, isOptional: true) is { } spanCtorArray)
+            {
+                asSpanMethod = spanCtorArray.AsMember(spanCtorArray.ContainingType.Construct(arrayType.ElementType));
+                return true;
+            }
+
+            if (type is not NamedTypeSymbol namedType)
+            {
+                asSpanMethod = null;
+                return false;
+            }
+
+            if (namedType.OriginalDefinition.Equals(_compilation.GetWellKnownType(WellKnownType.System_Span_T), TypeCompareKind.ConsiderEverything)
+                || namedType.OriginalDefinition.Equals(_compilation.GetWellKnownType(WellKnownType.System_ReadOnlySpan_T), TypeCompareKind.ConsiderEverything))
+            {
+                asSpanMethod = null;
+                return true;
+            }
+
+            if (namedType.OriginalDefinition.Equals(_compilation.GetWellKnownType(WellKnownType.System_Collections_Immutable_ImmutableArray_T), TypeCompareKind.ConsiderEverything)
+                && _factory.WellKnownMethod(WellKnownMember.System_Collections_Immutable_ImmutableArray_T__AsSpan, isOptional: true) is { } immutableArrayAsSpanMethod)
+            {
+                asSpanMethod = immutableArrayAsSpanMethod.AsMember(namedType);
+                return true;
+            }
+
+            if (namedType.OriginalDefinition.Equals(_compilation.GetWellKnownType(WellKnownType.System_Collections_Generic_List_T), TypeCompareKind.ConsiderEverything)
+                && _factory.WellKnownMethod(WellKnownMember.System_Runtime_InteropServices_CollectionsMarshal__AsSpan_T, isOptional: true) is { } collectionsMarshalAsSpanMethod)
+            {
+                asSpanMethod = collectionsMarshalAsSpanMethod.Construct(namedType.TypeArgumentsWithAnnotationsNoUseSiteDiagnostics[0].Type);
+                return true;
+            }
+
+            asSpanMethod = null;
+            return false;
+        }
+
+        private bool TryConvertToSpanOrReadOnlySpan(BoundExpression expression, [NotNullWhen(true)] out BoundExpression? span)
+        {
+            var type = expression.Type;
+            Debug.Assert(type is not null);
+
+            if (!TryGetSpanConversion(type, out var asSpanMethod))
+            {
+                span = null;
+                return false;
+            }
+
+            span = CallAsSpanMethod(expression, asSpanMethod);
+            return true;
+        }
+
+        private BoundExpression CallAsSpanMethod(BoundExpression spreadExpression, MethodSymbol? asSpanMethod)
+        {
+            if (asSpanMethod is null)
+            {
+                return spreadExpression;
+            }
+            if (asSpanMethod is MethodSymbol { MethodKind: MethodKind.Constructor } constructor)
+            {
+                return _factory.New(constructor, spreadExpression);
+            }
+            else if (asSpanMethod is MethodSymbol { IsStatic: true, ParameterCount: 1 })
+            {
+                return _factory.Call(receiver: null, asSpanMethod, spreadExpression);
+            }
+            else
+            {
+                return _factory.Call(spreadExpression, asSpanMethod);
+            }
+        }
+
+        /// <summary>
+        /// Verifies presence of methods necessary for the CopyTo optimization
+        /// without performing mutating actions e.g. appending to side effects or locals builders.
+        /// </summary>
+        private (MethodSymbol spanSliceMethod, BoundExpression spreadElementAsSpan, MethodSymbol getLengthMethod, MethodSymbol copyToMethod)? PrepareCopyToOptimization(
+            BoundCollectionExpressionSpreadElement spreadElement,
+            BoundExpression rewrittenSpreadOperand)
+        {
+            // Cannot use CopyTo when spread element has non-identity conversion to target element type.
+            // Could do a covariant conversion of ReadOnlySpan in future: https://github.com/dotnet/roslyn/issues/71106
+            if (spreadElement.IteratorBody is not BoundExpressionStatement expressionStatement || expressionStatement.Expression is BoundConversion)
+                return null;
+
+            if (_factory.WellKnownMethod(WellKnownMember.System_Span_T__Slice_Int_Int, isOptional: true) is not { } spanSliceMethod)
+                return null;
+
+            if (!TryConvertToSpanOrReadOnlySpan(rewrittenSpreadOperand, out var spreadOperandAsSpan))
+                return null;
+
+            if ((getSpanMethodsForSpread(WellKnownType.System_ReadOnlySpan_T, WellKnownMember.System_ReadOnlySpan_T__get_Length, WellKnownMember.System_ReadOnlySpan_T__CopyTo_Span_T)
+                    ?? getSpanMethodsForSpread(WellKnownType.System_Span_T, WellKnownMember.System_Span_T__get_Length, WellKnownMember.System_Span_T__CopyTo_Span_T))
+                is not (var getLengthMethod, var copyToMethod))
+            {
+                return null;
+            }
+
+            return (spanSliceMethod, spreadOperandAsSpan, getLengthMethod, copyToMethod);
+
+            // gets either Span or ReadOnlySpan methods for operating on the source spread element.
+            (MethodSymbol getLengthMethod, MethodSymbol copyToMethod)? getSpanMethodsForSpread(
+                WellKnownType wellKnownSpanType,
+                WellKnownMember getLengthMember,
+                WellKnownMember copyToMember)
+            {
+                if (spreadOperandAsSpan.Type!.OriginalDefinition.Equals(this._compilation.GetWellKnownType(wellKnownSpanType))
+                    && _factory.WellKnownMethod(getLengthMember, isOptional: true) is { } getLengthMethod
+                    && _factory.WellKnownMethod(copyToMember, isOptional: true) is { } copyToMethod)
+                {
+                    return (getLengthMethod, copyToMethod);
+                }
+
+                return null;
+            }
+        }
+
+        private void PerformCopyToOptimization(
+            ArrayBuilder<BoundExpression> sideEffects,
+            ArrayBuilder<BoundLocal> localsBuilder,
+            BoundLocal indexTemp,
+            BoundExpression spanTemp,
+            BoundExpression rewrittenSpreadOperand,
+            MethodSymbol spanSliceMethod,
+            BoundExpression spreadOperandAsSpan,
+            MethodSymbol getLengthMethod,
+            MethodSymbol copyToMethod)
+        {
+            // before:
+            // ..e1 // in [e0, ..e1]
+            //
+            // after (roughly):
+            // var e1Span = e1.AsSpan();
+            // e1Span.CopyTo(destinationSpan.Slice(indexTemp, e1Span.Length);
+            // indexTemp += e1Span.Length;
+
+            Debug.Assert((object)spreadOperandAsSpan != rewrittenSpreadOperand || spreadOperandAsSpan is BoundLocal { LocalSymbol.SynthesizedKind: SynthesizedLocalKind.LoweringTemp });
+            if ((object)spreadOperandAsSpan != rewrittenSpreadOperand)
+            {
+                spreadOperandAsSpan = _factory.StoreToTemp(spreadOperandAsSpan, out var assignmentToTemp);
+                sideEffects.Add(assignmentToTemp);
+                localsBuilder.Add((BoundLocal)spreadOperandAsSpan);
+            }
+
+            // e1Span.CopyTo(destinationSpan.Slice(indexTemp, e1Span.Length);
+            var spreadLength = _factory.Call(spreadOperandAsSpan, getLengthMethod.AsMember((NamedTypeSymbol)spreadOperandAsSpan.Type!));
+            var targetSlice = _factory.Call(spanTemp, spanSliceMethod.AsMember((NamedTypeSymbol)spanTemp.Type!), indexTemp, spreadLength);
+            sideEffects.Add(_factory.Call(spreadOperandAsSpan, copyToMethod.AsMember((NamedTypeSymbol)spreadOperandAsSpan.Type!), targetSlice));
+
+            // indexTemp += e1Span.Length;
+            sideEffects.Add(new BoundAssignmentOperator(rewrittenSpreadOperand.Syntax, indexTemp, _factory.Binary(BinaryOperatorKind.Addition, indexTemp.Type, indexTemp, spreadLength), isRef: false, indexTemp.Type));
+        }
+
         /// <summary>
         /// Create and populate an list from a collection expression.
         /// The collection may or may not have a known length.
@@ -722,11 +933,20 @@ private BoundExpression CreateAndPopulateList(BoundCollectionExpression node, Ty
                                 _factory.Binary(BinaryOperatorKind.Addition, indexTemp.Type, indexTemp, _factory.Literal(1)),
                                 isRef: false,
                                 indexTemp.Type));
+                    },
+                    (ArrayBuilder<BoundExpression> sideEffects, BoundExpression spanTemp, BoundCollectionExpressionSpreadElement spreadElement, BoundExpression rewrittenSpreadOperand) =>
+                    {
+                        if (PrepareCopyToOptimization(spreadElement, rewrittenSpreadOperand) is not var (spanSliceMethod, spreadElementAsSpan, getLengthMethod, copyToMethod))
+                            return false;
+
+                        PerformCopyToOptimization(sideEffects, localsBuilder, indexTemp, spanTemp, rewrittenSpreadOperand, spanSliceMethod, spreadElementAsSpan, getLengthMethod, copyToMethod);
+                        return true;
                     });
             }
             else
             {
-                var addMethod = ((MethodSymbol)_factory.WellKnownMember(WellKnownMember.System_Collections_Generic_List_T__Add)).AsMember(collectionType);
+                var addMethod = _factory.WellKnownMethod(WellKnownMember.System_Collections_Generic_List_T__Add).AsMember(collectionType);
+                var addRangeMethod = _factory.WellKnownMethod(WellKnownMember.System_Collections_Generic_List_T__AddRange, isOptional: true)?.AsMember(collectionType);
                 AddCollectionExpressionElements(
                     elements,
                     listTemp,
@@ -738,6 +958,25 @@ private BoundExpression CreateAndPopulateList(BoundCollectionExpression node, Ty
                         // list.Add(element);
                         expressions.Add(
                             _factory.Call(listTemp, addMethod, rewrittenValue));
+                    },
+                    (ArrayBuilder<BoundExpression> sideEffects, BoundExpression listTemp, BoundCollectionExpressionSpreadElement spreadElement, BoundExpression rewrittenSpreadOperand) =>
+                    {
+                        if (addRangeMethod is null)
+                            return false;
+
+                        var type = rewrittenSpreadOperand.Type!;
+
+                        var useSiteInfo = GetNewCompoundUseSiteInfo();
+                        var conversion = _compilation.Conversions.ClassifyConversionFromType(type, addRangeMethod.Parameters[0].Type, isChecked: false, ref useSiteInfo);
+                        _diagnostics.Add(rewrittenSpreadOperand.Syntax, useSiteInfo);
+                        if (conversion.IsIdentity || (conversion.IsImplicit && conversion.IsReference))
+                        {
+                            conversion.MarkUnderlyingConversionsCheckedRecursive();
+                            sideEffects.Add(_factory.Call(listTemp, addRangeMethod, rewrittenSpreadOperand));
+                            return true;
+                        }
+
+                        return false;
                     });
             }
 
@@ -782,7 +1021,8 @@ private void AddCollectionExpressionElements(
             ArrayBuilder<BoundLocal> rewrittenExpressions,
             int numberIncludingLastSpread,
             ArrayBuilder<BoundExpression> sideEffects,
-            Action<ArrayBuilder<BoundExpression>, BoundExpression, BoundExpression> addElement)
+            Action<ArrayBuilder<BoundExpression>, BoundExpression, BoundExpression> addElement,
+            Func<ArrayBuilder<BoundExpression>, BoundExpression, BoundCollectionExpressionSpreadElement, BoundExpression, bool> tryOptimizeSpreadElement)
         {
             for (int i = 0; i < elements.Length; i++)
             {
@@ -793,6 +1033,9 @@ private void AddCollectionExpressionElements(
 
                 if (element is BoundCollectionExpressionSpreadElement spreadElement)
                 {
+                    if (tryOptimizeSpreadElement(sideEffects, rewrittenReceiver, spreadElement, rewrittenExpression))
+                        continue;
+
                     var rewrittenElement = MakeCollectionExpressionSpreadElement(
                         spreadElement,
                         rewrittenExpression,