<commit_before>// Ҹ boost asio http://neive.tistory.com/22

#pragma once
#include"stdafx.h"

boostAsioServer::boostAsioServer() : m_acceptor(g_io_service, tcp::endpoint(tcp::v4(), SERVERPORT)), m_socket(g_io_service)
{
	getMyServerIP();
	CheckThisCPUcoreCount();

	acceptThread();
	start_io_service();
}

boostAsioServer::~boostAsioServer()
{
	// make_shared   ,  ʿ䰡 پ ?
	for (auto ptr : g_clients) { delete ptr; }
}

void boostAsioServer::getMyServerIP()
{
	tcp::resolver			m_resolver(g_io_service);
	tcp::resolver::query	m_query(boost::asio::ip::host_name(), "");
	tcp::resolver::iterator m_resolver_iterator = m_resolver.resolve(m_query);

	while (m_resolver_iterator != tcp::resolver::iterator()) {
		using boost::asio::ip::address;
		address addr = (m_resolver_iterator++)->endpoint().address();
		if (!addr.is_v6()) { cout << "This Server's IPv4 address: " << addr.to_string() << endl; }
		//else if (addr.is_v6()) { cout << "This Server's IPv6 address: " << addr.to_string() << endl; }
	}
}

void boostAsioServer::CheckThisCPUcoreCount()
{
	SYSTEM_INFO si;
	GetSystemInfo(&si);
	m_cpuCore = static_cast<int>(si.dwNumberOfProcessors) * 2;
	printf("CPU Core Count = %d, threads = %d\n", m_cpuCore / 2, m_cpuCore);
}

void boostAsioServer::start_io_service()
{
	m_worker_threads.reserve(m_cpuCore);

	for (int i = 0; i < m_cpuCore; ++i) { m_worker_threads.emplace_back(new thread{ [&]() -> void { g_io_service.run(); } }); }

	while (m_ServerShutdown) { Sleep(1000); }

	// workerThread ߵ
	for (auto thread : m_worker_threads) {
		thread->join();
		delete thread;
	}
}

void boostAsioServer::acceptThread()
{
	m_acceptor.async_accept(m_socket, [&](boost::system::error_code error_code) {
		if (true == (!error_code)) {
			cout << "Client No. [ " << ++m_playerIndex << " ] Connected \t:: IP = " << m_socket.remote_endpoint().address().to_string() << ", Port = " << m_socket.remote_endpoint().port() << "\n";
			g_clients.emplace_back(new player_session (std::move(m_socket), m_playerIndex));
			g_clients[m_playerIndex]->Init();
		}
		if (false == m_ServerShutdown) { acceptThread(); }
	});
}

// player_session class ------------------------------------------------------------------------------------------------------------------------

void player_session::Init()
{
	m_connect_state = true;

	// ⺻  ʱȭ
	Packet init_this_player_buf[MAX_BUF_SIZE];

	init_this_player_buf[0] = sizeof(player_data) + 2;
	init_this_player_buf[1] = INIT_CLIENT;

	m_player_data.id = m_id;
	m_player_data.pos.x = 400;
	m_player_data.pos.y = 300;

	memcpy(&init_this_player_buf[2], g_clients[m_id]->get_player_data(), init_this_player_buf[0]);
	g_clients[m_id]->send_packet(init_this_player_buf);

	// ʱȭ   2 -   ٸ ֵ , ٸ ֵ
	/*Packet other_info_to_me_buf[MAX_BUF_SIZE];*/
	Packet my_info_to_other_buf[MAX_BUF_SIZE];

	/*other_info_to_me_buf[0] =*/ my_info_to_other_buf[0] = sizeof(player_data) + 2;
	/*other_info_to_me_buf[1] =*/ my_info_to_other_buf[1] = KEYINPUT;
	my_info_to_other_buf[1] = INIT_CLIENT;

	//    ٸ ÷̾
	memcpy(&my_info_to_other_buf[2], &m_player_data, my_info_to_other_buf[0]);

	for (auto players : g_clients)
	{
		if (DISCONNECTED == players->get_current_connect_state()) { continue; }
		if (m_id == players->get_id()) { continue; }

		// ٸ ֵ  ؼ ְ,   ... (   ޴° ˼  ü... )
		/*memcpy(&other_info_to_me_buf[2], players->get_player_data(), other_info_to_me_buf[0] - 2);
		send_packet(other_info_to_me_buf);*/

		//    ٸ ֵ  Ǵµ..
		players->send_packet(my_info_to_other_buf);
	}

	/*
		ó ÷̾,  ÷̾  ˸
		view list   ߰ ~ !!
	*/

	m_recv_packet();
}

void player_session::m_recv_packet()
{
	//auto self(shared_from_this());
	m_socket.async_read_some(boost::asio::buffer(m_recv_buf, MAX_BUF_SIZE), [&](boost::system::error_code error_code, std::size_t length) -> void {
		if (error_code) {
			if (error_code.value() == boost::asio::error::operation_aborted) { return; }
			// client was disconnected
			if (false == g_clients[m_id]->get_current_connect_state()) { return; }

			cout << "Client No. [ " << m_id << " ] Disonnected \t:: IP = " << m_socket.remote_endpoint().address().to_string() << ", Port = " << m_socket.remote_endpoint().port() << "\n";
			m_socket.shutdown(m_socket.shutdown_both);
			m_socket.close();

			m_connect_state = false;

			/*
				ó ÷̾,  ÷̾  ˸
				view list    ~ !!
			*/

			Packet temp_buf[MAX_BUF_SIZE] = { sizeof(player_data) + 2, PLAYER_DISCONNECTED };
			memcpy(&temp_buf[2], &m_player_data, temp_buf[0]);

			for (auto players : g_clients)
			{
				if (DISCONNECTED == players->m_connect_state) { continue; }
				if (m_id == players->m_id) { continue; }

				players->send_packet(temp_buf);
			}

			return;
		}

		int current_data_processing = static_cast<int>(length);
		Packet *buf = m_recv_buf;
		while (0 < current_data_processing) {
			if (0 == m_packet_size_current) {
				m_packet_size_current = buf[0];
				if (buf[0] > MAX_BUF_SIZE) {
					cout << "player_session::m_recv_packet() Error, Client No. [ " << m_id << " ] recv buf[0] is out of MAX_BUF_SIZE\n";
					exit(-1);
				}
			}
			int need_to_build = m_packet_size_current - m_packet_size_previous;
			if (need_to_build <= current_data_processing) {
				// Packet building Complete & Process Packet
				memcpy(m_data_buf + m_packet_size_previous, buf, need_to_build);

				m_process_packet(m_id, m_data_buf);

				m_packet_size_current = 0;
				m_packet_size_previous = 0;
				current_data_processing -= need_to_build;
				buf += need_to_build;
			}
			else {
				// Packet build continue
				memcpy(m_data_buf + m_packet_size_previous, buf, current_data_processing);
				m_packet_size_previous += current_data_processing;
				current_data_processing = 0;
				buf += current_data_processing;
			}
		}
		m_recv_packet();
	});
}

void player_session::send_packet(Packet *packet)
{
	int packet_size = packet[0];
	Packet *sendBuf = new Packet[packet_size];
	memcpy(sendBuf, packet, packet_size);

	//auto self(shared_from_this());
	m_socket.async_write_some(boost::asio::buffer(sendBuf, packet_size), [=](boost::system::error_code error_code, std::size_t bytes_transferred) -> void {
		if (!error_code) {
			if (packet_size != bytes_transferred) { cout << "Client No. [ " << m_id << " ] async_write_some packet bytes was NOT SAME !!\n"; }
			delete[] sendBuf;
		}
	});
}

void player_session::m_process_packet(const unsigned int& id, Packet buf[])
{
	// packet[0] = packet size		> 0° ڸ , Ŷ ũⰡ ߸ Ѵ.
	// packet[1] = type				> 1° ڸ  Ŷ  Ŷ Ӽ ִ ̴.
	// packet[...] = data			> 2° ʹ Ӽ ´  ó ش.

	// buf[1] ° Ӽ з  ڿ, ο 2°  ͸ óϱ Ѵ.

	{
		switch (buf[1])
		{
		case TEST:
			//  Ŷ ״ ش.
			cout << "Client No. [ " << m_id << " ] TEST Packet Recived !!\n";
			printf("buf[0] = %d, buf[1] = %d, buf[2] = %d\n\n", buf[0], buf[1], buf[2]);
			g_clients[id]->send_packet(buf);
			break;

		case KEYINPUT:
			g_clients[id]->m_player_data.pos = reinterpret_cast<player_data*>(&buf[2])->pos;

			// ʿ ֵ ̵  ѷ
			for (auto players : g_clients)
			{
				if (DISCONNECTED == players->m_connect_state) { continue; }
				if (id == players->m_id) { continue; }

				players->send_packet(buf);
			}

			break;
		default:
			break;
		}
	}
}<commit_msg>[김형준] 100 100<commit_after>// Ҹ boost asio http://neive.tistory.com/22

#pragma once
#include"stdafx.h"

boostAsioServer::boostAsioServer() : m_acceptor(g_io_service, tcp::endpoint(tcp::v4(), SERVERPORT)), m_socket(g_io_service)
{
	getMyServerIP();
	CheckThisCPUcoreCount();

	acceptThread();
	start_io_service();
}

boostAsioServer::~boostAsioServer()
{
	// make_shared   ,  ʿ䰡 پ ?
	for (auto ptr : g_clients) { delete ptr; }
}

void boostAsioServer::getMyServerIP()
{
	tcp::resolver			m_resolver(g_io_service);
	tcp::resolver::query	m_query(boost::asio::ip::host_name(), "");
	tcp::resolver::iterator m_resolver_iterator = m_resolver.resolve(m_query);

	while (m_resolver_iterator != tcp::resolver::iterator()) {
		using boost::asio::ip::address;
		address addr = (m_resolver_iterator++)->endpoint().address();
		if (!addr.is_v6()) { cout << "This Server's IPv4 address: " << addr.to_string() << endl; }
		//else if (addr.is_v6()) { cout << "This Server's IPv6 address: " << addr.to_string() << endl; }
	}
}

void boostAsioServer::CheckThisCPUcoreCount()
{
	SYSTEM_INFO si;
	GetSystemInfo(&si);
	m_cpuCore = static_cast<int>(si.dwNumberOfProcessors) * 2;
	printf("CPU Core Count = %d, threads = %d\n", m_cpuCore / 2, m_cpuCore);
}

void boostAsioServer::start_io_service()
{
	m_worker_threads.reserve(m_cpuCore);

	for (int i = 0; i < m_cpuCore; ++i) { m_worker_threads.emplace_back(new thread{ [&]() -> void { g_io_service.run(); } }); }

	while (m_ServerShutdown) { Sleep(1000); }

	// workerThread ߵ
	for (auto thread : m_worker_threads) {
		thread->join();
		delete thread;
	}
}

void boostAsioServer::acceptThread()
{
	m_acceptor.async_accept(m_socket, [&](boost::system::error_code error_code) {
		if (true == (!error_code)) {
			cout << "Client No. [ " << ++m_playerIndex << " ] Connected \t:: IP = " << m_socket.remote_endpoint().address().to_string() << ", Port = " << m_socket.remote_endpoint().port() << "\n";
			g_clients.emplace_back(new player_session (std::move(m_socket), m_playerIndex));
			g_clients[m_playerIndex]->Init();
		}
		if (false == m_ServerShutdown) { acceptThread(); }
	});
}

// player_session class ------------------------------------------------------------------------------------------------------------------------

