pslang/libs/parser/source/finilize.cpp

#include <pslang/parser/indented_statement.hpp>
#include <pslang/parser/error.hpp>
#include <pslang/ast/statement.hpp>

#include <stdexcept>
#include <vector>

namespace pslang::parser
{

	namespace
	{

		ast::statement_list * get_statement_list(ast::expression_ptr &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::assignment &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::variable_declaration &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::if_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::else_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::else_if_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		// NB: if chain merging happens after retrieving statement list
		ast::statement_list * get_statement_list(ast::if_chain & node)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::while_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::function_definition & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::return_statement &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::field_definition &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::struct_definition &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::statement & statement)
		{
			return std::visit([](auto & value){ return get_statement_list(value); }, statement);
		}

	}

	ast::statement_list_ptr finilize(indented_statement_list statements)
	{
		ast::statement_list_ptr result = std::make_unique<ast::statement_list>();

		using stack_entry = std::variant<ast::statement_list *, ast::struct_definition *>;

		std::vector<stack_entry> stack;
		stack.push_back(result.get());
		std::size_t current_indent = 0;

		auto current_statement_list = [&](ast::location const & location) -> ast::statement_list *
		{
			if (stack.empty())
				throw internal_error("empty finilization stack");

			if (auto list = std::get_if<ast::statement_list *>(&stack.back()))
				return *list;

			throw parse_error("unexpected statement inside struct definition", location);
		};

		auto current_struct_definition = [&](ast::location const & location) -> ast::struct_definition *
		{
			if (stack.empty())
				throw std::runtime_error("Internal error: empty finilization stack");

			if (auto list = std::get_if<ast::struct_definition *>(&stack.back()))
				return *list;

			throw parse_error("unexpected statement outside struct definition", location);
		};

		for (auto & statement : statements.statements)
		{
			auto location = ast::get_location(*statement.statement);

			if (statement.indentation > current_indent)
				throw parse_error("unexpected indent", location);

			while (statement.indentation < current_indent)
			{
				stack.pop_back();
				--current_indent;
			}

			 // Now statement.indentation == current_indent

			auto list = get_statement_list(*statement.statement);

			if (auto if_block = std::get_if<ast::if_block>(statement.statement.get()))
			{
				ast::if_chain chain;
				chain.blocks.push_back({.condition = std::move(if_block->condition), .statements = std::move(if_block->statements)});
				current_statement_list(location)->statements.push_back(std::make_unique<ast::statement>(std::move(chain)));
			}
			else if (auto else_block = std::get_if<ast::else_block>(statement.statement.get()))
			{
				if (current_statement_list(location)->statements.empty())
					throw parse_error("unexpected else block", location);
				auto chain = std::get_if<ast::if_chain>(current_statement_list(location)->statements.back().get());
				if (!chain || chain->blocks.empty() || !chain->blocks.back().condition)
					throw parse_error("unexpected else block", location);

				chain->blocks.push_back({.condition = nullptr, .statements = std::move(else_block->statements)});
			}
			else if (auto else_if_block = std::get_if<ast::else_if_block>(statement.statement.get()))
			{
				if (current_statement_list(location)->statements.empty())
					throw parse_error("unexpected else if block", location);
				auto chain = std::get_if<ast::if_chain>(current_statement_list(location)->statements.back().get());
				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::move(else_if_block->statements)});
			}
			else if (auto field_definition = std::get_if<ast::field_definition>(statement.statement.get()))
			{
				current_struct_definition(location)->fields.push_back(*field_definition);
			}
			else if (std::get_if<ast::struct_definition>(statement.statement.get()))
			{
				current_statement_list(location)->statements.push_back(std::move(statement.statement));
				stack.push_back(std::get_if<ast::struct_definition>(current_statement_list(location)->statements.back().get()));
				++current_indent;
			}
			else
			{
				current_statement_list(location)->statements.push_back(std::move(statement.statement));
			}

			if (list)
			{
				stack.push_back(list);
				++current_indent;
			}
		}

		return result;
	}

}