Fix finding variable definitions in workspace

src/PowerShellEditorServices/Language/AstOperations.cs

@@ -24,7 +24,7 @@ namespace Microsoft.PowerShell.EditorServices
     internal static class AstOperations
     {
         /// <summary>
-        /// Gets completions for the symbol found in the Ast at 
+        /// Gets completions for the symbol found in the Ast at
         /// the given file offset.
         /// </summary>
         /// <param name="scriptAst">
@@ -47,14 +47,14 @@ internal static class AstOperations
         /// symbol at the given offset.
         /// </returns>
         static public async Task<CommandCompletion> GetCompletions(
-            Ast scriptAst, 
-            Token[] currentTokens, 
+            Ast scriptAst,
+            Token[] currentTokens,
             int fileOffset,
             PowerShellContext powerShellContext,
             CancellationToken cancellationToken)
         {
             var type = scriptAst.Extent.StartScriptPosition.GetType();
-            var method = 
+            var method =
 #if CoreCLR
                 type.GetMethod(
                     "CloneWithNewOffset",
@@ -67,9 +67,9 @@ static public async Task<CommandCompletion> GetCompletions(
                     new[] { typeof(int) }, null);
 #endif
 
-            IScriptPosition cursorPosition = 
+            IScriptPosition cursorPosition =
                 (IScriptPosition)method.Invoke(
-                    scriptAst.Extent.StartScriptPosition, 
+                    scriptAst.Extent.StartScriptPosition,
                     new object[] { fileOffset });
 
             Logger.Write(
@@ -138,7 +138,7 @@ static public async Task<CommandCompletion> GetCompletions(
         }
 
         /// <summary>
-        /// Finds the symbol at a given file location 
+        /// Finds the symbol at a given file location
         /// </summary>
         /// <param name="scriptAst">The abstract syntax tree of the given script</param>
         /// <param name="lineNumber">The line number of the cursor for the given script</param>
@@ -176,15 +176,15 @@ static public SymbolReference FindCommandAtPosition(Ast scriptAst, int lineNumbe
         /// <param name="AliasToCmdletDictionary">Dictionary maping aliases to cmdlets for finding alias references</param>
         /// <returns></returns>
         static public IEnumerable<SymbolReference> FindReferencesOfSymbol(
-            Ast scriptAst, 
-            SymbolReference symbolReference, 
+            Ast scriptAst,
+            SymbolReference symbolReference,
             Dictionary<String, List<String>> CmdletToAliasDictionary,
             Dictionary<String, String> AliasToCmdletDictionary)
         {
             // find the symbol evaluators for the node types we are handling
-            FindReferencesVisitor referencesVisitor = 
+            FindReferencesVisitor referencesVisitor =
                 new FindReferencesVisitor(
-                    symbolReference, 
+                    symbolReference,
                     CmdletToAliasDictionary,
                     AliasToCmdletDictionary);
             scriptAst.Visit(referencesVisitor);
@@ -202,8 +202,8 @@ static public IEnumerable<SymbolReference> FindReferencesOfSymbol(
         /// <returns>A collection of SymbolReference objects that are refrences to the symbolRefrence
         /// not including aliases</returns>
         static public IEnumerable<SymbolReference> FindReferencesOfSymbol(
-            ScriptBlockAst scriptAst, 
-            SymbolReference foundSymbol, 
+            ScriptBlockAst scriptAst,
+            SymbolReference foundSymbol,
             bool needsAliases)
         {
             FindReferencesVisitor referencesVisitor =
@@ -214,7 +214,7 @@ static public IEnumerable<SymbolReference> FindReferencesOfSymbol(
         }
 
         /// <summary>
-        /// Finds the definition of the symbol 
+        /// Finds the definition of the symbol
         /// </summary>
         /// <param name="scriptAst">The abstract syntax tree of the given script</param>
         /// <param name="symbolReference">The symbol that we are looking for the definition of</param>
@@ -223,12 +223,12 @@ static public SymbolReference FindDefinitionOfSymbol(
             Ast scriptAst,
             SymbolReference symbolReference)
         {
-            FindDeclartionVisitor declarationVisitor = 
-                new FindDeclartionVisitor(
+            FindDeclarationVisitor declarationVisitor =
+                new FindDeclarationVisitor(
                     symbolReference);
             scriptAst.Visit(declarationVisitor);
 
-            return declarationVisitor.FoundDeclartion;
+            return declarationVisitor.FoundDeclaration;
         }
 
         /// <summary>

src/PowerShellEditorServices/Language/FindDeclartionVisitor.cs → src/PowerShellEditorServices/Language/FindDeclarationVisitor.cs

@@ -9,30 +9,36 @@
 namespace Microsoft.PowerShell.EditorServices
 {
     /// <summary>
-    /// The vistor used to find the defintion of a symbol
+    /// The visitor used to find the definition of a symbol
     /// </summary>
-    internal class FindDeclartionVisitor : AstVisitor
+    internal class FindDeclarationVisitor : AstVisitor
     {
         private SymbolReference symbolRef;
+        private string variableName;
 
-        public SymbolReference FoundDeclartion{ get; private set; }
+        public SymbolReference FoundDeclaration{ get; private set; }
 
-        public FindDeclartionVisitor(SymbolReference symbolRef)
+        public FindDeclarationVisitor(SymbolReference symbolRef)
         {
             this.symbolRef = symbolRef;
+            if (this.symbolRef.SymbolType == SymbolType.Variable)
+            {
+                // converts `$varName` to `varName` or of the form ${varName} to varName
+                variableName = symbolRef.SymbolName.TrimStart('$').Trim('{', '}');
+            }
         }
 
         /// <summary>
-        /// Decides if the current function defintion is the right defition
-        /// for the symbol being searched for. The defintion of the symbol will be a of type 
+        /// Decides if the current function definition is the right definition
+        /// for the symbol being searched for. The definition of the symbol will be a of type
         /// SymbolType.Function and have the same name as the symbol
         /// </summary>
         /// <param name="functionDefinitionAst">A FunctionDefinitionAst in the script's AST</param>
-        /// <returns>A descion to stop searching if the right FunctionDefinitionAst was found, 
+        /// <returns>A descion to stop searching if the right FunctionDefinitionAst was found,
         /// or a decision to continue if it wasn't found</returns>
         public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
         {
-            // Get the start column number of the function name, 
+            // Get the start column number of the function name,
             // instead of the the start column of 'function' and create new extent for the functionName
             int startColumnNumber =
                 functionDefinitionAst.Extent.Text.IndexOf(
@@ -50,7 +56,7 @@ public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst fun
             if (symbolRef.SymbolType.Equals(SymbolType.Function) &&
                  nameExtent.Text.Equals(symbolRef.ScriptRegion.Text, StringComparison.CurrentCultureIgnoreCase))
             {
-                this.FoundDeclartion =
+                this.FoundDeclaration =
                     new SymbolReference(
                         SymbolType.Function,
                         nameExtent);
@@ -62,27 +68,27 @@ public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst fun
         }
 
         /// <summary>
-        /// Decides if the current variable expression is the right defition for 
-        /// the symbol being searched for. The defintion of the symbol will be a of type 
-        /// SymbolType.Variable and have the same name as the symbol
+        /// Check if the left hand side of an assignmentStatementAst is a VariableExpressionAst
+        /// with the same name as that of symbolRef.
         /// </summary>
-        /// <param name="variableExpressionAst">A FunctionDefinitionAst in the script's AST</param>
-        /// <returns>A descion to stop searching if the right VariableExpressionAst was found, 
+        /// <param name="assignmentStatementAst">An AssignmentStatementAst/param>
+        /// <returns>A decision to stop searching if the right VariableExpressionAst was found,
         /// or a decision to continue if it wasn't found</returns>
-        public override AstVisitAction VisitVariableExpression(VariableExpressionAst variableExpressionAst)
+        public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst assignmentStatementAst)
         {
-            if(symbolRef.SymbolType.Equals(SymbolType.Variable) &&
-                variableExpressionAst.Extent.Text.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
+            var variableExprAst = assignmentStatementAst.Left as VariableExpressionAst;
+            if (variableExprAst == null ||
+                variableName == null ||
+                !variableExprAst.VariablePath.UserPath.Equals(
+                    variableName,
+                    StringComparison.OrdinalIgnoreCase))
             {
-                this.FoundDeclartion =
-                    new SymbolReference(
-                        SymbolType.Variable,
-                        variableExpressionAst.Extent);
-
-                return AstVisitAction.StopVisit;
+                return AstVisitAction.Continue;
             }
 
-            return AstVisitAction.Continue;
+            // TODO also find instances of set-variable
+            FoundDeclaration = new SymbolReference(SymbolType.Variable, variableExprAst.Extent);
+            return AstVisitAction.StopVisit;
         }
     }
 }

src/PowerShellEditorServices/Language/SymbolReference.cs

@@ -3,6 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Diagnostics;
 using System.Management.Automation.Language;
 

test/PowerShellEditorServices.Test.Host/LanguageServerTests.cs

@@ -455,6 +455,66 @@ await this.SendRequest(
             Assert.Equal(7, locations[0].Range.End.Character);
         }
 
+        [Fact]
+        public async Task FindsDefinitionOfVariableInOtherFile()
+        {
+            await this.SendOpenFileEvent("TestFiles\\FindReferences.ps1");
+
+            Location[] locations =
+                await this.SendRequest(
+                    DefinitionRequest.Type,
+                    new TextDocumentPositionParams
+                    {
+                        TextDocument = new TextDocumentIdentifier
+                        {
+                            Uri = "TestFiles\\FindReferences.ps1"
+                        },
+                        Position = new Position
+                        {
+                            Line = 15,
+                            Character = 20,
+                        }
+                    });
+
+            Assert.NotNull(locations);
+            Assert.Equal(1, locations.Length);
+            Assert.EndsWith("VariableDefinition.ps1", locations[0].Uri);
+            Assert.Equal(0, locations[0].Range.Start.Line);
+            Assert.Equal(0, locations[0].Range.Start.Character);
+            Assert.Equal(0, locations[0].Range.End.Line);
+            Assert.Equal(20, locations[0].Range.End.Character);
+        }
+
+        [Fact]
+        public async Task FindDefinitionOfVariableWithSpecialChars()
+        {
+            await this.SendOpenFileEvent("TestFiles\\FindReferences.ps1");
+
+            Location[] locations =
+                await this.SendRequest(
+                    DefinitionRequest.Type,
+                    new TextDocumentPositionParams
+                    {
+                        TextDocument = new TextDocumentIdentifier
+                        {
+                            Uri = "TestFiles\\FindReferences.ps1"
+                        },
+                        Position = new Position
+                        {
+                            Line = 18,
+                            Character = 24,
+                        }
+                    });
+
+            Assert.NotNull(locations);
+            Assert.Equal(1, locations.Length);
+            Assert.EndsWith("FindReferences.ps1", locations[0].Uri);
+            Assert.Equal(17, locations[0].Range.Start.Line);
+            Assert.Equal(0, locations[0].Range.Start.Character);
+            Assert.Equal(17, locations[0].Range.End.Line);
+            Assert.Equal(27, locations[0].Range.End.Character);
+        }
+
         [Fact]
         public async Task FindsOccurencesOnFunctionDefinition()
         {

test/PowerShellEditorServices.Test.Host/TestFiles/FindReferences.ps1

@@ -10,4 +10,10 @@ My-Function $things
 
 Write-Output "Hi";
 
-Write-Output ""
+Write-Output ""
+
+. .\VariableDefinition.ps1
+Write-Output $variableInOtherFile
+
+${variable-with-weird-name} = "this variable has special characters"
+Write-Output ${variable-with-weird-name}

test/PowerShellEditorServices.Test.Host/TestFiles/VariableDefinition.ps1

@@ -0,0 +1 @@
+$variableInOtherFile = "some value"