void player_session::Init()
{
	m_connect_state = true;

	// ⺻  ʱȭ
	Packet init_this_player_buf[MAX_BUF_SIZE];

	init_this_player_buf[0] = sizeof(player_data) + 2;
	init_this_player_buf[1] = INIT_CLIENT;

	m_player_data.id = m_id;
	m_player_data.pos.x = 100;
	m_player_data.pos.y = 100;

	memcpy(&init_this_player_buf[2], g_clients[m_id]->get_player_data(), init_this_player_buf[0]);
	g_clients[m_id]->send_packet(init_this_player_buf);

	// ʱȭ   2 -   ٸ ֵ , ٸ ֵ
	/*Packet other_info_to_me_buf[MAX_BUF_SIZE];*/
	Packet my_info_to_other_buf[MAX_BUF_SIZE];

	/*other_info_to_me_buf[0] =*/ my_info_to_other_buf[0] = sizeof(player_data) + 2;
	/*other_info_to_me_buf[1] =*/ my_info_to_other_buf[1] = KEYINPUT;
	my_info_to_other_buf[1] = INIT_CLIENT;

	//    ٸ ÷̾
	memcpy(&my_info_to_other_buf[2], &m_player_data, my_info_to_other_buf[0]);

	for (auto players : g_clients)
	{
		if (DISCONNECTED == players->get_current_connect_state()) { continue; }
		if (m_id == players->get_id()) { continue; }

		// ٸ ֵ  ؼ ְ,   ... (   ޴° ˼  ü... )
		/*memcpy(&other_info_to_me_buf[2], players->get_player_data(), other_info_to_me_buf[0] - 2);
		send_packet(other_info_to_me_buf);*/

		//    ٸ ֵ  Ǵµ..
		players->send_packet(my_info_to_other_buf);
	}

	/*
		ó ÷̾,  ÷̾  ˸
		view list   ߰ ~ !!
	*/

	m_recv_packet();
}

void player_session::m_recv_packet()
{
	//auto self(shared_from_this());
	m_socket.async_read_some(boost::asio::buffer(m_recv_buf, MAX_BUF_SIZE), [&](boost::system::error_code error_code, std::size_t length) -> void {
		if (error_code) {
			if (error_code.value() == boost::asio::error::operation_aborted) { return; }
			// client was disconnected
			if (false == g_clients[m_id]->get_current_connect_state()) { return; }

			cout << "Client No. [ " << m_id << " ] Disonnected \t:: IP = " << m_socket.remote_endpoint().address().to_string() << ", Port = " << m_socket.remote_endpoint().port() << "\n";
			m_socket.shutdown(m_socket.shutdown_both);
			m_socket.close();

			m_connect_state = false;

			/*
				ó ÷̾,  ÷̾  ˸
				view list    ~ !!
			*/

			Packet temp_buf[MAX_BUF_SIZE] = { sizeof(player_data) + 2, PLAYER_DISCONNECTED };
			memcpy(&temp_buf[2], &m_player_data, temp_buf[0]);

			for (auto players : g_clients)
			{
				if (DISCONNECTED == players->m_connect_state) { continue; }
				if (m_id == players->m_id) { continue; }

				players->send_packet(temp_buf);
			}

			return;
		}

		int current_data_processing = static_cast<int>(length);
		Packet *buf = m_recv_buf;
		while (0 < current_data_processing) {
			if (0 == m_packet_size_current) {
				m_packet_size_current = buf[0];
				if (buf[0] > MAX_BUF_SIZE) {
					cout << "player_session::m_recv_packet() Error, Client No. [ " << m_id << " ] recv buf[0] is out of MAX_BUF_SIZE\n";
					exit(-1);
				}
			}
			int need_to_build = m_packet_size_current - m_packet_size_previous;
			if (need_to_build <= current_data_processing) {
				// Packet building Complete & Process Packet
				memcpy(m_data_buf + m_packet_size_previous, buf, need_to_build);

				m_process_packet(m_id, m_data_buf);

				m_packet_size_current = 0;
				m_packet_size_previous = 0;
				current_data_processing -= need_to_build;
				buf += need_to_build;
			}
			else {
				// Packet build continue
				memcpy(m_data_buf + m_packet_size_previous, buf, current_data_processing);
				m_packet_size_previous += current_data_processing;
				current_data_processing = 0;
				buf += current_data_processing;
			}
		}
		m_recv_packet();
	});
}

void player_session::send_packet(Packet *packet)
{
	int packet_size = packet[0];
	Packet *sendBuf = new Packet[packet_size];
	memcpy(sendBuf, packet, packet_size);

	//auto self(shared_from_this());
	m_socket.async_write_some(boost::asio::buffer(sendBuf, packet_size), [=](boost::system::error_code error_code, std::size_t bytes_transferred) -> void {
		if (!error_code) {
			if (packet_size != bytes_transferred) { cout << "Client No. [ " << m_id << " ] async_write_some packet bytes was NOT SAME !!\n"; }
			delete[] sendBuf;
		}
	});
}

void player_session::m_process_packet(const unsigned int& id, Packet buf[])
{
	// packet[0] = packet size		> 0° ڸ , Ŷ ũⰡ ߸ Ѵ.
	// packet[1] = type				> 1° ڸ  Ŷ  Ŷ Ӽ ִ ̴.
	// packet[...] = data			> 2° ʹ Ӽ ´  ó ش.

	// buf[1] ° Ӽ з  ڿ, ο 2°  ͸ óϱ Ѵ.

	{
		switch (buf[1])
		{
		case TEST:
			//  Ŷ ״ ش.
			cout << "Client No. [ " << m_id << " ] TEST Packet Recived !!\n";
			printf("buf[0] = %d, buf[1] = %d, buf[2] = %d\n\n", buf[0], buf[1], buf[2]);
			g_clients[id]->send_packet(buf);
			break;

		case KEYINPUT:
			g_clients[id]->m_player_data.pos = reinterpret_cast<player_data*>(&buf[2])->pos;

			// ʿ ֵ ̵  ѷ
			for (auto players : g_clients)
			{
				if (DISCONNECTED == players->m_connect_state) { continue; }
				if (id == players->m_id) { continue; }

				players->send_packet(buf);
			}

			break;
		default:
			break;
		}
	}
}<|endoftext|>


<commit_before>/*
 *
 * Copyright (C) 2000 Frans Kaashoek (kaashoek@lcs.mit.edu)
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
 * USA
 *
 */

#include "chord.h"
#include "dhash.h"
#include "parseopt.h"

EXITFN (cleanup);

#define MAX_VNODES 1024

ptr<chord> chordnode;
static str p2psocket;
int do_cache;
str db_name;
dhash *dh[MAX_VNODES + 1];
int ndhash = 0;
int nreplica;

void stats ();
void print ();

void
client_accept (int fd)
{
  if (fd < 0)
    fatal ("EOF\n");
  ref<axprt_stream> x = axprt_stream::alloc (fd);
  dhashclient *c = New dhashclient (x, nreplica, chordnode);
  if (do_cache) c->set_caching (1);
  else c->set_caching (0);
}

static void
client_accept_socket (int lfd)
{
  sockaddr_un sun;
  bzero (&sun, sizeof (sun));
  socklen_t sunlen = sizeof (sun);
  int fd = accept (lfd, reinterpret_cast<sockaddr *> (&sun), &sunlen);
  if (fd >= 0)
    client_accept (fd);
}

static void
client_listening_cb (int fd, int status)
{
  if (status)
    close (fd);
  else if (listen (fd, 5) < 0) {
    warn ("not listening for sfskey: listen: %m\n");
    close (fd);
  }
  else {
    fdcb (fd, selread, wrap (client_accept_socket, fd));
  }
}

static void
cleanup ()
{
   unlink (p2psocket);
}

static void
startclntd()
{
  sockaddr_un sun;
  if (p2psocket.len () >= sizeof (sun.sun_path)) {
    warn ("not listening on socket: path too long: %s\n", p2psocket.cstr ());
    return;
  }
  int clntfd = socket (AF_UNIX, SOCK_STREAM, 0);
  if (clntfd < 0) {
    warn ("not listening on socket: socket: %m\n");
    return;
  }
  make_async (clntfd);

  bzero (&sun, sizeof (sun));
  sun.sun_family = AF_UNIX;
  strcpy (sun.sun_path, p2psocket);

#if 0
  pid_t pid = afork ();
  if (pid == -1) {
    warn ("not listening on socket: fork: %m\n");
    return;
  }
  else if (!pid) {
#endif
    umask (077);
    if (bind (clntfd, (sockaddr *) &sun, sizeof (sun)) < 0) {
      if (errno == EADDRINUSE)
	unlink (sun.sun_path);
      if (bind (clntfd, (sockaddr *) &sun, sizeof (sun)) < 0) {
	warn ("not listening on socket: %s: %m\n", sun.sun_path);
	err_flush ();
	_exit (1);
      }
    }
#if 0
    err_flush ();
    _exit (0);
  }
  chldcb (pid, wrap (start_listening_cb, clntfd));
#endif
  client_listening_cb (clntfd, 0);
}

static void
newvnode_cb (int nreplica, int n, vnode *my)
{
  str db_name_prime = strbuf () << db_name << "-" << n;
  if (ndhash == MAX_VNODES) fatal << "Too many virtual nodes (1024)\n";
  dh[ndhash++] = New dhash (db_name_prime, my, nreplica);
  if (n > 0) chordnode->newvnode (wrap (newvnode_cb, nreplica, n-1));
}

static void
initID (str s, chordID *ID)
{
  chordID n (s);
  chordID b (1);
  b = b << NBIT;
  b = b - 1;
  *ID = n & b;
}

static void
parseconfigfile (str cf, int nvnode, int set_rpcdelay)
{
  parseargs pa (cf);
  bool errors = false;
  int line;
  vec<str> av;
  int myport;
  str myhost;
  str wellknownhost;
  int wellknownport;
  bool myid = false;
  chordID myID;
  chordID wellknownID;
  int ss = 10000;
  int cs = 1000;
  myport = 0;
  int max_connections = 400;
  int max_loccache = 250;

  while (pa.getline (&av, &line)) {
    if (!strcasecmp (av[0], "#")) {
    } else if (!strcasecmp (av[0], "myport")) {
      if (av.size () != 2 || !convertint (av[1], &myport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: myport <number>\n";
      }
    } else if (!strcasecmp (av[0], "myname")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: myname <string>\n";
      }
      myhost = av[1];
    } else if (!strcasecmp (av[0], "wellknownport")) {
      if (av.size () != 2 || !convertint (av[1], &wellknownport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownport <number>\n";
      }
    } else if (!strcasecmp (av[0], "myID")) {
      if (av.size () != 2) {
        errors = true;
       warn << cf << ":" << line << ": usage: myID <number>\n";
      } else {
	initID (av[1], &myID);
	myid = true;
      }
    } else if (!strcasecmp (av[0], "wellknownID")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownID <number>\n";
      } else {
	initID (av[1], &wellknownID);
      }
    } else if (!strcasecmp (av[0], "wellknownport")) {
      if (av.size () != 2 || !convertint (av[1], &wellknownport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownport <number>\n";
      }
    } else if (!strcasecmp (av[0], "nreplica")) {
      if (av.size () != 2 || !convertint (av[1], &nreplica)) {
        errors = true;
        warn << cf << ":" << line << ": usage: nreplica <number>\n";
      }
   } else if (!strcasecmp (av[0], "wellknownhost")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownhost <hostname>\n";
      }
      else
        wellknownhost = av[1];
   } else if (!strcasecmp (av[0], "storesize")) {
     if (av.size () != 2 || !convertint (av[1], &ss)) {
       errors = true;
       warn << cf << ":" << line << ": usage: storesize <size in elements>\n";
     }
   } else if (!strcasecmp (av[0], "cachesize")) {
     if (av.size () != 2 || !convertint (av[1], &cs)) {
       errors = true;
       warn << cf << ":" << line << ": usage: cachesize <size in elements>\n";
     }
   } else if (!strcasecmp (av[0], "maxopenconnections")) {
     if (av.size () != 2 || !convertint (av[1], &max_connections)) {
       errors = true;
       warn << cf << ":" << line << ": usage: maxopenconnections <number>\n";
     }
   } else if (!strcasecmp (av[0], "maxlocationcache")) {
     if (av.size () != 2 || !convertint (av[1], &max_loccache)) {
       errors = true;
       warn << cf << ":" << line << ": usage: maxlocationcache <number>\n";
     }
   }
  }
  if (errors) {
    fatal ("errors in config file\n");
  }
  if (!myhost) {
    myhost = myname ();
  }
  chordnode = New refcounted<chord> (wellknownhost, wellknownport,
				     wellknownID, myport, myhost, set_rpcdelay,
				     max_loccache, max_connections);
  if (myid) chordnode->newvnode (myID, wrap (newvnode_cb, nreplica, nvnode-1));
  else chordnode->newvnode (wrap (newvnode_cb, nreplica, nvnode-1));
  sigcb(SIGUSR1, wrap (&stats));
  sigcb(SIGUSR2, wrap (&print));
}

