Fix the word wrapping in formatting to handle escape sequences properly (PowerShell#17316)

Author: daxian-dbw

Settings.StyleCop

@@ -162,6 +162,7 @@
           <Value>op</Value>
           <Value>my</Value>
           <Value>sb</Value>
+          <Value>vt</Value>
         </CollectionProperty>
       </AnalyzerSettings>
     </Analyzer>

src/System.Management.Automation/FormatAndOutput/common/ComplexWriter.cs

@@ -6,8 +6,10 @@
 using System.Collections.ObjectModel;
 using System.Collections.Specialized;
 using System.Globalization;
+using System.Management.Automation;
 using System.Management.Automation.Internal;
 using System.Text;
+using System.Text.RegularExpressions;
 
 namespace Microsoft.PowerShell.Commands.Internal.Format
 {
@@ -146,7 +148,7 @@ private void WriteToScreen()
             int indentationAbsoluteValue = (firstLineIndentation > 0) ? firstLineIndentation : -firstLineIndentation;
             if (indentationAbsoluteValue >= usefulWidth)
             {
-                // valu too big, we reset it to zero
+                // value too big, we reset it to zero
                 firstLineIndentation = 0;
             }
 
@@ -353,27 +355,58 @@ static StringManipulationHelper()
         private static IEnumerable<GetWordsResult> GetWords(string s)
         {
             StringBuilder sb = new StringBuilder();
-            GetWordsResult result = new GetWordsResult();
+            StringBuilder vtSeqs = null;
+            Dictionary<int, int> vtRanges = null;
 
+            var valueStrDec = new ValueStringDecorated(s);
+            if (valueStrDec.IsDecorated)
+            {
+                vtSeqs = new StringBuilder();
+                vtRanges = valueStrDec.EscapeSequenceRanges;
+            }
+
+            bool wordHasVtSeqs = false;
             for (int i = 0; i < s.Length; i++)
             {
-                // Soft hyphen = \u00AD - Should break, and add a hyphen if needed. If not needed for a break, hyphen should be absent
-                if (s[i] == ' ' || s[i] == '\t' || s[i] == s_softHyphen)
+                if (vtRanges?.TryGetValue(i, out int len) == true)
                 {
-                    result.Word = sb.ToString();
-                    sb.Clear();
-                    result.Delim = new string(s[i], 1);
+                    var vtSpan = s.AsSpan(i, len);
+                    sb.Append(vtSpan);
+                    vtSeqs.Append(vtSpan);
 
-                    yield return result;
+                    wordHasVtSeqs = true;
+                    i += len - 1;
+                    continue;
+                }
+
+                string delimiter = null;
+                if (s[i] == ' ' || s[i] == '\t' || s[i] == s_softHyphen)
+                {
+                    // Soft hyphen = \u00AD - Should break, and add a hyphen if needed.
+                    // If not needed for a break, hyphen should be absent.
+                    delimiter = new string(s[i], 1);
                 }
-                // Non-breaking space = \u00A0 - ideally shouldn't wrap
-                // Hard hyphen = \u2011 - Should not break
                 else if (s[i] == s_hardHyphen || s[i] == s_nonBreakingSpace)
                 {
-                    result.Word = sb.ToString();
-                    sb.Clear();
-                    result.Delim = string.Empty;
+                    // Non-breaking space = \u00A0 - ideally shouldn't wrap.
+                    // Hard hyphen = \u2011 - Should not break.
+                    delimiter = string.Empty;
+                }
+
+                if (delimiter is not null)
+                {
+                    if (wordHasVtSeqs && !sb.EndsWith(PSStyle.Instance.Reset))
+                    {
+                        sb.Append(PSStyle.Instance.Reset);
+                    }
 
+                    var result = new GetWordsResult()
+                    {
+                        Word = sb.ToString(),
+                        Delim = delimiter
+                    };
+
+                    sb.Clear().Append(vtSeqs);
                     yield return result;
                 }
                 else
@@ -382,10 +415,23 @@ private static IEnumerable<GetWordsResult> GetWords(string s)
                 }
             }
 
-            result.Word = sb.ToString();
-            result.Delim = string.Empty;
+            if (wordHasVtSeqs)
+            {
+                if (sb.Length == vtSeqs.Length)
+                {
+                    // This indicates 'sb' only contains all VT sequences, which may happen when the string ends with a word delimiter.
+                    // For a word that contains VT sequence only, it's the same as an empty string to the formatting system,
+                    // because nothing will actually be rendered.
+                    // So, we use an empty string in this case to avoid unneeded string allocations.
+                    sb.Clear();
+                }
+                else if (!sb.EndsWith(PSStyle.Instance.Reset))
+                {
+                    sb.Append(PSStyle.Instance.Reset);
+                }
+            }
 
-            yield return result;
+            yield return new GetWordsResult() { Word = sb.ToString(), Delim = string.Empty };
         }
 
         internal static StringCollection GenerateLines(DisplayCells displayCells, string val, int firstLineLen, int followingLinesLen)
@@ -412,9 +458,9 @@ private static StringCollection GenerateLinesWithoutWordWrap(DisplayCells displa
             }
 
             // break string on newlines and process each line separately
-            string[] lines = SplitLines(val);
+            List<string> lines = SplitLines(val);
 
-            for (int k = 0; k < lines.Length; k++)
+            for (int k = 0; k < lines.Count; k++)
             {
                 string currentLine = lines[k];
 
@@ -530,9 +576,9 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
             }
 
             // break string on newlines and process each line separately
-            string[] lines = SplitLines(val);
+            List<string> lines = SplitLines(val);
 
-            for (int k = 0; k < lines.Length; k++)
+            for (int k = 0; k < lines.Count; k++)
             {
                 if (lines[k] == null || displayCells.Length(lines[k]) <= firstLineLen)
                 {
@@ -545,28 +591,34 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
                 int lineWidth = firstLineLen;
                 bool firstLine = true;
                 StringBuilder singleLine = new StringBuilder();
+                string resetStr = PSStyle.Instance.Reset;
 
                 foreach (GetWordsResult word in GetWords(lines[k]))
                 {
                     string wordToAdd = word.Word;
+                    string suffix = null;
 
                     // Handle soft hyphen
-                    if (word.Delim == s_softHyphen.ToString())
+                    if (word.Delim.Length == 1 && word.Delim[0] == s_softHyphen)
                     {
                         int wordWidthWithHyphen = displayCells.Length(wordToAdd) + displayCells.Length(s_softHyphen);
 
                         // Add hyphen only if necessary
                         if (wordWidthWithHyphen == spacesLeft)
                         {
-                            wordToAdd += "-";
+                            suffix = "-";
                         }
                     }
-                    else
+                    else if (!string.IsNullOrEmpty(word.Delim))
                     {
-                        if (!string.IsNullOrEmpty(word.Delim))
-                        {
-                            wordToAdd += word.Delim;
-                        }
+                        suffix = word.Delim;
+                    }
+
+                    if (suffix is not null)
+                    {
+                        wordToAdd = wordToAdd.EndsWith(resetStr)
+                            ? wordToAdd.Insert(wordToAdd.Length - resetStr.Length, suffix)
+                            : wordToAdd + suffix;
                     }
 
                     int wordWidth = displayCells.Length(wordToAdd);
@@ -591,15 +643,35 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
                     // Word is wider than a single line
                     if (wordWidth > lineWidth)
                     {
-                        foreach (char c in wordToAdd)
+                        Dictionary<int, int> vtRanges = null;
+                        StringBuilder vtSeqs = null;
+
+                        var valueStrDec = new ValueStringDecorated(wordToAdd);
+                        if (valueStrDec.IsDecorated)
                         {
-                            char charToAdd = c;
-                            int charWidth = displayCells.Length(c);
+                            vtSeqs = new StringBuilder();
+                            vtRanges = valueStrDec.EscapeSequenceRanges;
+                        }
 
-                            // corner case: we have a two cell character and the current
-                            // display length is one.
-                            // add a single cell arbitrary character instead of the original
-                            // one and keep going
+                        bool hasEscSeqs = false;
+                        for (int i = 0; i < wordToAdd.Length; i++)
+                        {
+                            if (vtRanges?.TryGetValue(i, out int len) == true)
+                            {
+                                var vtSpan = wordToAdd.AsSpan(i, len);
+                                singleLine.Append(vtSpan);
+                                vtSeqs.Append(vtSpan);
+
+                                hasEscSeqs = true;
+                                i += len - 1;
+                                continue;
+                            }
+
+                            char charToAdd = wordToAdd[i];
+                            int charWidth = displayCells.Length(charToAdd);
+
+                            // Corner case: we have a two cell character and the current display length is one.
+                            // Add a single cell arbitrary character instead of the original one and keep going.
                             if (charWidth > lineWidth)
                             {
                                 charToAdd = '?';
@@ -608,9 +680,13 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
 
                             if (charWidth > spacesLeft)
                             {
+                                if (hasEscSeqs && !singleLine.EndsWith(resetStr))
+                                {
+                                    singleLine.Append(resetStr);
+                                }
+
                                 retVal.Add(singleLine.ToString());
-                                singleLine.Clear();
-                                singleLine.Append(charToAdd);
+                                singleLine.Clear().Append(vtSeqs).Append(charToAdd);
 
                                 if (firstLine)
                                 {
@@ -632,8 +708,7 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
                         if (wordWidth > spacesLeft)
                         {
                             retVal.Add(singleLine.ToString());
-                            singleLine.Clear();
-                            singleLine.Append(wordToAdd);
+                            singleLine.Clear().Append(wordToAdd);
 
                             if (firstLine)
                             {
@@ -663,49 +738,77 @@ private static StringCollection GenerateLinesWithWordWrap(DisplayCells displayCe
         /// </summary>
         /// <param name="s">String to split.</param>
         /// <returns>String array with the values.</returns>
-        internal static string[] SplitLines(string s)
+        internal static List<string> SplitLines(string s)
         {
-            if (string.IsNullOrEmpty(s))
-                return new string[1] { s };
+            if (string.IsNullOrEmpty(s) || !s.Contains('\n'))
+            {
+                return new List<string>(capacity: 1) { s?.Replace("\r", string.Empty) };
+            }
 
             StringBuilder sb = new StringBuilder();
+            List<string> list = new List<string>();
+
+            StringBuilder vtSeqs = null;
+            Dictionary<int, int> vtRanges = null;
 
-            foreach (char c in s)
+            var valueStrDec = new ValueStringDecorated(s);
+            if (valueStrDec.IsDecorated)
             {
-                if (c != '\r')
-                    sb.Append(c);
+                vtSeqs = new StringBuilder();
+                vtRanges = valueStrDec.EscapeSequenceRanges;
             }
 
-            return sb.ToString().Split(s_newLineChar);
-        }
-
-#if false
-        internal static string StripNewLines (string s)
-        {
-            if (string.IsNullOrEmpty (s))
-                return s;
-
-            string[] lines = SplitLines (s);
+            bool hasVtSeqs = false;
+            for (int i = 0; i < s.Length; i++)
+            {
+                if (vtRanges?.TryGetValue(i, out int len) == true)
+                {
+                    var vtSpan = s.AsSpan(i, len);
+                    sb.Append(vtSpan);
+                    vtSeqs.Append(vtSpan);
 
-            if (lines.Length == 0)
-                return null;
+                    hasVtSeqs = true;
+                    i += len - 1;
+                    continue;
+                }
 
-            if (lines.Length == 1)
-                return lines[0];
+                char c = s[i];
+                if (c == '\n')
+                {
+                    if (hasVtSeqs && !sb.EndsWith(PSStyle.Instance.Reset))
+                    {
+                        sb.Append(PSStyle.Instance.Reset);
+                    }
 
-            StringBuilder sb = new StringBuilder ();
+                    list.Add(sb.ToString());
+                    sb.Clear().Append(vtSeqs);
+                }
+                else if (c != '\r')
+                {
+                    sb.Append(c);
+                }
+            }
 
-            for (int k = 0; k < lines.Length; k++)
+            if (hasVtSeqs)
             {
-                if (k == 0)
-                    sb.Append (lines[k]);
-                else
-                    sb.Append (" " + lines[k]);
+                if (sb.Length == vtSeqs.Length)
+                {
+                    // This indicates 'sb' only contains all VT sequences, which may happen when the string ends with '\n'.
+                    // For a sub-string that contains VT sequence only, it's the same as an empty string to the formatting
+                    // system, because nothing will actually be rendered.
+                    // So, we use an empty string in this case to avoid unneeded string allocations.
+                    sb.Clear();
+                }
+                else if (!sb.EndsWith(PSStyle.Instance.Reset))
+                {
+                    sb.Append(PSStyle.Instance.Reset);
+                }
             }
 
-            return sb.ToString ();
+            list.Add(sb.ToString());
+            return list;
         }
-#endif
+
         internal static string TruncateAtNewLine(string s)
         {
             if (string.IsNullOrEmpty(s))

src/System.Management.Automation/FormatAndOutput/common/ILineOutput.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Collections.Generic;
 using System.Globalization;
 using System.IO;
 using System.Management.Automation;
@@ -379,10 +380,10 @@ private void WriteLineInternal(string val, int cols)
             }
 
             // check for line breaks
-            string[] lines = StringManipulationHelper.SplitLines(val);
+            List<string> lines = StringManipulationHelper.SplitLines(val);
 
             // process the substrings as separate lines
-            for (int k = 0; k < lines.Length; k++)
+            for (int k = 0; k < lines.Count; k++)
             {
                 // compute the display length of the string
                 int displayLength = _displayCells.Length(lines[k]);