print.cpp

#include "sema/lookup.hpp"
#define NOMINMAX
#define WIN32_LEAN_AND_MEAN

#include "module.hpp"
#include "print.hpp"
#include "interface.hpp"
#include "sema/function_arity.hpp"

#include <sstream>

std::ostream& print_impl(std::ostream& out, const doir::module& mod, ecrs::entity_t subtree, bool pretty, bool debug, size_t indent = 0);
std::string print_lookup_name(const doir::module& mod, doir::lookup::lookup lookup, bool debug);

std::string print_function_type(const doir::module& mod, ecrs::entity_t type, bool debug) {
	auto inputs = mod.has_component<doir::function_inputs>(type)
		? doir::lookup::function_inputs::to_lookup(mod.get_component<doir::function_inputs>(type))
		: mod.get_component<doir::lookup::function_inputs>(type);
	std::optional<doir::lookup::function_return_type> return_type = {};
	if(mod.has_component<doir::function_return_type>(type))
		return_type = {mod.get_component<doir::function_return_type>(type).related[0]};
	else if(mod.has_component<doir::lookup::function_return_type>(type))
		return_type = mod.get_component<doir::lookup::function_return_type>(type);

	std::vector<std::string> names;
	if(mod.has_component<doir::function_parameter_names>(type)) {
		auto& param_names = mod.get_component<doir::function_parameter_names>(type);
		names = {param_names.begin(), param_names.end()};
	} else {
		names.resize(inputs.size());
		for(size_t i = 0; i < names.size(); ++i)
			names[i] = "a" + std::to_string(i);
	}

	std::ostringstream out;
	out << "(";
	for(size_t i = 0; i < inputs.size(); ++i)
		out << (i > 0 ? "," : "") << names[i] << ":" << print_lookup_name(mod, inputs[i], debug);
	out << ")";

	if(return_type)
		out << "->" << print_lookup_name(mod, *return_type, debug);
	if(debug)
		out << "[" << type << "]";
	return out.str();
}

std::string print_lookup_name(const doir::module& mod, doir::lookup::lookup lookup, bool debug) {
	if(!lookup.resolved())
		return debug ? "lookup(" + std::string(lookup.name()) + ")" : std::string(lookup.name());

	if(mod.has_component<doir::name>(lookup.entity())) {
		auto out = std::string(mod.get_component<doir::name>(lookup.entity()));
		auto parent = doir::find_parent(mod, lookup.entity());
		while(
			mod.has_component<doir::flags>(parent) && mod.get_component<doir::flags>(parent).namespace_set()
			&& mod.has_component<doir::name>(parent)
		) {
			out = std::string(mod.get_component<doir::name>(parent)) + "." + out;
			parent = doir::find_parent(mod, parent);
		}
		if(debug) out += "[" + std::to_string(lookup.entity()) + "]";
		return out;
	}

	if(mod.has_component<doir::type_definition>(lookup.entity()) && (mod.has_component<doir::function_inputs>(lookup.entity()) || mod.has_component<doir::lookup::function_inputs>(lookup.entity())))
		return print_function_type(mod, lookup.entity(), debug);

	return "%" + std::to_string(lookup.entity());
}


std::tuple<std::string, std::string, std::optional<diagnose::source_location::detailed>, std::string> get_common_assignment_elements(const doir::module& mod, ecrs::entity_t subtree, bool debug) {
	std::string ident;
	if(mod.has_component<doir::name>(subtree)) {
		ident = std::string(mod.get_component<doir::name>(subtree));
		if(debug) ident += "[" + std::to_string(subtree) + "]";
	} else ident = "%" + std::to_string(subtree);

	std::string type = "<error>";
	if(mod.has_component<doir::type_of>(subtree))
		type = print_lookup_name(mod, mod.get_component<doir::type_of>(subtree).related[0], debug);
	else if(mod.has_component<doir::lookup::type_of>(subtree))
		type = print_lookup_name(mod, mod.get_component<doir::lookup::type_of>(subtree), debug);
	else if(mod.has_component<doir::type_definition>(subtree))
		type = "type";
	else if(mod.has_component<doir::flags>(subtree) && mod.get_component<doir::flags>(subtree).namespace_set())
		type = "namespace";

	std::optional<diagnose::source_location::detailed> location = {};
	if(mod.has_component<diagnose::source_location::detailed>(subtree))
		location = mod.get_component<diagnose::source_location::detailed>(subtree);
	else if(mod.has_component<diagnose::source_location>(subtree))
		location = mod.get_component<diagnose::source_location>(subtree).to_detailed(mod.source);

	std::string Export = "";
	if(mod.has_component<doir::flags>(subtree))
		Export = mod.get_component<doir::flags>(subtree).export_set() ? "export " : "";

	return {ident, type, location, Export};
}