void
stats ()
{
  chordnode->stats ();
  for (int i = 0 ; i < ndhash; i++)
    dh[i]->print_stats ();
}

void
print ()
{
  chordnode->print ();
}
static void
usage ()
{
  warnx << "Usage: " << progname
	<< "-d <dbfile> -S <sock> -v <nvnode> -c <cache?> -f <conffile>\n";
  exit (1);
}

int
main (int argc, char **argv)
{
  int vnode = 1;
  setprogname (argv[0]);
  sfsconst_init ();
  random_init ();
  int ch;
  do_cache = 0;
  int set_name = 0;
  int set_rpcdelay = 0;

  while ((ch = getopt (argc, argv, "d:S:v:f:c")) != -1)
    switch (ch) {
    case 'S':
      p2psocket = optarg;
      break;
    case 'v':
      vnode = atoi (optarg);
      break;
    case 'r':
      set_rpcdelay = atoi(optarg);
      break;
    case 'f':
      if (!set_name) fatal("must specify db name\n");
      parseconfigfile (optarg, vnode, set_rpcdelay);
      break;
    case 'c':
      do_cache = 1;
      break;
    case 'd':
      db_name = optarg;
      set_name = 1;
      break;
    default:
      usage ();
      break;
    }

  if (!chordnode)
    fatal ("Specify config file\n");
  if (p2psocket)
    startclntd();
  amain ();
}

<commit_msg>stupid--scale location table with number of virtual nodes<commit_after>/*
 *
 * Copyright (C) 2000 Frans Kaashoek (kaashoek@lcs.mit.edu)
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
 * USA
 *
 */

#include "chord.h"
#include "dhash.h"
#include "parseopt.h"

EXITFN (cleanup);

#define MAX_VNODES 1024

ptr<chord> chordnode;
static str p2psocket;
int do_cache;
str db_name;
dhash *dh[MAX_VNODES + 1];
int ndhash = 0;
int nreplica;

void stats ();
void print ();

void
client_accept (int fd)
{
  if (fd < 0)
    fatal ("EOF\n");
  ref<axprt_stream> x = axprt_stream::alloc (fd);
  dhashclient *c = New dhashclient (x, nreplica, chordnode);
  if (do_cache) c->set_caching (1);
  else c->set_caching (0);
}

static void
client_accept_socket (int lfd)
{
  sockaddr_un sun;
  bzero (&sun, sizeof (sun));
  socklen_t sunlen = sizeof (sun);
  int fd = accept (lfd, reinterpret_cast<sockaddr *> (&sun), &sunlen);
  if (fd >= 0)
    client_accept (fd);
}

static void
client_listening_cb (int fd, int status)
{
  if (status)
    close (fd);
  else if (listen (fd, 5) < 0) {
    warn ("not listening for sfskey: listen: %m\n");
    close (fd);
  }
  else {
    fdcb (fd, selread, wrap (client_accept_socket, fd));
  }
}

static void
cleanup ()
{
   unlink (p2psocket);
}

static void
startclntd()
{
  sockaddr_un sun;
  if (p2psocket.len () >= sizeof (sun.sun_path)) {
    warn ("not listening on socket: path too long: %s\n", p2psocket.cstr ());
    return;
  }
  int clntfd = socket (AF_UNIX, SOCK_STREAM, 0);
  if (clntfd < 0) {
    warn ("not listening on socket: socket: %m\n");
    return;
  }
  make_async (clntfd);

  bzero (&sun, sizeof (sun));
  sun.sun_family = AF_UNIX;
  strcpy (sun.sun_path, p2psocket);

#if 0
  pid_t pid = afork ();
  if (pid == -1) {
    warn ("not listening on socket: fork: %m\n");
    return;
  }
  else if (!pid) {
#endif
    umask (077);
    if (bind (clntfd, (sockaddr *) &sun, sizeof (sun)) < 0) {
      if (errno == EADDRINUSE)
	unlink (sun.sun_path);
      if (bind (clntfd, (sockaddr *) &sun, sizeof (sun)) < 0) {
	warn ("not listening on socket: %s: %m\n", sun.sun_path);
	err_flush ();
	_exit (1);
      }
    }
#if 0
    err_flush ();
    _exit (0);
  }
  chldcb (pid, wrap (start_listening_cb, clntfd));
#endif
  client_listening_cb (clntfd, 0);
}

static void
newvnode_cb (int nreplica, int n, vnode *my)
{
  str db_name_prime = strbuf () << db_name << "-" << n;
  if (ndhash == MAX_VNODES) fatal << "Too many virtual nodes (1024)\n";
  dh[ndhash++] = New dhash (db_name_prime, my, nreplica);
  if (n > 0) chordnode->newvnode (wrap (newvnode_cb, nreplica, n-1));
}

static void
initID (str s, chordID *ID)
{
  chordID n (s);
  chordID b (1);
  b = b << NBIT;
  b = b - 1;
  *ID = n & b;
}

static void
parseconfigfile (str cf, int nvnode, int set_rpcdelay)
{
  parseargs pa (cf);
  bool errors = false;
  int line;
  vec<str> av;
  int myport;
  str myhost;
  str wellknownhost;
  int wellknownport;
  bool myid = false;
  chordID myID;
  chordID wellknownID;
  int ss = 10000;
  int cs = 1000;
  myport = 0;
  int max_connections = 400;
  int max_loccache = 100;

  while (pa.getline (&av, &line)) {
    if (!strcasecmp (av[0], "#")) {
    } else if (!strcasecmp (av[0], "myport")) {
      if (av.size () != 2 || !convertint (av[1], &myport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: myport <number>\n";
      }
    } else if (!strcasecmp (av[0], "myname")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: myname <string>\n";
      }
      myhost = av[1];
    } else if (!strcasecmp (av[0], "wellknownport")) {
      if (av.size () != 2 || !convertint (av[1], &wellknownport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownport <number>\n";
      }
    } else if (!strcasecmp (av[0], "myID")) {
      if (av.size () != 2) {
        errors = true;
       warn << cf << ":" << line << ": usage: myID <number>\n";
      } else {
	initID (av[1], &myID);
	myid = true;
      }
    } else if (!strcasecmp (av[0], "wellknownID")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownID <number>\n";
      } else {
	initID (av[1], &wellknownID);
      }
    } else if (!strcasecmp (av[0], "wellknownport")) {
      if (av.size () != 2 || !convertint (av[1], &wellknownport)) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownport <number>\n";
      }
    } else if (!strcasecmp (av[0], "nreplica")) {
      if (av.size () != 2 || !convertint (av[1], &nreplica)) {
        errors = true;
        warn << cf << ":" << line << ": usage: nreplica <number>\n";
      }
   } else if (!strcasecmp (av[0], "wellknownhost")) {
      if (av.size () != 2) {
        errors = true;
        warn << cf << ":" << line << ": usage: wellknownhost <hostname>\n";
      }
      else
        wellknownhost = av[1];
   } else if (!strcasecmp (av[0], "storesize")) {
     if (av.size () != 2 || !convertint (av[1], &ss)) {
       errors = true;
       warn << cf << ":" << line << ": usage: storesize <size in elements>\n";
     }
   } else if (!strcasecmp (av[0], "cachesize")) {
     if (av.size () != 2 || !convertint (av[1], &cs)) {
       errors = true;
       warn << cf << ":" << line << ": usage: cachesize <size in elements>\n";
     }
   } else if (!strcasecmp (av[0], "maxopenconnections")) {
     if (av.size () != 2 || !convertint (av[1], &max_connections)) {
       errors = true;
       warn << cf << ":" << line << ": usage: maxopenconnections <number>\n";
     }
   } else if (!strcasecmp (av[0], "maxlocationcache")) {
     if (av.size () != 2 || !convertint (av[1], &max_loccache)) {
       errors = true;
       warn << cf << ":" << line << ": usage: maxlocationcache <number>\n";
     }
   }
  }
  if (errors) {
    fatal ("errors in config file\n");
  }
  if (!myhost) {
    myhost = myname ();
  }
  max_loccache = max_loccache * nvnode;
  chordnode = New refcounted<chord> (wellknownhost, wellknownport,
				     wellknownID, myport, myhost, set_rpcdelay,
				     max_loccache, max_connections);
  if (myid) chordnode->newvnode (myID, wrap (newvnode_cb, nreplica, nvnode-1));
  else chordnode->newvnode (wrap (newvnode_cb, nreplica, nvnode-1));
  sigcb(SIGUSR1, wrap (&stats));
  sigcb(SIGUSR2, wrap (&print));
}

void
stats ()
{
  chordnode->stats ();
  for (int i = 0 ; i < ndhash; i++)
    dh[i]->print_stats ();
}

void
print ()
{
  chordnode->print ();
}
static void
usage ()
{
  warnx << "Usage: " << progname
	<< "-d <dbfile> -S <sock> -v <nvnode> -c <cache?> -f <conffile>\n";
  exit (1);
}

int
main (int argc, char **argv)
{
  int vnode = 1;
  setprogname (argv[0]);
  sfsconst_init ();
  random_init ();
  int ch;
  do_cache = 0;
  int set_name = 0;
  int set_rpcdelay = 0;

  while ((ch = getopt (argc, argv, "d:S:v:f:c")) != -1)
    switch (ch) {
    case 'S':
      p2psocket = optarg;
      break;
    case 'v':
      vnode = atoi (optarg);
      break;
    case 'r':
      set_rpcdelay = atoi(optarg);
      break;
    case 'f':
      if (!set_name) fatal("must specify db name\n");
      parseconfigfile (optarg, vnode, set_rpcdelay);
      break;
    case 'c':
      do_cache = 1;
      break;
    case 'd':
      db_name = optarg;
      set_name = 1;
      break;
    default:
      usage ();
      break;
    }

  if (!chordnode)
    fatal ("Specify config file\n");
  if (p2psocket)
    startclntd();
  amain ();
}

<|endoftext|>


<commit_before>
/**
 * @file /home/ryan/programming/atl/test/lists.cpp
 * @author Ryan Domigan <ryan_domigan@sutdents@uml.edu>
 * Created on Feb 14, 2015
 */

#include "../atl.hpp"
#include "../helpers.hpp"

#include "../debug.hpp"

#include <iostream>
#include <vector>

#include <gtest/gtest.h>

using namespace atl;
using namespace std;

struct ListTest : public ::testing::Test
{
    Atl atl;

    // Assert that two simple lists (nested lists, Symbols, and
    // Fixnums) are equivalent.
    template<class Range0, class Range1>
    void assert_equiv(Range0 const& aa, Range1 const& bb)
    {
        auto aitr = aa.begin(), bitr = bb.begin();
        for(;
            (aitr != aa.end()) && (bitr != bb.end());
            ++aitr, ++bitr)
            {
                ASSERT_EQ(((*aitr)._tag),
                          ((*bitr)._tag));
                switch((*aitr)._tag)
                    {
                    case tag<Symbol>::value:
                        ASSERT_EQ((unwrap<std::string>(*aitr)),
                                  (unwrap<std::string>(*bitr)));
                        break;
                    case tag<Ast>::value:
                        assert_equiv((unwrap<Ast>(*aitr)),
                                     (unwrap<Ast>(*bitr)));
                        break;
                    default:
                        ASSERT_EQ((*aitr),
                                  (*bitr));
                        break;
                    }
            }
        ASSERT_EQ((aitr == aa.end()),
                  (bitr == bb.end()));
    }

    bool _assert_not_equiv(Ast const& aa, Ast const& bb)
    {
        auto aitr = aa.begin(), bitr = bb.begin();
        for(;
            (aitr != aa.end()) || (bitr != bb.end());
            ++aitr, ++bitr)
            {
                if((*aitr)._tag != (*bitr)._tag)
                    return true;

                switch((*aitr)._tag)
                    {
                    case tag<Symbol>::value:
                        if(unwrap<std::string>(*aitr) != unwrap<std::string>(*bitr))
                            return true;
                        break;
                    case tag<Ast>::value:
                        if(_assert_not_equiv((unwrap<Ast>(*aitr)),
                                             (unwrap<Ast>(*bitr))))
                            return true;
                        break;
                    default:
                        if((*aitr).value != (*bitr).value)
                            return true;
                        break;
                    }
            }
        return (aitr != aa.end()) && (bitr != bb.end());
    }

