Correctly infer return type enum fields for completions

2 files changed, 56 insertions, 2 deletions

diff --git a/src/server/completion.odin b/src/server/completion.odin
index 01a7d32..e04963f 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -938,7 +938,7 @@ get_implicit_completion :: proc(
 
 	//infer bitset and enums based on the identifier comp_lit, i.e. a := My_Struct { my_ident = . } 
 	if position_context.comp_lit != nil {
-		if position_context.parent_comp_lit.type != nil {
+		if position_context.parent_comp_lit != nil {
 			field_name: string
 
 			if position_context.field_value != nil {
@@ -949,7 +949,7 @@ get_implicit_completion :: proc(
 				}
 			}
 
-			if symbol, ok := resolve_type_expression(ast_context, position_context.parent_comp_lit.type); ok {
+			if symbol, ok := resolve_comp_literal(ast_context, position_context); ok {
 				if comp_symbol, comp_lit, ok := resolve_type_comp_literal(
 					ast_context,
 					position_context,
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index bae9476..21f10a7 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -3973,3 +3973,57 @@ ast_completion_bitset_named_proc_arg_should_remove_already_used :: proc(t: ^test
 	}
 	test.expect_completion_details(t, &source, "", {"B"}, {"A"})
 }
+
+@(test)
+ast_completion_return_comp_lit_enum :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foo: Foo,
+		}
+
+		foo :: proc() -> Bar {
+			return {
+				foo = .{*}
+			}
+		}
+		`,
+	}
+	test.expect_completion_details(t, &source, "", {"A", "B"})
+}
+
+@(test)
+ast_completion_return_nested_comp_lit_enum :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foo: Foo,
+		}
+
+		Bazz :: struct {
+			bar: Bar,
+		}
+
+		foo :: proc() -> Bazz {
+			return {
+				bar = {
+					foo = .{*}
+				}
+			}
+		}
+		`,
+	}
+	test.expect_completion_details(t, &source, "", {"A", "B"})
+}