std::ostream& print_block(std::ostream& out, const doir::module& mod, const doir::block& block, bool pretty, bool debug, bool skip_parameters, size_t indent) {
	out << "{" << (pretty ? "\n" : "");
	for(auto& elem: block.related) {
		if(skip_parameters && mod.has_component<doir::function_parameter>(elem))
			continue;
		print_impl(out, mod, elem, pretty, debug, indent + 1);
		if(pretty) out << std::endl;
		else out << ";";
	}
	out << (pretty ? std::string(indent, '\t') : "") << "}";
	return out;
}

std::ostream& print_type_of(std::ostream& out, const doir::module& mod, ecrs::entity_t subtree, bool pretty, bool debug, size_t indent) {
	auto indent_string = pretty ? std::string(indent, '\t') : "";
	if(mod.has_component<doir::flags>(subtree) && mod.get_component<doir::flags>(subtree).valueless_set()) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type;
		if(location) out << *location;

	} else if(mod.has_component<doir::call>(subtree) || mod.has_component<doir::lookup::call>(subtree) || (mod.has_component<doir::print_as_call>(subtree) && !debug)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		doir::lookup::lookup call = mod.has_component<doir::call>(subtree)
			? doir::lookup::lookup(mod.get_component<doir::call>(subtree).related[0])
			: mod.has_component<doir::print_as_call>(subtree)
				? doir::lookup::lookup(mod.get_component<doir::print_as_call>(subtree).related[0])
			 	: doir::lookup::lookup(mod.get_component<doir::lookup::call>(subtree));
		auto inputs = mod.has_component<doir::function_inputs>(subtree)
			? doir::lookup::function_inputs::to_lookup(mod.get_component<doir::function_inputs>(subtree))
			: mod.get_component<doir::lookup::function_inputs>(subtree);
		std::string flags = "";
		bool comptime = false;
		if(mod.has_component<doir::flags>(subtree)) {
			auto f = mod.get_component<doir::flags>(subtree);
			if(f.inline_set()) flags += "inline ";
			if(f.flatten_set()) flags += "flatten ";
			if(f.tail_set()) flags += "tail ";
			if(debug && f.comptime_set()) flags += "comptime ";
		}
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=")
			<< flags << print_lookup_name(mod, call, debug) << "(";

		for(size_t i = 0, size = inputs.size(); i < size; ++i) {
			if(i > 0) out << (pretty ? ", " : ",");
			out << print_lookup_name(mod, inputs[i], debug);
		}
		out << ")";

	} else if(mod.has_component<doir::number>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		auto value = mod.get_component<doir::number>(subtree).value;
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=") << value;
		if(location) out << *location;

	} else if(mod.has_component<doir::string>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		auto value = mod.get_component<doir::string>(subtree).value;
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=")
			<< '"' << escape_python_string(value) << '"';
		if(location) out << *location;

	} else if(mod.has_component<doir::function_return_type>(subtree) || mod.has_component<doir::lookup::function_return_type>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		out << indent_string << Export << (debug ? "f:" : "") << ident << (pretty ? ": " : ":");

		doir::lookup::lookup lookup = mod.has_component<doir::type_of>(subtree)
			? doir::lookup::lookup(mod.get_component<doir::type_of>(subtree).related[0])
			: doir::lookup::lookup(mod.get_component<doir::lookup::type_of>(subtree));
		if(!lookup.resolved()) {
			out << type << (pretty ? " = " : "=");
			print_block(out, mod, mod.get_component<doir::block>(subtree), pretty, debug, true, indent);
		} else {
			auto ft = lookup.entity();
			auto resolved = doir::sema::validate::resolve_type_modifications(mod, ft);
			bool ft_is_modification = ft != resolved;

			if(mod.has_component<doir::block>(subtree)) {
				auto inputs = mod.has_component<doir::function_inputs>(ft)
					? doir::lookup::function_inputs::to_lookup(mod.get_component<doir::function_inputs>(ft))
					: mod.get_component<doir::lookup::function_inputs>(ft);
				std::optional<doir::lookup::function_return_type> return_type = {};
				if(mod.has_component<doir::function_return_type>(ft))
					return_type = {mod.get_component<doir::function_return_type>(ft).related[0]};
				else if(mod.has_component<doir::lookup::function_return_type>(ft))
					return_type = mod.get_component<doir::lookup::function_return_type>(ft);

				auto parameters = inputs.associated_parameters(mod, mod.get_component<doir::block>(subtree));

				if(ft_is_modification)
					out << type;
				else {
					out << "(";
					for(size_t i = 0; i < parameters.size(); ++i)
						print_impl(out, mod, parameters[i], pretty, debug) << (i < parameters.size() - 1 ? (pretty ? ", " : ",") : "");
					out << ")";

					if(return_type)
						out << (pretty ? " -> " : "->") << print_lookup_name(mod, *return_type, debug);
					if(debug)
						out << "[" << ft << "]";
				}

				out << (pretty ? " = " : "=");
				print_block(out, mod, mod.get_component<doir::block>(subtree), pretty, debug, true, indent);
			} else { // Valueless function
				out << (ft_is_modification ? type : print_function_type(mod, ft, debug));
			}
		}

	} else if(mod.has_component<doir::block>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=");
		print_block(out, mod, mod.get_component<doir::block>(subtree), pretty, debug, false, indent);

	} else out << "<type_of error>";
	return out;
}