    void assert_not_equiv(Ast const& aa, Ast const& bb)
    { ASSERT_TRUE(_assert_not_equiv(aa, bb)); }
};

TEST_F(ListTest, test_equiv_ast)
{
    // Test manually constructed asts and my test harness's equiv
    // function.
    auto seq = atl.gc.dynamic_seq();
    auto car = seq->push_seq<Ast>();

    seq->push_back(wrap<Fixnum>(1));
    seq->push_back(wrap<Fixnum>(2));
    car->end_at(seq->end());

    auto seq3 = atl.gc.dynamic_seq();
    car = seq3->push_seq<Ast>();

    seq3->push_back(wrap<Fixnum>(1));
    seq3->push_back(wrap<Fixnum>(2));
    car->end_at(seq3->end());

    auto seq2 = atl.gc.dynamic_seq();
    car = seq2->push_seq<Ast>();

    seq2->push_back(wrap<Fixnum>(1));
    seq2->push_back(wrap<Fixnum>(2));
    seq2->push_back(wrap<Fixnum>(3));
    car->end_at(seq2->end());

    assert_not_equiv((unwrap<Ast>(*seq->begin())),
                     (unwrap<Ast>(*seq2->begin())));

    assert_equiv((unwrap<Ast>(*seq->begin())),
                 (unwrap<Ast>(*seq3->begin())));
}

TEST_F(ListTest, test_quote)
{
    // Constructing a list from function calls is complicated by the
    // need to tag the list memebers.  Make sure quoting works first.
    {
        auto rval = atl.string_("'(1 2 (a b))");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 2 (a b))"))));
    }

    {
        auto rval = atl.string_("'(1)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1)"))));
    }
}

TEST_F(ListTest, test_indexing)
{
    {
        auto rval = atl.string_("(nth '(1 2 3) 0)");
        ASSERT_EQ((unwrap<Fixnum>(*unwrap<Pointer>(rval).value).value),
                  1);
    }

    {
        auto rval = atl.string_("(nth '(1 2 3) 1)");
        ASSERT_EQ((unwrap<Fixnum>(*unwrap<Pointer>(rval).value).value),
                  2);
    }
}

TEST_F(ListTest, test_slice)
{
    auto result = atl.string_("(slice '(1 2 3 4) 2)");

    assert_equiv((unwrap<Slice>(result)),
                 (unwrap<Ast>(atl.parse.string_("(3 4)"))));
}

TEST_F(ListTest, test_make_ast)
{
	using namespace make_ast;

	{
		auto ast = make(lift(1), lift(2), lift(3))
			(*atl.env.gc.dynamic_seq());

		assert_equiv((*ast),
		             unwrap<Ast>(atl.parse.string_("(1 2 3)")));
	}

	{
		auto ast = make(lift(1),
		                make(lift(2), lift(3)),
		                lift(4))
			(*atl.env.gc.dynamic_seq());

		assert_equiv((*ast),
		             unwrap<Ast>(atl.parse.string_("(1 (2 3) 4)")));
	}
}

TEST_F(ListTest, test_list)
{
    {
        auto rval = atl.string_("(Ast 1 2 3)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 2 3)"))));
    }

    {
        auto rval = atl.string_("(Ast)");
        assert_equiv((unwrap<Ast>(rval)),
                     (*make_ast::make()(*atl.env.gc.dynamic_seq())));
    }

    {
        auto rval = atl.string_("(Ast 1 (Ast 2 3))");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 (2 3))"))));
    }

    {
        auto rval = atl.string_("(Ast 1 (Ast 2 3) 4)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 (2 3) 4)"))));
    }
}

TEST_F(ListTest, test_cons)
{
    auto result = atl.string_("(cons 0 '(1))");
    cout << printer::any(result) << endl;
}

<commit_msg>Test embedded Asts with quote and nth<commit_after>
/**
 * @file /home/ryan/programming/atl/test/lists.cpp
 * @author Ryan Domigan <ryan_domigan@sutdents@uml.edu>
 * Created on Feb 14, 2015
 */

#include "../atl.hpp"
#include "../helpers.hpp"

#include "../debug.hpp"

#include <iostream>
#include <vector>

#include <gtest/gtest.h>

using namespace atl;
using namespace std;

struct ListTest : public ::testing::Test
{
    Atl atl;

    // Assert that two simple lists (nested lists, Symbols, and
    // Fixnums) are equivalent.
    template<class Range0, class Range1>
    void assert_equiv(Range0 const& aa, Range1 const& bb)
    {
        auto aitr = aa.begin(), bitr = bb.begin();
        for(;
            (aitr != aa.end()) && (bitr != bb.end());
            ++aitr, ++bitr)
            {
                ASSERT_EQ(((*aitr)._tag),
                          ((*bitr)._tag));
                switch((*aitr)._tag)
                    {
                    case tag<Symbol>::value:
                        ASSERT_EQ((unwrap<std::string>(*aitr)),
                                  (unwrap<std::string>(*bitr)));
                        break;
                    case tag<Ast>::value:
                        assert_equiv((unwrap<Ast>(*aitr)),
                                     (unwrap<Ast>(*bitr)));
                        break;
                    default:
                        ASSERT_EQ((*aitr),
                                  (*bitr));
                        break;
                    }
            }
        ASSERT_EQ((aitr == aa.end()),
                  (bitr == bb.end()));
    }

    bool _assert_not_equiv(Ast const& aa, Ast const& bb)
    {
        auto aitr = aa.begin(), bitr = bb.begin();
        for(;
            (aitr != aa.end()) || (bitr != bb.end());
            ++aitr, ++bitr)
            {
                if((*aitr)._tag != (*bitr)._tag)
                    return true;

                switch((*aitr)._tag)
                    {
                    case tag<Symbol>::value:
                        if(unwrap<std::string>(*aitr) != unwrap<std::string>(*bitr))
                            return true;
                        break;
                    case tag<Ast>::value:
                        if(_assert_not_equiv((unwrap<Ast>(*aitr)),
                                             (unwrap<Ast>(*bitr))))
                            return true;
                        break;
                    default:
                        if((*aitr).value != (*bitr).value)
                            return true;
                        break;
                    }
            }
        return (aitr != aa.end()) && (bitr != bb.end());
    }

    void assert_not_equiv(Ast const& aa, Ast const& bb)
    { ASSERT_TRUE(_assert_not_equiv(aa, bb)); }
};

TEST_F(ListTest, test_equiv_ast)
{
    // Test manually constructed asts and my test harness's equiv
    // function.
    auto seq = atl.gc.dynamic_seq();
    auto car = seq->push_seq<Ast>();

    seq->push_back(wrap<Fixnum>(1));
    seq->push_back(wrap<Fixnum>(2));
    car->end_at(seq->end());

    auto seq3 = atl.gc.dynamic_seq();
    car = seq3->push_seq<Ast>();

    seq3->push_back(wrap<Fixnum>(1));
    seq3->push_back(wrap<Fixnum>(2));
    car->end_at(seq3->end());

    auto seq2 = atl.gc.dynamic_seq();
    car = seq2->push_seq<Ast>();

    seq2->push_back(wrap<Fixnum>(1));
    seq2->push_back(wrap<Fixnum>(2));
    seq2->push_back(wrap<Fixnum>(3));
    car->end_at(seq2->end());

    assert_not_equiv((unwrap<Ast>(*seq->begin())),
                     (unwrap<Ast>(*seq2->begin())));

    assert_equiv((unwrap<Ast>(*seq->begin())),
                 (unwrap<Ast>(*seq3->begin())));
}

TEST_F(ListTest, test_quote)
{
    // Constructing a list from function calls is complicated by the
    // need to tag the list memebers.  Make sure quoting works first.
    {
        auto rval = atl.string_("'(1 2 (a b))");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 2 (a b))"))));
    }

    {
        auto rval = atl.string_("'(1)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1)"))));
    }
}

// Check that I can quote an embedded Ast
TEST_F(ListTest, test_quote_embedded)
{
	using namespace make_ast;
	Arena arena;

	auto inner = make
		(lift(1), lift(2), lift(3))
		(*arena.dynamic_seq());

	auto expr = make
		(lift<Quote>(),
		 lift(*inner))
		(*arena.dynamic_seq());

	assert_equiv(*expr,
	             unwrap<Ast>(atl.parse.string_("'(1 2 3)")));
}

TEST_F(ListTest, test_indexing)
{
    {
        auto rval = atl.string_("(nth '(1 2 3) 0)");
        ASSERT_EQ((unwrap<Fixnum>(*unwrap<Pointer>(rval).value).value),
                  1);
    }

    {
        auto rval = atl.string_("(nth '(1 2 3) 1)");
        ASSERT_EQ((unwrap<Fixnum>(*unwrap<Pointer>(rval).value).value),
                  2);
    }
}

// Check that I can get the index of an embedded Ast
TEST_F(ListTest, test_index_embedded)
{
	using namespace make_ast;
	Arena arena;

	auto inner = make
		(lift<Quote>(),
		 make(lift(1), lift(2), lift(3)))
		(*arena.dynamic_seq());

	auto expr = make
		(lift(arena.amake<Symbol>("nth")),
		 lift(*inner),
		 lift(1))
		(*arena.dynamic_seq());

	auto rval = atl.eval(wrap(*expr));
	ASSERT_EQ(unwrap<Fixnum>(*unwrap<Pointer>(rval).value).value,
	          2);
}

TEST_F(ListTest, test_slice)
{
	auto result = atl.string_("(slice '(1 2 3 4) 2)");

    assert_equiv((unwrap<Slice>(result)),
                 (unwrap<Ast>(atl.parse.string_("(3 4)"))));
}

TEST_F(ListTest, test_make_ast)
{
	using namespace make_ast;

	{
		auto ast = make(lift(1), lift(2), lift(3))
			(*atl.env.gc.dynamic_seq());

		assert_equiv((*ast),
		             unwrap<Ast>(atl.parse.string_("(1 2 3)")));
	}

	{
		auto ast = make(lift(1),
		                make(lift(2), lift(3)),
		                lift(4))
			(*atl.env.gc.dynamic_seq());

		assert_equiv((*ast),
		             unwrap<Ast>(atl.parse.string_("(1 (2 3) 4)")));
	}
}

TEST_F(ListTest, test_list)
{
    {
        auto rval = atl.string_("(Ast 1 2 3)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 2 3)"))));
    }

    {
        auto rval = atl.string_("(Ast)");
        assert_equiv((unwrap<Ast>(rval)),
                     (*make_ast::make()(*atl.env.gc.dynamic_seq())));
    }

    {
        auto rval = atl.string_("(Ast 1 (Ast 2 3))");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 (2 3))"))));
    }

    {
        auto rval = atl.string_("(Ast 1 (Ast 2 3) 4)");
        assert_equiv((unwrap<Ast>(rval)),
                     (unwrap<Ast>(atl.parse.string_("(1 (2 3) 4)"))));
    }
}

TEST_F(ListTest, test_cons)
{
    auto result = atl.string_("(cons 0 '(1))");
    cout << printer::any(result) << endl;
}

<|endoftext|>


<commit_before><commit_msg>Fix a typo pointed out by wtc after I checked in.<commit_after><|endoftext|>


<commit_before>//===- ReduceArguments.cpp - Specialized Delta Pass -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a function which calls the Generic Delta pass in order
// to reduce uninteresting Arguments from defined functions.
//
//===----------------------------------------------------------------------===//

#include "ReduceBasicBlocks.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Value.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include <vector>

using namespace llvm;

