Merge pull request #1226 from BradLewis/fix/implicit-selector-completion-assignment-proc-arg

Fix implicit selector completions for proc call args within assignments

3 files changed, 133 insertions, 82 deletions

diff --git a/src/server/completion.odin b/src/server/completion.odin
index 7c30f78..282b7d7 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -1206,35 +1206,6 @@ get_implicit_completion :: proc(
 		}
 	}
 
-	if position_context.assign != nil &&
-	   position_context.assign.lhs != nil &&
-	   len(position_context.assign.lhs) == 1 &&
-	   is_bitset_assignment_operator(position_context.assign.op.text) {
-		//bitsets
-		if symbol, ok := resolve_type_expression(ast_context, position_context.assign.lhs[0]); ok {
-			set_ast_package_set_scoped(ast_context, symbol.pkg)
-			if value, ok := unwrap_bitset(ast_context, symbol); ok {
-				for name in value.names {
-					if position_context.comp_lit != nil && field_exists_in_comp_lit(position_context.comp_lit, name) {
-						continue
-					}
-
-					item := CompletionItem {
-						label  = name,
-						kind   = .EnumMember,
-						detail = name,
-					}
-
-					append(results, CompletionResult{completion_item = item})
-				}
-
-				return is_incomplete
-			}
-		}
-
-		reset_ast_context(ast_context)
-	}
-
 	if position_context.comp_lit != nil &&
 	   position_context.parent_binary != nil &&
 	   is_bitset_binary_operator(position_context.binary.op.text) {
@@ -1332,53 +1303,6 @@ get_implicit_completion :: proc(
 		}
 	}
 
-	if position_context.assign != nil && position_context.assign.rhs != nil && position_context.assign.lhs != nil {
-		rhs_index: int
-
-		for elem in position_context.assign.rhs {
-			if position_in_node(elem, position_context.position) {
-				break
-			} else {
-				//procedures are the only types that can return more than one value
-				if symbol, ok := resolve_type_expression(ast_context, elem); ok {
-					if procedure, ok := symbol.value.(SymbolProcedureValue); ok {
-						if procedure.return_types == nil {
-							return is_incomplete
-						}
-
-						rhs_index += len(procedure.return_types)
-					} else {
-						rhs_index += 1
-					}
-				}
-			}
-		}
-
-		if len(position_context.assign.lhs) > rhs_index {
-			if enum_value, unwrapped_super_enum, ok := unwrap_enum(
-				ast_context,
-				position_context.assign.lhs[rhs_index],
-			); ok {
-				for name in enum_value.names {
-					item := CompletionItem {
-						label  = name,
-						kind   = .EnumMember,
-						detail = name,
-					}
-					if unwrapped_super_enum {
-						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
-					}
-
-					append(results, CompletionResult{completion_item = item})
-				}
-
-				return is_incomplete
-			}
-		}
-
-		reset_ast_context(ast_context)
-	}
-
 	if position_context.returns != nil && position_context.function != nil {
 		return_index: int
 
@@ -1551,6 +1475,83 @@ get_implicit_completion :: proc(
 
 		reset_ast_context(ast_context)
 	}
+
+	if position_context.assign != nil &&
+	   position_context.assign.lhs != nil &&
+	   len(position_context.assign.lhs) == 1 &&
+	   is_bitset_assignment_operator(position_context.assign.op.text) {
+		//bitsets
+		if symbol, ok := resolve_type_expression(ast_context, position_context.assign.lhs[0]); ok {
+			set_ast_package_set_scoped(ast_context, symbol.pkg)
+			if value, ok := unwrap_bitset(ast_context, symbol); ok {
+				for name in value.names {
+					if position_context.comp_lit != nil && field_exists_in_comp_lit(position_context.comp_lit, name) {
+						continue
+					}
+
+					item := CompletionItem {
+						label  = name,
+						kind   = .EnumMember,
+						detail = name,
+					}
+
+					append(results, CompletionResult{completion_item = item})
+				}
+
+				return is_incomplete
+			}
+		}
+
+		reset_ast_context(ast_context)
+	}
+
+	if position_context.assign != nil && position_context.assign.rhs != nil && position_context.assign.lhs != nil {
+		rhs_index: int
+
+		for elem in position_context.assign.rhs {
+			if position_in_node(elem, position_context.position) {
+				break
+			} else {
+				//procedures are the only types that can return more than one value
+				if symbol, ok := resolve_type_expression(ast_context, elem); ok {
+					if procedure, ok := symbol.value.(SymbolProcedureValue); ok {
+						if procedure.return_types == nil {
+							return is_incomplete
+						}
+
+						rhs_index += len(procedure.return_types)
+					} else {
+						rhs_index += 1
+					}
+				}
+			}
+		}
+
+		if len(position_context.assign.lhs) > rhs_index {
+			if enum_value, unwrapped_super_enum, ok := unwrap_enum(
+				ast_context,
+				position_context.assign.lhs[rhs_index],
+			); ok {
+				for name in enum_value.names {
+					item := CompletionItem {
+						label  = name,
+						kind   = .EnumMember,
+						detail = name,
+					}
+					if unwrapped_super_enum {
+						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+					}
+
+					append(results, CompletionResult{completion_item = item})
+				}
+
+				return is_incomplete
+			}
+		}
+
+		reset_ast_context(ast_context)
+	}
+
 	return is_incomplete
 }
 
