Merge pull request #1139 from BradLewis/fix/resolve-local-whens

Correctly resolve local when blocks

2 files changed, 23 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
+			}
 		}
 	}