/// Replaces BB Terminator with one that only contains Chunk BBs
static void replaceBranchTerminator(BasicBlock &BB,
                                    std::set<BasicBlock *> BBsToKeep) {
  auto Term = BB.getTerminator();
  std::vector<BasicBlock *> ChunkSucessors;
  for (auto Succ : successors(&BB))
    if (BBsToKeep.count(Succ))
      ChunkSucessors.push_back(Succ);

  // BB only references Chunk BBs
  if (ChunkSucessors.size() == Term->getNumSuccessors())
    return;

  Term->eraseFromParent();

  if (ChunkSucessors.empty()) {
    ReturnInst::Create(BB.getContext(), nullptr, &BB);
    return;
  }

  if (isa<BranchInst>(Term))
    BranchInst::Create(ChunkSucessors[0], &BB);

  if (auto IndBI = dyn_cast<IndirectBrInst>(Term)) {
    auto NewIndBI =
        IndirectBrInst::Create(IndBI->getAddress(), ChunkSucessors.size(), &BB);
    for (auto Dest : ChunkSucessors)
      NewIndBI->addDestination(Dest);
  }
}

/// Removes uninteresting BBs from switch, if the default case ends up being
/// uninteresting, the switch is replaced with a void return (since it has to be
/// replace with something)
static void removeUninterestingBBsFromSwitch(SwitchInst &SwInst,
                                             std::set<BasicBlock *> BBsToKeep) {
  if (!BBsToKeep.count(SwInst.getDefaultDest())) {
    ReturnInst::Create(SwInst.getContext(), nullptr, SwInst.getParent());
    SwInst.eraseFromParent();
  } else
    for (int I = 0, E = SwInst.getNumCases(); I != E; ++I) {
      auto Case = SwInst.case_begin() + I;
      if (!BBsToKeep.count(Case->getCaseSuccessor())) {
        SwInst.removeCase(Case);
        --I;
        --E;
      }
    }
}

/// Removes out-of-chunk arguments from functions, and modifies their calls
/// accordingly. It also removes allocations of out-of-chunk arguments.
/// @returns the Module stripped of out-of-chunk functions
static void extractBasicBlocksFromModule(std::vector<Chunk> ChunksToKeep,
                                         Module *Program) {
  int I = 0, BBCount = 0;
  std::set<BasicBlock *> BBsToKeep;

  for (auto &F : *Program)
    for (auto &BB : F)
      if (I < (int)ChunksToKeep.size()) {
        if (ChunksToKeep[I].contains(++BBCount))
          BBsToKeep.insert(&BB);
        if (ChunksToKeep[I].end == BBCount)
          ++I;
      }

  std::vector<BasicBlock *> BBsToDelete;
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (!BBsToKeep.count(&BB)) {
        BBsToDelete.push_back(&BB);
        // Remove out-of-chunk BB from successor phi nodes
        for (auto *Succ : successors(&BB))
          Succ->removePredecessor(&BB);
      }
    }

  // Replace terminators that reference out-of-chunk BBs
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (auto *SwInst = dyn_cast<SwitchInst>(BB.getTerminator()))
        removeUninterestingBBsFromSwitch(*SwInst, BBsToKeep);
      else
        replaceBranchTerminator(BB, BBsToKeep);
    }

  // Replace out-of-chunk switch uses
  for (auto &BB : BBsToDelete) {
    // Instructions might be referenced in other BBs
    for (auto &I : *BB)
      I.replaceAllUsesWith(UndefValue::get(I.getType()));
    BB->eraseFromParent();
  }
}

/// Counts the amount of basic blocks and prints their name & respective index
static int countBasicBlocks(Module *Program) {
  // TODO: Silence index with --quiet flag
  outs() << "----------------------------\n";
  int BBCount = 0;
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (BB.hasName())
        outs() << "\t" << ++BBCount << ": " << BB.getName() << "\n";
      else
        outs() << "\t" << ++BBCount << ": Unnamed\n";
    }

  return BBCount;
}

void llvm::reduceBasicBlocksDeltaPass(TestRunner &Test) {
  outs() << "*** Reducing Basic Blocks...\n";
  int BBCount = countBasicBlocks(Test.getProgram());
  runDeltaPass(Test, BBCount, extractBasicBlocksFromModule);
}
<commit_msg>llvm-reduce: Remove inaccurate doxy comment about a return that isn't returned<commit_after>//===- ReduceArguments.cpp - Specialized Delta Pass -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a function which calls the Generic Delta pass in order
// to reduce uninteresting Arguments from defined functions.
//
//===----------------------------------------------------------------------===//

#include "ReduceBasicBlocks.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Value.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include <vector>

using namespace llvm;

/// Replaces BB Terminator with one that only contains Chunk BBs
static void replaceBranchTerminator(BasicBlock &BB,
                                    std::set<BasicBlock *> BBsToKeep) {
  auto Term = BB.getTerminator();
  std::vector<BasicBlock *> ChunkSucessors;
  for (auto Succ : successors(&BB))
    if (BBsToKeep.count(Succ))
      ChunkSucessors.push_back(Succ);

  // BB only references Chunk BBs
  if (ChunkSucessors.size() == Term->getNumSuccessors())
    return;

  Term->eraseFromParent();

  if (ChunkSucessors.empty()) {
    ReturnInst::Create(BB.getContext(), nullptr, &BB);
    return;
  }

  if (isa<BranchInst>(Term))
    BranchInst::Create(ChunkSucessors[0], &BB);

  if (auto IndBI = dyn_cast<IndirectBrInst>(Term)) {
    auto NewIndBI =
        IndirectBrInst::Create(IndBI->getAddress(), ChunkSucessors.size(), &BB);
    for (auto Dest : ChunkSucessors)
      NewIndBI->addDestination(Dest);
  }
}

/// Removes uninteresting BBs from switch, if the default case ends up being
/// uninteresting, the switch is replaced with a void return (since it has to be
/// replace with something)
static void removeUninterestingBBsFromSwitch(SwitchInst &SwInst,
                                             std::set<BasicBlock *> BBsToKeep) {
  if (!BBsToKeep.count(SwInst.getDefaultDest())) {
    ReturnInst::Create(SwInst.getContext(), nullptr, SwInst.getParent());
    SwInst.eraseFromParent();
  } else
    for (int I = 0, E = SwInst.getNumCases(); I != E; ++I) {
      auto Case = SwInst.case_begin() + I;
      if (!BBsToKeep.count(Case->getCaseSuccessor())) {
        SwInst.removeCase(Case);
        --I;
        --E;
      }
    }
}

/// Removes out-of-chunk arguments from functions, and modifies their calls
/// accordingly. It also removes allocations of out-of-chunk arguments.
static void extractBasicBlocksFromModule(std::vector<Chunk> ChunksToKeep,
                                         Module *Program) {
  int I = 0, BBCount = 0;
  std::set<BasicBlock *> BBsToKeep;

  for (auto &F : *Program)
    for (auto &BB : F)
      if (I < (int)ChunksToKeep.size()) {
        if (ChunksToKeep[I].contains(++BBCount))
          BBsToKeep.insert(&BB);
        if (ChunksToKeep[I].end == BBCount)
          ++I;
      }

  std::vector<BasicBlock *> BBsToDelete;
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (!BBsToKeep.count(&BB)) {
        BBsToDelete.push_back(&BB);
        // Remove out-of-chunk BB from successor phi nodes
        for (auto *Succ : successors(&BB))
          Succ->removePredecessor(&BB);
      }
    }

  // Replace terminators that reference out-of-chunk BBs
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (auto *SwInst = dyn_cast<SwitchInst>(BB.getTerminator()))
        removeUninterestingBBsFromSwitch(*SwInst, BBsToKeep);
      else
        replaceBranchTerminator(BB, BBsToKeep);
    }

  // Replace out-of-chunk switch uses
  for (auto &BB : BBsToDelete) {
    // Instructions might be referenced in other BBs
    for (auto &I : *BB)
      I.replaceAllUsesWith(UndefValue::get(I.getType()));
    BB->eraseFromParent();
  }
}

/// Counts the amount of basic blocks and prints their name & respective index
static int countBasicBlocks(Module *Program) {
  // TODO: Silence index with --quiet flag
  outs() << "----------------------------\n";
  int BBCount = 0;
  for (auto &F : *Program)
    for (auto &BB : F) {
      if (BB.hasName())
        outs() << "\t" << ++BBCount << ": " << BB.getName() << "\n";
      else
        outs() << "\t" << ++BBCount << ": Unnamed\n";
    }

  return BBCount;
}

void llvm::reduceBasicBlocksDeltaPass(TestRunner &Test) {
  outs() << "*** Reducing Basic Blocks...\n";
  int BBCount = countBasicBlocks(Test.getProgram());
  runDeltaPass(Test, BBCount, extractBasicBlocksFromModule);
}
<|endoftext|>


<commit_before>#include "util/testing.h"

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

namespace cert_trans {
namespace test {

void InitTesting(const char* name, int* argc, char*** argv,
                 bool remove_flags) {
  // Change the defaults. Can be overridden on command line.
  // Log to stderr instead of log files.
  FLAGS_logtostderr = true;
  // Only log fatal messages by default.
  FLAGS_minloglevel = 3;
  ::testing::InitGoogleTest(argc, *argv);
  google::ParseCommandLineFlags(argc, argv, remove_flags);
  google::InitGoogleLogging(name);
}

}  // namespace test
}  // namespace cert_trans
<commit_msg>Call evthread_use_pthreads() in all tests.<commit_after>#include "util/testing.h"

#include <event2/thread.h>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

namespace cert_trans {
namespace test {

void InitTesting(const char* name, int* argc, char*** argv,
                 bool remove_flags) {
  // Change the defaults. Can be overridden on command line.
  // Log to stderr instead of log files.
  FLAGS_logtostderr = true;
  // Only log fatal messages by default.
  FLAGS_minloglevel = 3;
  ::testing::InitGoogleTest(argc, *argv);
  google::ParseCommandLineFlags(argc, argv, remove_flags);
  google::InitGoogleLogging(name);
  evthread_use_pthreads();
}

}  // namespace test
}  // namespace cert_trans
<|endoftext|>


<commit_before>//  Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/fluid/framework/tensor.h"
#include <gtest/gtest.h>
#include <string>
#include "paddle/fluid/platform/float16.h"

namespace framework = paddle::framework;
namespace platform = paddle::platform;

TEST(Tensor, Dims) {
  framework::Tensor tt;
  tt.Resize({2, 3, 4});
  framework::DDim dims = tt.dims();
  ASSERT_EQ(arity(dims), 3);
  for (int i = 0; i < 3; ++i) {
    EXPECT_EQ(i + 2, dims[i]);
  }
}

TEST(Tensor, DataAssert) {
  framework::Tensor src_tensor;

  bool caught = false;
  try {
    src_tensor.data<double>();
  } catch (platform::EnforceNotMet err) {
    caught = true;
    std::string msg =
        "holder_ should not be null\nTensor holds no memory. Call "
        "Tensor::mutable_data first.";
    const char* what = err.what();
    for (size_t i = 0; i < msg.length(); ++i) {
      ASSERT_EQ(what[i], msg[i]);
    }
  }
  ASSERT_TRUE(caught);
}

TEST(Tensor, MutableData) {
  {
    framework::Tensor src_tensor;
    float* p1 = nullptr;
    float* p2 = nullptr;
    // initialization
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                        platform::CPUPlace());
    EXPECT_NE(p1, nullptr);
    // set src_tensor a new dim with large size
    // momery is supposed to be re-allocated
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 4}),
                                        platform::CPUPlace());
    EXPECT_NE(p2, nullptr);
    EXPECT_NE(p1, p2);
    // set src_tensor a new dim with same size
    // momery block is supposed to be unchanged
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
                                        platform::CPUPlace());
    EXPECT_EQ(p1, p2);
    // set src_tensor a new dim with smaller size
    // momery block is supposed to be unchanged
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2}),
                                        platform::CPUPlace());
    EXPECT_EQ(p1, p2);
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    float* p1 = nullptr;
    float* p2 = nullptr;
    // initialization
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                        platform::CUDAPlace());
    EXPECT_NE(p1, nullptr);
    // set src_tensor a new dim with large size
    // momery is supposed to be re-allocated
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 4}),
                                        platform::CUDAPlace());
    EXPECT_NE(p2, nullptr);
    EXPECT_NE(p1, p2);
    // set src_tensor a new dim with same size
    // momery block is supposed to be unchanged
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
                                        platform::CUDAPlace());
    EXPECT_EQ(p1, p2);
    // set src_tensor a new dim with smaller size
    // momery block is supposed to be unchanged
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2}),
                                        platform::CUDAPlace());
    EXPECT_EQ(p1, p2);
  }
