Correctly resolve semantic tokens for basic literals and casts

2 files changed, 34 insertions, 13 deletions

diff --git a/src/server/semantic_tokens.odin b/src/server/semantic_tokens.odin
index 8259df2..ac1beed 100644
--- a/src/server/semantic_tokens.odin
+++ b/src/server/semantic_tokens.odin
@@ -390,6 +390,15 @@ visit_node :: proc(node: ^ast.Node, builder: ^SemanticTokenBuilder) {
 	}
 }
 
+is_variable_declaration :: proc(expr: ^ast.Expr) -> bool {
+	#partial switch v in expr.derived {
+	case ^ast.Comp_Lit, ^ast.Basic_Lit, ^ast.Type_Cast:
+		return true
+	case:
+		return false
+	}
+}
+
 visit_value_decl :: proc(value_decl: ast.Value_Decl, builder: ^SemanticTokenBuilder) {
 	using ast
 
@@ -402,22 +411,15 @@ visit_value_decl :: proc(value_decl: ast.Value_Decl, builder: ^SemanticTokenBuil
 	if len(value_decl.values) == len(value_decl.names) {
 		for name, i in value_decl.names {
 			ident := name.derived.(^Ident) or_continue
-			if _, ok := value_decl.values[i].derived.(^ast.Comp_Lit); ok {
-				visit_ident(ident, ident, modifiers, builder, true)
-			} else {
-				visit_ident(ident, ident, modifiers, builder)
-			}
+			is_var_decl := is_variable_declaration(value_decl.values[i])
+			visit_ident(ident, ident, modifiers, builder, is_var_decl)
 		}
 	} else if len(value_decl.values) > 0 {
 		// a, b: int
-		_, ok := value_decl.values[0].derived.(^ast.Comp_Lit)
+		is_var_decl := is_variable_declaration(value_decl.values[0])
 		for name in value_decl.names {
 			ident := name.derived.(^Ident) or_continue
-			if ok {
-				visit_ident(ident, ident, modifiers, builder, true)
-			} else {
-				visit_ident(ident, ident, modifiers, builder)
-			}
+			visit_ident(ident, ident, modifiers, builder, is_var_decl)
 		}
 	} else {
 		for name in value_decl.names {
@@ -555,7 +557,7 @@ visit_ident :: proc(
 	symbol_ptr: rawptr,
 	modifiers: SemanticTokenModifiers,
 	builder: ^SemanticTokenBuilder,
-	is_comp_lit := false,
+	is_variable_decl := false,
 ) {
 	using ast
 
@@ -571,7 +573,7 @@ visit_ident :: proc(
 		modifiers += {.ReadOnly}
 	}
 
-	if is_comp_lit {
+	if is_variable_decl {
 		write_semantic_node(builder, ident, .Variable, modifiers)
 		return
 	}
diff --git a/tests/semantic_tokens_test.odin b/tests/semantic_tokens_test.odin
index 099ffde..88ad3c9 100644
--- a/tests/semantic_tokens_test.odin
+++ b/tests/semantic_tokens_test.odin
@@ -50,3 +50,22 @@ semantic_tokens_global_consts :: proc(t: ^testing.T) {
 		{0, 11, 3,  .Type,      {.ReadOnly}}, // [3]  f32
 	})
 }
+
+@(test)
+semantic_tokens_literals_with_explicit_types :: proc(t: ^testing.T) {
+	src := test.Source {
+		main = `package test
+		Foo :: 1
+		Foo2 : int : 1
+		Foo3 :: cast(string) "hello"
+		`
+	}
+
+	test.expect_semantic_tokens(t, &src, {
+		{1, 2,  3, .Variable, {.ReadOnly}}, // [0]  Foo
+		{1, 2,  4, .Variable, {.ReadOnly}}, // [1]  Foo2
+		{0, 7,  3, .Type,     {.ReadOnly}}, // [2]  int
+		{1, 2,  4, .Variable, {.ReadOnly}}, // [3]  Foo3
+		{0, 13, 6, .Type,     {.ReadOnly}}, // [4]  string
+	})
+}