checker: fix comptime main fn (#25747)

Author: kbkpbot

vlib/v/checker/checker.v

@@ -476,11 +476,27 @@ pub fn (mut c Checker) check_files(ast_files []&ast.File) {
 	}
 }
 
-// do checks specific to files in main module
-// returns `true` if a main function is in the file
-fn (mut c Checker) file_has_main_fn(file &ast.File) bool {
+fn (mut c Checker) stmts_has_main_fn(stmts []ast.Stmt) bool {
 	mut has_main_fn := false
-	for stmt in file.stmts {
+	for stmt in stmts {
+		if stmt is ast.ExprStmt {
+			// top level comptime main fn
+			if stmt.expr is ast.IfExpr && stmt.expr.is_comptime {
+				// $if a ? { fn main(){} } $else { fn main() {} }
+				for branch in stmt.expr.branches {
+					if c.stmts_has_main_fn(branch.stmts) {
+						has_main_fn = true
+					}
+				}
+			} else if stmt.expr is ast.MatchExpr && stmt.expr.is_comptime {
+				// $match os { 'windows' { fn main() {} } ...
+				for branch in stmt.expr.branches {
+					if c.stmts_has_main_fn(branch.stmts) {
+						has_main_fn = true
+					}
+				}
+			}
+		}
 		if stmt is ast.FnDecl {
 			if stmt.name == 'main.main' {
 				if has_main_fn {
@@ -504,6 +520,12 @@ fn (mut c Checker) file_has_main_fn(file &ast.File) bool {
 	return has_main_fn
 }
 
+// do checks specific to files in main module
+// returns `true` if a main function is in the file
+fn (mut c Checker) file_has_main_fn(file &ast.File) bool {
+	return c.stmts_has_main_fn(file.stmts)
+}
+
 @[direct_array_access]
 fn (mut c Checker) check_valid_snake_case(name string, identifier string, pos token.Pos) {
 	if c.pref.translated || c.file.is_translated {

vlib/v/checker/tests/run/comptime_main_fn.run.out

@@ -0,0 +1 @@
+main2

vlib/v/checker/tests/run/comptime_main_fn.vv

@@ -0,0 +1,9 @@
+$if usemain1 ? {
+	fn main() {
+		println('main1')
+	}
+} $else {
+	fn main() {
+		println('main2')
+	}
+}