Fix procedure's scope

4 files changed, 47 insertions, 25 deletions

diff --git a/src/checker/checker.cpp b/src/checker/checker.cpp
index 2a084c543..22d734af6 100644
--- a/src/checker/checker.cpp
+++ b/src/checker/checker.cpp
@@ -89,14 +89,15 @@ ExpressionInfo make_expression_info(b32 is_lhs, AddressingMode mode, Type *type,
 }
 
 struct ProcedureInfo {
-	AstFile *file;
-	Token            token;
+	AstFile * file;
+	Token     token;
 	DeclInfo *decl;
-	Type *           type; // Type_Procedure
-	AstNode *        body; // AstNode_BlockStatement
+	Type *    type; // Type_Procedure
+	AstNode * body; // AstNode_BlockStatement
 };
 
 struct Scope {
+	b32 is_proc;
 	Scope *parent;
 	Scope *prev, *next;
 	Scope *first_child, *last_child;
@@ -231,6 +232,9 @@ void check_open_scope(Checker *c, AstNode *stmt) {
 	GB_ASSERT(is_ast_node_stmt(stmt) || stmt->kind == AstNode_ProcType);
 	Scope *scope = make_scope(c->context.scope, c->allocator);
 	add_scope(c, stmt, scope);
+	if (stmt->kind == AstNode_ProcType) {
+		scope->is_proc = true;
+	}
 	c->context.scope = scope;
 }
 
@@ -238,23 +242,36 @@ void check_close_scope(Checker *c) {
 	c->context.scope = c->context.scope->parent;
 }
 
-void scope_lookup_parent_entity(Scope *s, String name, Scope **scope, Entity **entity) {
+void scope_lookup_parent_entity(Checker *c, Scope *s, String name, Scope **scope, Entity **entity) {
+	b32 gone_thru_proc = false;
 	HashKey key = hash_string(name);
 	for (; s != NULL; s = s->parent) {
+
 		Entity **found = map_get(&s->elements, key);
 		if (found) {
-			if (entity) *entity = *found;
+			Entity *e = *found;
+			if (gone_thru_proc) {
+				if (e->kind == Entity_Variable && e->parent != c->global_scope) {
+					continue;
+				}
+			}
+
+			if (entity) *entity = e;
 			if (scope) *scope = s;
 			return;
 		}
+
+		if (s->is_proc) {
+			gone_thru_proc = true;
+		}
 	}
 	if (entity) *entity = NULL;
 	if (scope) *scope = NULL;
 }
 
-Entity *scope_lookup_entity(Scope *s, String name) {
+Entity *scope_lookup_entity(Checker *c, Scope *s, String name) {
 	Entity *entity = NULL;
-	scope_lookup_parent_entity(s, name, NULL, &entity);
+	scope_lookup_parent_entity(c, s, name, NULL, &entity);
 	return entity;
 }
 
@@ -500,7 +517,6 @@ void add_file_entity(Checker *c, AstNode *identifier, Entity *e, DeclInfo *d) {
 
 	add_entity(c, c->global_scope, identifier, e);
 	map_set(&c->info.entities, hash_pointer(e), d);
-	e->order = gb_array_count(c->info.entities.entries);
 }
 
 
@@ -637,8 +653,6 @@ void check_parsed_files(Checker *c) {
 				DeclInfo *d = make_declaration_info(c->allocator, e->parent);
 				d->proc_decl = decl;
 				map_set(&c->info.entities, hash_pointer(e), d);
-				e->order = gb_array_count(c->info.entities.entries);
-
 			case_end;
 
 			case_ast_node(ld, LoadDecl, decl);
diff --git a/src/checker/entity.cpp b/src/checker/entity.cpp
index 182bc1916..272d028d0 100644
--- a/src/checker/entity.cpp
+++ b/src/checker/entity.cpp
@@ -34,7 +34,6 @@ struct Entity {
 	Scope *parent;
 	Token token;
 	Type *type;
-	isize order;
 
 	union {
 		struct { ExactValue value; } Constant;
diff --git a/src/checker/expr.cpp b/src/checker/expr.cpp
index 128d9e6e0..bd317c08d 100644
--- a/src/checker/expr.cpp
+++ b/src/checker/expr.cpp
@@ -202,7 +202,7 @@ void check_identifier(Checker *c, Operand *o, AstNode *n, Type *named_type) {
 	o->mode = Addressing_Invalid;
 	o->expr = n;
 	ast_node(i, Ident, n);
-	Entity *e = scope_lookup_entity(c->context.scope, i->token.string);
+	Entity *e = scope_lookup_entity(c, c->context.scope, i->token.string);
 	if (e == NULL) {
 		error(&c->error_collector, i->token,
 		    "Undeclared type or identifier `%.*s`", LIT(i->token.string));
@@ -1959,14 +1959,9 @@ ExpressionKind check__expr_base(Checker *c, Operand *o, AstNode *node, Type *typ
 	case_end;
 
 	case_ast_node(pl, ProcLit, node);
-		auto curr_context = c->context;
-		c->context.scope = c->global_scope;
 		check_open_scope(c, pl->type);
 		c->context.decl = make_declaration_info(c->allocator, c->context.scope);
-		defer ({
-			check_close_scope(c);
-			c->context = curr_context;
-		});
+		defer (check_close_scope(c));
 		Type *proc_type = check_type(c, pl->type);
 		if (proc_type != NULL) {
 			check_proc_body(c, empty_token, c->context.decl, proc_type, pl->body);
diff --git a/src/checker/stmt.cpp b/src/checker/stmt.cpp
index 197ab1bb0..a78db4623 100644
--- a/src/checker/stmt.cpp
+++ b/src/checker/stmt.cpp
@@ -190,7 +190,7 @@ Type *check_assignment_variable(Checker *c, Operand *op_a, AstNode *lhs) {
 	b32 used = false;
 	if (node->kind == AstNode_Ident) {
 		ast_node(i, Ident, node);
-		e = scope_lookup_entity(c->context.scope, i->token.string);
+		e = scope_lookup_entity(c, c->context.scope, i->token.string);
 		if (e != NULL && e->kind == Entity_Variable) {
 			used = e->Variable.used; // TODO(bill): Make backup just in case
 		}
@@ -398,20 +398,34 @@ void check_proc_decl(Checker *c, Entity *e, DeclInfo *d, b32 check_body_later) {
 	e->type = proc_type;
 	ast_node(pd, ProcDecl, d->proc_decl);
 
-	Scope *original_curr_scope = c->context.scope;
-	c->context.scope = c->global_scope;
 	check_open_scope(c, pd->type);
 	defer ({
 		check_close_scope(c);
-		c->context.scope = original_curr_scope;
 	});
 
 
 	check_procedure_type(c, proc_type, pd->type);
-	b32 is_foreign   = (pd->tags & ProcTag_foreign) != 0;
-	b32 is_inline    = (pd->tags & ProcTag_inline) != 0;
+	b32 is_foreign   = (pd->tags & ProcTag_foreign)   != 0;
+	b32 is_inline    = (pd->tags & ProcTag_inline)    != 0;
 	b32 is_no_inline = (pd->tags & ProcTag_no_inline) != 0;
 
+
+
+	if (d->scope == c->global_scope &&
+	    are_strings_equal(e->token.string, make_string("main"))) {
+		if (proc_type != NULL) {
+			auto *pt = &proc_type->proc;
+			if (pt->param_count != 0 ||
+			    pt->result_count) {
+				gbString str = type_to_string(proc_type);
+				defer (gb_string_free(str));
+
+				error(&c->error_collector, e->token,
+				      "Procedure type of `main` was expected to be `proc()`, got %s", str);
+			}
+		}
+	}
+
 	if (is_inline && is_no_inline) {
 		error(&c->error_collector, ast_node_token(pd->type),
 		      "You cannot apply both `inline` and `no_inline` to a procedure");