fixed peek definition tokenizing (#281)

* fixed peek definition tokenizing

* fixed signature help test

* added new heuristic for PeekDefinition behaviour

* added tests for new heuristic

* fixed code according to Kevin's CR

* fixed failing test due to shared connection

* changed uri for procedure test

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs

@@ -774,40 +774,18 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             return completionItem;
         }
 
+        
         /// <summary>
-        /// Get definition for a selected sql object using SMO Scripting
+        /// Queue a task to the binding queue
         /// </summary>
         /// <param name="textDocumentPosition"></param>
+        /// <param name="scriptParseInfo"></param>
+        /// <param name="connectionInfo"></param>
         /// <param name="scriptFile"></param>
-        /// <param name="connInfo"></param>
-        /// <returns> Location with the URI of the script file</returns>
-        internal DefinitionResult GetDefinition(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile, ConnectionInfo connInfo)
-        {
-            // Parse sql
-            ScriptParseInfo scriptParseInfo = GetScriptParseInfo(textDocumentPosition.TextDocument.Uri);
-            if (scriptParseInfo == null)
-            {
-                return null;
-            }
-
-            if (RequiresReparse(scriptParseInfo, scriptFile))
-            {
-                scriptParseInfo.ParseResult = ParseAndBind(scriptFile, connInfo);
-            }
-
-            // Get token from selected text
-            Token selectedToken = ScriptDocumentInfo.GetToken(scriptParseInfo, textDocumentPosition.Position.Line + 1, textDocumentPosition.Position.Character);
-            if (selectedToken == null)
-            {
-                return null;
-            }
-            // Strip "[" and "]"(if present) from the token text to enable matching with the suggestions.
-            // The suggestion title does not contain any sql punctuation
-            string tokenText = TextUtilities.RemoveSquareBracketSyntax(selectedToken.Text);
-
-            if (scriptParseInfo.IsConnected && Monitor.TryEnter(scriptParseInfo.BuildingMetadataLock))
-            {
-                try
+        /// <param name="tokenText"></param>
+        /// <returns> Returns the result of the task as a DefinitionResult </returns>
+        private DefinitionResult QueueTask(TextDocumentPosition textDocumentPosition, ScriptParseInfo scriptParseInfo, 
+                                            ConnectionInfo connInfo, ScriptFile scriptFile, string tokenText)
         {
             // Queue the task with the binding queue
             QueueItem queueItem = this.BindingQueue.QueueBindingOperation(
@@ -838,7 +816,33 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                            
             // wait for the queue item
             queueItem.ItemProcessed.WaitOne();
-                    return queueItem.GetResultAsT<DefinitionResult>();
+            var result = queueItem.GetResultAsT<DefinitionResult>();
+            return result;
+        }
+
+        private DefinitionResult GetDefinitionFromTokenList(TextDocumentPosition textDocumentPosition, List<Token> tokenList,
+                ScriptParseInfo scriptParseInfo, ScriptFile scriptFile, ConnectionInfo connInfo)
+        {
+
+            DefinitionResult lastResult = null;
+            foreach (var token in tokenList)
+            {
+
+                // Strip "[" and "]"(if present) from the token text to enable matching with the suggestions.
+                // The suggestion title does not contain any sql punctuation
+                string tokenText = TextUtilities.RemoveSquareBracketSyntax(token.Text);
+                textDocumentPosition.Position.Line = token.StartLocation.LineNumber;
+                textDocumentPosition.Position.Character = token.StartLocation.ColumnNumber;
+                if (Monitor.TryEnter(scriptParseInfo.BuildingMetadataLock))
+                {
+                    try
+                    {
+                        var result = QueueTask(textDocumentPosition, scriptParseInfo, connInfo, scriptFile, tokenText);
+                        lastResult = result;
+                        if (!result.IsErrorResult)
+                        {
+                            return result;
+                        }
                     }
                     catch (Exception ex)
                     {
@@ -857,6 +861,65 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     }
                 }
                 else
+                {
+                    Logger.Write(LogLevel.Error, "Timeout waiting to query metadata from server");
+                }
+            }
+            return (lastResult != null) ? lastResult : null;
+        }
+
+        /// <summary>
+        /// Get definition for a selected sql object using SMO Scripting
+        /// </summary>
+        /// <param name="textDocumentPosition"></param>
+        /// <param name="scriptFile"></param>
+        /// <param name="connInfo"></param>
+        /// <returns> Location with the URI of the script file</returns>
+        internal DefinitionResult GetDefinition(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile, ConnectionInfo connInfo)
+        {
+            // Parse sql
+            ScriptParseInfo scriptParseInfo = GetScriptParseInfo(textDocumentPosition.TextDocument.Uri);
+            if (scriptParseInfo == null)
+            {
+                return null;
+            }
+
+            if (RequiresReparse(scriptParseInfo, scriptFile))
+            {
+                scriptParseInfo.ParseResult = ParseAndBind(scriptFile, connInfo);
+            }
+
+            // Get token from selected text
+            Tuple<Stack<Token>, Queue<Token>> selectedToken = ScriptDocumentInfo.GetPeekDefinitionTokens(scriptParseInfo, 
+                textDocumentPosition.Position.Line + 1, textDocumentPosition.Position.Character + 1);
+
+            if (selectedToken == null)
+            {
+                return null;
+            }
+
+            if (scriptParseInfo.IsConnected)
+            {
+                //try children tokens first
+                Stack<Token> childrenTokens = selectedToken.Item1;
+                List<Token> tokenList = childrenTokens.ToList();
+                DefinitionResult childrenResult = GetDefinitionFromTokenList(textDocumentPosition, tokenList, scriptParseInfo, scriptFile, connInfo);
+
+                // if the children peak definition returned null then 
+                // try the parents
+                if (childrenResult == null || childrenResult.IsErrorResult)
+                {
+                    Queue<Token> parentTokens = selectedToken.Item2;
+                    tokenList = parentTokens.ToList();
+                    DefinitionResult parentResult = GetDefinitionFromTokenList(textDocumentPosition, tokenList, scriptParseInfo, scriptFile, connInfo);
+                    return (parentResult == null) ? null : parentResult;
+                }
+                else
+                {
+                    return childrenResult;
+                }    
+            }
+            else
             {
                 // User is not connected.
                 return new DefinitionResult
@@ -868,6 +931,24 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             }
         }
    
+        /// <summary>
+        /// Wrapper around find token method 
+        /// </summary>
+        /// <param name="scriptParseInfo"></param>
+        /// <param name="startLine"></param>
+        /// <param name="startColumn"></param>
+        /// <returns> token index</returns>
+        private int FindTokenWithCorrectOffset(ScriptParseInfo scriptParseInfo, int startLine, int startColumn)
+        {
+            var tokenIndex = scriptParseInfo.ParseResult.Script.TokenManager.FindToken(startLine, startColumn);
+            var end = scriptParseInfo.ParseResult.Script.TokenManager.GetToken(tokenIndex).EndLocation;
+            if (end.LineNumber == startLine && end.ColumnNumber == startColumn)
+            {
+                return tokenIndex + 1;
+            }
+            return tokenIndex;
+        }
+
         /// <summary>
         /// Extract schema name for a token, if present
         /// </summary>
@@ -878,13 +959,13 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         private string GetSchemaName(ScriptParseInfo scriptParseInfo, Position position, ScriptFile scriptFile)
         {
             // Offset index by 1 for sql parser
-            int startLine = position.Line + 1;
-            int startColumn = position.Character + 1;
+            int startLine = position.Line;
+            int startColumn = position.Character;
 
             // Get schema name
             if (scriptParseInfo != null && scriptParseInfo.ParseResult != null && scriptParseInfo.ParseResult.Script != null && scriptParseInfo.ParseResult.Script.Tokens != null)
             {
-                var tokenIndex = scriptParseInfo.ParseResult.Script.TokenManager.FindToken(startLine, startColumn);
+                var tokenIndex = FindTokenWithCorrectOffset(scriptParseInfo, startLine, startColumn);
                 var prevTokenIndex = scriptParseInfo.ParseResult.Script.TokenManager.GetPreviousSignificantTokenIndex(tokenIndex);
                 var prevTokenText = scriptParseInfo.ParseResult.Script.TokenManager.GetText(prevTokenIndex);
                 if (prevTokenText != null && prevTokenText.Equals("."))

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/ScriptDocumentInfo.cs

@@ -3,9 +3,11 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlTools.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
+using System.Collections.Generic;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion
 {
@@ -129,5 +131,62 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion
             }
             return null;
         }
+
+        /// <summary>
+        /// Returns the token that is used for Peek Definition objects
+        /// </summary>
+        internal static Tuple<Stack<Token>, Queue<Token>> GetPeekDefinitionTokens(ScriptParseInfo scriptParseInfo, int startLine, int startColumn)
+        {
+            Stack<Token> childrenTokens = new Stack<Token>();
+            Queue<Token> parentTokens = new Queue<Token>();
+            if (scriptParseInfo != null 
+                && scriptParseInfo.ParseResult != null 
+                && scriptParseInfo.ParseResult.Script != null 
+                && scriptParseInfo.ParseResult.Script.Tokens != null)
+            {
+                var tokenIndex = scriptParseInfo.ParseResult.Script.TokenManager.FindToken(startLine, startColumn);
+                if (tokenIndex >= 0)
+                {
+                    // return the current token and the ones to its right
+                    // until we hit a whitespace token
+                    int currentIndex = 0;
+                    foreach (var token in scriptParseInfo.ParseResult.Script.Tokens)
+                    {
+                        if (currentIndex == tokenIndex)
+                        {
+                            // push all parent tokens until we hit whitespace
+                            int parentIndex = currentIndex;
+                            while (true)
+                            {
+                                if (scriptParseInfo.ParseResult.Script.TokenManager.GetToken(parentIndex).Type != "LEX_WHITE")
+                                {
+                                    parentTokens.Enqueue(scriptParseInfo.ParseResult.Script.TokenManager.GetToken(parentIndex));
+                                    parentIndex--;
+                                }
+                                else
+                                {
+                                    break;
+                                }
+                            }
+                        }
+                        else if (currentIndex > tokenIndex)
+                        {
+                            // push all children tokens until we hit whitespace
+                            if (scriptParseInfo.ParseResult.Script.TokenManager.GetToken(currentIndex).Type != "LEX_WHITE")
+                            {
+                                childrenTokens.Push(token);
+                            }
+                            else 
+                            {
+                                break;
+                            }
+                        }
+                        ++currentIndex;
+                    }
+                    return Tuple.Create(childrenTokens, parentTokens);
+                }
+            }
+          return null;
+        }
     }
 }

src/Microsoft.SqlTools.ServiceLayer/Scripting/ScripterCore.cs

@@ -133,8 +133,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Scripting
         /// <returns>Location object of the script file</returns>
         internal DefinitionResult GetScript(ParseResult parseResult, Position position, IMetadataDisplayInfoProvider metadataDisplayInfoProvider, string tokenText, string schemaName)
         {
-            int parserLine = position.Line + 1;
-            int parserColumn = position.Character + 1;
+            int parserLine = position.Line;
+            int parserColumn = position.Character;
             // Get DeclarationItems from The Intellisense Resolver for the selected token. The type of the selected token is extracted from the declarationItem.
             IEnumerable<Declaration> declarationItems = GetCompletionsForToken(parseResult, parserLine, parserColumn, metadataDisplayInfoProvider);
             if (declarationItems != null && declarationItems.Count() > 0)

src/Microsoft.SqlTools.ServiceLayer/Workspace/Contracts/TextDocument.cs

@@ -157,6 +157,35 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
         /// Gets or sets the zero-based column number.
         /// </summary>
         public int Character { get; set; }
+
+        /// <summary>
+        /// Overrides the base equality method
+        /// </summary>
+        /// <param name="obj"></param>
+        /// <returns></returns>
+        public override bool Equals(object obj)
+        {
+            if (obj == null ||  (obj as Position == null))
+            {
+                return false;
+            }
+            Position p = (Position) obj;
+            bool result = (Line == p.Line) && (Character == p.Character);
+            return result;
+        }
+
+
+        /// <summary>
+        /// Overrides the base GetHashCode method
+        /// </summary>
+        /// <returns></returns>
+        public override int GetHashCode()
+        {
+            int hash = 17;
+            hash = hash * 23 + Line.GetHashCode();
+            hash = hash * 23 + Character.GetHashCode();
+            return hash;
+        }
     }
 
     [DebuggerDisplay("Start = {Start.Line}:{Start.Character}, End = {End.Line}:{End.Character}")]
@@ -171,6 +200,37 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
         /// Gets or sets the ending position of the range.
         /// </summary>
         public Position End { get; set; }
+
+        /// <summary>
+        /// Overrides the base equality method
+        /// </summary>
+        /// <param name="obj"></param>
+        /// <returns></returns>
+        public override bool Equals(object obj)
+        {
+           
+
+            if (obj == null || !(obj is Range))
+            {
+                return false;
+            }
+            Range range = (Range) obj;
+            bool sameStart = range.Start.Equals(Start);
+            bool sameEnd = range.End.Equals(End);
+            return (sameStart && sameEnd);
+        }
+
+        /// <summary>
+        /// Overrides the base GetHashCode method
+        /// </summary>
+        /// <returns></returns>
+        public override int GetHashCode()
+        {
+            int hash = 17;
+            hash = hash * 23 + Start.GetHashCode();
+            hash = hash * 23 + End.GetHashCode();
+            return hash;
+        }
     }
 
     [DebuggerDisplay("Range = {Range.Start.Line}:{Range.Start.Character} - {Range.End.Line}:{Range.End.Character}, Uri = {Uri}")]
@@ -185,6 +245,35 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
         /// Gets or sets the Range indicating the range in which location refers.
         /// </summary>
         public Range Range { get; set; }
+
+        /// <summary>
+        /// Overrides the base equality method
+        /// </summary>
+        /// <param name="obj"></param>
+        /// <returns></returns>
+        public override bool Equals(object obj)
+        {
+            if (obj == null || (obj as Location == null))
+            {
+                return false;
+            }
+            Location loc = (Location)obj;
+            bool sameUri = string.Equals(loc.Uri, Uri);
+            bool sameRange = loc.Range.Equals(Range);
+            return (sameUri && sameRange);
+        }
+
+        /// <summary>
+        /// Overrides the base GetHashCode method
+        /// </summary>
+        /// <returns></returns>
+        public override int GetHashCode()
+        {
+            int hash = 17;
+            hash = hash * 23 + Uri.GetHashCode();
+            hash = hash * 23 + Range.GetHashCode();
+            return hash;
+        }
     }
 
     public enum FileChangeType

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/LanguageServer/PeekDefinitionTests.cs