#endif
}

TEST(Tensor, ShareDataWith) {
  {
    framework::Tensor src_tensor;
    framework::Tensor dst_tensor;
    // Try to share data form uninitialized tensor
    bool caught = false;
    try {
      dst_tensor.ShareDataWith(src_tensor);
    } catch (paddle::platform::EnforceNotMet err) {
      caught = true;
      std::string msg =
          "holder_ should not be null\nTensor holds no memory. Call "
          "Tensor::mutable_data first.";
      const char* what = err.what();
      for (size_t i = 0; i < msg.length(); ++i) {
        ASSERT_EQ(what[i], msg[i]);
      }
    }
    ASSERT_TRUE(caught);

    src_tensor.mutable_data<int>(framework::make_ddim({2, 3, 4}),
                                 platform::CPUPlace());
    dst_tensor.ShareDataWith(src_tensor);
    ASSERT_EQ(src_tensor.data<int>(), dst_tensor.data<int>());
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    framework::Tensor dst_tensor;
    src_tensor.mutable_data<int>(framework::make_ddim({2, 3, 4}),
                                 platform::CUDAPlace());
    dst_tensor.ShareDataWith(src_tensor);
    ASSERT_EQ(src_tensor.data<int>(), dst_tensor.data<int>());
  }
#endif
}

TEST(Tensor, Slice) {
  {
    framework::Tensor src_tensor;
    src_tensor.mutable_data<int>(framework::make_ddim({5, 3, 4}),
                                 platform::CPUPlace());
    framework::Tensor slice_tensor = src_tensor.Slice(1, 3);
    framework::DDim slice_dims = slice_tensor.dims();
    ASSERT_EQ(arity(slice_dims), 3);
    EXPECT_EQ(slice_dims[0], 2);
    EXPECT_EQ(slice_dims[1], 3);
    EXPECT_EQ(slice_dims[2], 4);

    uintptr_t src_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.data<int>());
    uintptr_t src_mutable_data_address = reinterpret_cast<uintptr_t>(
        src_tensor.mutable_data<int>(src_tensor.dims(), platform::CPUPlace()));
    uintptr_t slice_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.data<int>());
    uintptr_t slice_mutable_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.mutable_data<int>(
            slice_tensor.dims(), platform::CPUPlace()));
    EXPECT_EQ(src_data_address, src_mutable_data_address);
    EXPECT_EQ(slice_data_address, slice_mutable_data_address);
    EXPECT_EQ(src_data_address + 3 * 4 * 1 * sizeof(int), slice_data_address);
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    src_tensor.mutable_data<double>(framework::make_ddim({6, 9}),
                                    platform::CUDAPlace());
    framework::Tensor slice_tensor = src_tensor.Slice(2, 6);
    framework::DDim slice_dims = slice_tensor.dims();
    ASSERT_EQ(arity(slice_dims), 2);
    EXPECT_EQ(slice_dims[0], 4);
    EXPECT_EQ(slice_dims[1], 9);

    uintptr_t src_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.data<double>());
    uintptr_t src_mutable_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.mutable_data<double>(
            src_tensor.dims(), platform::CUDAPlace()));
    uintptr_t slice_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.data<double>());
    uintptr_t slice_mutable_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.mutable_data<double>(
            slice_tensor.dims(), platform::CUDAPlace()));
    EXPECT_EQ(src_data_address, src_mutable_data_address);
    EXPECT_EQ(slice_data_address, slice_mutable_data_address);
    EXPECT_EQ(src_data_address + 9 * 2 * sizeof(double), slice_data_address);
  }
#endif
}

TEST(Tensor, ReshapeToMatrix) {
  framework::Tensor src;
  int* src_ptr = src.mutable_data<int>({2, 3, 4, 9}, platform::CPUPlace());
  for (int i = 0; i < 2 * 3 * 4 * 9; ++i) {
    src_ptr[i] = i;
  }
  framework::Tensor res = framework::ReshapeToMatrix(src, 2);
  ASSERT_EQ(res.dims()[0], 2 * 3);
  ASSERT_EQ(res.dims()[1], 4 * 9);
}

TEST(Tensor, Layout) {
  framework::Tensor src;
  ASSERT_EQ(src.layout(), framework::DataLayout::kNCHW);
  src.set_layout(framework::DataLayout::kAnyLayout);
  ASSERT_EQ(src.layout(), framework::DataLayout::kAnyLayout);
}

TEST(Tensor, FP16) {
  using platform::float16;
  framework::Tensor src;
  float16* src_ptr = src.mutable_data<float16>({2, 3}, platform::CPUPlace());
  for (int i = 0; i < 2 * 3; ++i) {
    src_ptr[i] = static_cast<float16>(i);
  }
  EXPECT_EQ(src.memory_size(), 2 * 3 * sizeof(float16));
  // EXPECT a human readable error message
  // src.data<uint8_t>();
  // Tensor holds the wrong type, it holds N6paddle8platform7float16E at
  // [/paddle/Paddle/paddle/fluid/framework/tensor_impl.h:43]
}
<commit_msg>add a small test to verify tensor type<commit_after>//  Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/fluid/framework/tensor.h"
#include <gtest/gtest.h>
#include <string>
#include "paddle/fluid/platform/float16.h"

namespace framework = paddle::framework;
namespace platform = paddle::platform;

TEST(Tensor, Dims) {
  framework::Tensor tt;
  tt.Resize({2, 3, 4});
  framework::DDim dims = tt.dims();
  ASSERT_EQ(arity(dims), 3);
  for (int i = 0; i < 3; ++i) {
    EXPECT_EQ(i + 2, dims[i]);
  }
}

TEST(Tensor, DataAssert) {
  framework::Tensor src_tensor;

  bool caught = false;
  try {
    src_tensor.data<double>();
  } catch (platform::EnforceNotMet err) {
    caught = true;
    std::string msg =
        "holder_ should not be null\nTensor holds no memory. Call "
        "Tensor::mutable_data first.";
    const char* what = err.what();
    for (size_t i = 0; i < msg.length(); ++i) {
      ASSERT_EQ(what[i], msg[i]);
    }
  }
  ASSERT_TRUE(caught);
}

TEST(Tensor, MutableData) {
  {
    framework::Tensor src_tensor;
    float* p1 = nullptr;
    float* p2 = nullptr;
    // initialization
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                        platform::CPUPlace());
    EXPECT_NE(p1, nullptr);
    // set src_tensor a new dim with large size
    // momery is supposed to be re-allocated
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 4}),
                                        platform::CPUPlace());
    EXPECT_NE(p2, nullptr);
    EXPECT_NE(p1, p2);
    // set src_tensor a new dim with same size
    // momery block is supposed to be unchanged
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
                                        platform::CPUPlace());
    EXPECT_EQ(p1, p2);
    // set src_tensor a new dim with smaller size
    // momery block is supposed to be unchanged
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2}),
                                        platform::CPUPlace());
    EXPECT_EQ(p1, p2);
  }
  // Not sure if it's desired, but currently, Tensor type can be changed.
  {
    framework::Tensor src_tensor;
    int8_t* p1 = src_tensor.mutable_data<int8_t>(framework::make_ddim({1}),
                                                 platform::CPUPlace());
    EXPECT_NE(p1, nullptr);
    *p1 = 1;

    uint8_t* p2 = src_tensor.mutable_data<uint8_t>(framework::make_ddim({1}),
                                                   platform::CPUPlace());
    EXPECT_NE(p2, nullptr);
    EXPECT_EQ(static_cast<int>(p2[0]), 1);
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    float* p1 = nullptr;
    float* p2 = nullptr;
    // initialization
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({1, 2, 3}),
                                        platform::CUDAPlace());
    EXPECT_NE(p1, nullptr);
    // set src_tensor a new dim with large size
    // momery is supposed to be re-allocated
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({3, 4}),
                                        platform::CUDAPlace());
    EXPECT_NE(p2, nullptr);
    EXPECT_NE(p1, p2);
    // set src_tensor a new dim with same size
    // momery block is supposed to be unchanged
    p1 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2, 3}),
                                        platform::CUDAPlace());
    EXPECT_EQ(p1, p2);
    // set src_tensor a new dim with smaller size
    // momery block is supposed to be unchanged
    p2 = src_tensor.mutable_data<float>(framework::make_ddim({2, 2}),
                                        platform::CUDAPlace());
    EXPECT_EQ(p1, p2);
  }
#endif
}

TEST(Tensor, ShareDataWith) {
  {
    framework::Tensor src_tensor;
    framework::Tensor dst_tensor;
    // Try to share data form uninitialized tensor
    bool caught = false;
    try {
      dst_tensor.ShareDataWith(src_tensor);
    } catch (paddle::platform::EnforceNotMet err) {
      caught = true;
      std::string msg =
          "holder_ should not be null\nTensor holds no memory. Call "
          "Tensor::mutable_data first.";
      const char* what = err.what();
      for (size_t i = 0; i < msg.length(); ++i) {
        ASSERT_EQ(what[i], msg[i]);
      }
    }
    ASSERT_TRUE(caught);

    src_tensor.mutable_data<int>(framework::make_ddim({2, 3, 4}),
                                 platform::CPUPlace());
    dst_tensor.ShareDataWith(src_tensor);
    ASSERT_EQ(src_tensor.data<int>(), dst_tensor.data<int>());
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    framework::Tensor dst_tensor;
    src_tensor.mutable_data<int>(framework::make_ddim({2, 3, 4}),
                                 platform::CUDAPlace());
    dst_tensor.ShareDataWith(src_tensor);
    ASSERT_EQ(src_tensor.data<int>(), dst_tensor.data<int>());
  }
#endif
}

TEST(Tensor, Slice) {
  {
    framework::Tensor src_tensor;
    src_tensor.mutable_data<int>(framework::make_ddim({5, 3, 4}),
                                 platform::CPUPlace());
    framework::Tensor slice_tensor = src_tensor.Slice(1, 3);
    framework::DDim slice_dims = slice_tensor.dims();
    ASSERT_EQ(arity(slice_dims), 3);
    EXPECT_EQ(slice_dims[0], 2);
    EXPECT_EQ(slice_dims[1], 3);
    EXPECT_EQ(slice_dims[2], 4);

    uintptr_t src_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.data<int>());
    uintptr_t src_mutable_data_address = reinterpret_cast<uintptr_t>(
        src_tensor.mutable_data<int>(src_tensor.dims(), platform::CPUPlace()));
    uintptr_t slice_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.data<int>());
    uintptr_t slice_mutable_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.mutable_data<int>(
            slice_tensor.dims(), platform::CPUPlace()));
    EXPECT_EQ(src_data_address, src_mutable_data_address);
    EXPECT_EQ(slice_data_address, slice_mutable_data_address);
    EXPECT_EQ(src_data_address + 3 * 4 * 1 * sizeof(int), slice_data_address);
  }

#ifdef PADDLE_WITH_CUDA
  {
    framework::Tensor src_tensor;
    src_tensor.mutable_data<double>(framework::make_ddim({6, 9}),
                                    platform::CUDAPlace());
    framework::Tensor slice_tensor = src_tensor.Slice(2, 6);
    framework::DDim slice_dims = slice_tensor.dims();
    ASSERT_EQ(arity(slice_dims), 2);
    EXPECT_EQ(slice_dims[0], 4);
    EXPECT_EQ(slice_dims[1], 9);

    uintptr_t src_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.data<double>());
    uintptr_t src_mutable_data_address =
        reinterpret_cast<uintptr_t>(src_tensor.mutable_data<double>(
            src_tensor.dims(), platform::CUDAPlace()));
    uintptr_t slice_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.data<double>());
    uintptr_t slice_mutable_data_address =
        reinterpret_cast<uintptr_t>(slice_tensor.mutable_data<double>(
            slice_tensor.dims(), platform::CUDAPlace()));
    EXPECT_EQ(src_data_address, src_mutable_data_address);
    EXPECT_EQ(slice_data_address, slice_mutable_data_address);
    EXPECT_EQ(src_data_address + 9 * 2 * sizeof(double), slice_data_address);
  }
