Store location of module members in the database

src/Analysis/Ast/Impl/Analyzer/PythonAnalyzerEntry.cs

@@ -264,7 +264,7 @@ private HashSet<AnalysisModuleKey> FindDependencies(IPythonModule module, Python
         }
 
         private static bool Ignore(IModuleManagement moduleResolution, string fullName, string modulePath)
-            => moduleResolution.BuiltinModuleName.EqualsOrdinal(fullName) || moduleResolution.GetSpecializedModule(fullName, modulePath) != null;
+            => moduleResolution.BuiltinModuleName.EqualsOrdinal(fullName) || moduleResolution.IsSpecializedModule(fullName, modulePath);
 
         private void UpdateAnalysisTcs(int analysisVersion) {
             _analysisVersion = analysisVersion;

src/Analysis/Ast/Impl/Caching/Definitions/IModuleDatabaseService.cs

@@ -33,5 +33,10 @@ internal interface IModuleDatabaseService {
         /// Writes module data to the database.
         /// </summary>
         Task StoreModuleAnalysisAsync(IDocumentAnalysis analysis, CancellationToken cancellationToken = default);
+
+        /// <summary>
+        /// Determines if module analysis exists in the storage.
+        /// </summary>
+        bool ModuleExistsInStorage(string moduleName, string filePath);
     }
 }

src/Analysis/Ast/Impl/Extensions/PythonModuleExtensions.cs

