From 599a5981fb347000b68320c09e3ed44af92d30d6 Mon Sep 17 00:00:00 2001
From: Andy Jordan <andy.jordan@microsoft.com>
Date: Wed, 8 Mar 2023 16:33:00 -0800
Subject: [PATCH] Fix declaration detection for variables with type constraints

---
 .../Symbols/Visitors/SymbolVisitor.cs         | 23 ++++++++++++++++++-
 .../FindsTypedVariableDefinition.cs           | 20 ++++++++++++++++
 .../References/SimpleFile.ps1                 |  3 +++
 .../Language/SymbolsServiceTests.cs           | 12 ++++++++++
 4 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 test/PowerShellEditorServices.Test.Shared/Definition/FindsTypedVariableDefinition.cs

diff --git a/src/PowerShellEditorServices/Services/Symbols/Visitors/SymbolVisitor.cs b/src/PowerShellEditorServices/Services/Symbols/Visitors/SymbolVisitor.cs
index fda5a53dc..58d7f581a 100644
--- a/src/PowerShellEditorServices/Services/Symbols/Visitors/SymbolVisitor.cs
+++ b/src/PowerShellEditorServices/Services/Symbols/Visitors/SymbolVisitor.cs
@@ -110,6 +110,27 @@ public override AstVisitAction VisitVariableExpression(VariableExpressionAst var
             ));
         }
 
+        // Traverse the parents to determine if this is a declaration.
+        static bool isDeclaration(Ast current)
+        {
+            Ast next = current.Parent;
+            while (true)
+            {
+                // Should come from an assignment statement.
+                if (next is AssignmentStatementAst assignment)
+                {
+                    return assignment.Left == current;
+                }
+                // Or we might have type constraints or attributes to traverse first.
+                if (next is not ConvertExpressionAst or not AttributedExpressionAst)
+                {
+                    return false;
+                }
+                current = next;
+                next = next.Parent;
+            }
+        }
+
         // TODO: Consider tracking unscoped variable references only when they're declared within
         // the same function definition.
         return _action(new SymbolReference(
@@ -119,7 +140,7 @@ public override AstVisitAction VisitVariableExpression(VariableExpressionAst var
             variableExpressionAst.Extent,
             variableExpressionAst.Extent, // TODO: Maybe parent?
             _file,
-            isDeclaration: variableExpressionAst.Parent is AssignmentStatementAst or ParameterAst));
+            isDeclaration(variableExpressionAst)));
     }
 
     public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst typeDefinitionAst)
diff --git a/test/PowerShellEditorServices.Test.Shared/Definition/FindsTypedVariableDefinition.cs b/test/PowerShellEditorServices.Test.Shared/Definition/FindsTypedVariableDefinition.cs
new file mode 100644
index 000000000..98892dfab
--- /dev/null
+++ b/test/PowerShellEditorServices.Test.Shared/Definition/FindsTypedVariableDefinition.cs
@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Shared.Definition
+{
+    public static class FindsTypedVariableDefinitionData
+    {
+        public static readonly ScriptRegion SourceDetails = new(
+            file: TestUtilities.NormalizePath("References/SimpleFile.ps1"),
+            text: string.Empty,
+            startLineNumber: 25,
+            startColumnNumber: 13,
+            startOffset: 0,
+            endLineNumber: 0,
+            endColumnNumber: 0,
+            endOffset: 0);
+    }
+}
diff --git a/test/PowerShellEditorServices.Test.Shared/References/SimpleFile.ps1 b/test/PowerShellEditorServices.Test.Shared/References/SimpleFile.ps1
index b5c1ef7ce..819f2a2c6 100644
--- a/test/PowerShellEditorServices.Test.Shared/References/SimpleFile.ps1
+++ b/test/PowerShellEditorServices.Test.Shared/References/SimpleFile.ps1
@@ -20,3 +20,6 @@ Get-ChildItem
 My-Alias
 
 Invoke-Command -ScriptBlock ${Function:My-Function}
+
+[string]$hello = "test"
+Write-Host $hello
diff --git a/test/PowerShellEditorServices.Test/Language/SymbolsServiceTests.cs b/test/PowerShellEditorServices.Test/Language/SymbolsServiceTests.cs
index 109a3667f..100b1ee8e 100644
--- a/test/PowerShellEditorServices.Test/Language/SymbolsServiceTests.cs
+++ b/test/PowerShellEditorServices.Test/Language/SymbolsServiceTests.cs
@@ -300,6 +300,18 @@ public async Task FindsVariableDefinition()
             AssertIsRegion(symbol.NameRegion, 6, 1, 6, 8);
         }
 
+        [Fact]
+        public async Task FindsTypedVariableDefinition()
+        {
+            IEnumerable<SymbolReference> definitions = await GetDefinitions(FindsTypedVariableDefinitionData.SourceDetails).ConfigureAwait(true);
+            SymbolReference symbol = Assert.Single(definitions);
+            Assert.Equal("var hello", symbol.Id);
+            Assert.Equal("$hello", symbol.Name);
+            Assert.Equal(SymbolType.Variable, symbol.Type);
+            Assert.True(symbol.IsDeclaration);
+            AssertIsRegion(symbol.NameRegion, 24, 9, 24, 15);
+        }
+
         [Fact]
         public async Task FindsReferencesOnVariable()
         {