Optimize IEnumerable variant of replace operator (#14221)

iSazonov · web-flow · commit ba49eb225f82 · 2020-12-09T11:46:27.000-08:00
diff --git a/src/System.Management.Automation/engine/lang/parserutils.cs b/src/System.Management.Automation/engine/lang/parserutils.cs
@@ -965,6 +965,7 @@ internal static object ReplaceOperator(ExecutionContext context, IScriptExtent e
                 }
             }
 
+            var replacer = ReplaceOperatorImpl.Create(context, rr, substitute);
             IEnumerator list = LanguagePrimitives.GetEnumerator(lval);
             if (list == null)
             {
@@ -978,53 +979,63 @@ internal static object ReplaceOperator(ExecutionContext context, IScriptExtent e
                     lvalString = lval?.ToString() ?? string.Empty;
                 }
 
-                return ReplaceOperatorImpl(context, lvalString, rr, substitute);
+                return replacer.Replace(lvalString);
             }
             else
             {
                 List<object> resultList = new List<object>();
                 while (ParserOps.MoveNext(context, errorPosition, list))
                 {
                     string lvalString = PSObject.ToStringParser(context, ParserOps.Current(errorPosition, list));
-                    resultList.Add(ReplaceOperatorImpl(context, lvalString, rr, substitute));
+                    resultList.Add(replacer.Replace(lvalString));
                 }
 
                 return resultList.ToArray();
             }
         }
 
-        /// <summary>
-        /// ReplaceOperator implementation.
-        /// Abstracts away conversion of the optional substitute parameter to either a string or a MatchEvaluator delegate
-        /// and finally returns the result of the final Regex.Replace operation.
-        /// </summary>
-        /// <param name="context">The execution context in which to evaluate the expression.</param>
-        /// <param name="input">The input string.</param>
-        /// <param name="regex">A Regex instance.</param>
-        /// <param name="substitute">The substitute value.</param>
-        /// <returns>The result of the regex.Replace operation.</returns>
-        private static object ReplaceOperatorImpl(ExecutionContext context, string input, Regex regex, object substitute)
+        private struct ReplaceOperatorImpl
         {
-            switch (substitute)
+            public static ReplaceOperatorImpl Create(ExecutionContext context, Regex regex, object substitute)
+            {
+                return new ReplaceOperatorImpl(context, regex, substitute);
+            }
+
+            private readonly Regex _regex;
+            private readonly string _cachedReplacementString;
+            private readonly MatchEvaluator _cachedMatchEvaluator;
+
+            private ReplaceOperatorImpl(
+                ExecutionContext context,
+                Regex regex,
+                object substitute)
             {
-                case string replacementString:
-                    return regex.Replace(input, replacementString);
+                _regex = regex;
+                _cachedReplacementString = null;
+                _cachedMatchEvaluator = null;
+
+                switch (substitute)
+                {
+                    case string replacement:
+                        _cachedReplacementString = replacement;
+                        break;
 
-                case ScriptBlock sb:
-                    MatchEvaluator me = GetMatchEvaluator(context, sb);
-                    return regex.Replace(input, me);
+                    case ScriptBlock sb:
+                        _cachedMatchEvaluator = GetMatchEvaluator(context, sb);
+                        break;
 
-                case object val when LanguagePrimitives.TryConvertTo(val, out MatchEvaluator matchEvaluator):
-                    return regex.Replace(input, matchEvaluator);
+                    case object val when LanguagePrimitives.TryConvertTo(val, out _cachedMatchEvaluator):
+                        break;
 
-                default:
-                    string replacement = PSObject.ToStringParser(context, substitute);
-                    return regex.Replace(input, replacement);
+                    default:
+                        _cachedReplacementString = PSObject.ToStringParser(context, substitute);
+                        break;
+                }
             }
 
             // Local helper function to avoid creating an instance of the generated delegate helper class
             // every time 'ReplaceOperatorImpl' is invoked.
-            static MatchEvaluator GetMatchEvaluator(ExecutionContext context, ScriptBlock sb)
+            private static MatchEvaluator GetMatchEvaluator(ExecutionContext context, ScriptBlock sb)
             {
                 return match =>
                 {
@@ -1039,6 +1050,22 @@ static MatchEvaluator GetMatchEvaluator(ExecutionContext context, ScriptBlock sb
                     return PSObject.ToStringParser(context, result);
                 };
             }
+
+            /// <summary>
+            /// ReplaceOperator implementation.
+            /// Abstracts away conversion of the optional substitute parameter to either a string or a MatchEvaluator delegate
+            /// and finally returns the result of the final Regex.Replace operation.
+            /// </summary>
+            public object Replace(string input)
+            {
+                if (_cachedReplacementString is not null)
+                {
+                    return _regex.Replace(input, _cachedReplacementString);
+                }
+
+                Dbg.Assert(_cachedMatchEvaluator is not null, "_cachedMatchEvaluator should be not null when code reach here.");
+                return _regex.Replace(input, _cachedMatchEvaluator);
+            }
         }
 
         /// <summary>
diff --git a/test/powershell/Language/Operators/ReplaceOperator.Tests.ps1 b/test/powershell/Language/Operators/ReplaceOperator.Tests.ps1
@@ -11,6 +11,12 @@ Describe "Replace Operator" -Tags CI {
             $res | Should -BeExactly "image.jpg"
         }
 
+        It "Replace operator can convert an substitution object to string" {
+            $substitution = Get-Process -Id $pid
+            $res = "!3!" -replace "3",$substitution
+            $res | Should -BeExactly "!System.Diagnostics.Process (pwsh)!"
+        }
+
         It "Replace operator can be case-insensitive and case-sensitive" {
             $res = "book" -replace "B","C"
             $res | Should -BeExactly "Cook"
@@ -29,6 +35,18 @@ Describe "Replace Operator" -Tags CI {
             $res = "PowerPoint" -replace "Point"
             $res | Should -BeExactly "Power"
         }
+
+        It "Replace operator can take an enumerable as first argument, a mandatory pattern, and an optional substitution" {
+            $res = "PowerPoint1","PowerPoint2" -replace "Point","Shell"
+            $res.Count | Should -Be 2
+            $res[0] | Should -BeExactly "PowerShell1"
+            $res[1] | Should -BeExactly "PowerShell2"
+
+            $res = "PowerPoint1","PowerPoint2" -replace "Point"
+            $res.Count | Should -Be 2
+            $res[0] | Should -BeExactly "Power1"
+            $res[1] | Should -BeExactly "Power2"
+        }
     }
 
     Context "Replace operator substitutions" {
