Fix filename tracking in locations & separate prelude location and full block location for AST nodes containing statement lists

libs/ast/include/pslang/ast/control.hpp

@@ -36,6 +36,8 @@ namespace pslang::ast
 		{
 			expression_ptr condition;
 			statement_list_ptr statements;
+			location prelude_location;
+			location location;
 		};
 
 		std::vector<block> blocks;
@@ -46,7 +48,8 @@ namespace pslang::ast
 	{
 		expression_ptr condition;
 		statement_list_ptr statements;
-		ast::location location;
+		location prelude_location;
+		location location;
 	};
 
 }

libs/ast/include/pslang/ast/error.hpp

@@ -13,10 +13,11 @@ namespace pslang::ast
 	{
 		error(std::string message, location location)
 			: message_(std::move(message))
-			, filename_(location.filename)
+			, filename_(location.begin.filename)
 			, location_(location)
 		{
-			location_.filename = filename_;
+			location_.begin.filename = filename_;
+			location_.end.filename = filename_;
 		}
 
 		char const * what() const noexcept

libs/ast/include/pslang/ast/function.hpp

@@ -24,18 +24,21 @@ namespace pslang::ast
 		std::string name;
 		std::vector<argument> arguments;
 		type_ptr return_type;
-		ast::location location;
+		ast::location prelude_location;
 	};
 
 	struct function_definition
 		: function_declaration
 	{
 		statement_list_ptr statements;
+		ast::location location;
 	};
 
 	struct foreign_function_declaration
 		: function_declaration
