analyzer: support selective imports and improve module resolution (resolves #17, #98, #103, #105, #107, #149, #151, #153, #159) (#177)

Author: wenxuanjun

src/analyzer/psi/Identifier.v

@@ -30,5 +30,14 @@ pub fn (i Identifier) reference() PsiReference {
 }
 
 pub fn (i Identifier) resolve() ?PsiElement {
+	if parent := i.parent() {
+		if parent is PsiNamedElement {
+			if ident := parent.identifier() {
+				if ident.is_equal(i) {
+					return parent as PsiElement
+				}
+			}
+		}
+	}
 	return i.reference().resolve()
 }

src/analyzer/psi/ImportSpec.v

@@ -77,3 +77,11 @@ pub fn (n ImportSpec) resolve_directory() string {
 
 	return stubs_index.get_module_root(fqn)
 }
+
+pub fn (n &ImportSpec) selective_list() ?&SelectiveImportList {
+	list := n.find_child_by_type_or_stub(.selective_import_list)?
+	if list is SelectiveImportList {
+		return list
+	}
+	return none
+}

src/analyzer/psi/PsiFile.v

@@ -217,3 +217,54 @@ pub fn (p &PsiFile) process_declarations(mut processor PsiScopeProcessor) bool {
 
 	return true
 }
+
+pub fn (p &PsiFile) resolve_selective_import_symbol(name string) ?PsiElement {
+	imports := p.get_imports()
+
+	for spec in imports {
+		list := spec.selective_list() or { continue }
+		symbols := list.symbols()
+
+		for ref in symbols {
+			if ref.name() == name {
+				if found := p.resolve_symbol_in_import_spec(spec, name) {
+					return found
+				}
+			}
+		}
+	}
+
+	return none
+}
+
+pub fn (p &PsiFile) resolve_symbol_in_import_spec(spec ImportSpec, name string) ?PsiElement {
+	import_name := spec.qualified_name()
+	real_module_fqn := stubs_index.find_real_module_fqn(import_name)
+
+	if found := p.find_in_module(real_module_fqn, name) {
+		return found
+	}
+
+	return none
+}
+
+fn (p &PsiFile) find_in_module(module_fqn string, name string) ?PsiElement {
+	elements := stubs_index.get_all_declarations_from_module(module_fqn, false)
+	for element in elements {
+		if element is PsiNamedElement {
+			if element.name() == name {
+				return element as PsiElement
+			}
+		}
+	}
+
+	types := stubs_index.get_all_declarations_from_module(module_fqn, true)
+	for type_element in types {
+		if type_element is PsiNamedElement {
+			if type_element.name() == name {
+				return type_element as PsiElement
+			}
+		}
+	}
+	return none
+}

src/analyzer/psi/ReferenceImpl.v

@@ -92,8 +92,10 @@ 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 {
-			elements := stubs_index.get_all_declarations_from_module(resolved.qualified_name(),
-				r.for_types)
+			import_name := resolved.qualified_name()
+			real_fqn := stubs_index.find_real_module_fqn(import_name)
+
+			elements := stubs_index.get_all_declarations_from_module(real_fqn, r.for_types)
 			for element in elements {
 				if !processor.execute(element) {
 					return false
@@ -351,7 +353,7 @@ pub fn (r &SubResolver) process_unqualified_resolve(mut processor PsiScopeProces
 	if !r.process_imported_modules(mut processor) {
 		return false
 	}
-	if !r.process_module_clause(mut processor) {
+	if !r.process_selective_imports(mut processor) {
 		return false
 	}
 	if !r.process_owner_generic_ts(mut processor) {
@@ -441,6 +443,9 @@ pub fn (r &SubResolver) process_unqualified_resolve(mut processor PsiScopeProces
 	if !r.process_elements(mod_decls, mut processor) {
 		return false
 	}
+	if !r.process_module_clause(mut processor) {
+		return false
+	}
 
 	return true
 }
@@ -570,6 +575,30 @@ pub fn (r &SubResolver) process_imported_modules(mut processor PsiScopeProcessor
 	return true
 }
 
+pub fn (r &SubResolver) process_selective_imports(mut processor PsiScopeProcessor) bool {
+	element := r.element as PsiElement
+	name := r.element.name()
+	file := r.containing_file or { return true }
+
+	if parent_list := element.parent_of_type(.selective_import_list) {
+		if spec := parent_list.parent_of_type(.import_spec) {
+			if spec is ImportSpec {
+				if found := file.resolve_symbol_in_import_spec(spec, name) {
+					return processor.execute(found)
+				}
+				return true
+			}
+		}
+	}
+
+	if resolved := file.resolve_selective_import_symbol(name) {
+		if !processor.execute(resolved) {
+			return false
+		}
+	}
+	return true
+}
+
 pub fn (r &SubResolver) process_enum_fetch(parent PsiElement, mut processor PsiScopeProcessor) bool {
 	context_type := TypeInferer{}.infer_context_type(parent)
 	return r.process_type(context_type, mut processor)

src/analyzer/psi/SelectiveImportList.v

@@ -0,0 +1,18 @@
+module psi
+
+pub struct SelectiveImportList {
+	PsiElementImpl
+}
+
+pub fn (n &SelectiveImportList) symbols() []ReferenceExpression {
+	children := n.find_children_by_type_or_stub(.reference_expression)
+	mut res := []ReferenceExpression{cap: children.len}
+	for child in children {
+		if child is ReferenceExpression {
+			res << child
+		}
+	}
+	return res
+}
+
+fn (n &SelectiveImportList) stub() {}

src/analyzer/psi/StubBase.v

@@ -125,6 +125,7 @@ pub fn (s &StubBase) element_type() bindings.NodeType {
 		.import_path { .import_path }
 		.import_name { .import_name }
 		.import_alias { .import_alias }
+		.selective_import_list { .selective_import_list }
 		.module_clause { .module_clause }
 		.reference_expression { .reference_expression }
 		.generic_parameters { .generic_parameters }

src/analyzer/psi/StubIndex.v

@@ -268,9 +268,11 @@ pub fn (s &StubIndex) get_all_sink_depends_on(module_fqn string) []StubIndexSink
 		if sink.kind != .workspace {
 			continue
 		}
-
-		if module_fqn in sink.imported_modules {
-			sinks << sink
+		for imported in sink.imported_modules {
+			if s.find_real_module_fqn(imported) == module_fqn {
+				sinks << sink
+				break
+			}
 		}
 	}
 	return sinks
@@ -377,3 +379,29 @@ mut:
 	stubs []&StubBase
 	psis  []PsiElement
 }
+
+pub fn (s &StubIndex) find_real_module_fqn(name string) string {
+	workspace_idx := int(StubIndexLocationKind.workspace)
+	workspace_modules := s.all_elements_by_modules[workspace_idx]
+
+	if name in workspace_modules {
+		return name
+	}
+
+	suffix := '.' + name
+	for mod_fqn, _ in workspace_modules {
+		if mod_fqn.ends_with(suffix) {
+			return mod_fqn
+		}
+	}
+
+	$for kind in StubIndexLocationKind.values {
+		if kind.value != StubIndexLocationKind.workspace {
+			modules := s.all_elements_by_modules[kind.value]
+			if name in modules {
+				return name
+			}
+		}
+	}
+	return name
+}

src/analyzer/psi/StubbedElementTypeImpl.v

@@ -52,6 +52,7 @@ pub enum StubType as u8 {
 	import_path
 	import_name
 	import_alias
+	selective_import_list
 	module_clause
 	reference_expression
 	generic_parameters
@@ -108,6 +109,8 @@ pub fn node_type_to_stub_type(typ bindings.NodeType) StubType {
 		.import_path { .import_path }
 		.import_name { .import_name }
 		.import_alias { .import_alias }
+		.selective_import_list { .selective_import_list }
+		.element_list { .selective_import_list }
 		.module_clause { .module_clause }
 		.reference_expression { .reference_expression }
 		.generic_parameters { .generic_parameters }
@@ -132,7 +135,7 @@ pub fn (_ &StubbedElementType) index_stub(stub &StubBase, mut sink IndexSink) {
 
 	if stub.stub_type == .function_declaration {
 		name := stub.name()
-		if name == 'main' || name.starts_with('test_') {
+		if name.starts_with('test_') {
 			return
 		}
 
@@ -346,6 +349,11 @@ pub fn (_ &StubbedElementType) create_psi(stub &StubBase) ?PsiElement {
 			PsiElementImpl: base_psi
 		}
 	}
+	if stub_type == .selective_import_list {
+		return SelectiveImportList{
+			PsiElementImpl: base_psi
+		}
+	}
 	if stub_type == .module_clause {
 		return ModuleClause{
 			PsiElementImpl: base_psi
@@ -571,6 +579,7 @@ pub fn (s &StubbedElementType) create_stub(psi PsiElement, parent_stub &StubBase
 			include_text: psi.node().type_name !in [
 				.import_list,
 				.import_declaration,
+				.selective_import_list,
 			]
 		)
 	}

src/analyzer/psi/TypeAliasDeclaration.v

@@ -7,7 +7,13 @@ pub struct TypeAliasDeclaration {
 }
 
 pub fn (a &TypeAliasDeclaration) get_type() types.Type {
-	return types.new_interface_type(a.name(), a.module_name())
+	types_list := a.types()
+	inner_type := if types_list.len > 0 {
+		convert_type(types_list[0])
+	} else {
+		types.Type(types.unknown_type)
+	}
+	return types.new_alias_type(a.name(), a.module_name(), inner_type)
 }
 
 pub fn (a &TypeAliasDeclaration) generic_parameters() ?&GenericParameters {

src/analyzer/psi/TypeInferer.v

@@ -340,6 +340,10 @@ pub fn (t &TypeInferer) infer_embedded_definition_type(element EmbeddedDefinitio
 }
 
 pub fn (t &TypeInferer) infer_reference_expression_type(element ReferenceExpression) types.Type {
+	if resolved := element.resolve() {
+		return t.infer_type(resolved)
+	}
+
 	if element.text_matches('it') {
 		call := get_it_call(element) or { return types.unknown_type }
 		caller_type := call.caller_type()
@@ -349,8 +353,7 @@ pub fn (t &TypeInferer) infer_reference_expression_type(element ReferenceExpress
 		return types.unknown_type
 	}
 
-	resolved := element.resolve() or { return types.unknown_type }
-	return t.infer_type(resolved)
+	return types.unknown_type
 }
 
 pub fn (t &TypeInferer) infer_if_expression_type(element IfExpression) types.Type {