From 847d511c4c132ca773d6f29c705a0ad12823af63 Mon Sep 17 00:00:00 2001
From: Brad Lewis <22850972+BradLewis@users.noreply.github.com>
Date: Thu, 30 Oct 2025 17:04:50 -0400
Subject: Correctly resolve local when blocks

---
 src/server/ast.odin         | 14 +++++++++-----
 src/server/locals.odin      | 24 ++++++++++++++----------
 tests/completions_test.odin | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/src/server/ast.odin b/src/server/ast.odin
index 70c568d..bee9ca9 100644
--- a/src/server/ast.odin
+++ b/src/server/ast.odin
@@ -476,10 +476,15 @@ collect_when_stmt :: proc(
 	if when_decl.body == nil {
 		return
 	}
+	if stmt, ok := get_when_block_stmt(when_decl); ok {
+		collect_when_body(exprs, file, file_tags, stmt)
+	}
+}
 
+get_when_block_stmt :: proc(when_decl: ^ast.When_Stmt) -> (^ast.Block_Stmt, bool) {
 	if resolve_when_condition(when_decl.cond) {
 		if block, ok := when_decl.body.derived.(^ast.Block_Stmt); ok {
-			collect_when_body(exprs, file, file_tags, block)
+			return block, true
 		}
 	} else {
 		else_stmt := when_decl.else_stmt
@@ -488,19 +493,18 @@ collect_when_stmt :: proc(
 			if else_when, ok := else_stmt.derived.(^ast.When_Stmt); ok {
 				if resolve_when_condition(else_when.cond) {
 					if block, ok := else_when.body.derived.(^ast.Block_Stmt); ok {
-						collect_when_body(exprs, file, file_tags, block)
+						return block, true
 					}
-					return
 				}
 				else_stmt = else_when.else_stmt
 			} else {
 				if block, ok := else_stmt.derived.(^ast.Block_Stmt); ok {
-					collect_when_body(exprs, file, file_tags, block)
+					return block, true
 				}
-				return
 			}
 		}
 	}
+	return nil, false
 }
 
 collect_when_body :: proc(
diff --git a/src/server/locals.odin b/src/server/locals.odin
index 337e557..4a542bf 100644
--- a/src/server/locals.odin
+++ b/src/server/locals.odin
@@ -438,8 +438,9 @@ get_locals_stmt :: proc(
 	case ^Using_Stmt:
 		get_locals_using_stmt(v^, ast_context)
 	case ^When_Stmt:
-		get_locals_stmt(file, v.else_stmt, ast_context, document_position)
-		get_locals_stmt(file, v.body, ast_context, document_position)
+		if stmt, ok := get_when_block_stmt(v); ok {
+			get_locals_block_stmt(file, stmt^, ast_context, document_position, true)
+		}
 	case ^Case_Clause:
 		get_locals_case_clause(file, v, ast_context, document_position)
 	case ^ast.Defer_Stmt:
@@ -470,6 +471,7 @@ get_locals_block_stmt :: proc(
 	block: ast.Block_Stmt,
 	ast_context: ^AstContext,
 	document_position: ^DocumentPositionContext,
+	skip_position_check := false,
 ) {
 	/*
 	   We need to handle blocks for non mutable and mutable: non mutable has no order for their value declarations, except for nested blocks where they are hidden by scope
@@ -491,14 +493,16 @@ get_locals_block_stmt :: proc(
 	   } <-- document_position.position
 	*/
 
-	if ast_context.non_mutable_only {
-		if !(block.pos.offset <= document_position.nested_position &&
-			   document_position.nested_position <= block.end.offset) {
-			return
-		}
-	} else {
-		if !(block.pos.offset <= document_position.position && document_position.position <= block.end.offset) {
-			return
+	if !skip_position_check {
+		if ast_context.non_mutable_only {
+			if !(block.pos.offset <= document_position.nested_position &&
+			document_position.nested_position <= block.end.offset) {
+				return
+			}
+		} else {
+			if !(block.pos.offset <= document_position.position && document_position.position <= block.end.offset) {
+				return
+			}
 		}
 	}
 
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index 628107a..ce30589 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -4889,3 +4889,39 @@ ast_completion_cast_rawptr_selector :: proc(t: ^testing.T) {
 	}
 	test.expect_completion_docs(t, &source, "", {"Foo.bar: bool"})
 }
+
+@(test)
+ast_completion_local_when_condition :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		main :: proc() {
+			when true {
+				foo : i32 = 5
+			} else {
+				foo : i64 = 6
+			}
+			fo{*}
+		}
+		`,
+	}
+	test.expect_completion_docs(t, &source, "", {"test.foo: i32"})
+}
+
+@(test)
+ast_completion_local_when_condition_false :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		main :: proc() {
+			when false {
+				foo : i32 = 5
+			} else {
+				foo : i64 = 6
+			}
+			fo{*}
+		}
+		`,
+	}
+	test.expect_completion_docs(t, &source, "", {"test.foo: i64"})
+}
-- 
cgit v1.2.3