-	{};
+	{
+		ast::location location;
+	};
 
 	struct return_statement
 	{

libs/ast/include/pslang/ast/location.hpp

@@ -10,11 +10,11 @@ namespace pslang::ast
 	{
 		struct position
 		{
+			std::string_view filename = {};
 			std::size_t line = 1;
 			std::size_t column = 1;
 		};
 
-		std::string_view filename;
 		position begin = {};
 		position end = {};
 
@@ -46,7 +46,6 @@ namespace pslang::ast
 	inline location merge(location const & l1,location const & l2)
 	{
 		return location {
-			.filename = l1.filename,
 			.begin = l1.begin,
 			.end = l2.end,
 		};
@@ -54,7 +53,7 @@ namespace pslang::ast
 
 	inline std::ostream & operator << (std::ostream & out, location const & location)
 	{
-		out << location.filename << ':' << location.begin.line << '.' << location.begin.column;
+		out << location.begin.filename << ':' << location.begin.line << '.' << location.begin.column;
 
 		std::size_t end_column = std::max<std::size_t>(1, location.end.column) - 1;
 		if (location.begin.line < location.end.line)

libs/ast/include/pslang/ast/struct.hpp

@@ -22,6 +22,7 @@ namespace pslang::ast
 	{
 		std::string name;
 		std::vector<field_definition> fields;
+		ast::location prelude_location;
 		ast::location location;
 	};
 

libs/ast/source/validate.cpp

@@ -3,8 +3,6 @@
 #include <pslang/ast/error.hpp>
 #include <pslang/types/type.hpp>
 
-#include <iostream>
-
 namespace pslang::ast
 {
 
@@ -41,6 +39,7 @@ namespace pslang::ast
 
 			void apply(function_definition const & node)
 			{
+				apply(*node.statements);
 				if (!types::equal(*get_type(*node.return_type), types::unit_type{}))
 					if (node.statements->statements.empty() || (true
 						&& !std::get_if<return_statement>(node.statements->statements.back().get())

libs/parser/rules/pslang.l

@@ -53,7 +53,7 @@ f64  { return bp::make_f64(ctx.location); }
 
 [a-zA-Z_]+[a-zA-Z0-9_]* { return bp::make_name(yytext, ctx.location); }
 
-"\n" { ctx.location.move_lines(1); return bp::make_newline(ctx.location); }
+"\n" { auto old_location = ctx.location; ctx.location.move_lines(1); return bp::make_newline(old_location); }
 "\t" { return bp::make_indent(ctx.location); }
 "="  { return bp::make_assignment(ctx.location); }
 ":"  { return bp::make_colon(ctx.location); }

libs/parser/rules/pslang.y

@@ -211,12 +211,12 @@ statement
 | if expression colon { $$ = ast::if_block{std::make_unique<ast::expression>($2), @$}; }
 | else colon { $$ = ast::else_block{@$}; }
 | else if expression colon { $$ = ast::else_if_block{std::make_unique<ast::expression>($3), @$}; }
-| while expression colon { $$ = ast::while_block{std::make_unique<ast::expression>($2), {}, @$}; }
+| while expression colon { $$ = ast::while_block{std::make_unique<ast::expression>($2), {}, @$, @$}; }
 | func name lparen function_declaration_argument_list rparen function_return_type colon { $$ = ast::function_definition{{$2, $4, $6, @$}, {}}; }
 | foreign func name lparen function_declaration_argument_list rparen function_return_type { $$ = ast::foreign_function_declaration{{$3, $5, $7, @$}}; }
 | return expression { $$ = ast::return_statement{std::make_unique<ast::expression>($2), @$}; }
 | return { $$ = ast::return_statement{nullptr, @$}; }
-| struct name colon { $$ = ast::struct_definition{$2, {}, @$}; }
+| struct name colon { $$ = ast::struct_definition{$2, {}, @$, @$}; }
 | name colon type_expression { $$ = ast::field_definition{$1, std::make_unique<ast::type>($3), @$}; }
 ;
 

libs/parser/source/finalize.cpp

@@ -1,12 +1,118 @@
 #include <pslang/parser/indented_statement.hpp>
 #include <pslang/parser/error.hpp>
 #include <pslang/ast/statement.hpp>
+#include <pslang/ast/statement_visitor.hpp>
 
 #include <vector>
 
 namespace pslang::parser
 {
 
+	namespace
+	{
+
+		struct fill_location_visitor
+			: ast::statement_visitor<fill_location_visitor>
+		{
+			using statement_visitor::apply;
+
+			ast::location apply(ast::expression_ptr const & node)
+			{
+				return ast::get_location(node);
+			}
+
+			ast::location apply(ast::assignment const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::variable_declaration const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::if_block const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::else_block const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::else_if_block const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::if_chain & node)
+			{
+				bool first = true;
+				for (auto & block : node.blocks)
+				{
+					block.location = apply(*block.statements);
+					if (first)
+						node.location = block.location;
+					else
+						node.location = ast::merge(node.location, block.location);
+					first = false;
+				}
+				return node.location;
+			}
+
+			ast::location apply(ast::while_block & node)
+			{
+				return node.location = ast::merge(node.prelude_location, apply(*node.statements));
+			}
+
+			ast::location apply(ast::function_definition & node)
+			{
+				return node.location = ast::merge(node.prelude_location, apply(*node.statements));
+			}
+
+			ast::location apply(ast::foreign_function_declaration & node)
+			{
+				return node.location = node.prelude_location;
+			}
+
+			ast::location apply(ast::return_statement const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::field_definition const & node)
+			{
+				return node.location;
+			}
+
+			ast::location apply(ast::struct_definition & node)
+			{
+				node.location = node.prelude_location;
+				for (auto const & field : node.fields)
+					node.location = ast::merge(node.location, apply(field));
+				return node.location;
+			}
+
+			ast::location apply(ast::statement_list & list)
+			{
+				ast::location result;
+				bool first = true;
+				for (auto & statement : list.statements)
+				{
+					auto statement_location = apply(*statement);
+					if (first)
+						result = statement_location;
+					else
+						result = ast::merge(result, statement_location);
+					first = false;
+				}
+				return result;
+			}
+		};
+
+	}
+
 	ast::statement_list_ptr finalize(indented_statement_list statements)
 	{
 		ast::statement_list_ptr result = std::make_unique<ast::statement_list>();
@@ -64,7 +170,7 @@ namespace pslang::parser
 			{
 				ast::if_chain chain;
 				chain.location = if_block->location;
-				chain.blocks.push_back({.condition = std::move(if_block->condition), .statements = std::make_unique<ast::statement_list>()});
+				chain.blocks.push_back({.condition = std::move(if_block->condition), .statements = std::make_unique<ast::statement_list>(), .prelude_location = if_block->location});
 				list = chain.blocks.back().statements.get();
 				current_statement_list(location)->statements.push_back(std::make_unique<ast::statement>(std::move(chain)));
 			}
@@ -76,7 +182,7 @@ namespace pslang::parser
 				if (!chain || chain->blocks.empty() || !chain->blocks.back().condition)
 					throw parse_error("Unexpected else block", location);
 
-				chain->blocks.push_back({.condition = nullptr, .statements = std::make_unique<ast::statement_list>()});
+				chain->blocks.push_back({.condition = nullptr, .statements = std::make_unique<ast::statement_list>(), .prelude_location = else_block->location});
 				list = chain->blocks.back().statements.get();
 			}
 			else if (auto else_if_block = std::get_if<ast::else_if_block>(statement.statement.get()))
@@ -87,7 +193,7 @@ namespace pslang::parser
 				if (!chain || chain->blocks.empty() || !chain->blocks.back().condition)
 					throw parse_error("Unexpected else if block", location);
 
-				chain->blocks.push_back({.condition = std::move(else_if_block->condition), .statements = std::make_unique<ast::statement_list>()});
+				chain->blocks.push_back({.condition = std::move(else_if_block->condition), .statements = std::make_unique<ast::statement_list>(), .prelude_location = else_if_block->location});
 				list = chain->blocks.back().statements.get();
 			}
 			else if (auto while_block = std::get_if<ast::while_block>(statement.statement.get()))
@@ -140,6 +246,8 @@ namespace pslang::parser
 			}
 		}
 
+		fill_location_visitor{}.apply(*result);
+
 		return result;
 	}
 

libs/parser/source/parser.cpp

@@ -14,7 +14,7 @@ namespace pslang::parser
 		if (!yyin)
 			throw std::system_error(std::make_error_code(static_cast<std::errc>(errno)));
 
-		ast::location location{path};
+		ast::location location{.begin = {.filename = path}, .end = {.filename = path}};
 		indented_statement_list result;
 		context ctx{location, result};