analyzer: refactor folding ranges, support multi-resolve and grammar updates

src/analyzer/psi/PsiFile.v

@@ -191,18 +191,25 @@ pub fn (p &PsiFile) get_imports() []ImportSpec {
 }
 
 pub fn (p &PsiFile) resolve_import_spec(name string) ?ImportSpec {
+	specs := p.resolve_import_specs(name)
+	if specs.len > 0 {
+		return specs.first()
+	}
+	return none
+}
+
+pub fn (p &PsiFile) resolve_import_specs(name string) []ImportSpec {
 	imports := p.get_imports()
 	if imports.len == 0 {
-		return none
+		return []
 	}
-
+	mut result := []ImportSpec{cap: 2}
 	for imp in imports {
 		if imp.import_name() == name {
-			return imp
+			result << imp
 		}
 	}
-
-	return none
+	return result
 }
 
 pub fn (p &PsiFile) process_declarations(mut processor PsiScopeProcessor) bool {

src/analyzer/psi/PsiReference.v

@@ -3,4 +3,5 @@ module psi
 pub interface PsiReference {
 	element() PsiElement
 	resolve() ?PsiElement
+	multi_resolve() []PsiElement
 }

src/analyzer/psi/ReferenceImpl.v

@@ -59,6 +59,49 @@ pub fn (r &ReferenceImpl) resolve() ?PsiElement {
 	return result
 }
 
+pub fn (r &ReferenceImpl) multi_resolve() []PsiElement {
+	file := r.file or { return [] }
+
+	if res := r.resolve_as_import_spec() {
+		return res
+	}
+
+	sub := SubResolver{
+		containing_file: file
+		element:         r.element
+		for_types:       r.for_types
+		for_attributes:  r.for_attributes
+	}
+	mut processor := ResolveProcessor{
+		containing_file: file
+		ref:             r.element
+		ref_name:        r.element.name()
+		collect_all:     true
+	}
+
+	sub.process_resolve_variants(mut processor)
+
+	return processor.result
+}
+
+fn (r &ReferenceImpl) resolve_as_import_spec() ?[]PsiElement {
+	if r.element is Identifier {
+		parent := r.element.parent()?
+		if parent !is ImportName {
+			return none
+		}
+		spec := parent.parent()?.parent()?
+		if spec is ImportSpec {
+			if ident := spec.identifier() {
+				if ident.is_equal(parent) {
+					return [spec]
+				}
+			}
+		}
+	}
+	return none
+}
+
 pub struct SubResolver {
 pub:
 	containing_file ?&PsiFile
@@ -91,16 +134,22 @@ pub fn (r &SubResolver) process_qualifier_expression(qualifier PsiElement, mut p
 
 	if qualifier is ReferenceExpressionBase {
 		resolved := qualifier.resolve() or { return true }
+
 		if resolved is ImportSpec {
-			import_name := resolved.qualified_name()
-			real_fqn := stubs_index.find_real_module_fqn(import_name)
+			import_name := resolved.import_name()
+			file := r.containing_file or { return true }
+			specs := file.resolve_import_specs(import_name)
 
-			elements := stubs_index.get_all_declarations_from_module(real_fqn, r.for_types)
-			for element in elements {
-				if !processor.execute(element) {
+			for _, spec in specs {
+				target_fqn := spec.qualified_name()
+				real_fqn := stubs_index.find_real_module_fqn(target_fqn)
+				elements := stubs_index.get_all_declarations_from_module(real_fqn, r.for_types)
+
+				if !r.process_elements(elements, mut processor) {
 					return false
 				}
 			}
+			return true
 		}
 
 		if resolved is ModuleClause {
@@ -821,7 +870,8 @@ pub struct ResolveProcessor {
 	ref             ReferenceExpressionBase
 	ref_name        string
 mut:
-	result []PsiElement
+	result      []PsiElement
+	collect_all bool
 }
 
 fn (mut r ResolveProcessor) execute(element PsiElement) bool {
@@ -836,6 +886,9 @@ fn (mut r ResolveProcessor) execute(element PsiElement) bool {
 		}
 		if name == r.ref_name {
 			r.result << element as PsiElement
+			if r.collect_all {
+				return true
+			}
 			return false
 		}
 	}

src/analyzer/psi/StubbedElementTypeImpl.v

@@ -572,7 +572,14 @@ pub fn (s &StubbedElementType) create_stub(psi PsiElement, parent_stub &StubBase
 		return declaration_stub(*psi, parent_stub, .import_spec, include_text: true)
 	}
 
-	if node_type in [.import_list, .import_declaration, .import_path, .import_name, .import_alias] {
+	if node_type in [
+		.import_list,
+		.import_declaration,
+		.import_path,
+		.import_name,
+		.import_alias,
+		.selective_import_list,
+	] {
 		stub_type := node_type_to_stub_type(node_type)
 		return text_based_stub(psi, parent_stub, stub_type,
 			include_text: node_type !in [

src/analyzer/psi/search/ReferencesSearch.v

@@ -274,6 +274,17 @@ pub fn (r &ReferencesSearch) search_in(element psi.PsiNamedElement, search_root
 						result << node
 					}
 				}
+				if element is psi.ImportSpec && resolved is psi.ImportSpec {
+					if element.import_name() == resolved.import_name() {
+						if element_file := element.containing_file() {
+							if resolved_file := resolved.containing_file() {
+								if element_file.path == resolved_file.path {
+									result << node
+								}
+							}
+						}
+					}
+				}
 			}
 		}
 	}

src/server/documentation/provider.v

@@ -267,6 +267,11 @@ fn (mut p Provider) const_documentation(element psi.ConstantDefinition) ? {
 	p.sb.write_string(' = ')
 	if value := element.expression() {
 		p.sb.write_string(value.get_text())
+		if element.stub_based() {
+			if mut file := value.containing_file() {
+				file.free()
+			}
+		}
 	}
 	p.sb.write_string('\n')
 	p.sb.write_string('```')

src/server/features_definition.v

@@ -17,19 +17,21 @@ pub fn (mut ls LanguageServer) definition(params lsp.TextDocumentPositionParams)
 		return none
 	}
 
-	resolved := element.resolve() or {
-		loglib.with_fields({
-			'caller': @METHOD
-			'name':   element.name()
-		}).warn('Cannot resolve reference')
+	resolved_elements := element.reference().multi_resolve()
+	if resolved_elements.len == 0 {
 		return none
 	}
 
-	containing_file := resolved.containing_file() or { return [] }
-	data := new_resolve_result(containing_file, resolved) or { return [] }
-	return [
-		data.to_location_link(element.text_range()),
-	]
+	mut links := []lsp.LocationLink{cap: resolved_elements.len}
+
+	for resolved in resolved_elements {
+		containing_file := resolved.containing_file() or { continue }
+		if data := new_resolve_result(containing_file, resolved) {
+			links << data.to_location_link(element.text_range())
+		}
+	}
+
+	return links
 }
 
 struct ResolveResult {

src/server/features_folding_range.v

@@ -1,45 +1,20 @@
 module server
 
 import lsp
-import analyzer.psi
+import loglib
+import server.folding
 
 pub fn (mut ls LanguageServer) folding_range(params lsp.FoldingRangeParams) ?[]lsp.FoldingRange {
 	uri := params.text_document.uri.normalize()
-	file := ls.get_file(uri)?
-
-	mut result := []lsp.FoldingRange{}
-	mut walker := psi.new_tree_walker(file.psi_file.root().node())
-	defer { walker.free() }
-
-	for {
-		node := walker.next() or { break }
-
-		if node.type_name == .import_list {
-			element := psi.create_element(node, file.psi_file)
-			decls := element.find_children_by_type(.import_declaration)
-			if decls.len < 2 {
-				continue
-			}
-
-			first := decls.first()
-			range := element.text_range()
-
-			first_range := if first is psi.ImportDeclaration {
-				spec := first.spec() or { continue }
-				spec.text_range()
-			} else {
-				continue
-			}
-
-			result << lsp.FoldingRange{
-				start_line:      first_range.line
-				start_character: first_range.column
-				end_line:        range.end_line - 2
-				end_character:   range.end_column
-				kind:            lsp.folding_range_kind_imports
-			}
-		}
+	file := ls.get_file(uri) or {
+		loglib.with_fields({
+			'uri': uri.str()
+		}).warn('Folding range requested for unopened file')
+		return []
 	}
 
-	return result
+	mut visitor := folding.FoldingVisitor.new(file.psi_file)
+	ranges := visitor.accept(file.psi_file.root())
+
+	return ranges
 }