Fix issues with union switch and also now using the renamed package name

3 files changed, 46 insertions, 29 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index a114b92..e6de07e 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1,19 +1,19 @@
 package server
 
-import "core:odin/parser"
-import "core:odin/ast"
-import "core:odin/tokenizer"
 import "core:fmt"
 import "core:log"
-import "core:strings"
-import path "core:path/slashpath"
 import "core:mem"
-import "core:strconv"
+import "core:odin/ast"
+import "core:odin/parser"
+import "core:odin/tokenizer"
 import "core:path/filepath"
-import "core:sort"
+import path "core:path/slashpath"
+import "core:reflect"
 import "core:slice"
+import "core:sort"
+import "core:strconv"
+import "core:strings"
 import "core:unicode/utf8"
-import "core:reflect"
 
 import "shared:common"
 
@@ -2750,6 +2750,23 @@ get_using_packages :: proc(ast_context: ^AstContext) -> []string {
 	return usings
 }
 
+get_symbol_pkg_name :: proc(
+	ast_context: ^AstContext,
+	symbol: Symbol,
+) -> string {
+
+	name := path.base(symbol.pkg, false, context.temp_allocator)
+
+	for imp in ast_context.imports {
+		if imp.base_original == name {
+			return imp.base
+		}
+	}
+
+	return name
+}
+
+
 make_symbol_procedure_from_ast :: proc(
 	ast_context: ^AstContext,
 	n: ^ast.Node,
@@ -4439,6 +4456,7 @@ unwrap_union :: proc(
 	bool,
 ) {
 	if union_symbol, ok := resolve_type_expression(ast_context, node); ok {
+		ast_context.current_package = union_symbol.pkg
 		if union_value, ok := union_symbol.value.(SymbolUnionValue); ok {
 			return union_value, true
 		}
diff --git a/src/server/completion.odin b/src/server/completion.odin
index 044483a..00e7b0f 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -1,18 +1,18 @@
 package server
 
-import "core:odin/parser"
-import "core:odin/ast"
-import "core:odin/tokenizer"
 import "core:fmt"
 import "core:log"
-import "core:strings"
-import path "core:path/slashpath"
 import "core:mem"
-import "core:strconv"
+import "core:odin/ast"
+import "core:odin/parser"
+import "core:odin/tokenizer"
+import "core:os"
 import "core:path/filepath"
-import "core:sort"
+import path "core:path/slashpath"
 import "core:slice"
-import "core:os"
+import "core:sort"
+import "core:strconv"
+import "core:strings"
 
 
 import "shared:common"
@@ -464,7 +464,7 @@ get_selector_completion :: proc(
 
 		for type in v.types {
 			if symbol, ok := resolve_type_expression(ast_context, type); ok {
-				base := path.base(symbol.pkg, false, context.temp_allocator)
+				base := get_symbol_pkg_name(ast_context, symbol)
 
 				item := CompletionItem {
 					kind          = .EnumMember,
@@ -496,7 +496,7 @@ get_selector_completion :: proc(
 							symbol.pointers,
 							context.temp_allocator,
 						),
-						path.base(symbol.pkg, false, context.temp_allocator),
+						get_symbol_pkg_name(ast_context, symbol),
 						common.node_to_string(type, true),
 					)
 				}
@@ -1516,6 +1516,9 @@ get_type_switch_completion :: proc(
 				for name in case_clause.list {
 					if ident, ok := name.derived.(^ast.Ident); ok {
 						used_unions[ident.name] = true
+					} else if selector, ok := name.derived.(^ast.Selector_Expr);
+					   ok {
+						used_unions[selector.field.name] = true
 					}
 				}
 			}
@@ -1533,9 +1536,6 @@ get_type_switch_completion :: proc(
 					union_value.types[i],
 				); ok {
 					name := symbol.name
-					if name in used_unions {
-						continue
-					}
 
 					item := CompletionItem {
 						kind = .EnumMember,
@@ -1559,11 +1559,7 @@ get_type_switch_completion :: proc(
 								symbol.pointers,
 								context.temp_allocator,
 							),
-							path.base(
-								symbol.pkg,
-								false,
-								context.temp_allocator,
-							),
+							get_symbol_pkg_name(ast_context, symbol),
 							name,
 						)
 					}
diff --git a/src/server/documents.odin b/src/server/documents.odin
index bbde1ca..5553fae 100644
--- a/src/server/documents.odin
+++ b/src/server/documents.odin
@@ -24,9 +24,10 @@ ParserError :: struct {
 }
 
 Package :: struct {
-	name:     string, //the entire absolute path to the directory
-	base:     string,
-	original: string,
+	name:          string, //the entire absolute path to the directory
+	base:          string,
+	base_original: string,
+	original:      string,
 }
 
 Document :: struct {
@@ -486,6 +487,7 @@ parse_imports :: proc(document: ^Document, config: ^common.Config) {
 
 			if imp.name.text != "" {
 				import_.base = imp.name.text
+				import_.base_original = path.base(import_.name, false)
 			} else {
 				import_.base = path.base(import_.name, false)
 			}
@@ -510,6 +512,7 @@ parse_imports :: proc(document: ^Document, config: ^common.Config) {
 
 			if imp.name.text != "" {
 				import_.base = imp.name.text
+				import_.base_original = path.base(import_.name, false)
 			} else {
 				import_.base = path.base(import_.name, false)
 			}