#endif
}

TEST(Tensor, ReshapeToMatrix) {
  framework::Tensor src;
  int* src_ptr = src.mutable_data<int>({2, 3, 4, 9}, platform::CPUPlace());
  for (int i = 0; i < 2 * 3 * 4 * 9; ++i) {
    src_ptr[i] = i;
  }
  framework::Tensor res = framework::ReshapeToMatrix(src, 2);
  ASSERT_EQ(res.dims()[0], 2 * 3);
  ASSERT_EQ(res.dims()[1], 4 * 9);
}

TEST(Tensor, Layout) {
  framework::Tensor src;
  ASSERT_EQ(src.layout(), framework::DataLayout::kNCHW);
  src.set_layout(framework::DataLayout::kAnyLayout);
  ASSERT_EQ(src.layout(), framework::DataLayout::kAnyLayout);
}

TEST(Tensor, FP16) {
  using platform::float16;
  framework::Tensor src;
  float16* src_ptr = src.mutable_data<float16>({2, 3}, platform::CPUPlace());
  for (int i = 0; i < 2 * 3; ++i) {
    src_ptr[i] = static_cast<float16>(i);
  }
  EXPECT_EQ(src.memory_size(), 2 * 3 * sizeof(float16));
  // EXPECT a human readable error message
  // src.data<uint8_t>();
  // Tensor holds the wrong type, it holds N6paddle8platform7float16E at
  // [/paddle/Paddle/paddle/fluid/framework/tensor_impl.h:43]
}
<|endoftext|>


<commit_before>#include <pandabase.h>
#ifndef HAVE_GETOPT
  #include <gnu_getopt.h>
#else
  #include <getopt.h>
#endif
#include <multifile.h>
#include <filename.h>

int
main(int argc, char *argv[]) {
  if (argc < 3) {
    cerr << "Usage: multify -[x,c|v] <dest_file> <src_file> ..." << endl;
    return 0;
  }

  bool verbose = true;
  bool extract = false;

  extern char *optarg;
  extern int optind;
  int flag = getopt(argc, argv, "xcvr:");
  Filename rel_path;
  while (flag != EOF) {
    switch (flag) {
      case 'x':
	extract = true;
        break;
      case 'c':
        break;
      case 'v':
	verbose = true;
 	break;
      case 'r':
	rel_path = optarg;
	break;
      default:
	cerr << "Unhandled switch: " << flag << endl;
	break;
    }
    flag = getopt(argc, argv, "xcvr:");
  }
  argc -= (optind - 1);
  argv += (optind - 1);

  if (argc <= 1) {
    cerr << "Usage: multify -[x,c|v] <dest_file> <src_file> ..." << endl;
    return 0;
  }

  Filename dest_file = argv[1];
  dest_file.set_binary();

  if (verbose == true) {
    if (extract == true)
      cerr << "Extracting from: " << dest_file << endl;
    else
      cerr << "Appending to: " << dest_file << endl;
  }

  Multifile mfile;

  if (extract == false) {
    for (int i = 2; i < argc; i++) {
      Filename in_file = argv[i];
      in_file.set_binary();
      mfile.add(in_file);
    }

    if (mfile.write(dest_file) == false)
      cerr << "Failed to write: " << dest_file << endl;
  } else {
    mfile.read(dest_file);
    mfile.extract_all(rel_path);
  }

  return 1;
}
<commit_msg>*** empty log message ***<commit_after>#include <pandabase.h>
#ifndef HAVE_GETOPT
  #include <gnu_getopt.h>
#else
  #include <getopt.h>
#endif
#include <multifile.h>
#include <filename.h>

int
main(int argc, char *argv[]) {
  if (argc < 3) {
    cerr << "Usage: multify -[x,c|vr] <dest_file> <src_file> ..." << endl;
    return 0;
  }

  bool verbose = true;
  bool extract = false;

  extern char *optarg;
  extern int optind;
  int flag = getopt(argc, argv, "xcvr:");
  Filename rel_path;
  while (flag != EOF) {
    switch (flag) {
      case 'x':
	extract = true;
        break;
      case 'c':
        break;
      case 'v':
	verbose = true;
 	break;
      case 'r':
	rel_path = optarg;
	break;
      default:
	cerr << "Unhandled switch: " << flag << endl;
	break;
    }
    flag = getopt(argc, argv, "xcvr:");
  }
  argc -= (optind - 1);
  argv += (optind - 1);

  if (argc <= 1) {
    cerr << "Usage: multify -[x,c|v] <dest_file> <src_file> ..." << endl;
    return 0;
  }

  Filename dest_file = argv[1];
  dest_file.set_binary();

  if (verbose == true) {
    if (extract == true)
      cerr << "Extracting from: " << dest_file << endl;
    else
      cerr << "Appending to: " << dest_file << endl;
  }

  Multifile mfile;

  if (extract == false) {
    for (int i = 2; i < argc; i++) {
      Filename in_file = argv[i];
      in_file.set_binary();
      mfile.add(in_file);
    }

    if (mfile.write(dest_file) == false)
      cerr << "Failed to write: " << dest_file << endl;
  } else {
    mfile.read(dest_file);
    mfile.extract_all(rel_path);
  }

  return 1;
}
<|endoftext|>


<commit_before>// Filename: windowsRegistry.cxx
// Created by:  drose (06Aug01)
//
////////////////////////////////////////////////////////////////////
//
// PANDA 3D SOFTWARE
// Copyright (c) 2001, Disney Enterprises, Inc.  All rights reserved
//
// All use of this software is subject to the terms of the Panda 3d
// Software license.  You should have received a copy of this license
// along with this source code; you will also find a current copy of
// the license at http://www.panda3d.org/license.txt .
//
// To contact the maintainers of this program write to
// panda3d@yahoogroups.com .
//
////////////////////////////////////////////////////////////////////

#include "windowsRegistry.h"
#include "config_express.h"

#ifdef WIN32_VC

#include <windows.h>

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::set_string_value
//       Access: Published, Static
//  Description: Sets the registry key to the indicated value as a
//               string.  The supplied string value is automatically
//               converted from whatever encoding is set by
//               TextEncoder::set_default_encoding() and written as a
//               Unicode string.  The registry key must already exist
//               prior to calling this function.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
set_string_value(const string &key, const string &name, const string &value) {
  TextEncoder encoder;
  wstring wvalue = encoder.decode_text(value);

  bool okflag = true;

  // Now convert the string to Windows' idea of the correct wide-char
  // encoding, so we can store it in the registry.  This might well be
  // the same string we just decoded from, but it might not.

  // Windows likes to have a null character trailing the string (even
  // though we also pass a length).
  wvalue += (wchar_t)0;
  int mb_result_len =
    WideCharToMultiByte(CP_ACP, 0,
                        wvalue.data(), wvalue.length(),
                        NULL, 0,
                        NULL, NULL);
  if (mb_result_len == 0) {
    express_cat.error()
      << "Unable to convert '" << value
      << "' from Unicode to MultiByte form.\n";
    return false;
  }

  char *mb_result = (char *)alloca(mb_result_len);
  WideCharToMultiByte(CP_ACP, 0,
                      wvalue.data(), wvalue.length(),
                      mb_result, mb_result_len,
                      NULL, NULL);

  if (express_cat.is_debug()) {
    express_cat.debug()
      << "Converted '" << value << "' to '" << mb_result
      << "' for storing in registry.\n";
  }

  return do_set(key, name, REG_SZ, mb_result, mb_result_len);
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::set_int_value
//       Access: Published, Static
//  Description: Sets the registry key to the indicated value as an
//               integer.  The registry key must already exist prior
//               to calling this function.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
set_int_value(const string &key, const string &name, int value) {
  DWORD dw = value;
  return do_set(key, name, REG_DWORD, &dw, sizeof(dw));
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_key_type
//       Access: Published, Static
//  Description: Returns the type of the indicated key, or T_none if
//               the key is not known or is some unsupported type.
////////////////////////////////////////////////////////////////////
WindowsRegistry::Type WindowsRegistry::
get_key_type(const string &key, const string &name) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return T_none;
  }

  switch (data_type) {
  case REG_SZ:
    return T_string;

  case REG_DWORD:
    return T_int;

  default:
    return T_none;
  }
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_string_value
//       Access: Published, Static
//  Description: Returns the value associated with the indicated
//               registry key, assuming it is a string value.  The
//               string value is automatically encoded using
//               TextEncoder::get_default_encoding().  If the key is
//               not defined or is not a string type value,
//               default_value is returned instead.
////////////////////////////////////////////////////////////////////
string WindowsRegistry::
get_string_value(const string &key, const string &name,
                 const string &default_value) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return default_value;
  }

  if (data_type != REG_SZ) {
    express_cat.warning()
      << "Registry key " << key << " does not contain a string value.\n";
    return default_value;
  }

  // Now we have to decode the MultiByte string to Unicode, and re-encode
  // it according to our own encoding.

  if (data.empty()) {
    return data;
  }

  int wide_result_len =
    MultiByteToWideChar(CP_ACP, 0,
                        data.data(), data.length(),
                        NULL, 0);
  if (wide_result_len == 0) {
    express_cat.error()
      << "Unable to convert '" << data
      << "' from MultiByte to Unicode form.\n";
    return data;
  }

  wchar_t *wide_result = (wchar_t *)alloca(wide_result_len * sizeof(wchar_t));
  MultiByteToWideChar(CP_ACP, 0,
                      data.data(), data.length(),
                      wide_result, wide_result_len);

  wstring wdata(wide_result, wide_result_len);

  TextEncoder encoder;
  string result = encoder.encode_wtext(wdata);

  if (express_cat.is_debug()) {
    express_cat.debug()
      << "Converted '" << data << "' from registry to '" << result << "'\n";
  }

  return result;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_int_value
//       Access: Published, Static
//  Description: Returns the value associated with the indicated
//               registry key, assuming it is an integer value.  If
//               the key is not defined or is not an integer type
//               value, default_value is returned instead.
////////////////////////////////////////////////////////////////////
int WindowsRegistry::
get_int_value(const string &key, const string &name, int default_value) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return default_value;
  }

  if (data_type != REG_DWORD) {
    express_cat.warning()
      << "Registry key " << key << " does not contain an integer value.\n";
    return default_value;
  }

  // Now we have a DWORD encoded in a string.
  nassertr(data.length() == sizeof(DWORD), default_value);
  DWORD dw = *(DWORD *)data.data();
  return dw;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::do_set
//       Access: Private, Static
//  Description: The internal function to actually make all of the
//               appropriate windows calls to set the registry value.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
do_set(const string &key, const string &name,
       int data_type, const void *data, int data_length) {
  HKEY hkey;
  LONG error;

  error =
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, key.c_str(), 0, KEY_SET_VALUE, &hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.error()
      << "Unable to open registry key " << key
      << ": " << format_message(error) << "\n";
    return false;
  }

  bool okflag = true;

  error =
    RegSetValueEx(hkey, name.c_str(), 0, data_type,
                  (CONST BYTE *)data, data_length);
  if (error != ERROR_SUCCESS) {
    express_cat.error()
      << "Unable to set registry key " << key << " name " << name
      << ": " << format_message(error) << "\n";
    okflag = false;
  }

  error = RegCloseKey(hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.warning()
      << "Unable to close opened registry key " << key
      << ": " << format_message(error) << "\n";
  }

  return okflag;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::do_get
