Add completion for loop labels after Break/Continue (#16438)

Adds completion for loop labels after labeled control flow statements such as Break and Continue.
The completion will look for any loops with a label that surrounds the Break/Continue keywords.
The tooltips for the completion entries shows the relevant loop.

src/System.Management.Automation/engine/CommandCompletion/CompletionAnalysis.cs

@@ -850,6 +850,16 @@ namespace System.Management.Automation
                                         result = GetResultForEnumPropertyValueOfDSCResource(completionContext, string.Empty, ref replacementIndex, ref replacementLength, out unused);
                                         break;
                                     }
+                                case TokenKind.Break:
+                                case TokenKind.Continue:
+                                    {
+                                        if ((lastAst is BreakStatementAst breakStatement && breakStatement.Label is null)
+                                            || (lastAst is ContinueStatementAst continueStatement && continueStatement.Label is null))
+                                        {
+                                            result = CompleteLoopLabel(completionContext);
+                                        }
+                                        break;
+                                    }
                                 default:
                                     break;
                             }
@@ -1759,6 +1769,11 @@ namespace System.Management.Automation
             var tokenAtCursorText = tokenAtCursor.Text;
             completionContext.WordToComplete = tokenAtCursorText;
 
+            if (lastAst.Parent is BreakStatementAst || lastAst.Parent is ContinueStatementAst)
+            {
+                return CompleteLoopLabel(completionContext);
+            }
+
             var strConst = lastAst as StringConstantExpressionAst;
             if (strConst != null)
             {
@@ -2213,5 +2228,40 @@ namespace System.Management.Automation
 
             return result;
         }
+
+        /// <summary>
+        /// Complete loop labels after labeled control flow statements such as Break and Continue.
+        /// </summary>
+        private static List<CompletionResult> CompleteLoopLabel(CompletionContext completionContext)
+        {
+            var result = new List<CompletionResult>();
+            foreach (Ast ast in completionContext.RelatedAsts)
+            {
+                if (ast is LabeledStatementAst labeledStatement
+                    && labeledStatement.Label is not null
+                    && (completionContext.WordToComplete is null || labeledStatement.Label.StartsWith(completionContext.WordToComplete, StringComparison.OrdinalIgnoreCase)))
+                {
+                    result.Add(new CompletionResult(labeledStatement.Label, labeledStatement.Label, CompletionResultType.Text, labeledStatement.Extent.Text));
+                }
+                else if (ast is ErrorStatementAst errorStatement)
+                {
+                    // Handles incomplete do/switch loops (other labeled statements do not need this special treatment)
+                    // The regex looks for the loopLabel of errorstatements that look like do/switch loops
+                    // For example in ":Label do " it will find "Label".
+                    var labelMatch = Regex.Match(errorStatement.Extent.Text, @"(?<=^:)\w+(?=\s+(do|switch)\b(?!-))", RegexOptions.IgnoreCase);
+                    if (labelMatch.Success)
+                    {
+                        result.Add(new CompletionResult(labelMatch.Value, labelMatch.Value, CompletionResultType.Text, errorStatement.Extent.Text));
+                    }
+                }
+            }
+
+            if (result.Count == 0)
+            {
+                return null;
+            }
+
+            return result;
+        }
     }
 }

test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1

@@ -1179,6 +1179,63 @@ dir -Recurse `
             $res.CompletionMatches | Should -HaveCount 2
             [string]::Join(',', ($res.CompletionMatches.completiontext | Sort-Object)) | Should -BeExactly "1.0,1.1"
         }
+        It '<Intent>' -TestCases @(
+            @{
+                Intent = 'Complete loop labels with no input'
+                Expected = 'Outer','Inner'
+                TestString = ':Outer while ($true){:Inner while ($true){ break ^ }}'
+            }
+            @{
+                Intent = 'Complete loop labels that are accessible'
+                Expected = 'Outer'
+                TestString = ':Outer do {:Inner while ($true){ break } continue ^ } until ($false)'
+            }
+            @{
+                Intent = 'Complete loop labels with partial input'
+                Expected = 'Outer'
+                TestString = ':Outer do {:Inner while ($true){ break } continue o^ut } while ($true)'
+            }
+            @{
+                Intent = 'Complete loop label for incomplete switch'
+                Expected = 'Outer'
+                TestString = ':Outer switch ($x){"randomValue"{ continue ^'
+            }
+            @{
+                Intent = 'Complete loop label for incomplete do loop'
+                Expected = 'Outer'
+                TestString = ':Outer do {:Inner while ($true){ break } continue ^'
+            }
+            @{
+                Intent = 'Complete loop label for incomplete for loop'
+                Expected = 'forLoop'
+                TestString = ':forLoop for ($i = 0; $i -lt $SomeCollection.Count; $i++) {continue ^'
+            }
+            @{
+                Intent = 'Complete loop label for incomplete while loop'
+                Expected = 'WhileLoop'
+                TestString = ':WhileLoop while ($true){ break ^'
+            }
+            @{
+                Intent = 'Complete loop label for incomplete foreach loop'
+                Expected = 'foreachLoop'
+                TestString = ':foreachLoop foreach ($x in $y) { break ^'
+            }
+            @{
+                Intent = 'Not Complete loop labels with colon'
+                Expected = $null
+                TestString = ':Outer foreach ($x in $y){:Inner for ($i = 0; $i -lt $X.Count; $i++){ break :O^}}'
+            }
+            @{
+                Intent = 'Not Complete loop labels if cursor is in front of existing label'
+                Expected = $null
+                TestString = ':Outer switch ($x){"Value1"{break ^ Outer}}'
+            }
+        ){
+            param($Expected, $TestString)
+            $CursorIndex = $TestString.IndexOf('^')
+            $res = TabExpansion2 -cursorColumn $CursorIndex -inputScript $TestString.Remove($CursorIndex, 1)
+            $res.CompletionMatches.CompletionText | Should -BeExactly $Expected
+        }
     }
 
     Context "Module completion for 'using module'" {