From 0f9322bf363f736ed9c79dd5dd5b66c408a39344 Mon Sep 17 00:00:00 2001
From: Enzo Baldisserri <enzo.baldisserri@outlook.fr>
Date: Sat, 18 Apr 2020 00:20:38 +0200
Subject: [PATCH] checker: print error and pos for lacking main or pub main

---
 vlib/v/checker/checker.v                   | 34 +++++++++++++---------
 vlib/v/checker/tests/inout/no_fn_main.out  |  5 ++++
 vlib/v/checker/tests/inout/no_fn_main.vv   |  3 ++
 vlib/v/checker/tests/inout/pub_fn_main.out |  5 ++++
 vlib/v/checker/tests/inout/pub_fn_main.vv  |  3 ++
 5 files changed, 37 insertions(+), 13 deletions(-)
 create mode 100644 vlib/v/checker/tests/inout/no_fn_main.out
 create mode 100644 vlib/v/checker/tests/inout/no_fn_main.vv
 create mode 100644 vlib/v/checker/tests/inout/pub_fn_main.out
 create mode 100644 vlib/v/checker/tests/inout/pub_fn_main.vv

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 3c14a7a0dc..0241fb26ce 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -62,10 +62,17 @@ pub fn (c mut Checker) check2(ast_file ast.File) []scanner.Error {
 }
 
 pub fn (c mut Checker) check_files(ast_files []ast.File) {
-	mut all_mods := map[string]int
+	mut has_main_fn := false
 	for file in ast_files {
 		c.check(file)
-		all_mods[ file.mod.name ] = all_mods[ file.mod.name ] + 1
+		if file.mod.name == 'main' {
+			if fn_decl := get_main_fn_decl(file) {
+				has_main_fn = true
+				if fn_decl.is_pub {
+					c.error('function `main` cannot be declared public', fn_decl.pos)
+				}
+			}
+		}
 	}
 	// Make sure fn main is defined in non lib builds
 	if c.pref.build_mode == .build_module || c.pref.is_test {
@@ -75,20 +82,21 @@ pub fn (c mut Checker) check_files(ast_files []ast.File) {
 		// shared libs do not need to have a main
 		return
 	}
-	// check that a main program has a `fn main(){}` function:
-	if all_mods['main'] > 0 {
-		for i, f in c.table.fns {
-			if f.name == 'main' {
-				if f.is_pub {
-					c.error('function `main` cannot be declared public', token.Position{})
-					exit(1)
-				}
-				return
+	if !has_main_fn {
+		c.error('function `main` must be declared in the main module', token.Position{})
+	}
+}
+
+fn get_main_fn_decl(file ast.File) ?ast.FnDecl {
+	for stmt in file.stmts {
+		if stmt is ast.FnDecl {
+			fn_decl := stmt as ast.FnDecl
+			if fn_decl.name == 'main' {
+				return fn_decl
 			}
 		}
-		c.error('function `main` is undeclared in the main module', token.Position{})
-		exit(1)
 	}
+	return none
 }
 
 pub fn (c mut Checker) struct_decl(decl ast.StructDecl) {
diff --git a/vlib/v/checker/tests/inout/no_fn_main.out b/vlib/v/checker/tests/inout/no_fn_main.out
new file mode 100644
index 0000000000..efaa176f82
--- /dev/null
+++ b/vlib/v/checker/tests/inout/no_fn_main.out
@@ -0,0 +1,5 @@
+vlib/v/checker/tests/inout/no_fn_main.v:1:1: error: function `main` must be declared in the main module
+    1| fn no_main() {
+       ^
+    2|     println('Hello world !')
+    3| }
diff --git a/vlib/v/checker/tests/inout/no_fn_main.vv b/vlib/v/checker/tests/inout/no_fn_main.vv
new file mode 100644
index 0000000000..d9a0b24dae
--- /dev/null
+++ b/vlib/v/checker/tests/inout/no_fn_main.vv
@@ -0,0 +1,3 @@
+fn no_main() {
+	println('Hello world !')
+}
diff --git a/vlib/v/checker/tests/inout/pub_fn_main.out b/vlib/v/checker/tests/inout/pub_fn_main.out
new file mode 100644
index 0000000000..a0d7830153
--- /dev/null
+++ b/vlib/v/checker/tests/inout/pub_fn_main.out
@@ -0,0 +1,5 @@
+vlib/v/checker/tests/inout/pub_fn_main.v:1:1: error: function `main` cannot be declared public
+    1| pub fn main() {
+       ~~~
+    2|     println('Hello world !')
+    3| }
diff --git a/vlib/v/checker/tests/inout/pub_fn_main.vv b/vlib/v/checker/tests/inout/pub_fn_main.vv
new file mode 100644
index 0000000000..0f47c312c0
--- /dev/null
+++ b/vlib/v/checker/tests/inout/pub_fn_main.vv
@@ -0,0 +1,3 @@
+pub fn main() {
+	println('Hello world !')
+}