Merge pull request #822 from BradLewis/feat/add-docs-for-enums-unions

Adds hover documentation for enums and unions

6 files changed, 257 insertions, 38 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 581a326..ab2efd7 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1274,11 +1274,8 @@ resolve_type_assertion_expr :: proc(ast_context: ^AstContext, v: ^ast.Type_Asser
 						return {}, false
 					}
 
-					if ok := internal_resolve_type_expression(
-						ast_context,
-						proc_value.return_types[0].type,
-						&symbol,
-					); ok {
+					if ok := internal_resolve_type_expression(ast_context, proc_value.return_types[0].type, &symbol);
+					   ok {
 						if union_value, ok := symbol.value.(SymbolUnionValue); ok {
 							if len(union_value.types) != 1 {
 								return {}, false
@@ -1711,8 +1708,7 @@ resolve_local_identifier :: proc(ast_context: ^AstContext, node: ast.Ident, loca
 
 			if !ok && !ast_context.overloading {
 				return_symbol, ok =
-					make_symbol_procedure_from_ast(ast_context, local.rhs, v.type^, node, {}, false, v.inlining),
-					true
+					make_symbol_procedure_from_ast(ast_context, local.rhs, v.type^, node, {}, false, v.inlining), true
 			}
 		} else {
 			return_symbol, ok =
@@ -3129,9 +3125,13 @@ make_symbol_union_from_ast :: proc(
 		}
 	}
 
+	docs, comments := get_field_docs_and_comments(ast_context.file, v.variants, ast_context.allocator)
+
 	symbol.value = SymbolUnionValue {
-		types = types[:],
-		poly  = v.poly_params,
+		types    = types[:],
+		poly     = v.poly_params,
+		docs     = docs[:],
+		comments = comments[:],
 	}
 
 	if v.poly_params != nil {
@@ -3172,11 +3172,15 @@ make_symbol_enum_from_ast :: proc(
 		append(&values, value)
 	}
 
+	docs, comments := get_field_docs_and_comments(ast_context.file, v.fields, ast_context.allocator)
+
 	symbol.value = SymbolEnumValue {
 		names     = names[:],
 		ranges    = ranges[:],
 		base_type = v.base_type,
 		values    = values[:],
+		docs      = docs[:],
+		comments  = comments[:],
 	}
 
 	return symbol
diff --git a/src/server/ast.odin b/src/server/ast.odin
index 71e17f6..326fd2d 100644
--- a/src/server/ast.odin
+++ b/src/server/ast.odin
@@ -305,11 +305,12 @@ collect_value_decl :: proc(
 	if !is_value_decl {
 		return
 	}
+	comment, _ := get_file_comment(file, value_decl.pos.line)
 
 	global_expr := GlobalExpr {
 		mutable    = value_decl.is_mutable,
 		docs       = value_decl.docs,
-		comment    = get_file_comment(file, value_decl.pos.line),
+		comment    = comment,
 		attributes = value_decl.attributes[:],
 		private    = file_tags.private,
 	}
@@ -1214,7 +1215,7 @@ repeat :: proc(value: string, count: int, allocator := context.allocator) -> str
 	return strings.repeat(value, count, allocator)
 }
 
-construct_struct_field_docs :: proc(file: ast.File, v: ^ast.Struct_Type) {
+construct_struct_field_docs :: proc(file: ast.File, v: ^ast.Struct_Type, allocator := context.temp_allocator) {
 	for field, i in v.fields.list {
 		// There is currently a bug in the odin parser where it adds line comments for a field to the
 		// docs of the following field, we address this problem here.
@@ -1225,7 +1226,7 @@ construct_struct_field_docs :: proc(file: ast.File, v: ^ast.Struct_Type) {
 		// skipped on the last line) eg
 		// Foo :: struct {
 		//     foo: int // my int <-- skipped as no ',' after 'int'
-	    // }
+		// }
 
 		// remove any unwanted docs
 		if i != len(v.fields.list) - 1 {
@@ -1239,11 +1240,7 @@ construct_struct_field_docs :: proc(file: ast.File, v: ^ast.Struct_Type) {
 						field.comment.list = list[:1]
 					}
 					if len(list) > 1 {
-						next_field.docs = ast.new(
-							ast.Comment_Group,
-							list[1].pos,
-							parser.end_pos(list[len(list) - 2]),
-						)
+						next_field.docs = ast.new(ast.Comment_Group, list[1].pos, parser.end_pos(list[len(list) - 2]))
 						next_field.docs.list = list[1:]
 					} else {
 						next_field.docs = nil
@@ -1252,31 +1249,32 @@ construct_struct_field_docs :: proc(file: ast.File, v: ^ast.Struct_Type) {
 			}
 		} else if field.comment == nil {
 			// We need to check the file to see if it contains a line comment as it might be skipped
-			field.comment = get_file_comment(file, field.pos.line)
+			field.comment, _ = get_file_comment(file, field.pos.line, allocator = allocator)
 		}
 	}
 }
 
-construct_bit_field_field_docs :: proc(file: ast.File, v: ^ast.Bit_Field_Type) {
+construct_bit_field_field_docs :: proc(file: ast.File, v: ^ast.Bit_Field_Type, allocator := context.temp_allocator) {
 	for field, i in v.fields {
 		// There is currently a bug in the odin parser where it adds line comments for a field to the
 		// docs of the following field, we address this problem here.
 		// see https://github.com/odin-lang/Odin/issues/5353
-			// We check if the comment is at the start of the next field
+		// We check if the comment is at the start of the next field
 		if i != len(v.fields) - 1 {
 			next_field := v.fields[i + 1]
 			if next_field.docs != nil && len(next_field.docs.list) > 0 {
 				list := next_field.docs.list
 				if list[0].pos.line == field.pos.line {
 					if field.comments == nil {
-						field.comments = ast.new(ast.Comment_Group, list[0].pos, parser.end_pos(list[0]))
+						field.comments = new_type(ast.Comment_Group, list[0].pos, parser.end_pos(list[0]), allocator)
 						field.comments.list = list[:1]
 					}
 					if len(list) > 1 {
-						next_field.docs = ast.new(
+						next_field.docs = new_type(
 							ast.Comment_Group,
 							list[1].pos,
 							parser.end_pos(list[len(list) - 2]),
+							allocator,
 						)
 						next_field.docs.list = list[1:]
 					} else {
@@ -1286,26 +1284,120 @@ construct_bit_field_field_docs :: proc(file: ast.File, v: ^ast.Bit_Field_Type) {
 			}
 		} else if field.comments == nil {
 			// We need to check the file to see if it contains a line comment as there is no next field
-			field.comments = get_file_comment(file, field.pos.line)
+			field.comments, _ = get_file_comment(file, field.pos.line, allocator = allocator)
 		}
 	}
 }
 
 // Retrives the comment group from the specified line of the file
-get_file_comment :: proc(file: ast.File, line: int) -> ^ast.Comment_Group {
+// Returns the index where the comment was found
+get_file_comment :: proc(
+	file: ast.File,
+	line: int,
+	start_index := 0,
+	allocator := context.temp_allocator,
+) -> (
+	^ast.Comment_Group,
+	int,
+) {
 	// TODO: linear scan might be a bit slow for files with lots of comments?
-	for c in file.comments {
+	for i := start_index; i < len(file.comments); i += 1 {
+		c := file.comments[i]
 		if c.pos.line == line {
 			for item, j in c.list {
-				comment := ast.new(ast.Comment_Group, item.pos, parser.end_pos(item))
+				comment := new_type(ast.Comment_Group, item.pos, parser.end_pos(item), allocator)
 				if j == len(c.list) - 1 {
 					comment.list = c.list[j:]
 				} else {
 					comment.list = c.list[j:j + 1]
 				}
-				return comment
+				return comment, i
 			}
 		}
 	}
-	return nil
+	return nil, -1
+}
+
+// Retrieves the comment group that ends on the specified line of the file
+// If start_line is specified, it will only add the docs that on that line and beyond
+get_file_doc :: proc(
+	file: ast.File,
+	end_line: int,
+	start_line := -1,
+	start_index := 0,
+	allocator := context.temp_allocator,
+) -> (
+	^ast.Comment_Group,
+	int,
+) {
+	for i := start_index; i < len(file.comments); i += 1 {
+		c := file.comments[i]
+		if c.end.line == end_line {
+			docs := new_type(ast.Comment_Group, c.pos, c.end, allocator)
+			if start_line != -1 {
+				for item, j in c.list {
+					if item.pos.line >= start_line {
+						docs.list = c.list[j:]
+						return docs, i
+					}
+				}
+			}
+			docs.list = c.list
+			return docs, i
+		}
+	}
+	return nil, -1
+}
+
+// Returns the docs and comments for a list of field types
+//
+// We use this as the odin parser does not include comments and docs on enum and union fields
+get_field_docs_and_comments :: proc(
+	file: ast.File,
+	fields: []^ast.Expr,
+	allocator := context.temp_allocator,
+) -> (
+	[dynamic]^ast.Comment_Group,
+	[dynamic]^ast.Comment_Group,
+) {
+	docs := make([dynamic]^ast.Comment_Group, allocator)
+	comments := make([dynamic]^ast.Comment_Group, allocator)
+	prev_line := -1
+	last_comment := -1
+	last_doc := -1
+	for n, i in fields {
+		doc: ^ast.Comment_Group
+		comment: ^ast.Comment_Group
+
+		if n.pos.line == prev_line {
+			// if we're on the same line, just add the previous docs and comments
+			doc = docs[i - 1]
+			comment = comments[i - 1]
+		} else {
+			// Check to see if there's space below the previous field for a comment
+			if n.pos.line - 1 > prev_line {
+				doc, last_doc = get_file_doc(
+					file,
+					n.pos.line - 1,
+					start_line = prev_line + 1,
+					start_index = last_doc + 1,
+					allocator = allocator,
+				)
+			}
+
+			comment, last_comment = get_file_comment(
+				file,
+				n.pos.line,
+				start_index = last_comment + 1,
+				allocator = allocator,
+			)
+		}
+
+		append(&docs, doc)
+		append(&comments, comment)
+		prev_line = n.pos.line
+	}
+
+
+	return docs, comments
 }
diff --git a/src/server/collector.odin b/src/server/collector.odin
index 6bad491..c06879d 100644
--- a/src/server/collector.odin
+++ b/src/server/collector.odin
@@ -135,7 +135,7 @@ collect_struct_fields :: proc(
 	file: ast.File,
 ) -> SymbolStructValue {
 	b := symbol_struct_value_builder_make(collection.allocator)
-	construct_struct_field_docs(file, struct_type)
+	construct_struct_field_docs(file, struct_type, collection.allocator)
 
 	for field in struct_type.fields.list {
 		for n in field.names {
@@ -174,7 +174,7 @@ collect_bit_field_fields :: proc(
 	package_map: map[string]string,
 	file: ast.File,
 ) -> SymbolBitFieldValue {
-	construct_bit_field_field_docs(file, bit_field_type)
+	construct_bit_field_field_docs(file, bit_field_type, collection.allocator)
 	names := make([dynamic]string, 0, len(bit_field_type.fields), collection.allocator)
 	types := make([dynamic]^ast.Expr, 0, len(bit_field_type.fields), collection.allocator)
 	ranges := make([dynamic]common.Range, 0, len(bit_field_type.fields), collection.allocator)
@@ -227,11 +227,17 @@ collect_enum_fields :: proc(
 		append(&values, clone_type(value, collection.allocator, &collection.unique_strings))
 	}
 
+	temp_docs, temp_comments := get_field_docs_and_comments(file, enum_type.fields, context.temp_allocator)
+	docs := clone_dynamic_array(temp_docs, collection.allocator, &collection.unique_strings)
+	comments := clone_dynamic_array(temp_comments, collection.allocator, &collection.unique_strings)
+
 	value := SymbolEnumValue {
 		names     = names[:],
 		ranges    = ranges[:],
 		values    = values[:],
 		base_type = clone_type(enum_type.base_type, collection.allocator, &collection.unique_strings),
+		comments  = comments[:],
+		docs      = docs[:],
 	}
 
 	return value
@@ -241,6 +247,7 @@ collect_union_fields :: proc(
 	collection: ^SymbolCollection,
 	union_type: ast.Union_Type,
 	package_map: map[string]string,
+	file: ast.File,
 ) -> SymbolUnionValue {
 	types := make([dynamic]^ast.Expr, 0, collection.allocator)
 
@@ -250,9 +257,15 @@ collect_union_fields :: proc(
 		append(&types, cloned)
 	}
 
+	temp_docs, temp_comments := get_field_docs_and_comments(file, union_type.variants, context.temp_allocator)
+	docs := clone_dynamic_array(temp_docs, collection.allocator, &collection.unique_strings)
+	comments := clone_dynamic_array(temp_comments, collection.allocator, &collection.unique_strings)
+
 	value := SymbolUnionValue {
-		types = types[:],
-		poly  = cast(^ast.Field_List)clone_type(union_type.poly_params, collection.allocator, &collection.unique_strings),
+		types    = types[:],
+		poly     = cast(^ast.Field_List)clone_type(union_type.poly_params, collection.allocator, &collection.unique_strings),
+		comments = comments[:],
+		docs     = docs[:],
 	}
 
 	return value
@@ -561,7 +574,7 @@ collect_symbols :: proc(collection: ^SymbolCollection, file: ast.File, uri: stri
 		case ^ast.Union_Type:
 			token = v^
 			token_type = .Union
-			symbol.value = collect_union_fields(collection, v^, package_map)
+			symbol.value = collect_union_fields(collection, v^, package_map, file)
 			symbol.signature = "union"
 		case ^ast.Bit_Set_Type:
 			token = v^
@@ -633,7 +646,7 @@ collect_symbols :: proc(collection: ^SymbolCollection, file: ast.File, uri: stri
 		symbol.type = token_type
 		symbol.doc = get_doc(expr.docs, collection.allocator)
 		symbol.uri = get_index_unique_string(collection, uri)
-		comment := get_file_comment(file, symbol.range.start.line + 1)
+		comment, _ := get_file_comment(file, symbol.range.start.line + 1)
 		symbol.comment = strings.clone(get_comment(comment), collection.allocator)
 		symbol.flags |= {.Distinct}
 
diff --git a/src/server/documentation.odin b/src/server/documentation.odin
index ab5b827..c22bbb4 100644
--- a/src/server/documentation.odin
+++ b/src/server/documentation.odin
@@ -174,13 +174,16 @@ get_signature :: proc(ast_context: ^AstContext, symbol: Symbol) -> string {
 		}
 		strings.write_string(&sb, "{\n")
 		for i in 0 ..< len(v.names) {
+			append_docs(&sb, v.docs, i)
 			strings.write_string(&sb, "\t")
 			strings.write_string(&sb, v.names[i])
 			if i < len(v.values) && v.values[i] != nil {
 				fmt.sbprintf(&sb, "%*s= ", longestNameLen - len(v.names[i]) + 1, "")
 				build_string_node(v.values[i], &sb, false)
 			}
-			strings.write_string(&sb, ",\n")
+			strings.write_string(&sb, ",")
+			append_comments(&sb, v.comments, i)
+			strings.write_string(&sb, "\n")
 		}
 		strings.write_string(&sb, "}")
 		return strings.to_string(sb)
@@ -213,9 +216,12 @@ get_signature :: proc(ast_context: ^AstContext, symbol: Symbol) -> string {
 		}
 		strings.write_string(&sb, " {\n")
 		for i in 0 ..< len(v.types) {
+			append_docs(&sb, v.docs, i)
 			strings.write_string(&sb, "\t")
 			build_string_node(v.types[i], &sb, false)
-			strings.write_string(&sb, ",\n")
+			strings.write_string(&sb, ",")
+			append_comments(&sb, v.comments, i)
+			strings.write_string(&sb, "\n")
 		}
 		strings.write_string(&sb, "}")
 		return strings.to_string(sb)
diff --git a/src/server/symbol.odin b/src/server/symbol.odin
index ad16339..7352425 100644
--- a/src/server/symbol.odin
+++ b/src/server/symbol.odin
@@ -77,12 +77,16 @@ SymbolEnumValue :: struct {
 	values:    []^ast.Expr,
 	base_type: ^ast.Expr,
 	ranges:    []common.Range,
+	docs:      []^ast.Comment_Group,
+	comments:  []^ast.Comment_Group,
 }
 
 SymbolUnionValue :: struct {
 	types:      []^ast.Expr,
 	poly:       ^ast.Field_List,
 	poly_names: []string,
+	docs:       []^ast.Comment_Group,
+	comments:   []^ast.Comment_Group,
 }
 
 SymbolDynamicArrayValue :: struct {
@@ -793,7 +797,7 @@ construct_bit_field_field_symbol :: proc(symbol: ^Symbol, parent_name: string, v
 
 construct_enum_field_symbol :: proc(symbol: ^Symbol, value: SymbolEnumValue, index: int) {
 	symbol.type = .Field
-	//symbol.doc = get_doc(value.docs[index], context.temp_allocator)
-	//symbol.comment = get_comment(value.comments[index])
+	symbol.doc = get_doc(value.docs[index], context.temp_allocator)
+	symbol.comment = get_comment(value.comments[index])
 	symbol.signature = get_enum_field_signature(value, index)
 }
diff --git a/tests/hover_test.odin b/tests/hover_test.odin
index 2f36538..923095b 100644
--- a/tests/hover_test.odin
+++ b/tests/hover_test.odin
@@ -3472,6 +3472,106 @@ ast_hover_multiple_chained_call_expr  :: proc(t: ^testing.T) {
 	}
 	test.expect_hover(t, &source, "test.a: int")
 }
+
+@(test)
+ast_hover_enum_field_documentation :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		F{*}oo :: enum {
+			A = 1, // this is a comment for A
+			// This is a doc for B
+			// across many lines
+			B,
+			C,
+			// D Doc
+			D,
+			E, // E comment
+		}
+		`,
+	}
+	test.expect_hover(
+		t,
+		&source,
+		"test.Foo: enum {\n\tA = 1, // this is a comment for A\n\t// This is a doc for B\n\t// across many lines\n\tB,\n\tC,\n\t// D Doc\n\tD,\n\tE, // E comment\n}"
+	)
+}
+
+@(test)
+ast_hover_enum_field_documentation_same_line :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		F{*}oo :: enum {
+			// Doc for A and B
+			// Mulitple lines!
+			A, B, // comment for A and B
+		}
+		`,
+	}
+	test.expect_hover(
+		t,
+		&source,
+		"test.Foo: enum {\n\t// Doc for A and B\n\t// Mulitple lines!\n\tA, // comment for A and B\n\t// Doc for A and B\n\t// Mulitple lines!\n\tB, // comment for A and B\n}"
+	)
+}
+
+@(test)
+ast_hover_enum_field_directly :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		Foo :: enum {
+			// Doc for A and B
+			// Mulitple lines!
+			A{*}, B, // comment for A and B
+		}
+		`,
+	}
+	test.expect_hover(
+		t,
+		&source,
+		"test.Foo: .A\n Doc for A and B\n Mulitple lines!\n\n// comment for A and B"
+	)
+}
+
+@(test)
+ast_hover_union_field_documentation :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		F{*}oo :: union {
+			int, // this is a comment for int
+			// This is a doc for string
+			// across many lines
+			string,
+			i16,
+			// i32 Doc
+			i32,
+			i64, // i64 comment
+		}
+		`,
+	}
+	test.expect_hover(
+		t,
+		&source,
+		"test.Foo: union {\n\tint, // this is a comment for int\n\t// This is a doc for string\n\t// across many lines\n\tstring,\n\ti16,\n\t// i32 Doc\n\ti32,\n\ti64, // i64 comment\n}"
+	)
+}
+
+@(test)
+ast_hover_union_field_documentation_same_line :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		F{*}oo :: union {
+			// Doc for int and string
+			// Mulitple lines!
+			int, string, // comment for int and string
+		}
+		`,
+	}
+	test.expect_hover(
+		t,
+		&source,
+		"test.Foo: union {\n\t// Doc for int and string\n\t// Mulitple lines!\n\tint, // comment for int and string\n\t// Doc for int and string\n\t// Mulitple lines!\n\tstring, // comment for int and string\n}"
+	)
+}
 /*
 
 Waiting for odin fix