From 0cbcf7b4d6583cc2f976c93c038e3fc9e1ed1561 Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Mon, 19 Jun 2023 22:30:39 +0200
Subject: [PATCH] Enable completion of variables across ScriptBlock scopes
 (#19819)

---
 .../CommandCompletion/CompletionCompleters.cs | 43 ++++++++++++++++---
 .../TabCompletion/TabCompletion.Tests.ps1     | 15 +++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
index 1d69925f4e..9dec8bcb3a 100644
--- a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
+++ b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
@@ -5375,6 +5375,20 @@ namespace System.Management.Automation
             "pv"
         };
 
+        private static readonly HashSet<string> s_localScopeCommandNames = new(StringComparer.OrdinalIgnoreCase)
+        {
+            "Microsoft.PowerShell.Core\\ForEach-Object",
+            "ForEach-Object",
+            "foreach",
+            "%",
+            "Microsoft.PowerShell.Core\\Where-Object",
+            "Where-Object",
+            "where",
+            "?",
+            "BeforeAll",
+            "BeforeEach"
+        };
+
         private sealed class VariableInfo
         {
             internal Type LastDeclaredConstraint;
@@ -5616,12 +5630,31 @@ namespace System.Management.Automation
 
             public override AstVisitAction VisitScriptBlockExpression(ScriptBlockExpressionAst scriptBlockExpressionAst)
             {
-                return scriptBlockExpressionAst != Top ? AstVisitAction.SkipChildren : AstVisitAction.Continue;
-            }
+                if (scriptBlockExpressionAst == Top)
+                {
+                    return AstVisitAction.Continue;
+                }
 
-            public override AstVisitAction VisitScriptBlock(ScriptBlockAst scriptBlockAst)
-            {
-                return scriptBlockAst != Top ? AstVisitAction.SkipChildren : AstVisitAction.Continue;
+                Ast parent = scriptBlockExpressionAst.Parent;
+                // This loop checks if the scriptblock is used as a command, or an argument for a command, eg: ForEach-Object -Process {$Var1 = "Hello"}, {Var2 = $true}
+                while (true)
+                {
+                    if (parent is CommandAst cmdAst)
+                    {
+                        string cmdName = cmdAst.GetCommandName();
+                        return s_localScopeCommandNames.Contains(cmdName)
+                            || (cmdAst.CommandElements[0] is ScriptBlockExpressionAst && cmdAst.InvocationOperator == TokenKind.Dot)
+                            ? AstVisitAction.Continue
+                            : AstVisitAction.SkipChildren;
+                    }
+
+                    if (parent is not CommandExpressionAst and not PipelineAst and not StatementBlockAst and not ArrayExpressionAst and not ArrayLiteralAst)
+                    {
+                        return AstVisitAction.SkipChildren;
+                    }
+
+                    parent = parent.Parent;
+                }
             }
         }
 
diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
index 2981f583eb..4113100100 100644
--- a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
+++ b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
@@ -274,6 +274,21 @@ switch ($x)
         $completionText -join ' ' | Should -BeExactly 'Ascending Descending Expression'
     }
 
+    It 'Should complete variable assigned in other scriptblock' {
+        $res = TabExpansion2 -inputScript 'ForEach-Object -Begin {$Test1 = "Hello"} -Process {$Test'
+        $res.CompletionMatches[0].CompletionText | Should -Be '$Test1'
+    }
+
+    It 'Should complete variable assigned in an array of scriptblocks' {
+        $res = TabExpansion2 -inputScript 'ForEach-Object -Process @({"Block1"},{$Test1="Hello"});$Test'
+        $res.CompletionMatches[0].CompletionText | Should -Be '$Test1'
+    }
+
+    It 'Should not complete variable assigned in an ampersand executed scriptblock' {
+        $res = TabExpansion2 -inputScript '& {$AmpeersandVarCompletionTest = "Hello"};$AmpeersandVarCompletionTes'
+        $res.CompletionMatches.Count | Should -Be 0
+    }
+
     context TypeConstructionWithHashtable {
         BeforeAll {
             class RandomTestType {