Fix #17: Add go to definition support for dot sourced file paths (#786)

* Implement go-to-definition for dot-sourced files

* Support $PSScriptRoot in dot-sourced files

* Add PathUtils for unified path normalization

Author: dee-see

src/PowerShellEditorServices/Language/AstOperations.cs

@@ -321,10 +321,11 @@ static private bool IsPowerShellDataFileAstNode(dynamic node, Type[] levelAstMap
         /// Finds all files dot sourced in a script
         /// </summary>
         /// <param name="scriptAst">The abstract syntax tree of the given script</param>
+        /// <param name="psScriptRoot">Pre-calculated value of $PSScriptRoot</param>
         /// <returns></returns>
-        static public string[] FindDotSourcedIncludes(Ast scriptAst)
+        static public string[] FindDotSourcedIncludes(Ast scriptAst, string psScriptRoot)
         {
-            FindDotSourcedVisitor dotSourcedVisitor = new FindDotSourcedVisitor();
+            FindDotSourcedVisitor dotSourcedVisitor = new FindDotSourcedVisitor(psScriptRoot);
             scriptAst.Visit(dotSourcedVisitor);
 
             return dotSourcedVisitor.DotSourcedFiles.ToArray();

src/PowerShellEditorServices/Language/FindDotSourcedVisitor.cs

@@ -3,8 +3,10 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Collections.Generic;
 using System.Management.Automation.Language;
+using Microsoft.PowerShell.EditorServices.Utility;
 
 namespace Microsoft.PowerShell.EditorServices
 {
@@ -13,14 +15,21 @@ namespace Microsoft.PowerShell.EditorServices
     /// </summary>
     internal class FindDotSourcedVisitor : AstVisitor
     {
-       /// <summary>
-       /// A hash set of the dot sourced files (because we don't want duplicates)
-       /// </summary>
+        private readonly string _psScriptRoot;
+
+        /// <summary>
+        /// A hash set of the dot sourced files (because we don't want duplicates)
+        /// </summary>
         public HashSet<string> DotSourcedFiles { get; private set; }
 
-        public FindDotSourcedVisitor()
+        /// <summary>
+        /// Creates a new instance of the FindDotSourcedVisitor class.
+        /// </summary>
+        /// <param name="psScriptRoot">Pre-calculated value of $PSScriptRoot</param>
+        public FindDotSourcedVisitor(string psScriptRoot)
         {
-            this.DotSourcedFiles = new HashSet<string>();
+            DotSourcedFiles = new HashSet<string>(StringComparer.CurrentCultureIgnoreCase);
+            _psScriptRoot = psScriptRoot;
         }
 
         /// <summary>
@@ -32,15 +41,50 @@ public FindDotSourcedVisitor()
         /// or a decision to continue if it wasn't found</returns>
         public override AstVisitAction VisitCommand(CommandAst commandAst)
         {
-            if (commandAst.InvocationOperator.Equals(TokenKind.Dot) &&
-                commandAst.CommandElements[0] is StringConstantExpressionAst)
+            CommandElementAst commandElementAst = commandAst.CommandElements[0];
+            if (commandAst.InvocationOperator.Equals(TokenKind.Dot))
             {
-                // Strip any quote characters off of the string
-                string fileName = commandAst.CommandElements[0].Extent.Text.Trim('\'', '"');
-                DotSourcedFiles.Add(fileName);
+                string path;
+                switch (commandElementAst)
+                {
+                    case StringConstantExpressionAst stringConstantExpressionAst:
+                        path = stringConstantExpressionAst.Value;
+                        break;
+
+                    case ExpandableStringExpressionAst expandableStringExpressionAst:
+                        path = GetPathFromExpandableStringExpression(expandableStringExpressionAst);
+                        break;
+
+                    default:
+                        path = null;
+                        break;
+                }
+
+                if (!string.IsNullOrWhiteSpace(path))
+                {
+                    DotSourcedFiles.Add(PathUtils.NormalizePathSeparators(path));
+                }
             }
 
             return base.VisitCommand(commandAst);
         }
+
+        private string GetPathFromExpandableStringExpression(ExpandableStringExpressionAst expandableStringExpressionAst)
+        {
+            var path = expandableStringExpressionAst.Value;
+            foreach (var nestedExpression in expandableStringExpressionAst.NestedExpressions)
+            {
+                // If the string contains the variable $PSScriptRoot, we replace it with the corresponding value.
+                if (!(nestedExpression is VariableExpressionAst variableAst
+                    && variableAst.VariablePath.UserPath.Equals("PSScriptRoot", StringComparison.OrdinalIgnoreCase)))
+                {
+                    return null; // We return null instead of a partially evaluated ExpandableStringExpression.
+                }
+
+                path = path.Replace(variableAst.ToString(), _psScriptRoot);
+            }
+
+            return path;
+        }
     }
 }

src/PowerShellEditorServices/Language/LanguageService.cs

@@ -14,6 +14,7 @@
 using System.Management.Automation.Language;
 using System.Runtime.InteropServices;
 using System.Security;
+using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -409,19 +410,30 @@ public async Task<GetDefinitionResult> GetDefinitionOfSymbol(
             // look through the referenced files until definition is found
             // or there are no more file to look through
             SymbolReference foundDefinition = null;
-            for (int i = 0; i < referencedFiles.Length; i++)
+            foreach (ScriptFile scriptFile in referencedFiles)
             {
                 foundDefinition =
                     AstOperations.FindDefinitionOfSymbol(
-                        referencedFiles[i].ScriptAst,
+                        scriptFile.ScriptAst,
                         foundSymbol);
 
-                filesSearched.Add(referencedFiles[i].FilePath);
+                filesSearched.Add(scriptFile.FilePath);
                 if (foundDefinition != null)
                 {
-                    foundDefinition.FilePath = referencedFiles[i].FilePath;
+                    foundDefinition.FilePath = scriptFile.FilePath;
                     break;
                 }
+
+                if (foundSymbol.SymbolType == SymbolType.Function)
+                {
+                    // Dot-sourcing is parsed as a "Function" Symbol.
+                    string dotSourcedPath = GetDotSourcedPath(foundSymbol, workspace, scriptFile);
+                    if (scriptFile.FilePath == dotSourcedPath)
+                    {
+                        foundDefinition = new SymbolReference(SymbolType.Function, foundSymbol.SymbolName, scriptFile.ScriptAst.Extent, scriptFile.FilePath);
+                        break;
+                    }
+                }
             }
 
             // if the definition the not found in referenced files
@@ -475,6 +487,21 @@ await CommandHelpers.GetCommandInfo(
                 null;
         }
 
+        /// <summary>
+        /// Gets a path from a dot-source symbol.
+        /// </summary>
+        /// <param name="symbol">The symbol representing the dot-source expression.</param>
+        /// <param name="workspace">The current workspace</param>
+        /// <param name="scriptFile">The script file containing the symbol</param>
+        /// <returns></returns>
+        private static string GetDotSourcedPath(SymbolReference symbol, Workspace workspace, ScriptFile scriptFile)
+        {
+            string cleanedUpSymbol = PathUtils.NormalizePathSeparators(symbol.SymbolName.Trim('\'', '"'));
+            string psScriptRoot = Path.GetDirectoryName(scriptFile.FilePath);
+            return workspace.ResolveRelativeScriptPath(psScriptRoot,
+                Regex.Replace(cleanedUpSymbol, @"\$PSScriptRoot|\${PSScriptRoot}", psScriptRoot, RegexOptions.IgnoreCase));
+        }
+
         /// <summary>
         /// Finds all the occurences of a symbol in the script given a file location
         /// </summary>
@@ -712,7 +739,7 @@ await _powerShellContext.GetRunspaceHandle(
                 {
                     if (!_cmdletToAliasDictionary.ContainsKey(aliasInfo.Definition))
                     {
-                        _cmdletToAliasDictionary.Add(aliasInfo.Definition, new List<String>{ aliasInfo.Name });
+                        _cmdletToAliasDictionary.Add(aliasInfo.Definition, new List<String> { aliasInfo.Name });
                     }
                     else
                     {

src/PowerShellEditorServices/Utility/PathUtils.cs

@@ -0,0 +1,49 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.IO;
+using System.Runtime.InteropServices;
+
+namespace Microsoft.PowerShell.EditorServices.Utility
+{
+    /// <summary>
+    /// Utility to help handling paths across different platforms.
+    /// </summary>
+    /// <remarks>
+    /// Some constants were copied from the internal System.Management.Automation.StringLiterals class.
+    /// </remarks>
+    internal static class PathUtils
+    {
+        /// <summary>
+        /// The default path separator used by the base implementation of the providers.
+        ///
+        /// Porting note: IO.Path.DirectorySeparatorChar is correct for all platforms. On Windows,
+        /// it is '\', and on Linux, it is '/', as expected.
+        /// </summary>
+        internal static readonly char DefaultPathSeparator = Path.DirectorySeparatorChar;
+        internal static readonly string DefaultPathSeparatorString = DefaultPathSeparator.ToString();
+
+        /// <summary>
+        /// The alternate path separator used by the base implementation of the providers.
+        ///
+        /// Porting note: we do not use .NET's AlternatePathSeparatorChar here because it correctly
+        /// states that both the default and alternate are '/' on Linux. However, for PowerShell to
+        /// be "slash agnostic", we need to use the assumption that a '\' is the alternate path
+        /// separator on Linux.
+        /// </summary>
+        internal static readonly char AlternatePathSeparator = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? '/' : '\\';
+        internal static readonly string AlternatePathSeparatorString = AlternatePathSeparator.ToString();
+
+        /// <summary>
+        /// Converts all alternate path separators to the current platform's main path separators.
+        /// </summary>
+        /// <param name="path">The path to normalize.</param>
+        /// <returns>The normalized path.</returns>
+        public static string NormalizePathSeparators(string path)
+        {
+            return string.IsNullOrWhiteSpace(path) ? path : path.Replace(AlternatePathSeparator, DefaultPathSeparator);
+        }
+    }
+}

src/PowerShellEditorServices/Workspace/ScriptFile.cs

@@ -20,7 +20,7 @@ public class ScriptFile
     {
         #region Private Fields
 
-        private static readonly string[] s_newlines = new []
+        private static readonly string[] s_newlines = new[]
         {
             "\r\n",
             "\n"
@@ -649,8 +649,7 @@ private void ParseFileContents()
                     .ToArray();
 
             //Get all dot sourced referenced files and store  them
-            this.ReferencedFiles =
-                AstOperations.FindDotSourcedIncludes(this.ScriptAst);
+            this.ReferencedFiles = AstOperations.FindDotSourcedIncludes(this.ScriptAst, Path.GetDirectoryName(this.FilePath));
         }
 
         #endregion

src/PowerShellEditorServices/Workspace/Workspace.cs

@@ -530,7 +530,7 @@ private string GetBaseFilePath(string filePath)
             return Path.GetDirectoryName(filePath);
         }
 
-        private string ResolveRelativeScriptPath(string baseFilePath, string relativePath)
+        internal string ResolveRelativeScriptPath(string baseFilePath, string relativePath)
         {
             string combinedPath = null;
             Exception resolveException = null;

test/PowerShellEditorServices.Test.Shared/Definition/FindsDotSourcedFile.cs

@@ -0,0 +1,20 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.PowerShell.EditorServices;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Shared.Definition
+{
+    public class FindsDotSourcedFile
+    {
+        public static readonly ScriptRegion SourceDetails =
+            new ScriptRegion
+            {
+                File = @"References\DotSources.ps1",
+                StartLineNumber = 1,
+                StartColumnNumber = 3
+            };
+    }
+}

test/PowerShellEditorServices.Test.Shared/References/DotSources.ps1

@@ -0,0 +1,15 @@
+. ./ReferenceFileE.ps1
+. "$PSScriptRoot/ReferenceFileE.ps1"
+. "${PSScriptRoot}/ReferenceFileE.ps1"
+. './ReferenceFileE.ps1'
+. "./ReferenceFileE.ps1"
+. .\ReferenceFileE.ps1
+. '.\ReferenceFileE.ps1'
+. ".\ReferenceFileE.ps1"
+. ReferenceFileE.ps1
+. 'ReferenceFileE.ps1'
+. "ReferenceFileE.ps1"
+. ./dir/../ReferenceFileE.ps1
+. ./invalidfile.ps1
+. ""
+. $someVar

test/PowerShellEditorServices.Test.Shared/References/ReferenceFileB.ps1

@@ -1,5 +1,5 @@
-. .\ReferenceFileC.ps1
+. "$PSScriptRoot\ReferenceFileC.ps1"
 
 Get-ChildItem
 
-My-Function "testb"
+My-Function "testb"

test/PowerShellEditorServices.Test/Language/LanguageServiceTests.cs

@@ -158,6 +158,23 @@ await this.GetDefinition(
             Assert.Equal("My-Function", definition.SymbolName);
         }
 
+        [Fact]
+        public async Task LanguageServiceFindsDotSourcedFile()
+        {
+            GetDefinitionResult definitionResult =
+                await this.GetDefinition(
+                    FindsDotSourcedFile.SourceDetails);
+
+            SymbolReference definition = definitionResult.FoundDefinition;
+            Assert.True(
+                definitionResult.FoundDefinition.FilePath.EndsWith(
+                    Path.Combine("References", "ReferenceFileE.ps1")),
+                "Unexpected reference file: " + definitionResult.FoundDefinition.FilePath);
+            Assert.Equal(1, definition.ScriptRegion.StartLineNumber);
+            Assert.Equal(1, definition.ScriptRegion.StartColumnNumber);
+            Assert.Equal("./ReferenceFileE.ps1", definition.SymbolName);
+        }
+
         [Fact]
         public async Task LanguageServiceFindsFunctionDefinitionInWorkspace()
         {