//       Access: Private, Static
//  Description: The internal function to actually make all of the
//               appropriate windows calls to retrieve the registry
//               value.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
do_get(const string &key, const string &name, int &data_type, string &data) {
  HKEY hkey;
  LONG error;

  error =
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, key.c_str(), 0, KEY_QUERY_VALUE, &hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.debug()
      << "Unable to open registry key " << key
      << ": " << format_message(error) << "\n";
    return false;
  }

  bool okflag = true;

  // We start with a 1K buffer; presumably that will be big enough
  // most of the time.
  static const size_t init_buffer_size = 1024;
  char init_buffer[init_buffer_size];
  DWORD buffer_size = init_buffer_size;
  DWORD dw_data_type;

  error =
    RegQueryValueEx(hkey, name.c_str(), 0, &dw_data_type,
                    (BYTE *)init_buffer, &buffer_size);
  if (error == ERROR_SUCCESS) {
    data_type = dw_data_type;
    if (data_type == REG_SZ ||
        data_type == REG_MULTI_SZ ||
        data_type == REG_EXPAND_SZ) {
      // Eliminate the trailing null character.
      buffer_size--;
    }
    data = string(init_buffer, buffer_size);

  } else if (error == ERROR_MORE_DATA) {
    // Huh, 1K wasn't big enough.  Ok, get a bigger buffer.

    // If we were querying HKEY_PERFORMANCE_DATA, we'd have to keep
    // guessing bigger and bigger until we got it.  Since we're
    // querying static data for now, we can just use the size Windows
    // tells us.
    char *new_buffer = new char[buffer_size];
    error =
      RegQueryValueEx(hkey, name.c_str(), 0, &dw_data_type,
                      (BYTE *)new_buffer, &buffer_size);
    if (error == ERROR_SUCCESS) {
      data_type = dw_data_type;
      if (data_type == REG_SZ ||
          data_type == REG_MULTI_SZ ||
          data_type == REG_EXPAND_SZ) {
        // Eliminate the trailing null character.
        buffer_size--;
      }
      data = string(new_buffer, buffer_size);
    }
    delete new_buffer;
  }

  if (error != ERROR_SUCCESS) {
    express_cat.info()
      << "Unable to get registry value " << name
      << ": " << format_message(error) << "\n";
    okflag = false;
  }

  error = RegCloseKey(hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.warning()
      << "Unable to close opened registry key " << key
      << ": " << format_message(error) << "\n";
  }

  return okflag;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::format_message
//       Access: Private, Static
//  Description: Returns the Windows error message associated with the
//               given error code.
////////////////////////////////////////////////////////////////////
string WindowsRegistry::
format_message(int error_code) {
  LPTSTR buffer;
  DWORD length =
    FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
                  NULL, error_code, 0, &buffer, 0, NULL);
  if (length == 0) {
    return "Unknown error message";
  }

  // Strip off \n's and \r's trailing the string.
  while (length > 0 &&
         (buffer[length - 1] == '\r' || buffer[length - 1] == '\n')) {
    length--;
  }

  string result((const char *)buffer, length);
  LocalFree(buffer);
  return result;
}

#endif  // WIN32_VC
<commit_msg>trim extra newlines from windows error message<commit_after>// Filename: windowsRegistry.cxx
// Created by:  drose (06Aug01)
//
////////////////////////////////////////////////////////////////////
//
// PANDA 3D SOFTWARE
// Copyright (c) 2001, Disney Enterprises, Inc.  All rights reserved
//
// All use of this software is subject to the terms of the Panda 3d
// Software license.  You should have received a copy of this license
// along with this source code; you will also find a current copy of
// the license at http://www.panda3d.org/license.txt .
//
// To contact the maintainers of this program write to
// panda3d@yahoogroups.com .
//
////////////////////////////////////////////////////////////////////

#include "windowsRegistry.h"
#include "config_express.h"

#ifdef WIN32_VC

#include <windows.h>

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::set_string_value
//       Access: Published, Static
//  Description: Sets the registry key to the indicated value as a
//               string.  The supplied string value is automatically
//               converted from whatever encoding is set by
//               TextEncoder::set_default_encoding() and written as a
//               Unicode string.  The registry key must already exist
//               prior to calling this function.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
set_string_value(const string &key, const string &name, const string &value) {
  TextEncoder encoder;
  wstring wvalue = encoder.decode_text(value);

  bool okflag = true;

  // Now convert the string to Windows' idea of the correct wide-char
  // encoding, so we can store it in the registry.  This might well be
  // the same string we just decoded from, but it might not.

  // Windows likes to have a null character trailing the string (even
  // though we also pass a length).
  wvalue += (wchar_t)0;
  int mb_result_len =
    WideCharToMultiByte(CP_ACP, 0,
                        wvalue.data(), wvalue.length(),
                        NULL, 0,
                        NULL, NULL);
  if (mb_result_len == 0) {
    express_cat.error()
      << "Unable to convert '" << value
      << "' from Unicode to MultiByte form.\n";
    return false;
  }

  char *mb_result = (char *)alloca(mb_result_len);
  WideCharToMultiByte(CP_ACP, 0,
                      wvalue.data(), wvalue.length(),
                      mb_result, mb_result_len,
                      NULL, NULL);

  if (express_cat.is_debug()) {
    express_cat.debug()
      << "Converted '" << value << "' to '" << mb_result
      << "' for storing in registry.\n";
  }

  return do_set(key, name, REG_SZ, mb_result, mb_result_len);
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::set_int_value
//       Access: Published, Static
//  Description: Sets the registry key to the indicated value as an
//               integer.  The registry key must already exist prior
//               to calling this function.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
set_int_value(const string &key, const string &name, int value) {
  DWORD dw = value;
  return do_set(key, name, REG_DWORD, &dw, sizeof(dw));
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_key_type
//       Access: Published, Static
//  Description: Returns the type of the indicated key, or T_none if
//               the key is not known or is some unsupported type.
////////////////////////////////////////////////////////////////////
WindowsRegistry::Type WindowsRegistry::
get_key_type(const string &key, const string &name) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return T_none;
  }

  switch (data_type) {
  case REG_SZ:
    return T_string;

  case REG_DWORD:
    return T_int;

  default:
    return T_none;
  }
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_string_value
//       Access: Published, Static
//  Description: Returns the value associated with the indicated
//               registry key, assuming it is a string value.  The
//               string value is automatically encoded using
//               TextEncoder::get_default_encoding().  If the key is
//               not defined or is not a string type value,
//               default_value is returned instead.
////////////////////////////////////////////////////////////////////
string WindowsRegistry::
get_string_value(const string &key, const string &name,
                 const string &default_value) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return default_value;
  }

  if (data_type != REG_SZ) {
    express_cat.warning()
      << "Registry key " << key << " does not contain a string value.\n";
    return default_value;
  }

  // Now we have to decode the MultiByte string to Unicode, and re-encode
  // it according to our own encoding.

  if (data.empty()) {
    return data;
  }

  int wide_result_len =
    MultiByteToWideChar(CP_ACP, 0,
                        data.data(), data.length(),
                        NULL, 0);
  if (wide_result_len == 0) {
    express_cat.error()
      << "Unable to convert '" << data
      << "' from MultiByte to Unicode form.\n";
    return data;
  }

  wchar_t *wide_result = (wchar_t *)alloca(wide_result_len * sizeof(wchar_t));
  MultiByteToWideChar(CP_ACP, 0,
                      data.data(), data.length(),
                      wide_result, wide_result_len);

  wstring wdata(wide_result, wide_result_len);

  TextEncoder encoder;
  string result = encoder.encode_wtext(wdata);

  if (express_cat.is_debug()) {
    express_cat.debug()
      << "Converted '" << data << "' from registry to '" << result << "'\n";
  }

  return result;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::get_int_value
//       Access: Published, Static
//  Description: Returns the value associated with the indicated
//               registry key, assuming it is an integer value.  If
//               the key is not defined or is not an integer type
//               value, default_value is returned instead.
////////////////////////////////////////////////////////////////////
int WindowsRegistry::
get_int_value(const string &key, const string &name, int default_value) {
  int data_type;
  string data;
  if (!do_get(key, name, data_type, data)) {
    return default_value;
  }

  if (data_type != REG_DWORD) {
    express_cat.warning()
      << "Registry key " << key << " does not contain an integer value.\n";
    return default_value;
  }

  // Now we have a DWORD encoded in a string.
  nassertr(data.length() == sizeof(DWORD), default_value);
  DWORD dw = *(DWORD *)data.data();
  return dw;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::do_set
//       Access: Private, Static
//  Description: The internal function to actually make all of the
//               appropriate windows calls to set the registry value.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
do_set(const string &key, const string &name,
       int data_type, const void *data, int data_length) {
  HKEY hkey;
  LONG error;

  error =
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, key.c_str(), 0, KEY_SET_VALUE, &hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.error()
      << "Unable to open registry key " << key
      << ": " << format_message(error) << "\n";
    return false;
  }

  bool okflag = true;

  error =
    RegSetValueEx(hkey, name.c_str(), 0, data_type,
                  (CONST BYTE *)data, data_length);
  if (error != ERROR_SUCCESS) {
    express_cat.error()
      << "Unable to set registry key " << key << " name " << name
      << ": " << format_message(error) << "\n";
    okflag = false;
  }

  error = RegCloseKey(hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.warning()
      << "Unable to close opened registry key " << key
      << ": " << format_message(error) << "\n";
  }

  return okflag;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::do_get
//       Access: Private, Static
//  Description: The internal function to actually make all of the
//               appropriate windows calls to retrieve the registry
//               value.
////////////////////////////////////////////////////////////////////
bool WindowsRegistry::
do_get(const string &key, const string &name, int &data_type, string &data) {
  HKEY hkey;
  LONG error;

  error =
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, key.c_str(), 0, KEY_QUERY_VALUE, &hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.debug()
      << "Unable to open registry key " << key
      << ": " << format_message(error) << "\n";
    return false;
  }

  bool okflag = true;

  // We start with a 1K buffer; presumably that will be big enough
  // most of the time.
  static const size_t init_buffer_size = 1024;
  char init_buffer[init_buffer_size];
  DWORD buffer_size = init_buffer_size;
  DWORD dw_data_type;

  error =
    RegQueryValueEx(hkey, name.c_str(), 0, &dw_data_type,
                    (BYTE *)init_buffer, &buffer_size);
  if (error == ERROR_SUCCESS) {
    data_type = dw_data_type;
    if (data_type == REG_SZ ||
        data_type == REG_MULTI_SZ ||
        data_type == REG_EXPAND_SZ) {
      // Eliminate the trailing null character.
      buffer_size--;
    }
    data = string(init_buffer, buffer_size);

  } else if (error == ERROR_MORE_DATA) {
    // Huh, 1K wasn't big enough.  Ok, get a bigger buffer.

    // If we were querying HKEY_PERFORMANCE_DATA, we'd have to keep
    // guessing bigger and bigger until we got it.  Since we're
    // querying static data for now, we can just use the size Windows
    // tells us.
    char *new_buffer = new char[buffer_size];
    error =
      RegQueryValueEx(hkey, name.c_str(), 0, &dw_data_type,
                      (BYTE *)new_buffer, &buffer_size);
    if (error == ERROR_SUCCESS) {
      data_type = dw_data_type;
      if (data_type == REG_SZ ||
          data_type == REG_MULTI_SZ ||
          data_type == REG_EXPAND_SZ) {
        // Eliminate the trailing null character.
        buffer_size--;
      }
      data = string(new_buffer, buffer_size);
    }
    delete new_buffer;
  }

  if (error != ERROR_SUCCESS) {
    express_cat.info()
      << "Unable to get registry value " << name
      << ": " << format_message(error) << "\n";
    okflag = false;
  }

  error = RegCloseKey(hkey);
  if (error != ERROR_SUCCESS) {
    express_cat.warning()
      << "Unable to close opened registry key " << key
      << ": " << format_message(error) << "\n";
  }

  return okflag;
}

////////////////////////////////////////////////////////////////////
//     Function: WindowsRegistry::format_message
//       Access: Private, Static
//  Description: Returns the Windows error message associated with the
//               given error code.
////////////////////////////////////////////////////////////////////
string WindowsRegistry::
format_message(int error_code) {
  PVOID buffer;
  DWORD length =
    FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
                  NULL, error_code, 0, (LPTSTR)&buffer, 0, NULL);
  if (length == 0) {
    return "Unknown error message";
  }

  const char *text = (const char *)buffer;

  // Strip off \n's and \r's trailing the string.
  while (length > 0 &&
         (text[length - 1] == '\r' || text[length - 1] == '\n')) {
    length--;
  }

  string result((const char *)text, length);
  LocalFree(buffer);
  return result;
}

#endif  // WIN32_VC
<|endoftext|>