From ac87d2ae0e5bc1778ffa3dfae4269aff08cbd5e4 Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Tue, 11 Jul 2023 02:15:05 +0200
Subject: [PATCH] Add completion for variables assigned by the `Data` statement
 (#19831)

---
 .../CommandCompletion/CompletionCompleters.cs     | 15 +++++++++++++++
 .../Host/TabCompletion/TabCompletion.Tests.ps1    |  6 ++++++
 2 files changed, 21 insertions(+)

diff --git a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
index 9dec8bcb3a..e4fb104485 100644
--- a/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
+++ b/src/System.Management.Automation/engine/CommandCompletion/CompletionCompleters.cs
@@ -5656,6 +5656,21 @@ namespace System.Management.Automation
                     parent = parent.Parent;
                 }
             }
+
+            public override AstVisitAction VisitDataStatement(DataStatementAst dataStatementAst)
+            {
+                if (dataStatementAst.Extent.StartOffset >= StopSearchOffset)
+                {
+                    return AstVisitAction.StopVisit;
+                }
+
+                if (dataStatementAst.Variable is not null)
+                {
+                    SaveVariableInfo(dataStatementAst.Variable, variableType: null, isConstraint: false);
+                }
+
+                return AstVisitAction.SkipChildren;
+            }
         }
 
         private static readonly Lazy<SortedSet<string>> s_specialVariablesCache = new Lazy<SortedSet<string>>(BuildSpecialVariablesCache);
diff --git a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1 b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
index ca159b798c..9e50e10f62 100644
--- a/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
+++ b/test/powershell/Host/TabCompletion/TabCompletion.Tests.ps1
@@ -657,6 +657,12 @@ ConstructorTestClass(int i, bool b)
         $res.CompletionMatches[0].CompletionText | Should -BeExactly Cat
     }
 
+    It 'Should complete variable assigned with Data statement' {
+        $TestString = 'data MyDataVar {"Hello"};$MyDatav'
+        $res = TabExpansion2 -inputScript $TestString
+        $res.CompletionMatches[0].CompletionText | Should -BeExactly '$MyDataVar'
+    }
+
     it 'Should complete "Value" parameter value in "Where-Object" for Enum property with no input' {
         $res = TabExpansion2 -inputScript 'Get-Command | where-Object CommandType -eq '
         $res.CompletionMatches[0].CompletionText | Should -BeExactly Alias