Merge pull request #703 from BradLewis/feat/reference-improvements

Improve resolution of symbols for references and renames

4 files changed, 249 insertions, 34 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 611bf4b..e7496da 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1717,7 +1717,7 @@ resolve_comp_literal :: proc(
 	symbol: Symbol,
 	ok: bool,
 ) {
-	if position_context.parent_comp_lit.type != nil {
+	if position_context.parent_comp_lit != nil && position_context.parent_comp_lit.type != nil {
 		symbol = resolve_type_expression(ast_context, position_context.parent_comp_lit.type) or_return
 	} else if position_context.call != nil {
 		if call_expr, ok := position_context.call.derived.(^ast.Call_Expr); ok {
@@ -1906,7 +1906,7 @@ resolve_implicit_selector :: proc(
 				if s, ok := comp_symbol.value.(SymbolStructValue); ok {
 					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
 
-					//We can either have the final 
+					//We can either have the final
 					elem_index := -1
 
 					for elem, i in comp_lit.elems {
@@ -1934,7 +1934,7 @@ resolve_implicit_selector :: proc(
 				} else if s, ok := comp_symbol.value.(SymbolBitFieldValue); ok {
 					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
 
-					//We can either have the final 
+					//We can either have the final
 					elem_index := -1
 
 					for elem, i in comp_lit.elems {
@@ -2756,17 +2756,9 @@ make_symbol_enum_from_ast :: proc(
 	ranges := make([dynamic]common.Range, ast_context.allocator)
 
 	for n in v.fields {
-		append(&ranges, common.get_token_range(n, ast_context.file.src))
-
-		if ident, ok := n.derived.(^ast.Ident); ok {
-			append(&names, ident.name)
-		} else if field, ok := n.derived.(^ast.Field_Value); ok {
-			if ident, ok := field.field.derived.(^ast.Ident); ok {
-				append(&names, ident.name)
-			} else if binary, ok := field.field.derived.(^ast.Binary_Expr); ok {
-				append(&names, binary.left.derived.(^ast.Ident).name)
-			}
-		}
+		name, range := get_enum_field_name_and_range(n, ast_context.file.src)
+		append(&names, name)
+		append(&ranges, range)
 	}
 
 	symbol.value = SymbolEnumValue {
@@ -2777,6 +2769,20 @@ make_symbol_enum_from_ast :: proc(
 	return symbol
 }
 
+get_enum_field_name_and_range :: proc(n: ^ast.Expr, document_text: string) -> (string, common.Range) {
+	if ident, ok := n.derived.(^ast.Ident); ok {
+		return ident.name, common.get_token_range(ident, document_text)
+	}
+	if field, ok := n.derived.(^ast.Field_Value); ok {
+		if ident, ok := field.field.derived.(^ast.Ident); ok {
+			return ident.name, common.get_token_range(ident, document_text)
+		} else if binary, ok := field.field.derived.(^ast.Binary_Expr); ok {
+			return binary.left.derived.(^ast.Ident).name, common.get_token_range(binary, document_text)
+		}
+	}
+	return "", {}
+}
+
 make_symbol_bitset_from_ast :: proc(
 	ast_context: ^AstContext,
 	v: ast.Bit_Set_Type,
diff --git a/src/server/file_resolve.odin b/src/server/file_resolve.odin
index b6c4a32..9b1ffa0 100644
--- a/src/server/file_resolve.odin
+++ b/src/server/file_resolve.odin
@@ -308,14 +308,20 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 	case ^Ellipsis:
 		resolve_node(n.expr, data)
 	case ^Comp_Lit:
+		// We only want to resolve the values, not the types
+		resolve_node(n.type, data)
+
 		//only set this for the parent comp literal, since we will need to walk through it to infer types.
+		set := false
 		if data.position_context.parent_comp_lit == nil {
+			set = true
 			data.position_context.parent_comp_lit = n
 		}
+		defer if set {
+			data.position_context.parent_comp_lit = nil
+		}
 
 		data.position_context.comp_lit = n
-
-		resolve_node(n.type, data)
 		resolve_nodes(n.elems, data)
 	case ^Tag_Expr:
 		resolve_node(n.expr, data)
@@ -409,6 +415,7 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 		resolve_nodes(n.rhs, data)
 	case ^Value_Decl:
 		data.position_context.value_decl = n
+
 		reset_position_context(data.position_context)
 		resolve_nodes(n.names, data)
 		resolve_node(n.type, data)
@@ -462,7 +469,10 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 				for name in field.names {
 					data.symbols[cast(uintptr)name] = SymbolAndNode {
 						node = name,
-						symbol = Symbol{range = common.get_token_range(name, string(data.document.text))},
+						symbol = Symbol{
+							range = common.get_token_range(name, string(data.document.text)),
+							uri = strings.clone(common.create_uri(field.pos.file, data.ast_context.allocator).uri, data.ast_context.allocator),
+						},
 					}
 				}
 			}
@@ -483,6 +493,20 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 					node = field,
 					symbol = Symbol{range = common.get_token_range(field, string(data.document.text))},
 				}
+				// In the case of a Field_Value, we explicitly add them so we can find the LHS correctly for things like renaming
+				if field, ok := field.derived.(^ast.Field_Value); ok {
+					if ident, ok := field.field.derived.(^ast.Ident); ok {
+						data.symbols[cast(uintptr)ident] = SymbolAndNode {
+							node = ident,
+							symbol = Symbol{name = ident.name, range = common.get_token_range(ident, string(data.document.text))},
+						}
+					} else if binary, ok := field.field.derived.(^ast.Binary_Expr); ok {
+						data.symbols[cast(uintptr)binary] = SymbolAndNode {
+							node = binary,
+							symbol = Symbol{name = "binary",range = common.get_token_range(binary, string(data.document.text))},
+						}
+					}
+				}
 			}
 		}
 	case ^Bit_Set_Type:
diff --git a/src/server/references.odin b/src/server/references.odin
index f6ab171..92cec09 100644
--- a/src/server/references.odin
+++ b/src/server/references.odin
@@ -48,24 +48,43 @@ prepare_references :: proc(
 	ok: bool,
 ) {
 	ok = false
-	reference := ""
 	pkg := ""
 
-	if position_context.label != nil {
-		return
-	} else if position_context.struct_type != nil {
+	if position_context.struct_type != nil {
 		found := false
 		done_struct: for field in position_context.struct_type.fields.list {
 			for name in field.names {
 				if position_in_node(name, position_context.position) {
 					symbol = Symbol {
-						range = common.get_token_range(name, string(document.text)),
+						range = common.get_token_range(name, ast_context.file.src),
+						pkg = ast_context.current_package,
 					}
 					found = true
 					resolve_flag = .Field
 					break done_struct
 				}
 			}
+			if position_in_node(field.type, position_context.position) {
+				if ident, ok := field.type.derived.(^ast.Ident); ok {
+					symbol, ok = resolve_location_identifier(ast_context, ident^)
+					if !ok {
+						return
+					}
+
+					found = true
+					resolve_flag = .Identifier
+					break done_struct
+				} else if selector, ok := field.type.derived.(^ast.Selector_Expr); ok {
+					symbol, ok = resolve_location_identifier(ast_context, ident^)
+					if !ok {
+						return
+					}
+
+					found = true
+					resolve_flag = .Identifier
+					break done_struct
+				}
+			}
 		}
 		if !found {
 			return
@@ -76,14 +95,35 @@ prepare_references :: proc(
 			if ident, ok := field.derived.(^ast.Ident); ok {
 				if position_in_node(ident, position_context.position) {
 					symbol = Symbol {
-						range = common.get_token_range(ident, string(document.text)),
+						pkg = ast_context.current_package,
+						range = common.get_token_range(ident, ast_context.file.src),
 					}
 					found = true
 					resolve_flag = .Field
 					break done_enum
 				}
-			}
+			} else if value, ok := field.derived.(^ast.Field_Value); ok {
+				if position_in_node(value.field, position_context.position) {
+					symbol = Symbol {
+						range = common.get_token_range(value.field, ast_context.file.src),
+						pkg = ast_context.current_package,
+					}
+					found = true
+					resolve_flag = .Field
+					break done_enum
+				} else if position_in_node(value.value, position_context.position) {
+					if ident, ok := value.value.derived.(^ast.Ident); ok {
+						symbol, ok = resolve_location_identifier(ast_context, ident^)
+						if !ok {
+							return
+						}
 
+						found = true
+						resolve_flag = .Identifier
+						break done_enum
+					}
+				}
+			}
 		}
 		if !found {
 			return
@@ -96,7 +136,6 @@ prepare_references :: proc(
 			if position_in_node(variant, position_context.position) {
 				if ident, ok := variant.derived.(^ast.Ident); ok {
 					symbol, ok = resolve_location_identifier(ast_context, ident^)
-					reference = ident.name
 					resolve_flag = .Identifier
 
 					if !ok {
@@ -140,6 +179,7 @@ prepare_references :: proc(
 			resolve_flag = .Identifier
 		} else {
 			symbol, ok = resolve_location_selector(ast_context, position_context.selector_expr)
+			symbol.flags -= {.Local}
 
 			resolve_flag = .Field
 		}
@@ -151,14 +191,13 @@ prepare_references :: proc(
 			position_context,
 			position_context.implicit_selector_expr,
 		)
+		symbol.flags -= {.Local}
 
 		if !ok {
 			return
 		}
 	} else if position_context.identifier != nil {
 		ident := position_context.identifier.derived.(^ast.Ident)
-
-		reference = ident.name
 		symbol, ok = resolve_location_identifier(ast_context, ident^)
 
 		resolve_flag = .Identifier
@@ -169,6 +208,9 @@ prepare_references :: proc(
 	} else {
 		return
 	}
+	if symbol.uri == "" {
+		symbol.uri = document.uri.uri
+	}
 
 	return symbol, resolve_flag, true
 }
@@ -195,6 +237,7 @@ resolve_references :: proc(
 			uri, _ := common.parse_uri(workspace.uri, context.temp_allocator)
 			filepath.walk(uri.path, walk_directories, document)
 		}
+	} else {
 	}
 
 	reset_ast_context(ast_context)
@@ -279,13 +322,18 @@ resolve_references :: proc(
 
 			if in_pkg || symbol.pkg == document.package_name {
 				symbols_and_nodes := resolve_entire_file(&document, resolve_flag, context.allocator)
-
 				for k, v in symbols_and_nodes {
-					if v.symbol.uri == symbol.uri && v.symbol.range == symbol.range {
+					// NOTE: the uri is sometimes empty for symbols used in the same file as they are derived
+					if (v.symbol.uri == symbol.uri || v.symbol.uri == "") && v.symbol.range == symbol.range {
 						node_uri := common.create_uri(v.node.pos.file, ast_context.allocator)
+						range := common.get_token_range(v.node^, string(document.text))
+						//We don't have to have the `.` with, otherwise it renames the dot.
+						if _, ok := v.node.derived.(^ast.Implicit_Selector_Expr); ok {
+							range.start.character += 1
+						}
 
 						location := common.Location {
-							range = common.get_token_range(v.node^, string(document.text)),
+							range = range,
 							uri   = strings.clone(node_uri.uri, ast_context.allocator),
 						}
 						append(&locations, location)
@@ -300,10 +348,11 @@ resolve_references :: proc(
 	symbols_and_nodes := resolve_entire_file(document, resolve_flag, context.allocator)
 
 	for k, v in symbols_and_nodes {
-		if v.symbol.uri == symbol.uri && v.symbol.range == symbol.range {
+		// NOTE: the uri is sometimes empty for symbols used in the same file as they are derived
+		if (v.symbol.uri == symbol.uri || v.symbol.uri == "") && v.symbol.range == symbol.range {
 			node_uri := common.create_uri(v.node.pos.file, ast_context.allocator)
 
-			range := common.get_token_range(v.node^, string(document.text))
+			range := common.get_token_range(v.node^, ast_context.file.src)
 
 			//We don't have to have the `.` with, otherwise it renames the dot.
 			if _, ok := v.node.derived.(^ast.Implicit_Selector_Expr); ok {
diff --git a/src/server/rename.odin b/src/server/rename.odin
index 786104a..95d53a5 100644
--- a/src/server/rename.odin
+++ b/src/server/rename.odin
@@ -89,8 +89,144 @@ get_prepare_rename :: proc(document: ^Document, position: common.Position) -> (c
 		get_locals(document.ast, position_context.function, &ast_context, &position_context)
 	}
 
-	symbol, _, ok2 := prepare_references(document, &ast_context, &position_context)
+	symbol, ok2 := prepare_rename(document, &ast_context, &position_context)
+	return symbol.range, ok2
+}
 
+// For preparing the rename, we want to position of the token within the current file,
+// not the position of the declaration
+prepare_rename :: proc(
+	document: ^Document,
+	ast_context: ^AstContext,
+	position_context: ^DocumentPositionContext,
+) -> (
+	symbol: Symbol,
+	ok: bool,
+) {
+	ok = false
+	pkg := ""
+
+	if position_context.struct_type != nil {
+		found := false
+		done_struct: for field in position_context.struct_type.fields.list {
+			for name in field.names {
+				if position_in_node(name, position_context.position) {
+					symbol = Symbol {
+						range = common.get_token_range(name, ast_context.file.src),
+					}
+					found = true
+					break done_struct
+				}
+			}
+			if position_in_node(field.type, position_context.position) {
+				if ident, ok := field.type.derived.(^ast.Ident); ok {
+					symbol = Symbol {
+						range = common.get_token_range(field.type, ast_context.file.src),
+					}
+
+					found = true
+					break done_struct
+				} else if selector, ok := field.type.derived.(^ast.Selector_Expr); ok {
+					symbol = Symbol {
+						range = common.get_token_range(selector.field, ast_context.file.src),
+					}
+
+					found = true
+					break done_struct
+				}
+			}
+		}
+		if !found {
+			return
+		}
+	} else if position_context.enum_type != nil {
+		found := false
+		done_enum: for field in position_context.enum_type.fields {
+			if ident, ok := field.derived.(^ast.Ident); ok {
+				if position_in_node(ident, position_context.position) {
+					symbol = Symbol {
+						range = common.get_token_range(ident, ast_context.file.src),
+					}
+					found = true
+					break done_enum
+				}
+			} else if value, ok := field.derived.(^ast.Field_Value); ok {
+				if position_in_node(value.field, position_context.position) {
+					symbol = Symbol {
+						range = common.get_token_range(value.field, ast_context.file.src),
+					}
+					found = true
+					break done_enum
+				} else if position_in_node(value.value, position_context.position) {
+					symbol = Symbol {
+						range = common.get_token_range(value.value, ast_context.file.src),
+					}
+					found = true
+					break done_enum
+				}
+			}
+		}
+		if !found {
+			return
+		}
+	} else if position_context.bitset_type != nil {
+		return
+	} else if position_context.union_type != nil {
+		found := false
+		for variant in position_context.union_type.variants {
+			if position_in_node(variant, position_context.position) {
+				symbol = Symbol {
+					range = common.get_token_range(variant, ast_context.file.src),
+				}
+				found = true
+				break
+			}
+		}
+		if !found {
+			return
+		}
 
-	return symbol.range, ok2
+	} else if position_context.field_value != nil &&
+	   position_context.comp_lit != nil &&
+	   !is_expr_basic_lit(position_context.field_value.field) &&
+	   position_in_node(position_context.field_value.field, position_context.position) {
+		symbol = Symbol {
+			range = common.get_token_range(position_context.field_value.field, ast_context.file.src)
+		}
+	} else if position_context.selector_expr != nil {
+		if position_in_node(position_context.selector, position_context.position) &&
+		   position_context.identifier != nil {
+			ident := position_context.identifier.derived.(^ast.Ident)
+
+			symbol, ok = resolve_location_identifier(ast_context, ident^)
+
+			if !ok {
+				return
+			}
+		} else {
+			symbol, ok = resolve_location_selector(ast_context, position_context.selector_expr)
+			if selector, ok := position_context.selector_expr.derived.(^ast.Selector_Expr); ok {
+				symbol.range = common.get_token_range(selector.field.expr_base, ast_context.file.src)
+			}
+
+		}
+	} else if position_context.implicit {
+		range := common.get_token_range(position_context.implicit_selector_expr, ast_context.file.src)
+		// Skip the `.`
+		range.start.character += 1
+		symbol = Symbol{
+			range = range,
+		}
+
+	} else if position_context.identifier != nil {
+		ident := position_context.identifier.derived.(^ast.Ident)
+
+		symbol = Symbol {
+			range = common.get_token_range(position_context.identifier^, ast_context.file.src)
+		}
+	} else {
+		return
+	}
+
+	return symbol, true
 }