Fix document highlight column

src/PowerShellEditorServices/Services/Symbols/Vistors/FindReferencesVisitor.cs

@@ -125,13 +125,13 @@ public override AstVisitAction VisitCommand(CommandAst commandAst)
         /// <returns>A visit action that continues the search for references</returns>
         public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
         {
-            int startColumnNumber = GetStartColumnNumberFromAst(functionDefinitionAst);
+            (int startColumnNumber, int startLineNumber) = GetStartColumnAndLineNumbersFromAst(functionDefinitionAst);
 
             IScriptExtent nameExtent = new ScriptExtent()
             {
                 Text = functionDefinitionAst.Name,
-                StartLineNumber = functionDefinitionAst.Extent.StartLineNumber,
-                EndLineNumber = functionDefinitionAst.Extent.StartLineNumber,
+                StartLineNumber = startLineNumber,
+                EndLineNumber = startLineNumber,
                 StartColumnNumber = startColumnNumber,
                 EndColumnNumber = startColumnNumber + functionDefinitionAst.Name.Length
             };
@@ -183,8 +183,10 @@ public override AstVisitAction VisitVariableExpression(VariableExpressionAst var
         }
 
         // Computes where the start of the actual function name is.
-        private static int GetStartColumnNumberFromAst(FunctionDefinitionAst ast)
+        private static (int, int) GetStartColumnAndLineNumbersFromAst(FunctionDefinitionAst ast)
         {
+            int startColumnNumber = ast.Extent.StartColumnNumber;
+            int startLineNumber = ast.Extent.StartLineNumber;
             int astOffset = 0;
 
             if (ast.IsFilter)
@@ -201,15 +203,31 @@ private static int GetStartColumnNumberFromAst(FunctionDefinitionAst ast)
             }
 
             string astText = ast.Extent.Text;
-            for (; astOffset < astText.Length; astOffset++)
+            // The line offset represents the offset on the line that we're on where as
+            // astOffset is the offset on the entire text of the AST.
+            int lineOffset = astOffset;
+            for (; astOffset < astText.Length; astOffset++, lineOffset++)
             {
-                if (!char.IsWhiteSpace(astText[astOffset]))
+                if (astText[astOffset] == '\n')
                 {
+                    // reset numbers since we are operating on a different line and increment the line number.
+                    startColumnNumber = 0;
+                    startLineNumber++;
+                    lineOffset = 0;
+                }
+                else if (astText[astOffset] == '\r')
+                {
+                    // Do nothing with carriage returns... we only look for line feeds since those
+                    // are used on every platform.
+                }
+                else if (!char.IsWhiteSpace(astText[astOffset]))
+                {
+                    // This is the start of the function name so we've found our start column and line number.
                     break;
                 }
             }
 
-            return ast.Extent.StartColumnNumber + astOffset;
+            return (startColumnNumber + lineOffset, startLineNumber);
         }
     }
 }

test/PowerShellEditorServices.Test/Services/Symbols/AstOperationsTests.cs

@@ -0,0 +1,86 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Collections.Generic;
+using System.Management.Automation;
+using System.Management.Automation.Language;
+using Microsoft.PowerShell.EditorServices.Services.Symbols;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Services.Symbols
+{
+    public class AstOperationsTests
+    {
+        private readonly ITestOutputHelper _output;
+        public AstOperationsTests(ITestOutputHelper output)
+        {
+            _output = output;
+        }
+        private static string s_scriptString = @"function BasicFunction {}
+BasicFunction
+
+function          FunctionWithExtraSpace
+{
+
+} FunctionWithExtraSpace
+
+function
+
+
+       FunctionNameOnDifferentLine
+
+
+
+
+
+
+           {}
+
+
+    FunctionNameOnDifferentLine
+";
+        private static ScriptBlockAst s_ast = (ScriptBlockAst) ScriptBlock.Create(s_scriptString).Ast;
+
+        [Trait("Category", "AstOperations")]
+        [Theory]
+        [InlineData(2, 3, "BasicFunction")]
+        [InlineData(7, 18, "FunctionWithExtraSpace")]
+        [InlineData(22, 13, "FunctionNameOnDifferentLine")]
+        public void CanFindSymbolAtPostion(int lineNumber, int columnNumber, string expectedName)
+        {
+            SymbolReference reference = AstOperations.FindSymbolAtPosition(s_ast, lineNumber, columnNumber);
+            Assert.NotNull(reference);
+            Assert.Equal(expectedName, reference.SymbolName);
+        }
+
+        [Trait("Category", "AstOperations")]
+        [Theory]
+        [MemberData(nameof(FindReferencesOfSymbolAtPostionData), parameters: 3)]
+        public void CanFindReferencesOfSymbolAtPostion(int lineNumber, int columnNumber, Position[] positions)
+        {
+            SymbolReference symbol = AstOperations.FindSymbolAtPosition(s_ast, lineNumber, columnNumber);
+
+            IEnumerable<SymbolReference> references = AstOperations.FindReferencesOfSymbol(s_ast, symbol, needsAliases: false);
+
+            int positionsIndex = 0;
+            foreach (SymbolReference reference in references)
+            {
+                Assert.Equal((int)positions[positionsIndex].Line, reference.ScriptRegion.StartLineNumber);
+                Assert.Equal((int)positions[positionsIndex].Character, reference.ScriptRegion.StartColumnNumber);
+
+                positionsIndex++;
+            }
+        }
+
+        public static object[][] FindReferencesOfSymbolAtPostionData = new object[][]
+        {
+            new object[] { 2, 3, new Position[] { new Position(1, 10), new Position(2, 1) } },
+            new object[] { 7, 18, new Position[] { new Position(4, 19), new Position(7, 3) } },
+            new object[] { 22, 13, new Position[] { new Position(12, 8), new Position(22, 5) } },
+        };
+    }
+}