Merge pull request #1279 from BradLewis/switch-code-actionnightly

Only append the missing cases, fix issues with parapoly unions

6 files changed, 128 insertions, 119 deletions

diff --git a/src/server/action.odin b/src/server/action.odin
index 64516a3..11e8f86 100644
--- a/src/server/action.odin
+++ b/src/server/action.odin
@@ -73,7 +73,6 @@ get_code_actions :: proc(document: ^Document, range: common.Range, config: ^comm
 
 	if position_context.switch_stmt != nil || position_context.switch_type_stmt != nil {
 		add_populate_switch_cases_action(
-			document,
 			&ast_context,
 			&position_context,
 			strings.clone(document.uri.uri),
diff --git a/src/server/action_populate_switch_cases.odin b/src/server/action_populate_switch_cases.odin
index f038e7c..1e60f15 100644
--- a/src/server/action_populate_switch_cases.odin
+++ b/src/server/action_populate_switch_cases.odin
@@ -4,6 +4,7 @@ package server
 
 import "core:fmt"
 import "core:odin/ast"
+import "core:odin/tokenizer"
 import "core:strings"
 
 import "src:common"
@@ -17,6 +18,7 @@ get_line_start_offset :: proc(src: string, offset: int) -> int {
 	}
 	return line_start
 }
+
 get_block_original_text :: proc(block: []^ast.Stmt, document_text: string) -> string {
 	if len(block) == 0 {
 		return ""
@@ -26,54 +28,57 @@ get_block_original_text :: proc(block: []^ast.Stmt, document_text: string) -> st
 	return string(document_text[start:end])
 }
 
-SwitchCaseInfo :: struct {
-	names:            []string,
-	body_indentation: string,
-	body:             string,
-}
 SwitchBlockInfo :: struct {
-	existing_cases:         []SwitchCaseInfo,
-	all_covered_case_names: []string,
-	all_case_names:         []string,
-	switch_indentation:     string,
-	is_enum:                bool,
+	names:              []string,
+	existing_cases:     map[string]struct{},
+	switch_indentation: string,
+	is_enum:            bool,
+	pos:                tokenizer.Pos,
 }
+
 get_switch_cases_info :: proc(
-	document: ^Document,
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
 ) -> (
-	info: SwitchBlockInfo,
-	ok: bool,
+	SwitchBlockInfo,
+	bool,
 ) {
-	if (position_context.switch_stmt == nil && position_context.switch_type_stmt == nil) ||
-	   (position_context.switch_stmt != nil && position_context.switch_stmt.cond == nil) {
+	if position_context.switch_stmt == nil && position_context.switch_type_stmt == nil {
 		return {}, false
 	}
+
+	if position_context.switch_stmt != nil && position_context.switch_stmt.cond == nil {
+		return {}, false
+	}
+
 	switch_block: ^ast.Block_Stmt
 	found_switch_block: bool
 	is_enum: bool
+	pos: tokenizer.Pos
 	if position_context.switch_stmt != nil {
 		switch_block, found_switch_block = position_context.switch_stmt.body.derived.(^ast.Block_Stmt)
 		is_enum = true
+		pos = position_context.switch_stmt.pos
 	}
+
 	if !found_switch_block && position_context.switch_type_stmt != nil {
 		switch_block, found_switch_block = position_context.switch_type_stmt.body.derived.(^ast.Block_Stmt)
+		pos = position_context.switch_type_stmt.pos
 	}
+
 	if !found_switch_block {
 		return {}, false
 	}
-	switch_indentation := get_line_indentation(string(document.text), switch_block.pos.offset)
-	existing_cases_in_order := make([dynamic]SwitchCaseInfo, context.temp_allocator)
-	all_covered_names := make([dynamic]string, context.temp_allocator)
+	switch_indentation := get_line_indentation(ast_context.file.src, switch_block.pos.offset)
+	existing_cases := make(map[string]struct{}, context.temp_allocator)
+
+
 	for stmt in switch_block.stmts {
 		if case_clause, ok := stmt.derived.(^ast.Case_Clause); ok {
-			case_names := make([dynamic]string, context.temp_allocator)
 			for clause in case_clause.list {
 				if is_enum {
 					if name, ok := get_used_switch_name(clause); ok && name != "" {
-						append(&case_names, name)
-						append(&all_covered_names, name)
+						existing_cases[name] = {}
 					}
 				} else {
 					reset_ast_context(ast_context)
@@ -85,19 +90,12 @@ get_switch_cases_info :: proc(
 							name = get_signature(ast_context, symbol)
 						}
 						if name != "" {
-							append(&case_names, name)
-							append(&all_covered_names, name)
+							existing_cases[name] = {}
 						}
 					}
 				}
 			}
-			if len(case_names) > 0 {
-				case_info := SwitchCaseInfo {
-					names = case_names[:],
-					body  = get_block_original_text(case_clause.body, string(document.text)),
-				}
-				append(&existing_cases_in_order, case_info)
-			}
+			pos = case_clause.stmt_base.end
 		}
 	}
 	if is_enum {
@@ -106,47 +104,47 @@ get_switch_cases_info :: proc(
 			return {}, false
 		}
 		return SwitchBlockInfo {
-				existing_cases = existing_cases_in_order[:],
-				all_covered_case_names = all_covered_names[:],
-				all_case_names = enum_value.names,
+				names = enum_value.names,
+				existing_cases = existing_cases,
 				switch_indentation = switch_indentation,
 				is_enum = !was_super_enum,
+				pos = pos,
 			},
 			true
-	} else {
-		st := position_context.switch_type_stmt
-		if st == nil {
-			return {}, false
-		}
+	}
+
+	st := position_context.switch_type_stmt
+	if st == nil {
+		return {}, false
+	}
+	reset_ast_context(ast_context)
+	union_value, unwrap_ok := unwrap_union(ast_context, st.tag.derived.(^ast.Assign_Stmt).rhs[0])
+	if !unwrap_ok {
+		return {}, false
+	}
+	all_case_names := make([]string, len(union_value.types), context.temp_allocator)
+	for t, i in union_value.types {
 		reset_ast_context(ast_context)
-		union_value, unwrap_ok := unwrap_union(ast_context, st.tag.derived.(^ast.Assign_Stmt).rhs[0])
-		if !unwrap_ok {
-			return {}, false
-		}
-		all_case_names := make([]string, len(union_value.types), context.temp_allocator)
-		for t, i in union_value.types {
-			reset_ast_context(ast_context)
-			if symbol, ok := resolve_type_expression(ast_context, t); ok {
-				case_name := get_qualified_union_case_name(&symbol, ast_context, position_context)
-				//TODO: this is wrong for anonymous enums and structs, where the name field is "enum" or "struct" respectively but we want to use the full signature
-				//we also can't use the signature all the time because type aliases need to use specifically the alias name here and not the signature
-				if case_name == "" {
-					case_name = get_signature(ast_context, symbol)
-				}
-				all_case_names[i] = case_name
-			} else {
-				all_case_names[i] = "invalid type expression"
+		if symbol, ok := resolve_type_expression(ast_context, t); ok {
+			case_name := get_qualified_union_case_name(&symbol, ast_context, position_context)
+			//TODO: this is wrong for anonymous enums and structs, where the name field is "enum" or "struct" respectively but we want to use the full signature
+			//we also can't use the signature all the time because type aliases need to use specifically the alias name here and not the signature
+			if case_name == "" {
+				case_name = get_signature(ast_context, symbol)
 			}
+			all_case_names[i] = case_name
+		} else {
+			all_case_names[i] = "invalid type expression"
 		}
-		return SwitchBlockInfo {
-				existing_cases = existing_cases_in_order[:],
-				all_covered_case_names = all_covered_names[:],
-				all_case_names = all_case_names,
-				switch_indentation = switch_indentation,
-				is_enum = false,
-			},
-			true
 	}
+	return SwitchBlockInfo {
+			names = all_case_names,
+			existing_cases = existing_cases,
+			switch_indentation = switch_indentation,
+			is_enum = false,
+			pos = pos,
+		},
+		true
 }
 
 create_populate_switch_cases_edit :: proc(
@@ -160,65 +158,39 @@ create_populate_switch_cases_edit :: proc(
 	if position_context.switch_stmt == nil && position_context.switch_type_stmt == nil {
 		return {}, false
 	}
-	//entirety of the switch block
-	range: common.Range
-	if info.is_enum {
-		range = common.get_token_range(position_context.switch_stmt.body.stmt_base, position_context.file.src)
-	} else {
-		range = common.get_token_range(position_context.switch_type_stmt.body.stmt_base, position_context.file.src)
+
+	pos := info.pos
+	pos.line += 1
+	pos.column = 1
+
+	position := common.token_pos_to_position(pos, position_context.file.src)
+
+	range := common.Range {
+		start = position,
+		end   = position,
 	}
+
 	replacement_builder := strings.builder_make()
 	dot := info.is_enum ? "." : ""
 	b := &replacement_builder
-	fmt.sbprintln(b, "{")
-	for case_info in info.existing_cases {
-		fmt.sbprint(b, info.switch_indentation, "case ", sep = "")
-		for name, i in case_info.names {
-			fmt.sbprint(b, dot, name, sep = "")
-			if i != len(case_info.names) - 1 {
-				fmt.sbprint(b, ", ", sep = "")
-			}
-		}
-		fmt.sbprintln(b, ":", sep = "")
-		case_body := case_info.body
-		if case_body != "" {
-			fmt.sbprintln(b, case_info.body)
-		}
-	}
-	existing_case_names := map[string]struct{}{}
-	for name in info.all_covered_case_names {
-		existing_case_names[name] = {}
-	}
-	for name in info.all_case_names {
-		if name in existing_case_names {continue} 	//covered by prev loop
+	for name in info.names {
+		if name in info.existing_cases {continue}
 		fmt.sbprintln(b, info.switch_indentation, "case ", dot, name, ":", sep = "")
 	}
-	fmt.sbprint(b, info.switch_indentation, "}", sep = "")
 	return TextEdit{range = range, newText = strings.to_string(replacement_builder)}, true
 }
+
 @(private = "package")
 add_populate_switch_cases_action :: proc(
-	document: ^Document,
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
 	uri: string,
 	actions: ^[dynamic]CodeAction,
 ) {
-	info, ok := get_switch_cases_info(document, ast_context, position_context)
+	info, ok := get_switch_cases_info(ast_context, position_context)
 	if !ok {return}
-	all_cases_covered := true
-	{
-		existing_case_names := map[string]struct{}{}
-		for name in info.all_covered_case_names {
-			existing_case_names[name] = {}
-		}
-		for name in info.all_case_names {
-			if name not_in existing_case_names {
-				all_cases_covered = false
-			}
-		}
-	}
-	if all_cases_covered {return} 	//action not needed
+
+	if len(info.existing_cases) == len(info.names) {return} 	//action not needed
 	edit, edit_ok := create_populate_switch_cases_edit(position_context, info)
 	if !edit_ok {return}
 	textEdits := make([dynamic]TextEdit, context.temp_allocator)
diff --git a/src/server/completion.odin b/src/server/completion.odin
index a7b1720..24b86a1 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -1957,18 +1957,37 @@ get_qualified_union_case_name :: proc(
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
 ) -> string {
-	if symbol.pkg == ast_context.document_package {
-		return fmt.aprintf("%v%v", repeat("^", symbol.pointers, context.temp_allocator), symbol.name)
-	} else {
-		return fmt.aprintf(
-			"%v%v.%v",
-			repeat("^", symbol.pointers, context.temp_allocator),
-			get_symbol_pkg_name(ast_context, symbol),
-			symbol.name,
-		)
+	sb := strings.builder_make(context.temp_allocator)
+	pointer_prefix := repeat("^", symbol.pointers, ast_context.allocator)
+	strings.write_string(&sb, pointer_prefix)
+	if symbol.pkg != ast_context.document_package {
+		strings.write_string(&sb, get_symbol_pkg_name(ast_context, symbol))
 	}
+	strings.write_string(&sb, symbol.name)
+	#partial switch v in symbol.value {
+	case SymbolUnionValue:
+		write_poly_names(&sb, v.poly_names)
+	case SymbolStructValue:
+		write_poly_names(&sb, v.poly_names)
+	}
+
+	return strings.to_string(sb)
 }
 
+write_poly_names :: proc(sb: ^strings.Builder, poly_names: []string) {
+	if len(poly_names) > 0 {
+		strings.write_string(sb, "(")
+		for name, i in poly_names {
+			strings.write_string(sb, name)
+			if i != len(poly_names) - 1 {
+				strings.write_string(sb, ", ")
+			}
+		}
+		strings.write_string(sb, ")")
+	}
+}
+
+
 get_type_switch_completion :: proc(
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
diff --git a/src/server/documents.odin b/src/server/documents.odin
index ed1fb52..b9a50dd 100644
--- a/src/server/documents.odin
+++ b/src/server/documents.odin
@@ -4,12 +4,10 @@ import "base:intrinsics"
 
 import "core:fmt"
 import "core:log"
-import "core:mem"
 import "core:mem/virtual"
 import "core:odin/ast"
 import "core:odin/parser"
 import "core:odin/tokenizer"
-import "core:os"
 import "core:path/filepath"
 import path "core:path/slashpath"
 import "core:strings"
diff --git a/src/server/generics.odin b/src/server/generics.odin
index 8aa111f..82f71c1 100644
--- a/src/server/generics.odin
+++ b/src/server/generics.odin
@@ -894,7 +894,7 @@ resolve_poly_union :: proc(ast_context: ^AstContext, poly_params: ^ast.Field_Lis
 			for arg, i in call_expr.args {
 				if ident, ok := arg.derived.(^ast.Ident); ok {
 					if expr, ok := poly_map[ident.name]; ok {
-						symbol_value.types[i] = expr
+						call_expr.args[i] = expr
 					}
 				}
 			}
diff --git a/tests/hover_test.odin b/tests/hover_test.odin
index 02cde4e..df65f3a 100644
--- a/tests/hover_test.odin
+++ b/tests/hover_test.odin
@@ -6034,6 +6034,27 @@ ast_hover_constant_unary_expr :: proc(t: ^testing.T) {
 	}
 	test.expect_hover(t, &source, "test.FOO :: ~u32(0)")
 }
+
+@(test)
+ast_hover_union_multiple_poly :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+		Foo :: struct($T: typeid) {}
+		Bar :: struct{}
+
+		Bazz :: union($T: typeid) {
+			Foo(T),
+			Bar,
+		}
+
+		main :: proc() {
+			T :: distinct int
+			bazz: Ba{*}zz(T)
+		}
+		`,
+	}
+	test.expect_hover(t, &source, "test.Bazz :: union(T) {\n\tFoo(T),\n\tBar,\n}")
+}
 /*
 
 Waiting for odin fix