protohackers

protohack-1.cpp (7349B)

---

 #include "inet.hpp"
 #include "fd.hpp"
 
 #include <poll.h>
 
 #include <iostream>
 #include <cmath>
 #include <map>
 #include <vector>
 #include <optional>
 #include <openssl/bn.h>
 
 using namespace std;
 
 struct bignum
 {
 	bignum()
 	:	p(BN_new(), BN_free)
 	{
 		BN_init(*this);
 	}
 
 	bignum(BIGNUM* bn)
 	:	p(bn, BN_free)
 	{}
 
 	bignum(unsigned long w)
 	:	bignum()
 	{
 		BN_set_word(*this, w);
 	}
 
 	bignum(bignum&&) = default;
 	bignum& operator=(bignum&&) = default;
 
 	bignum(bignum const& old)
 	:	bignum{BN_dup(old)}
 	{}
 	bignum& operator=(bignum const& old)
 	{
 		if (p)
 			BN_copy(*this, old);
 		else
 			p.reset(BN_dup(old));
 
 		return *this;
 	}
 
 	char const* to_str() const
 	{
 		return BN_bn2dec(*this);
 	}
 
 	bool is_prime() const
 	{
 		return BN_is_prime_fasttest_ex(*this, INT_MAX, nullptr, 1, nullptr);
 	}
 
 	bignum operator +(bignum const& rhs)
 	{
 		bignum r;
 		BN_add(r, *this, rhs);
 		return r;
 	}
 	bignum operator -(bignum const& rhs)
 	{
 		bignum r;
 		BN_sub(r, *this, rhs);
 		return r;
 	}
 
 private:
 	unique_ptr<BIGNUM, void(*)(BIGNUM*)> p;
 
 	operator BIGNUM const*() const
 	{
 		return p.get();
 	}
 	operator BIGNUM*()
 	{
 		return p.get();
 	}
 };
 
 ostream& operator <<(ostream& out, bignum const& rhs)
 {
 	return out << rhs.to_str();
 }
 
 struct json_object
 {
 	std::map<string_view, string_view> nv_pairs;
 	std::map<string_view, bignum> num_pairs;
 };
 
 template <typename T>
 struct parse_t
 {
 	T val;
 	string_view remaining;
 };
 
 std::optional<string_view> parse_char(string_view in, char c)
 {
 	if (in.empty())
 		return {};
 	if (in.front() != c)
 		return {};
 	return in.substr(1);
 }
 
 std::optional<string_view> parse_str(string_view in, string_view str)
 {
 	if (not in.starts_with(str))
 		return {};
 	return in.substr(str.size());
 }
 
 string_view parse_whitespace(string_view in)
 {
 	auto end = in.find_first_not_of(" \t\r\n");
 	return end != string::npos ? in.substr(end) : ""sv;
 }
 
 std::optional<parse_t<bool>> parse_bool(string_view in)
 {
 	if (auto next = parse_str(in, "true"))
 		return parse_t<bool>{true, *next};
 	if (auto next = parse_str(in, "false"))
 		return parse_t<bool>{false, *next};
 	return {};
 }
 
 size_t find_end_quote(string_view str)
 {
 	size_t start = 0;
 	for (;;)
 	{
 		auto pos = str.find_first_of("\"\\", start);
 		if ((pos == string::npos) or (str[pos] == '"'))
 			return pos;
 		start = pos + 2;
 		if (start > str.size())
 			return string::npos;
 	}
 }
 
 std::optional<parse_t<string_view>> parse_string(string_view in)
 {
 	if (auto next = parse_char(in, '"'))
 		in = *next;
 	else
 		return {};
 
 	if (in.empty())
 		return {};
 
 	auto end = find_end_quote(in);
 	auto str = in.substr(0, end);
 	in = end != string::npos ? in.substr(end) : ""sv;
 
 	if (auto next = parse_char(in, '"'))
 		in = *next;
 	else
 		return {};
 
 	return parse_t<string_view>{str, in};
 }
 
 std::optional<parse_t<bignum>> parse_bignum(string_view in)
 {
 	BIGNUM* p = nullptr;
 	auto end = in.find_first_not_of("-0123456789");
 	string str{in.substr(0, end)};
 	int r = BN_dec2bn(&p, str.c_str());
 	if (r == 0)
 		return {};
 	if (r < 0)
 		__builtin_trap();
 
 	in = end != string::npos ? in.substr(end) : ""sv;
 
 	if (auto next = parse_char(in, '.'))
 	{
 		in = *next;
 	}
 	else
 		return parse_t<bignum>{bignum{p}, in};
 
 	end = in.find_first_not_of("0123456789");
 	in = end != string::npos ? in.substr(end) : ""sv;
 	return parse_t<bignum>{bignum{0UL}, in};
 }
 
 std::optional<parse_t<json_object>> parse_object(string_view in)
 {
 	json_object o;
 
 	if (auto next = parse_char(in, '{'))
 		in = *next;
 	else
 		return {};
 
 	in = parse_whitespace(in);
 
 	if (auto next = parse_char(in, '}'))
 		return parse_t<json_object>{std::move(o), *next};
 
 	for (;;)
 	{
 		auto name = parse_string(in);
 		if (not name)
 			return {};
 
 		in = parse_whitespace(name->remaining);
 
 		if (auto next = parse_char(in, ':'))
 			in = parse_whitespace(*next);
 		else
 			return {};
 
 		if (auto value = parse_string(in))
 		{
 			auto& val = o.nv_pairs[name->val];
 			val = value->val;
 			in = value->remaining;
 		}
 		else if (auto value = parse_bignum(in))
 		{
 			auto& val = o.num_pairs[name->val];
 			val = value->val;
 			in = value->remaining;
 		}
 		else if (auto value = parse_bool(in))
 		{
 			//auto& val = o.bool_pairs[name->val];
 			//val = value->val;
 			in = value->remaining;
 		}
 		else if (auto value = parse_object(in))
 		{
 			//auto& val = o.obj_pairs[name->val];
 			//val = value->val;
 			in = value->remaining;
 		}
 		else
 			return {};
 
 		if (auto next = parse_char(in, ','))
 			in = parse_whitespace(*next);
 		else
 			break;
 	}
 	if (auto next = parse_char(in, '}'))
 		return parse_t<json_object>{std::move(o), *next};
 	else
 		return {};
 }
 
 ostream& operator <<(ostream& out, json_object const& rhs)
 {
 	out << '{';
 	for (auto const& nv : rhs.nv_pairs)
 		out << '"' << nv.first << "\": \"" << nv.second << "\",";
 	for (auto const& nv : rhs.num_pairs)
 		out << '"' << nv.first << "\": " << nv.second << ',';
 	return out << '}';
 }
 
 ostream& operator <<(ostream& out, std::optional<parse_t<json_object>> const& rhs)
 {
 	if (not rhs)
 		return out << "<failed parse>";
 	return out << rhs->val << " <" << rhs->remaining << ">";
 }
 
 struct client
 {
 	descriptor d;
 	std::string buf;
 
 	bool do_read()
 	{
 		string data = read(d);
 		if (data.empty())
 			return false;
 
 		buf.append(std::move(data));
 		return true;
 	}
 	optional<string> getline()
 	{
 		if (buf.empty())
 			return {};
 		auto n = buf.find('\n');
 		if (n == string::npos)
 			return {};
 		string line = buf.substr(0, n);
 		buf.erase(0, n+1);
 		return line;
 	}
 };
 
 std::optional<bignum> parse_json_request(string_view line)
 {
 	auto obj = parse_object(line);
 	if (not obj)
 	{
 		cout << "request is invalid JSON\n";
 		return {};
 	}
 	//std::cout << obj << std::flush;
 
 	auto srch = obj->val.nv_pairs.find("method");
 	if ((srch == obj->val.nv_pairs.end()) or (srch->second != "isPrime"))
 	{
 		cout << "'method' != isPrime or is missing\n";
 		return {};
 	}
 	auto srch2 = obj->val.num_pairs.find("number");
 	if (srch2 == obj->val.num_pairs.end())
 	{
 		cout << "'number' is not number type or is missing\n";
 		return {};
 	}
 	return srch2->second;
 }
 
 bool client_ready(client& c)
 {
 	if (not c.do_read())
 	{
 		cout << "* " << c.d << " disconnected\n";
 		return false;
 	}
 
 	while (auto line = c.getline())
 	{
 		if (auto number = parse_json_request(*line))
 		{
 			bool prime = number->is_prime();
 			std::cout << *number << ": " << prime << std::endl;
 
 			if (prime)
 				write(c.d, "{\"method\":\"isPrime\",\"prime\":true}\n"sv);
 			else
 				write(c.d, "{\"method\":\"isPrime\",\"prime\":false}\n"sv);
 		}
 		else
 		{
 			write(c.d, "xxx\n"sv);
 			return false;
 		}
 	}
 	return true;
 }
 
 int main()
 {
 	std::map<int, client> clients;
 	auto incoming = inet::listen(SOCK_STREAM);
 
 	std::cout << std::boolalpha << "Awaiting clients\n";
 
 	for (;;)
 	{
 		std::vector<pollfd> awaitables;
 		awaitables.push_back({incoming, POLLIN, 0});
 		for (auto const& c : clients)
 			awaitables.push_back({c.second.d, POLLIN, 0});
 
 		::poll(awaitables.data(), awaitables.size(), INFTIM);
 
 		for (auto& ev : awaitables)
 		{
 			if (~ev.revents & POLLIN)
 				continue;
 
 			if (ev.fd == incoming)
 			{
 				auto conn = inet::accept(incoming);
 
 				auto& x = clients[conn];
 				x.d = move(conn);
 				continue;
 			}
 
 			auto it = clients.find(ev.fd);
 			if (it == clients.end())
 			{
 				cerr << "Unknown fd " << ev.fd << endl;
 				return 1;
 			}
 
 			if (not client_ready(it->second))
 				clients.erase(it);
 		}
 	}
 }