Improve When resolving

1 files changed, 54 insertions, 37 deletions

diff --git a/src/server/ast.odin b/src/server/ast.odin
index d799e0d..78e037c 100644
--- a/src/server/ast.odin
+++ b/src/server/ast.odin
@@ -357,52 +357,69 @@ collect_value_decl :: proc(
 	}
 }
 
+collect_when_stmt :: proc(
+	exprs: ^[dynamic]GlobalExpr,
+	file: ast.File,
+	file_tags: parser.File_Tags,
+	when_decl: ^ast.When_Stmt,
+	skip_private: bool,
+) {
+	if when_decl.cond == nil {
+		return
+	}
+
+	if when_decl.body == nil {
+		return
+	}
+
+	if resolve_when_condition(when_decl.cond) {
+		if block, ok := when_decl.body.derived.(^ast.Block_Stmt); ok {
+			for stmt in block.stmts {
+				if when_stmt, ok := stmt.derived.(^ast.When_Stmt); ok {
+					collect_when_stmt(exprs, file, file_tags, when_stmt, skip_private)
+				} else {
+					collect_value_decl(exprs, file, file_tags, stmt, skip_private)
+				}
+			}
+		}
+	} else {
+		else_stmt := when_decl.else_stmt
+
+		for else_stmt != nil {
+			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 {
+						for stmt in block.stmts {
+							if when_stmt, ok := stmt.derived.(^ast.When_Stmt); ok {
+								collect_when_stmt(exprs, file, file_tags, when_stmt, skip_private)
+							} else {
+								collect_value_decl(exprs, file, file_tags, stmt, skip_private)
+							}
+						}						
+					} 
+					return
+				}
+				else_stmt = else_when.else_stmt
+			} else {
+				return
+			}
+		}
+	}
+
+
+}
+
 collect_globals :: proc(file: ast.File, skip_private := false) -> []GlobalExpr {
 	exprs := make([dynamic]GlobalExpr, context.temp_allocator)
 	defer shrink(&exprs)
 
 	file_tags := parser.parse_file_tags(file, context.temp_allocator)
 
-	_next_decl: for decl in file.decls {
+	for decl in file.decls {
 		if value_decl, ok := decl.derived.(^ast.Value_Decl); ok {
 			collect_value_decl(&exprs, file, file_tags, decl, skip_private)
 		} else if when_decl, ok := decl.derived.(^ast.When_Stmt); ok {
-			if when_decl.cond == nil {
-				continue
-			}
-
-			if when_decl.body == nil {
-				continue
-			}
-
-			if resolve_when_condition(when_decl.cond) {
-				if block, ok := when_decl.body.derived.(^ast.Block_Stmt); ok {
-					for stmt in block.stmts {
-						collect_value_decl(&exprs, file, file_tags, stmt, skip_private)
-					}
-				}
-				continue
-			} else {
-				else_stmt := when_decl.else_stmt
-
-				for else_stmt != nil {
-					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 {
-								for stmt in block.stmts {
-									collect_value_decl(&exprs, file, file_tags, stmt, skip_private)
-								}
-							}
-							continue _next_decl
-						}
-						else_stmt = else_when.else_stmt
-					} else {
-						continue _next_decl
-					}
-				}
-			}
-
-
+			collect_when_stmt(&exprs, file, file_tags, when_decl, skip_private)
 		} else if foreign_decl, ok := decl.derived.(^ast.Foreign_Block_Decl); ok {
 			if foreign_decl.body == nil {
 				continue