std::ostream& print_impl(std::ostream& out, const doir::module& mod, ecrs::entity_t subtree, bool pretty, bool debug, size_t indent /*= 0*/) {
	auto indent_string = pretty ? std::string(indent, '\t') : "";

	if(mod.has_component<doir::type_of>(subtree) || mod.has_component<doir::lookup::type_of>(subtree) || (mod.has_component<doir::print_as_call>(subtree) && !debug))
		print_type_of(out, mod, subtree, pretty, debug, indent);
	else if(mod.has_component<doir::type_definition>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=");
		if(mod.has_component<doir::block>(subtree))
			print_block(out, mod, mod.get_component<doir::block>(subtree), pretty, debug, false, indent);
		else if(mod.has_component<doir::pointer>(subtree)) {
			auto& p = mod.get_component<doir::pointer>(subtree);
			if(p.size == 0)
				out << "type.pointer(" + print_lookup_name(mod, p.related[0], debug) << ")";
			else out << "type.array(" + print_lookup_name(mod, p.related[0], debug) << ", " << p.size << ")";
		} else out << print_function_type(mod, subtree, debug);

	} else if(mod.has_component<doir::flags>(subtree) && mod.get_component<doir::flags>(subtree).namespace_set()) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		out << indent_string << Export << ident << (pretty ? ": " : ":") << type << (pretty ? " = " : "=");
		print_block(out, mod, mod.get_component<doir::block>(subtree), pretty, debug, false, indent);

	} else if(mod.has_component<doir::alias>(subtree) || mod.has_component<doir::lookup::alias>(subtree)) {
		auto [ident, type, location, Export] = get_common_assignment_elements(mod, subtree, debug);
		auto alias = mod.has_component<doir::alias>(subtree)
			? doir::lookup::alias(mod.get_component<doir::alias>(subtree).related[0], mod.get_component<doir::alias>(subtree).file)
			: mod.get_component<doir::lookup::alias>(subtree);

		out << indent_string << Export << ident << (pretty ? ": " : ":") << "alias" << (pretty ? " = " : "=")
			<< print_lookup_name(mod, alias, debug);
		if(debug && alias.file)
			out << "[" << *alias.file << "]";

	} else out << "<error>";
	return out;
}

std::ostream& doir::print(std::ostream& out, const doir::module& mod, ecrs::entity_t root, bool pretty /* = true */, bool debug /*= false */) {
	// If the root is a block, print the contents of the block
	if(mod.has_component<doir::block>(root))
		for(auto& elem: mod.get_component<doir::block>(root).related) {
			print_impl(out, mod, elem, pretty, debug);
			if(pretty) out << std::endl;
			else out << ";";
		}
	else print_impl(out, mod, root, pretty, debug);
	return out;
}