diff --git a/src/testing/testing.odin b/src/testing/testing.odin
index 6a6b80d..373ef62 100644
--- a/src/testing/testing.odin
+++ b/src/testing/testing.odin
@@ -165,7 +165,13 @@ expect_signature_parameter_position :: proc(t: ^testing.T, src: ^Source, positio
 	}
 }
 
-expect_completion_labels :: proc(t: ^testing.T, src: ^Source, trigger_character: string, expect_labels: []string) {
+expect_completion_labels :: proc(
+	t: ^testing.T,
+	src: ^Source,
+	trigger_character: string,
+	expect_labels: []string,
+	expect_excluded: []string = nil,
+) {
 	setup(src)
 	defer teardown(src)
 
@@ -198,6 +204,14 @@ expect_completion_labels :: proc(t: ^testing.T, src: ^Source, trigger_character:
 			log.errorf("Expected completion detail %v, but received %v", expect_labels[i], completion_list.items)
 		}
 	}
+
+	for expect_exclude in expect_excluded {
+		for completion in completion_list.items {
+			if expect_exclude == completion.label {
+				log.errorf("Expected completion label %v to not be included", expect_exclude)
+			}
+		}
+	}
 }
 
 expect_completion_docs :: proc(
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index 4403e02..1883032 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -5164,11 +5164,6 @@ ast_completion_empty_selector_with_ident_newline :: proc(t: ^testing.T) {
 	)
 	source := test.Source {
 		main = `package test
-
-		Foo :: struct{
-			x: int,
-		}
-
 		import "my_package"
 
 		main :: proc() {
@@ -5180,3 +5175,44 @@ ast_completion_empty_selector_with_ident_newline :: proc(t: ^testing.T) {
 	}
 	test.expect_completion_docs(t, &source, "", {"my_package.Foo :: struct{}"})
 }
+
+@(test)
+ast_completion_implicit_selector_binary_expr_proc_call :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+			Foo :: enum {
+				A,
+				B,
+				C,
+			}
+
+			Bar :: enum {
+				X,
+				Y,
+			}
+
+			foo :: proc(f: Foo) -> bit_set[Bar] {
+				return {.X}
+			}
+		`,
+		},
+	)
+	source := test.Source {
+		main = `package test
+		import "my_package"
+
+		main :: proc() {
+			results: bit_set[my_package.Bar]
+
+			results |= my_package.foo(.{*})
+		}
+		`,
+		packages = packages[:],
+	}
+	test.expect_completion_labels(t, &source, "", {"A", "B", "C"}, {"X", "Y"})
+}