@@ -20,7 +20,7 @@
 namespace Microsoft.Python.Analysis {
     public static class PythonModuleExtensions {
         internal static PythonAst GetAst(this IPythonModule module)
-            => (PythonAst)((IAstNodeContainer)module).GetAstNode(module);
+            => (PythonAst)(module as IAstNodeContainer)?.GetAstNode(module);
 
         internal static void SetAst(this IPythonModule module, PythonAst ast) {
             var contained = (IAstNodeContainer)module;

src/Analysis/Ast/Impl/Modules/Definitions/IModuleManagement.cs

@@ -48,15 +48,20 @@ public interface IModuleManagement: IModuleResolution {
         /// content is loaded and analyzed only for class/functions definitions
         /// so the original documentation can be extracted.
         /// </summary>
-        /// <param name="name">Module to specialize.</param>
+        /// <param name="fullName">Module to specialize.</param>
         /// <param name="specializationConstructor">Specialized module constructor.</param>
         /// <returns>Original (library) module loaded as stub, if any.</returns>
-        IPythonModule SpecializeModule(string name, Func<string, IPythonModule> specializationConstructor);
+        IPythonModule SpecializeModule(string fullName, Func<string, IPythonModule> specializationConstructor);
 
         /// <summary>
-        /// Returns specialized module, if any.
+        /// Returns specialized module, if any. Will attempt to load module from persistent state.
         /// </summary>
-        IPythonModule GetSpecializedModule(string name, string modulePath = null);
+        IPythonModule GetSpecializedModule(string fullName, bool allowCreation = false, string modulePath = null);
+
+        /// <summary>
+        /// Determines of module is specialized or exists in the database.
+        /// </summary>
+        bool IsSpecializedModule(string fullName, string modulePath = null);
 
         /// <summary>
         /// Root directory of the path resolver.

src/Analysis/Ast/Impl/Modules/PythonModule.cs

@@ -42,7 +42,7 @@ namespace Microsoft.Python.Analysis.Modules {
     /// to AST and the module analysis.
     /// </summary>
     [DebuggerDisplay("{Name} : {ModuleType}")]
-    internal class PythonModule : LocatedMember, IDocument, IAnalyzable, IEquatable<IPythonModule>, IAstNodeContainer {
+    internal class PythonModule : LocatedMember, IDocument, IAnalyzable, IEquatable<IPythonModule>, IAstNodeContainer, ILocationConverter {
         private enum State {
             None,
             Loading,
@@ -565,6 +565,11 @@ private string TryGetDocFromModuleInitFile() {
         }
         #endregion
 
+        #region ILocationConverter
+        public virtual SourceLocation IndexToLocation(int index) => this.GetAst()?.IndexToLocation(index) ?? default;
+        public virtual int LocationToIndex(SourceLocation location) => this.GetAst()?.LocationToIndex(location) ?? default;
+        #endregion
+
         private void RemoveReferencesInModule(IPythonModule module) {
             if (module.GlobalScope?.Variables != null) {
                 foreach (var v in module.GlobalScope.Variables) {

src/Analysis/Ast/Impl/Modules/PythonVariableModule.cs

@@ -19,14 +19,15 @@
 using Microsoft.Python.Analysis.Types;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Core;
+using Microsoft.Python.Core.Text;
 
 namespace Microsoft.Python.Analysis.Modules {
     /// <summary>
     /// Module-scoped representation of the module or implicit package
     /// Contains either module members, members + imported children of explicit package or imported implicit package children
     /// Instance is unique for each module analysis
     /// </summary>
-    internal sealed class PythonVariableModule : LocatedMember, IPythonModule, IEquatable<IPythonModule> {
+    internal sealed class PythonVariableModule : LocatedMember, IPythonModule, IEquatable<IPythonModule>, ILocationConverter {
         private readonly Dictionary<string, PythonVariableModule> _children = new Dictionary<string, PythonVariableModule>();
 
         public string Name { get; }
@@ -71,5 +72,10 @@ public PythonVariableModule(IPythonModule module): base(module) {
         public IMember CreateInstance(string typeName = null, IArgumentSet args = null) => this;
 
         public bool Equals(IPythonModule other) => other is PythonVariableModule module && Name.EqualsOrdinal(module.Name);
+
+        #region ILocationConverter
+        public SourceLocation IndexToLocation(int index) => (Module as ILocationConverter)?.IndexToLocation(index) ?? default;
+        public int LocationToIndex(SourceLocation location) => (Module as ILocationConverter)?.LocationToIndex(location) ?? default;
+        #endregion
     }
 }

src/Analysis/Ast/Impl/Modules/Resolution/MainModuleResolution.cs

@@ -69,14 +69,21 @@ protected override IPythonModule CreateModule(string name) {
                 return null;
             }
 
+            IPythonModule module;
             if (moduleImport.ModulePath != null) {
-                var module = GetRdt().GetDocument(new Uri(moduleImport.ModulePath));
+                module = GetRdt().GetDocument(new Uri(moduleImport.ModulePath));
                 if (module != null) {
                     GetRdt().LockDocument(module.Uri);
                     return module;
                 }
             }
 
+            var dbs = GetDbService();
+            if (dbs != null && dbs.TryCreateModule(name, moduleImport.ModulePath, out module) != ModuleStorageState.DoesNotExist && module != null) {
+                SpecializeModule(name, s => module);
+                return module;
+            }
+
             // If there is a stub, make sure it is loaded and attached
             // First check stub next to the module.
             if (!TryCreateModuleStub(name, moduleImport.ModulePath, out var stub)) {
@@ -152,17 +159,14 @@ public IPythonModule SpecializeModule(string name, Func<string, IPythonModule> s
         /// <summary>
         /// Returns specialized module, if any.
         /// </summary>
-        public IPythonModule GetSpecializedModule(string fullName, string modulePath = null) {
-            if (_specialized.TryGetValue(fullName, out var module)) {
-                return module;
-            }
-            var dbs = GetDbService();
-            if (dbs != null && dbs.TryCreateModule(fullName, modulePath, out module) != ModuleStorageState.DoesNotExist && module != null) {
-                SpecializeModule(fullName, s => module);
-                return module;
-            }
-            return null;
-        }
+        public IPythonModule GetSpecializedModule(string fullName, bool allowCreation = false, string modulePath = null) 
+            => _specialized.TryGetValue(fullName, out var module) ? module : null;
+
+        /// <summary>
+        /// Determines of module is specialized or exists in the database.
+        /// </summary>
+        public bool IsSpecializedModule(string fullName, string modulePath = null)
+            => _specialized.ContainsKey(fullName) || GetDbService()?.ModuleExistsInStorage(fullName, modulePath) == true;
 
         internal async Task LoadBuiltinTypesAsync(CancellationToken cancellationToken = default) {
             var analyzer = _services.GetService<IPythonAnalyzer>();

src/Analysis/Ast/Impl/Modules/Resolution/ModuleResolutionBase.cs

@@ -86,15 +86,25 @@ public IPythonModule GetImportedModule(string name)
             => Modules.TryGetValue(name, out var moduleRef) ? moduleRef.Value : _interpreter.ModuleResolution.GetSpecializedModule(name);
 
         public IPythonModule GetOrLoadModule(string name) {
-            if (Modules.TryGetValue(name, out var moduleRef)) {
-                return moduleRef.GetOrCreate(name, this);
+            // Specialized should always win. However, we don't want
+            // to allow loading from the database just yet since module
+            // may already exist in the analyzed state.
+            var module = GetImportedModule(name);
+            if (module != null) {
+                return module;
             }
 
-            var module = _interpreter.ModuleResolution.GetSpecializedModule(name);
+            // Now try restoring from the database.
+            module = _interpreter.ModuleResolution.GetSpecializedModule(name, true);
             if (module != null) {
                 return module;
             }
 
+            // Now try regular case.
+            if (Modules.TryGetValue(name, out var moduleRef)) {
+                return moduleRef.GetOrCreate(name, this);
+            }
+
             moduleRef = Modules.GetOrAdd(name, new ModuleRef());
             return moduleRef.GetOrCreate(name, this);
         }

src/Analysis/Ast/Impl/Types/ArgumentSet.cs

@@ -328,7 +328,7 @@ private IMember GetArgumentValue(Argument arg) {
             }
 
             if (arg.ValueIsDefault) {
-                using (Eval.OpenScope(DeclaringModule.Analysis.GlobalScope)) {
+                using (Eval.OpenScope(DeclaringModule.GlobalScope)) {
                     return Eval.GetValueFromExpression(arg.ValueExpression) ?? Eval.UnknownType;
                 }
             }

src/Analysis/Ast/Impl/Types/Location.cs

@@ -30,9 +30,8 @@ public Location(IPythonModule module, IndexSpan indexSpan) {
 
         public LocationInfo LocationInfo {
             get {
-                var ast = Module?.GetAst();
-                if (ast != null && !string.IsNullOrEmpty(Module?.FilePath) && Module?.Uri != null) {
-                    return new LocationInfo(Module.FilePath, Module.Uri, IndexSpan.ToSourceSpan(ast));
+                if (Module is ILocationConverter lc && !string.IsNullOrEmpty(Module?.FilePath) && Module?.Uri != null) {
+                    return new LocationInfo(Module.FilePath, Module.Uri, IndexSpan.ToSourceSpan(lc));
                 }
                 return LocationInfo.Empty;
             }

src/Analysis/Ast/Impl/Types/PythonFunctionOverload.cs

@@ -182,7 +182,7 @@ private IMember CreateSpecificReturnFromTypeVar(IPythonClassType selfClassType,
             // Try returning the constraint
             // TODO: improve this, the heuristic is pretty basic and tailored to simple func(_T) -> _T
             var name = StaticReturnValue.GetPythonType()?.Name;
-            var typeDefVar = DeclaringModule.Analysis.GlobalScope.Variables[name];
+            var typeDefVar = DeclaringModule.GlobalScope.Variables[name];
             if (typeDefVar?.Value is IGenericTypeDefinition gtp2) {
                 // See if the instance (self) type satisfies one of the constraints.
                 return selfClassType.Mro.Any(b => gtp2.Constraints.Any(c => c.Equals(b)))

src/Caching/Impl/Extensions/IndexSpanExtensions.cs

@@ -0,0 +1,27 @@
+// Python Tools for Visual Studio
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using Microsoft.Python.Analysis.Caching.Models;
+using Microsoft.Python.Core.Text;
+
+namespace Microsoft.Python.Analysis.Caching {
+    internal static class IndexSpanExtensions {
+        public static IndexSpanModel ToModel(this IndexSpan span) => new IndexSpanModel {
+            Start = span.Start,
+            Length = span.Length
+        };
+    }
+}

src/Caching/Impl/Factories/ClassFactory.cs

@@ -27,7 +27,7 @@ public ClassFactory(IEnumerable<ClassModel> classes, ModuleFactory mf)
         }
 
         protected override PythonClassType CreateMember(ClassModel cm, IPythonType declaringType) 
-            => new PythonClassType(cm.Name, ModuleFactory.DefaultLocation);
+            => new PythonClassType(cm.Name, new Location(ModuleFactory.Module, cm.IndexSpan.ToSpan()));
 
         protected override void CreateMemberParts(ClassModel cm, PythonClassType cls) {
             // In Python 3 exclude object since type creation will add it automatically.

src/Caching/Impl/Factories/FunctionFactory.cs

@@ -25,9 +25,9 @@ public FunctionFactory(IEnumerable<FunctionModel> classes, ModuleFactory mf)
         }
 
         protected override IPythonFunctionType CreateMember(FunctionModel fm, IPythonType declaringType) {
-            var f = new PythonFunctionType(fm.Name, ModuleFactory.DefaultLocation, declaringType, fm.Documentation);
+            var f = new PythonFunctionType(fm.Name, new Location(ModuleFactory.Module, fm.IndexSpan.ToSpan()), declaringType, fm.Documentation);
             foreach (var om in fm.Overloads) {
-                var o = new PythonFunctionOverload(fm.Name, ModuleFactory.DefaultLocation);
+                var o = new PythonFunctionOverload(fm.Name, new Location(ModuleFactory.Module, fm.IndexSpan.ToSpan()));
                 o.SetDocumentation(fm.Documentation);
                 o.SetReturnValue(ModuleFactory.ConstructMember(om.ReturnType), true);
                 o.SetParameters(om.Parameters.Select(ConstructParameter).ToArray());

src/Caching/Impl/Factories/PropertyFactory.cs

@@ -25,7 +25,7 @@ public PropertyFactory(ModuleFactory mf) {
         }
 
         public IPythonPropertyType Construct(PropertyModel pm, IPythonClassType cls) {
-            var prop = new PythonPropertyType(pm.Name, _mf.DefaultLocation, cls, (pm.Attributes & FunctionAttributes.Abstract) != 0);
+            var prop = new PythonPropertyType(pm.Name, new Location(_mf.Module, pm.IndexSpan.ToSpan()), cls, (pm.Attributes & FunctionAttributes.Abstract) != 0);
             prop.SetDocumentation(pm.Documentation);
             var o = new PythonFunctionOverload(pm.Name, _mf.DefaultLocation);
             o.SetDocumentation(pm.Documentation); // TODO: own documentation?

src/Caching/Impl/Factories/VariableFactory.cs

@@ -26,7 +26,7 @@ public VariableFactory(IEnumerable<VariableModel> models, ModuleFactory mf)
 
         protected override IVariable CreateMember(VariableModel vm, IPythonType declaringType) {
             var m = ModuleFactory.ConstructMember(vm.Value);
-            return new Variable(vm.Name, m, VariableSource.Declaration, ModuleFactory.DefaultLocation);
+            return new Variable(vm.Name, m, VariableSource.Declaration, new Location(ModuleFactory.Module, vm.IndexSpan.ToSpan()));
         }
     }
 }

src/Caching/Impl/Models/ClassModel.cs

@@ -85,6 +85,8 @@ private ClassModel(IPythonClassType cls) {
 
             Name = cls.TypeId == BuiltinTypeId.Ellipsis ? "ellipsis" : cls.Name;
             Id = Name.GetStableHash();
+            IndexSpan = cls.Location.IndexSpan.ToModel();
+
             Documentation = cls.Documentation;
             Bases = cls.Bases.OfType<IPythonClassType>().Select(t => t.GetQualifiedName()).ToArray();
             Methods = methods.ToArray();

src/Caching/Impl/Models/FunctionModel.cs

@@ -31,6 +31,7 @@ public static FunctionModel FromType(IPythonFunctionType ft) {
             return new FunctionModel {
                 Id = ft.Name.GetStableHash(),
                 Name = ft.Name,
+                IndexSpan = ft.Location.IndexSpan.ToModel(),
                 Documentation = ft.Documentation,
                 Overloads = ft.Overloads.Select(FromOverload).ToArray()
                 // TODO: attributes, inner functions and inner classes.

src/Caching/Impl/Models/IndexSpanModel.cs

@@ -0,0 +1,26 @@
+// Python Tools for Visual Studio
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using Microsoft.Python.Core.Text;
+
+namespace Microsoft.Python.Analysis.Caching.Models {
+    internal sealed class IndexSpanModel {
+        public int Start { get; set; }
+        public int Length { get; set; }
+
+        public IndexSpan ToSpan() => new IndexSpan(Start, Length);
+    }
+}

src/Caching/Impl/Models/MemberModel.cs

@@ -13,9 +13,12 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using Microsoft.Python.Core.Text;
+
 namespace Microsoft.Python.Analysis.Caching.Models {
     internal abstract class MemberModel {
         public int Id { get; set; }
         public string Name { get; set; }
+        public IndexSpanModel IndexSpan { get; set; }
     }
 }

src/Caching/Impl/Models/ModuleModel.cs

@@ -18,14 +18,31 @@
 using Microsoft.Python.Analysis.Types;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Core;
+using Microsoft.Python.Parsing;
 
 namespace Microsoft.Python.Analysis.Caching.Models {
     internal sealed class ModuleModel : MemberModel {
+        /// <summary>
+        /// Module unique id that includes version.
+        /// </summary>
         public string UniqueId { get; set; }
+
         public string Documentation { get; set; }
         public FunctionModel[] Functions { get; set; }
         public VariableModel[] Variables { get; set; }
         public ClassModel[] Classes { get; set; }
+
+        /// <summary>
+        /// Collection of new line information for conversion of linear spans
+        /// to line/columns in navigation to member definitions and references.
+        /// </summary>
+        public NewLineModel[] NewLines { get; set; }
+
+        /// <summary>
+        /// Length of the original module file. Used in conversion of indices to line/columns.
+        /// </summary>
+        public int FileSize { get; set; }
+
         // TODO: TypeVars, ...
 
         public static ModuleModel FromAnalysis(IDocumentAnalysis analysis, IServiceContainer services) {
@@ -42,15 +59,15 @@ public static ModuleModel FromAnalysis(IDocumentAnalysis analysis, IServiceConta
                 string typeName = null;
 
                 switch (v.Value) {
-                    case IPythonFunctionType ft 
+                    case IPythonFunctionType ft
                         when ft.DeclaringModule.Equals(analysis.Document) || ft.DeclaringModule.Equals(analysis.Document.Stub):
                         if (!functions.ContainsKey(ft.Name)) {
                             typeName = ft.Name;
                             functions[ft.Name] = FunctionModel.FromType(ft);
                         }
 
                         break;
-                    case IPythonClassType cls 
+                    case IPythonClassType cls
                         when cls.DeclaringModule.Equals(analysis.Document) || cls.DeclaringModule.Equals(analysis.Document.Stub):
                         if (!classes.ContainsKey(cls.Name)) {
                             typeName = cls.Name;
@@ -73,7 +90,12 @@ when cls.DeclaringModule.Equals(analysis.Document) || cls.DeclaringModule.Equals
                 Documentation = analysis.Document.Documentation,
                 Functions = functions.Values.ToArray(),
                 Variables = variables.Values.ToArray(),
-                Classes = classes.Values.ToArray()
+                Classes = classes.Values.ToArray(),
+                NewLines = analysis.Ast.NewLineLocations.Select(l => new NewLineModel {
+                    EndIndex = l.EndIndex,
+                    Kind = l.Kind
+                }).ToArray(),
+                FileSize = analysis.Ast.EndIndex
             };
         }
     }