@@ -18,6 +18,7 @@ using System.Threading;
 using Xunit;
 using ConnectionType = Microsoft.SqlTools.ServiceLayer.Connection.ConnectionType;
 using Location = Microsoft.SqlTools.ServiceLayer.Workspace.Contracts.Location;
+using System.Collections.Generic;
 
 namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServices
 {
@@ -27,6 +28,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServices
     public class PeekDefinitionTests
     {
         private const string OwnerUri = "testFile1";
+        private const string TestUri = "testFile2";
         private const string ReturnTableFunctionName = "pd_returnTable";
         private const string ReturnTableTableFunctionQuery = @"
 CREATE FUNCTION [dbo].[" + ReturnTableFunctionName + @"] ()
@@ -693,6 +695,114 @@ GO";
             Assert.True(connInfo.ConnectionTypeToConnectionMap.TryRemove(ConnectionType.Query, out connection));
         }
 
+        /// <summary>
+        /// Get Definition for a object with no definition. Expect a error result
+        /// </summary>
+        [Fact]
+        public async void GetDefinitionFromChildrenAndParents()
+        {
+            string queryString = "select * from master.sys.objects";
+
+            // place the cursor on every token
+
+            //cursor on objects
+            TextDocumentPosition objectDocument = CreateTextDocPositionWithCursor(26, OwnerUri);
+                
+            //cursor on sys
+            TextDocumentPosition sysDocument = CreateTextDocPositionWithCursor(22, OwnerUri);
+
+            //cursor on master
+            TextDocumentPosition masterDocument = CreateTextDocPositionWithCursor(15, OwnerUri);
+
+            LiveConnectionHelper.TestConnectionResult connectionResult = LiveConnectionHelper.InitLiveConnectionInfo();
+            ScriptFile scriptFile = connectionResult.ScriptFile;
+            ConnectionInfo connInfo = connectionResult.ConnectionInfo;
+            var bindingQueue = new ConnectedBindingQueue();
+            bindingQueue.AddConnectionContext(connInfo);
+            LanguageService.Instance.BindingQueue = bindingQueue;
+            scriptFile.Contents = queryString;
+
+            var service = LanguageService.Instance;
+            await service.UpdateLanguageServiceOnConnection(connectionResult.ConnectionInfo);
+            Thread.Sleep(2000);
+
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+            scriptInfo.ConnectionKey = bindingQueue.AddConnectionContext(connInfo);
+            LanguageService.Instance.ScriptParseInfoMap.Add(OwnerUri, scriptInfo);
+
+            // When I call the language service
+            var objectResult = LanguageService.Instance.GetDefinition(objectDocument, scriptFile, connInfo);
+            var sysResult = LanguageService.Instance.GetDefinition(sysDocument, scriptFile, connInfo);
+            var masterResult = LanguageService.Instance.GetDefinition(masterDocument, scriptFile, connInfo);
+
+            // Then I expect the results to be non-null
+            Assert.NotNull(objectResult);
+            Assert.NotNull(sysResult);
+            Assert.NotNull(masterResult);
+
+            // And I expect the all results to be the same
+            Assert.True(CompareLocations(objectResult.Locations, sysResult.Locations));
+            Assert.True(CompareLocations(objectResult.Locations, masterResult.Locations));
+
+            Cleanup(objectResult.Locations);
+            Cleanup(sysResult.Locations);
+            Cleanup(masterResult.Locations);
+            LanguageService.Instance.ScriptParseInfoMap.Remove(OwnerUri);
+        }
+
+        [Fact]
+        public async void GetDefinitionFromProcedures()
+        {
+
+            string queryString = "EXEC master.dbo.sp_MSrepl_startup";
+
+            // place the cursor on every token
+
+            //cursor on objects
+            TextDocumentPosition fnDocument = CreateTextDocPositionWithCursor(30, TestUri);
+
+            //cursor on sys
+            TextDocumentPosition dboDocument = CreateTextDocPositionWithCursor(14, TestUri);
+
+            //cursor on master
+            TextDocumentPosition masterDocument = CreateTextDocPositionWithCursor(10, TestUri);
+
+            LiveConnectionHelper.TestConnectionResult connectionResult = LiveConnectionHelper.InitLiveConnectionInfo();
+            ScriptFile scriptFile = connectionResult.ScriptFile;
+            ConnectionInfo connInfo = connectionResult.ConnectionInfo;
+            var bindingQueue = new ConnectedBindingQueue();
+            bindingQueue.AddConnectionContext(connInfo);
+            LanguageService.Instance.BindingQueue = bindingQueue;
+            scriptFile.Contents = queryString;
+
+            var service = LanguageService.Instance;
+            await service.UpdateLanguageServiceOnConnection(connectionResult.ConnectionInfo);
+            Thread.Sleep(2000);
+
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+            scriptInfo.ConnectionKey = bindingQueue.AddConnectionContext(connInfo);
+            LanguageService.Instance.ScriptParseInfoMap.Add(TestUri, scriptInfo);
+
+            // When I call the language service
+            var fnResult = LanguageService.Instance.GetDefinition(fnDocument, scriptFile, connInfo);
+            var sysResult = LanguageService.Instance.GetDefinition(dboDocument, scriptFile, connInfo);
+            var masterResult = LanguageService.Instance.GetDefinition(masterDocument, scriptFile, connInfo);
+
+            // Then I expect the results to be non-null
+            Assert.NotNull(fnResult);
+            Assert.NotNull(sysResult);
+            Assert.NotNull(masterResult);
+
+            // And I expect the all results to be the same
+            Assert.True(CompareLocations(fnResult.Locations, sysResult.Locations));
+            Assert.True(CompareLocations(fnResult.Locations, masterResult.Locations));
+
+            Cleanup(fnResult.Locations);
+            Cleanup(sysResult.Locations);
+            Cleanup(masterResult.Locations);
+            LanguageService.Instance.ScriptParseInfoMap.Remove(TestUri);
+        }
+
 
         /// <summary>
         /// Helper method to clean up script files
@@ -712,5 +822,42 @@ GO";
                 }
             }
         }
+
+        /// <summary>
+        /// Helper method to compare 2 Locations arrays
+        /// </summary>
+        /// <param name="locationsA"></param>
+        /// <param name="locationsB"></param>
+        /// <returns></returns>
+        private bool CompareLocations(Location[] locationsA, Location[] locationsB)
+        {
+            HashSet<Location> locationSet = new HashSet<Location>();
+            foreach (var location in locationsA)
+            {
+                locationSet.Add(location);
+            }
+            foreach (var location in locationsB)
+            {
+                if (!locationSet.Contains(location))
+                {
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        private TextDocumentPosition CreateTextDocPositionWithCursor(int column, string OwnerUri)
+        {
+            TextDocumentPosition textDocPos = new TextDocumentPosition
+            {
+                TextDocument = new TextDocumentIdentifier { Uri = OwnerUri },
+                Position = new Position
+                {
+                    Line = 0,
+                    Character = column
+                }
+            };
+            return textDocPos;
+        }
     }
 }