From c64a636d9bdcae3d00aab551685d87b48cc35e30 Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Wed, 26 Feb 2025 13:04:36 +0100
Subject: [PATCH] Fix completion of variables assigned inside Do loops (#25076)

---
 .../CommandCompletion/CompletionCompleters.cs      | 12 +++++++++---
 .../Host/TabCompletion/TabCompletion.Tests.ps1     | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
index 012b96d9a1..ebee01e8a0 100644
--- a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
+++ b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
@@ -5546,7 +5546,7 @@ namespace System.Management.Automation
             internal Type LastAssignedType;
         }
 
-        private sealed class FindVariablesVisitor : AstVisitor
+        private sealed class FindVariablesVisitor : AstVisitor2
         {
             internal Ast Top;
             internal Ast CompletionVariableAst;
@@ -5621,7 +5621,11 @@ namespace System.Management.Automation
             {
                 if (ast.Extent.StartOffset > StopSearchOffset)
                 {
-                    return AstVisitAction.StopVisit;
+                    // When visiting do while/until statements, the condition will be visited before the statement block.
+                    // The condition itself may not be interesting if it's after the cursor, but the statement block could be.
+                    return ast is PipelineBaseAst && ast.Parent is DoUntilStatementAst or DoWhileStatementAst
+                        ? AstVisitAction.SkipChildren
+                        : AstVisitAction.StopVisit;
                 }
 
                 return AstVisitAction.Continue;
@@ -5631,7 +5635,9 @@ namespace System.Management.Automation
             {
                 if (assignmentStatementAst.Extent.StartOffset > StopSearchOffset)
                 {
-                    return AstVisitAction.StopVisit;
+                    return assignmentStatementAst.Parent is DoUntilStatementAst or DoWhileStatementAst ?
+                        AstVisitAction.SkipChildren
+                        : AstVisitAction.StopVisit;
                 }
 
                 if (assignmentStatementAst.Left is AttributedExpressionAst attributedExpression)
diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
index daeae1c499..bfd9ef2cad 100644
--- a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
+++ b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
@@ -760,6 +760,20 @@ using `
         $completionText -join ' ' | Should -BeExactly 'Equals( new( ReferenceEquals('
     }
 
+    It 'Should complete variables assigned inside do while loop' {
+        $TestString =  'do{$Var1 = 1; $Var^ }while ($true)'
+        $CursorIndex = $TestString.IndexOf('^')
+        $res = TabExpansion2 -cursorColumn $CursorIndex -inputScript $TestString.Remove($CursorIndex, 1)
+        $res.CompletionMatches[0].CompletionText | Should -BeExactly '$Var1'
+    }
+
+    It 'Should complete variables assigned inside do until loop' {
+        $TestString =  'do{$Var1 = 1; $Var^ }until ($null = Get-ChildItem)'
+        $CursorIndex = $TestString.IndexOf('^')
+        $res = TabExpansion2 -cursorColumn $CursorIndex -inputScript $TestString.Remove($CursorIndex, 1)
+        $res.CompletionMatches[0].CompletionText | Should -BeExactly '$Var1'
+    }
+
     It 'Should show multiple constructors in the tooltip' {
         $res = TabExpansion2 -inputScript 'class ConstructorTestClass{ConstructorTestClass ([string] $s){}ConstructorTestClass ([int] $i){}ConstructorTestClass ([int] $i, [bool]$b){}};[ConstructorTestClass]::new'
         $res.CompletionMatches | Should -HaveCount 1