<commit_before>#include <iostream>
#include <fstream>

#include "mbedtls/sha512.h"
#include "mbedtls/aes.h"
#include "macros.h"

#define USE_SHA512 0

#define AES_CBC_BLOCK_SIZE 16
#define SHA512_LENGTH 64
#define IV_LENGTH 16
#define KEYBITS 128

void generate_random_string(unsigned char *iv, int len)
{
	static const char alphanum[] =
		"0123456789"
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
		"abcdefghijklmnopqrstuvwxyz";

	for (int i = 0; i < len; ++i) {
		iv[i] = alphanum[rand() % (sizeof(alphanum) - 1)];
	}
}

void pkcs7_pad(unsigned char *buf, unsigned int *count)
{
	unsigned int padding = AES_CBC_BLOCK_SIZE - *count;
	for(unsigned int i = 1; i <= padding; i++) {
		buf[AES_CBC_BLOCK_SIZE - i] = padding;
	}
	// We pad to the block size
	*count = AES_CBC_BLOCK_SIZE;
}

unsigned int pkcs7_detect(unsigned char *buf)
{
	unsigned int count = AES_CBC_BLOCK_SIZE;
	unsigned char last = buf[AES_CBC_BLOCK_SIZE - 1];
	if(last < 16) {
		for(int i = AES_CBC_BLOCK_SIZE - last; i < AES_CBC_BLOCK_SIZE; i++) {
			if(buf[i] != last) {
				return count;
			}
		}
		count -= last;
	}
	return count;
}

void set_aes_key(mbedtls_aes_context *aes_ctx, unsigned char *key, bool encrypt)
{
	int result;
	if(encrypt) {
		result = mbedtls_aes_setkey_enc(aes_ctx, key, KEYBITS);
	} else {
		result = mbedtls_aes_setkey_dec(aes_ctx, key, KEYBITS);
	}
	if(result == MBEDTLS_ERR_AES_INVALID_KEY_LENGTH) {
		DIE(2, "Invalid input key length");
	}
}

bool compare_hashes(unsigned char *h1, unsigned char *h2, int len) {
	for(int i = 0; i < len; i++) {
		if(h1[i] != h2[i]) return false;
	}
	return true;
}

void post_decrypt(unsigned char *sha512, std::ifstream *src)
{
	unsigned char read_sha512[SHA512_LENGTH];
	src->read(reinterpret_cast<char *>(&read_sha512), SHA512_LENGTH);
	printf("Comparing encrypted hashes ... ");
	if(!compare_hashes(sha512, read_sha512, SHA512_LENGTH)) {
		printf("failed\n");
	} else {
		printf("OK\n");
	}
}

void hash_n_stuff(std::ifstream *src, std::ofstream *dst, unsigned char *key, bool encrypt, unsigned char *sha512)
{
	// Initialization vector, this should be random
	unsigned char iv[16];
	unsigned char buf[AES_CBC_BLOCK_SIZE];
	unsigned char outbuf[AES_CBC_BLOCK_SIZE];
	unsigned int bytes;

	// Get how long the input is
	src->seekg(0, std::ios_base::end);
	size_t in_length = src->tellg();
	src->seekg(0, std::ios_base::beg);

	mbedtls_sha512_context sha_ctx;
	mbedtls_aes_context aes_ctx;

	mbedtls_sha512_init(&sha_ctx);
	mbedtls_sha512_starts(&sha_ctx, USE_SHA512);

	mbedtls_aes_init(&aes_ctx);
	set_aes_key(&aes_ctx, key, encrypt);

	if(!encrypt) {
		// There are two 64bit long hashes at the end of the file and IV at the beginning
		in_length -= SHA512_LENGTH + IV_LENGTH;
		src->read(reinterpret_cast<char *>(iv), IV_LENGTH);
	} else {
		// Generate initial IV and write it into the file
		generate_random_string(iv, IV_LENGTH);
		dst->write(reinterpret_cast<char *>(iv), IV_LENGTH);
	}

	// While there's something to read
	for(size_t i = 0; i < in_length; i += AES_CBC_BLOCK_SIZE) {
		// Read one block
		src->read(reinterpret_cast<char *>(buf), AES_CBC_BLOCK_SIZE);
		bytes = src->gcount();

		// Hash the block
		if(encrypt) {
			// Add padding when encrypting
			pkcs7_pad(buf, &bytes);
		} else {
			// Hash the read blocks when decrypting
			mbedtls_sha512_update(&sha_ctx, buf, bytes);
		}

		// Encrypt the block and write it to the output stream
		mbedtls_aes_crypt_cbc(&aes_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, bytes, iv, buf, outbuf);

		// Detect padding when decrypting
		if(!encrypt && i + AES_CBC_BLOCK_SIZE >= in_length) bytes = pkcs7_detect(outbuf);

		// Write the encrypted/decrypted block to disk
		dst->write(reinterpret_cast<char *>(outbuf), bytes);

		// Hash the encrypted block
		if(encrypt) mbedtls_sha512_update(&sha_ctx, outbuf, bytes);
	}

	// Finish generating checksums
	mbedtls_sha512_finish(&sha_ctx, sha512);

	// Clean things up
	mbedtls_sha512_free(&sha_ctx);
	mbedtls_aes_free(&aes_ctx);

	if(encrypt) {
		// Write the hash if we're encrypting
		dst->write(reinterpret_cast<char*>(sha512), SHA512_LENGTH);
	} else {
		// Run post-decrypt checks
		post_decrypt(sha512, src);
	}
}

int main(int argc, char **argv)
{
	if(argc != 5) {
		DIE(1, "Wrong argument count, expected 4, got %d\n", argc - 1);
	}
	bool encrypt;
	encrypt = argv[1][0] == 'e' ? true : false;
	std::ifstream src;
	std::ofstream dst;
	src.open(argv[2]);
	dst.open(argv[3]);
	unsigned char sha512[SHA512_LENGTH];
	hash_n_stuff(&src, &dst, reinterpret_cast<unsigned char *>(argv[4]), encrypt, sha512);
	src.close();
	dst.close();
	return 0;
}
<commit_msg>reworked the whole flow<commit_after>#include <iostream>
#include <fstream>
#include <cstring>

#include "mbedtls/sha512.h"
#include "mbedtls/aes.h"
#include "macros.h"

#define USE_SHA512 0

#define AES_CBC_BLOCK_SIZE 16
#define SHA512_LENGTH 64
#define IV_LENGTH 16
#define KEYBITS 128
#define HEADER_SIZE (SHA512_LENGTH + IV_LENGTH + 3)

void generate_random_string(unsigned char *iv, int len)
{
	static const char alphanum[] =
		"0123456789"
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
		"abcdefghijklmnopqrstuvwxyz";

	for (int i = 0; i < len; ++i) {
		iv[i] = alphanum[rand() % (sizeof(alphanum) - 1)];
	}
}

void pkcs7_pad(unsigned char *buf, unsigned int *count)
{
	unsigned int padding = AES_CBC_BLOCK_SIZE - *count;
	for(unsigned int i = 1; i <= padding; i++) {
		buf[AES_CBC_BLOCK_SIZE - i] = padding;
	}
	// We pad to the block size
	*count = AES_CBC_BLOCK_SIZE;
}

unsigned int pkcs7_detect(unsigned char *buf)
{
	unsigned int count = AES_CBC_BLOCK_SIZE;
	unsigned char last = buf[AES_CBC_BLOCK_SIZE - 1];
	if(last < 16) {
		for(int i = AES_CBC_BLOCK_SIZE - last; i < AES_CBC_BLOCK_SIZE; i++) {
			if(buf[i] != last) {
				return count;
			}
		}
		count -= last;
	}
	return count;
}

void set_aes_key(mbedtls_aes_context *aes_ctx, unsigned char *key, bool encrypt)
{
	int result;
	if(encrypt) {
		result = mbedtls_aes_setkey_enc(aes_ctx, key, KEYBITS);
	} else {
		result = mbedtls_aes_setkey_dec(aes_ctx, key, KEYBITS);
	}
	if(result == MBEDTLS_ERR_AES_INVALID_KEY_LENGTH) {
		DIE(2, "Invalid input key length");
	}
}

bool check_hashes(unsigned char *h1, unsigned char *h2, size_t length)
{
	for(size_t i = 0; i < length; i++) {
		if(h1[i] != h2[i]) return false;
	}
	return true;
}

void hash_n_stuff(std::ifstream *src, std::ofstream *dst, bool encrypt, unsigned char *key, unsigned char *sha512, unsigned char *iv, size_t length)
{
	unsigned char buf[AES_CBC_BLOCK_SIZE];
	unsigned char outbuf[AES_CBC_BLOCK_SIZE];
	unsigned int bytes;

	mbedtls_sha512_context sha_ctx;
	mbedtls_aes_context aes_ctx;

	mbedtls_sha512_init(&sha_ctx);
	mbedtls_sha512_starts(&sha_ctx, USE_SHA512);

	mbedtls_aes_init(&aes_ctx);
	set_aes_key(&aes_ctx, key, encrypt);
	size_t foo = src->tellg();

	// While there's something to read
	for(size_t i = 0; i < length; i += AES_CBC_BLOCK_SIZE) {
		// Read one block
		src->read(reinterpret_cast<char *>(buf), AES_CBC_BLOCK_SIZE);
		bytes = src->gcount();

		if(encrypt) {
			// Add padding when encrypting
			pkcs7_pad(buf, &bytes);
		}

		// Encrypt the block and write it to the output stream
		mbedtls_aes_crypt_cbc(&aes_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, bytes, iv, buf, outbuf);

		// Detect padding when decrypting
		if(!encrypt && i + AES_CBC_BLOCK_SIZE >= length) bytes = pkcs7_detect(outbuf);

		// Write the encrypted/decrypted block to disk
		dst->write(reinterpret_cast<char *>(outbuf), bytes);

		// Hash the block
		mbedtls_sha512_update(&sha_ctx, encrypt ? outbuf : buf, AES_CBC_BLOCK_SIZE);
	}

	// Finish generating checksums
	mbedtls_sha512_finish(&sha_ctx, sha512);

	// Clean things up
	mbedtls_sha512_free(&sha_ctx);
	mbedtls_aes_free(&aes_ctx);
}

size_t stream_length(std::fstream *stream)
{
	size_t pos = stream->tellg();
	stream->seekg(0, std::ios_base::end);
	size_t length = stream->tellg();
	stream->seekg(pos);
	return length;
}

void write_header(std::ofstream *out, unsigned char *sha512, unsigned char *iv)
{
	out->seekp(0, std::ios_base::beg);
	out->write("AES", 3);
	out->write(reinterpret_cast<char *>(iv), IV_LENGTH);
	out->write(reinterpret_cast<char *>(sha512), SHA512_LENGTH);
}

void read_header(std::ifstream *in, unsigned char *sha512, unsigned char *iv, size_t *length)
{
	char buf[3];
	in->read(buf, 3);
	if(std::strncmp(buf, "AES", 3) != 0) {
		DIE(5, "The encrpyted file misses the AES header.");
	}
	in->read(reinterpret_cast<char *>(iv), IV_LENGTH);
	if(in->gcount() != IV_LENGTH) {
		DIE(6, "The encrypted file does not contain the IV.");
	}
	in->read(reinterpret_cast<char *>(sha512), SHA512_LENGTH);
	if(in->gcount() != SHA512_LENGTH) {
		DIE(7, "The encrypted file does not contain the SHA512 hash.");
	}
	*length = stream_length((std::fstream *)in) - (HEADER_SIZE);
}

void perform(bool encrypt, char *src, char *dst, unsigned char *key)
{
	std::ifstream in;
	std::ofstream out;
	unsigned char sha512[SHA512_LENGTH], iv[IV_LENGTH], iv_backup[IV_LENGTH];
	size_t length;
	in.open(src);
	out.open(dst);

	if(encrypt) {
		length = stream_length((std::fstream *)&in);
		// Leave blank space for the header
		out.seekp(HEADER_SIZE);
		// Generate IV
		generate_random_string(iv, IV_LENGTH);
		memcpy(iv_backup, iv, IV_LENGTH*sizeof(unsigned char));
		// Encrypt the file
		hash_n_stuff(&in, &out, encrypt, key, sha512, iv, length);
		// Write the header
		write_header(&out, sha512, iv_backup);
	} else {
		unsigned char read_sha512[SHA512_LENGTH];
		read_header(&in, read_sha512, iv, &length);
		if(length % 16 != 0) {
			DIE(3, "The encrypted payload is of the wrong size.");
		}
		hash_n_stuff(&in, &out, encrypt, key, sha512, iv, length);
		if(!check_hashes(sha512, read_sha512, SHA512_LENGTH)) {
			DIE(4, "The hashes do not match.");
		}
	}
	in.close();
	out.close();
}

int main(int argc, char **argv)
{
	if(argc != 5) {
		DIE(1, "Wrong argument count, expected 4, got %d\n", argc - 1);
	}
	bool encrypt;
	encrypt = argv[1][0] == 'e' ? true : false;
	perform(encrypt, argv[2], argv[3], reinterpret_cast<unsigned char *>(argv[4]));
	return 0;
}
<|endoftext|>


<commit_before>#include <iostream>
#include <set>
#include <string>
#include <thread>

#include "src/eventloop.h"
#include "src/selector.h"
#include "src/tcpserver.h"
#include "src/util/concurrent.h"
#include "src/util/core.h"
#include "src/util/fd.h"
#include "src/util/net.h"
#include "src/buffer.h"

using namespace std;
using namespace happyntrain;
using namespace happyntrain::concurrent;
using namespace happyntrain::network;
using namespace happyntrain::fd;

struct A {
  int a;
  std::string b;
  A(int a_, std::string b_) : a(a_), b(b_) {}
  A() : A(0, "") {}
};

int main() {
  auto eventloop = newInstance<EventLoop>(0);
  TCPServer server(eventloop);

  server.OnConnect([](auto& socket) {
    socket->OnRead([](auto& self, auto& buffer) {
      INFO("%s", buffer.retrieve().c_str());
    });
  });

  const int port = 8888;
  server.Listen(port, [=]{
    INFO("Server is now listening on %d", port);
  });
  eventloop->Run();
  // auto server = TCPServer::Create([&](auto socket) {
  //                 socker.end("goodbye\n");
  //               })->OnError([&](auto error) {});
  // server.listen([&]() { DEBUG("opened server on ", server.Address()); });
  return 0;
}
<commit_msg>Update smoke test<commit_after>#include <iostream>
#include <string>

#include "src/eventloop.h"
#include "src/tcpserver.h"
#include "src/util/core.h"

using namespace std;
using namespace happyntrain;

int main() {
  auto eventloop = newInstance<EventLoop>(0);
  TCPServer server(eventloop);

  server.OnConnect([](auto& socket) {
    socket->OnRead([](auto& self, auto& buffer) {
      buffer.retrieve();
      self->Send("Hello World!");
    });
  });

  const int port = 8888;
  server.Listen(port, [=]{
    INFO("Server is now listening on %d", port);
  });
  eventloop->Run();
  return 0;
}
<|endoftext|>


<commit_before>86d5fff0-2749-11e6-9129-e0f84713e7b8<commit_msg>That didn't fix it<commit_after>86df4a4a-2749-11e6-b5e3-e0f84713e7b8<|endoftext|>


<commit_before>#ifndef MAIN_H

#define MAIN_H

#include <iostream>
#include <fstream>

// As per getopt example
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

using namespace std;

typedef struct {
  // Configuration to use for the game, in JSON
  string board_config;

  // The wordlist to use for validating wordplays
  string language_file;

  // Number of players playing and which go you are
  int players;
  int this_player_go;

  // Initialisation of options for defaults
  void init(void) {
    this->players = 2;
    this->this_player_go = 1;
    this->language_file = "./submodules/wordlists_en_mozilla/en_GB (Marco Pinto) - 2.32 - 2016-01-01/wordlist_marcoagpinto_20160101_153347w.txt";
  }

  // Validate options
  bool validate(void) {
    bool valid = true;

    // Must have a board config and player options
    valid &= !this->board_config.empty() || !this->players || !this->this_player_go;

    // Must have a language file / wordlist
    valid &= !this->language_file.empty();

    return valid;
  }
} options;

// Created Libraries
#include "Game.hpp"

void print_error(string err);
void usage(void);
bool check_options(options o);
bool check_file_exists(string s, bool print_err = false);

#endif
<commit_msg>Tidy validation code for options<commit_after>#ifndef MAIN_H

#define MAIN_H

#include <iostream>
#include <fstream>

// As per getopt example
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

using namespace std;

void print_error(string err);
void usage(void);

typedef struct {
  // Configuration to use for the game, in JSON
  string board_config;

  // The wordlist to use for validating wordplays
  string language_file;

  // Number of players playing and which go you are
  int players;
  int this_player_go;

  // Initialisation of options for defaults
  void init(void) {
    this->players = 2;
    this->this_player_go = 1;
    this->language_file = "./submodules/wordlists_en_mozilla/en_GB (Marco Pinto) - 2.32 - 2016-01-01/wordlist_marcoagpinto_20160101_153347w.txt";
  }

  // Validate options
  bool validate(void) {
    bool valid = true;

    // Must have a board config
    bool board_config_valid = !this->board_config.empty();
    valid &= board_config_valid;
    if (!board_config_valid) {
      print_error("No board config set. Please set a board configuration.");
      usage();
    }

    // Must have players and what player you are
    bool players_valid = !this->players || !this->this_player_go;
    valid &= players_valid;
    if (!players_valid) {
      print_error("You must set number of players and which player you are.");
      usage();
    }

    // Must have a language file / wordlist
    bool language_file_valid = !this->language_file.empty();
    valid &= language_file_valid;
    if (!language_file_valid) {
      print_error("No wordlist set. You must set a wordlist to use.");
      usage();
    }

    return valid;
  }
} options;

// Created Libraries
#include "Game.hpp"

bool check_options(options o);
bool check_file_exists(string s, bool print_err = false);

#endif
<|endoftext|>


<commit_before>#include "mcts.hpp"
#include "notation.hpp"

#include <stack>

#include <boost/range/algorithm/random_shuffle.hpp>

using namespace mcts;

void Node::initialize(State const& state) {
  hash = NodeTable::key(state.hash);
}

void Node::update(double result) {
  m_sample_size++;
  double mean0 = m_mean;
  double delta = result - mean0;
  m_mean += delta/m_sample_size;
  m_derivative = 0.9*m_derivative + 0.1*(m_mean - mean0);
}

void Node::adjoin_parent(Node* parent) {
  std::lock_guard<std::mutex> lock(parents_mutex);
  parents.insert(parent->hash);
}

std::string Node::format_statistics() {
  return str(boost::format("%1% %2% (d %3%) %4% (%5% parents)")
             % sample_size()
             % mean()
             % derivative()
             % selection_criterion()
             % parents.size());
}

Node* NodeTable::get(Hash hash) {
  size_t key = NodeTable::key(hash);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    return nullptr;
  return node;
}

Node* NodeTable::get_or_create(State const& state) {
  size_t key = NodeTable::key(state);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    node->initialize(state);
  return node;
}

void Graph::sample(State state, boost::mt19937& generator) {
  Node* node = nodes.get_or_create(state);
  sorted_vector<Hash> path;

  double result;

  // selection
  while (true) {
    path.insert(node->hash);

    Node* child = select_child(node, state, generator);

    if (!child) {
      // no legal successors; game over
      boost::optional<Color> winner = state.winner();
      result = !winner
        ? draw_value
        : (*winner == state.us
           ? win_value
           : loss_value);
      break;
    }

    if (state.drawn_by_50() || path.contains(child->hash)) {
      // draw per 50-move rule or by choosing repetition
      result = draw_value;
      break;
    }

    node = child;
  }

  backprop(node, result);
}

// returns nullptr if no legal successor states
// NOTE: state will be modified to be the corresponding successor
Node* Graph::select_child(Node* node, State& state, boost::mt19937& generator) {
  if (node->sample_size() < 30) {
    // don't waste time with unreliable statistics
    boost::optional<Move> move = moves::make_random_legal_move(state, generator);
    if (!move)
      return nullptr;
    Node* child = nodes.get_or_create(state);
    child->adjoin_parent(node);
    return child;
  } else {
    // use statistics to make selection
    std::vector<Move> moves;
    moves::moves(moves, state);

    // find a legal move with maximum selection criterion.
    boost::optional<Move> best_move;
    Node* best_child;
    double best_score;

    for (Move curr_move: moves) {
      Undo undo = state.make_move(curr_move);

      // if legal
      if (!state.their_king_attacked()) {
        Node* curr_child = nodes.get_or_create(state);

        curr_child->adjoin_parent(node);
        if (curr_child->sample_size() < 10)
          // select new nodes unconditionally
          return curr_child;

        double curr_score = curr_child->selection_criterion(generator);
        if (!best_move || curr_score > best_score) {
          best_move  = curr_move;
          best_child = curr_child;
          best_score = curr_score;
        }
      }

      state.unmake_move(undo);
    }

    if (!best_move)
      return nullptr;
    state.make_move(*best_move);
    return best_child;
  }
}

// update the node and all of its ancestors.  the graph is likely to be
// cyclic due to reversible moves.  we traverse it in a breadth-first
// manner so that the updates run along the shortest paths to each
// unique ancestor.
// NOTE: initial_result is from the perspective of the player who is to move in
// the state corresponding to node.  since the stored node values are from the
// perspective of the player who causes the node to be chosen, we have to invert
// it once.
void Graph::backprop(Node* node, double initial_result) {
  initial_result = invert_result(initial_result);

  std::unordered_set<Hash> encountered_nodes;
  auto encountered = [&](Hash hash) {
    if (encountered_nodes.count(hash) > 0) {
      return true;
    } else {
      encountered_nodes.insert(hash);
      return false;
    }
  };

  std::stack<std::pair<Node*, double> > backlog;
  backlog.emplace(node, initial_result);

#ifdef MC_EXPENSIVE_RUNTIME_TESTS
  State known_win_state("rn4nr/p4N1p/8/1p4Qk/1Pp4P/8/PP1PPP1P/RNB1KBR1 b Q - 0 0");
#endif

  while (!backlog.empty()) {
    std::pair<Node*, double> pair = backlog.top();
    backlog.pop();

    Node* node = pair.first;
    double result = pair.second;

#ifdef MC_EXPENSIVE_RUNTIME_TESTS
    //if (node->hash == NodeTable::key(known_win_state.hash)) {
      // assertion commented out because hash collisions can cause it to fail.
      // uncomment it to get a backtrace when the relevant test case fails.
//      assert(result == win_value);
    //}
#endif
    node->update(result);

    double parent_result = invert_result(result);
    node->do_parents([&](Hash parent) {
        if (!encountered(parent))
          backlog.emplace(nodes.get(parent), parent_result);
      });
  }
}

template <typename F>
boost::optional<Move> Graph::select_successor_by(State state, F f) {
  Node* node = nodes.get_or_create(state);
  boost::optional<Move> best_move;
  double best_score;
  node->do_successors(state, [&](State const& state, Move curr_move) {
      Node* curr_node = nodes.get_or_create(state);
      double curr_score = f(curr_node);
      if (!best_move || best_score < curr_score) {
        best_move  = curr_move;
        best_score = curr_score;
      }
    });
  return best_move;
}

void Graph::print_statistics(std::ostream& os, State state) {
  Node* node = nodes.get_or_create(state);
  os << node->format_statistics() << std::endl;
  node->do_successors(state, [&](State const& state, Move last_move) {
      Node* node = nodes.get_or_create(state);
      os << last_move << " " << node->format_statistics() << std::endl;
    });
}

boost::optional<Move> Graph::principal_move(State state) {
  return select_successor_by(state, [](Node const* node) {
      return node->mean();
    });
}

void Graph::print_principal_variation(std::ostream& os, State state) {
  sorted_vector<Hash> path;
  print_principal_variation(os, state, path);
}

void Graph::print_principal_variation(std::ostream& os, State state, sorted_vector<Hash>& path) {
  boost::optional<Move> move = principal_move(state);
  if (!move)
    return;
  state.make_move(*move);
  Node* child = nodes.get_or_create(state);
  if (child->sample_size() == 0 || path.contains(child->hash))
    return;
  path.insert(child->hash);
  os << *move << " " << child->format_statistics() << std::endl;
  print_principal_variation(os, state, path);
}

// FIXME: handle cycles
void Graph::graphviz(std::ostream& os, Node* node, State state, boost::optional<Move> last_move = boost::none) {
  auto hash_as_id = [this](Hash hash) {
    return str(boost::format("\"%1$#8x\"") % hash);
  };
  os << hash_as_id(state.hash)
     << "[label=\"{"
     << (last_move ? notation::coordinate::format(*last_move) : "-")
     << " | "
     << node->format_statistics()
     << "}\"];"
     << std::endl;
  node->do_successors(state, [&](State state, Move last_move) {
      // only include existing nodes (to keep the graph small)
      Node* child = nodes.get(state.hash);
      if (!child)
        return;
      graphviz(os, node, state, last_move);
      os << hash_as_id(node->hash) << " -> " << hash_as_id(child->hash) << ";" << std::endl;
    });
}
<commit_msg>no longer breadth-first; depth-first may be more efficient<commit_after>#include "mcts.hpp"
#include "notation.hpp"

#include <stack>

#include <boost/range/algorithm/random_shuffle.hpp>

using namespace mcts;

void Node::initialize(State const& state) {
  hash = NodeTable::key(state.hash);
}

void Node::update(double result) {
  m_sample_size++;
  double mean0 = m_mean;
  double delta = result - mean0;
  m_mean += delta/m_sample_size;
  m_derivative = 0.9*m_derivative + 0.1*(m_mean - mean0);
}

void Node::adjoin_parent(Node* parent) {
  std::lock_guard<std::mutex> lock(parents_mutex);
  parents.insert(parent->hash);
}

std::string Node::format_statistics() {
  return str(boost::format("%1% %2% (d %3%) %4% (%5% parents)")
             % sample_size()
             % mean()
             % derivative()
             % selection_criterion()
             % parents.size());
}

Node* NodeTable::get(Hash hash) {
  size_t key = NodeTable::key(hash);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    return nullptr;
  return node;
}

Node* NodeTable::get_or_create(State const& state) {
  size_t key = NodeTable::key(state);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    node->initialize(state);
  return node;
}

void Graph::sample(State state, boost::mt19937& generator) {
  Node* node = nodes.get_or_create(state);
  sorted_vector<Hash> path;

  double result;

  // selection
  while (true) {
    path.insert(node->hash);

    Node* child = select_child(node, state, generator);

    if (!child) {
      // no legal successors; game over
      boost::optional<Color> winner = state.winner();
      result = !winner
        ? draw_value
        : (*winner == state.us
           ? win_value
           : loss_value);
      break;
    }

    if (state.drawn_by_50() || path.contains(child->hash)) {
      // draw per 50-move rule or by choosing repetition
      result = draw_value;
      break;
    }

    node = child;
  }

  backprop(node, result);
}

// returns nullptr if no legal successor states
// NOTE: state will be modified to be the corresponding successor
Node* Graph::select_child(Node* node, State& state, boost::mt19937& generator) {
  if (node->sample_size() < 30) {
    // don't waste time with unreliable statistics
    boost::optional<Move> move = moves::make_random_legal_move(state, generator);
    if (!move)
      return nullptr;
    Node* child = nodes.get_or_create(state);
    child->adjoin_parent(node);
    return child;
  } else {
    // use statistics to make selection
    std::vector<Move> moves;
    moves::moves(moves, state);

    // find a legal move with maximum selection criterion.
    boost::optional<Move> best_move;
    Node* best_child;
    double best_score;

    for (Move curr_move: moves) {
      Undo undo = state.make_move(curr_move);

      // if legal
      if (!state.their_king_attacked()) {
        Node* curr_child = nodes.get_or_create(state);

        curr_child->adjoin_parent(node);
        if (curr_child->sample_size() < 10)
          // select new nodes unconditionally
          return curr_child;

        double curr_score = curr_child->selection_criterion(generator);
        if (!best_move || curr_score > best_score) {
          best_move  = curr_move;
          best_child = curr_child;
          best_score = curr_score;
        }
      }

      state.unmake_move(undo);
    }

    if (!best_move)
      return nullptr;
    state.make_move(*best_move);
    return best_child;
  }
}

// update the node and all of its ancestors.  the graph is likely to be
// cyclic due to reversible moves.
// NOTE: initial_result is from the perspective of the player who is to move in
// the state corresponding to node.  since the stored node values are from the
// perspective of the player who causes the node to be chosen, we have to invert
// it once.
void Graph::backprop(Node* node, double initial_result) {
  initial_result = invert_result(initial_result);

  std::unordered_set<Hash> encountered_nodes;
  auto encountered = [&](Hash hash) {
    if (encountered_nodes.count(hash) > 0) {
      return true;
    } else {
      encountered_nodes.insert(hash);
      return false;
    }
  };

  std::stack<std::pair<Node*, double> > backlog;
  backlog.emplace(node, initial_result);

#ifdef MC_EXPENSIVE_RUNTIME_TESTS
  State known_win_state("rn4nr/p4N1p/8/1p4Qk/1Pp4P/8/PP1PPP1P/RNB1KBR1 b Q - 0 0");
#endif

  while (!backlog.empty()) {
    std::pair<Node*, double> pair = backlog.top();
    backlog.pop();

    Node* node = pair.first;
    double result = pair.second;

#ifdef MC_EXPENSIVE_RUNTIME_TESTS
    //if (node->hash == NodeTable::key(known_win_state.hash)) {
      // assertion commented out because hash collisions can cause it to fail.
      // uncomment it to get a backtrace when the relevant test case fails.
//      assert(result == win_value);
    //}
#endif
    node->update(result);

    double parent_result = invert_result(result);
    node->do_parents([&](Hash parent) {
        if (!encountered(parent))
          backlog.emplace(nodes.get(parent), parent_result);
      });
  }
}

template <typename F>
boost::optional<Move> Graph::select_successor_by(State state, F f) {
  Node* node = nodes.get_or_create(state);
  boost::optional<Move> best_move;
  double best_score;
  node->do_successors(state, [&](State const& state, Move curr_move) {
      Node* curr_node = nodes.get_or_create(state);
      double curr_score = f(curr_node);
      if (!best_move || best_score < curr_score) {
        best_move  = curr_move;
        best_score = curr_score;
      }
    });
  return best_move;
}

void Graph::print_statistics(std::ostream& os, State state) {
  Node* node = nodes.get_or_create(state);
  os << node->format_statistics() << std::endl;
  node->do_successors(state, [&](State const& state, Move last_move) {
      Node* node = nodes.get_or_create(state);
      os << last_move << " " << node->format_statistics() << std::endl;
    });
}

boost::optional<Move> Graph::principal_move(State state) {
  return select_successor_by(state, [](Node const* node) {
      return node->mean();
    });
}

void Graph::print_principal_variation(std::ostream& os, State state) {
  sorted_vector<Hash> path;
  print_principal_variation(os, state, path);
}

void Graph::print_principal_variation(std::ostream& os, State state, sorted_vector<Hash>& path) {
  boost::optional<Move> move = principal_move(state);
  if (!move)
    return;
  state.make_move(*move);
  Node* child = nodes.get_or_create(state);
  if (child->sample_size() == 0 || path.contains(child->hash))
    return;
  path.insert(child->hash);
  os << *move << " " << child->format_statistics() << std::endl;
  print_principal_variation(os, state, path);
}

// FIXME: handle cycles
void Graph::graphviz(std::ostream& os, Node* node, State state, boost::optional<Move> last_move = boost::none) {
  auto hash_as_id = [this](Hash hash) {
    return str(boost::format("\"%1$#8x\"") % hash);
  };
  os << hash_as_id(state.hash)
     << "[label=\"{"
     << (last_move ? notation::coordinate::format(*last_move) : "-")
     << " | "
     << node->format_statistics()
     << "}\"];"
     << std::endl;
  node->do_successors(state, [&](State state, Move last_move) {
      // only include existing nodes (to keep the graph small)
      Node* child = nodes.get(state.hash);
      if (!child)
        return;
      graphviz(os, node, state, last_move);
      os << hash_as_id(node->hash) << " -> " << hash_as_id(child->hash) << ";" << std::endl;
    });
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Jan Christoph Uhde
////////////////////////////////////////////////////////////////////////////////

#include "RocksDBOptionFeature.h"
#include "Basics/process-utils.h"
#include "Logger/Logger.h"
#include "ProgramOptions/ProgramOptions.h"
#include "ProgramOptions/Section.h"

#include <rocksdb/options.h>
#include <rocksdb/table.h>
#include <rocksdb/utilities/transaction_db.h>

using namespace arangodb;
using namespace arangodb::application_features;
using namespace arangodb::options;

namespace {
rocksdb::TransactionDBOptions rocksDBTrxDefaults;
rocksdb::Options rocksDBDefaults;
rocksdb::BlockBasedTableOptions rocksDBTableOptionsDefaults;
}

RocksDBOptionFeature::RocksDBOptionFeature(
    application_features::ApplicationServer* server)
    : application_features::ApplicationFeature(server, "RocksDBOption"),
      _transactionLockTimeout(rocksDBTrxDefaults.transaction_lock_timeout),
      _writeBufferSize(rocksDBDefaults.write_buffer_size),
      _maxWriteBufferNumber(rocksDBDefaults.max_write_buffer_number),
      _maxTotalWalSize(80 << 20),
      _delayedWriteRate(rocksDBDefaults.delayed_write_rate),
      _minWriteBufferNumberToMerge(
          rocksDBDefaults.min_write_buffer_number_to_merge),
      _numLevels(rocksDBDefaults.num_levels),
      _numUncompressedLevels(2),
      _maxBytesForLevelBase(rocksDBDefaults.max_bytes_for_level_base),
      _maxBytesForLevelMultiplier(
          rocksDBDefaults.max_bytes_for_level_multiplier),
      _maxBackgroundJobs(rocksDBDefaults.max_background_jobs),
      _maxSubcompactions(rocksDBDefaults.max_subcompactions),
      _numThreadsHigh(0),
      _numThreadsLow(0),
      _blockCacheSize((TRI_PhysicalMemory >= (static_cast<uint64_t>(4) << 30))
        ? static_cast<uint64_t>(((TRI_PhysicalMemory - (static_cast<uint64_t>(2) << 30)) * 0.3))
        : (256 << 20)),
      _blockCacheShardBits(-1),
      _tableBlockSize(std::max(rocksDBTableOptionsDefaults.block_size, static_cast<decltype(rocksDBTableOptionsDefaults.block_size)>(16 * 1024))),
      _recycleLogFileNum(rocksDBDefaults.recycle_log_file_num),
      _compactionReadaheadSize(2 * 1024 * 1024),//rocksDBDefaults.compaction_readahead_size
      _level0CompactionTrigger(2),
      _level0SlowdownTrigger(rocksDBDefaults.level0_slowdown_writes_trigger),
      _level0StopTrigger(rocksDBDefaults.level0_stop_writes_trigger),
      _enablePipelinedWrite(rocksDBDefaults.enable_pipelined_write),
      _optimizeFiltersForHits(rocksDBDefaults.optimize_filters_for_hits),
      _useDirectReads(rocksDBDefaults.use_direct_reads),
      _useDirectIoForFlushAndCompaction(rocksDBDefaults.use_direct_io_for_flush_and_compaction),
      _useFSync(rocksDBDefaults.use_fsync),
      _skipCorrupted(false),
      _dynamicLevelBytes(true),
      _enableStatistics(false) {
  // setting the number of background jobs to
  _maxBackgroundJobs = static_cast<int32_t>(std::max((size_t)2,
                                                     std::min(TRI_numberProcessors(), (size_t)8)));
#ifdef _WIN32
  // Windows code does not (yet) support lowering thread priority of
  //  compactions.  Therefore it is possible for rocksdb to use all
  //  CPU time on compactions.  Essential network communications can be lost.
  //  Save one CPU for ArangoDB network and other activities.
  if (2 < _maxBackgroundJobs) {
    --_maxBackgroundJobs;
  } // if
#endif

  setOptional(true);
  startsAfter("Daemon");
  startsAfter("DatabasePath");
}

void RocksDBOptionFeature::collectOptions(
    std::shared_ptr<ProgramOptions> options) {
  options->addSection("rocksdb", "Configure the RocksDB engine");

  options->addObsoleteOption("--rocksdb.enabled",
                             "obsolete always active - Whether or not the "
                             "RocksDB engine is enabled for the persistent "
                             "index",
                             true);

  options->addOption("--rocksdb.wal-directory",
                     "optional path to the RocksDB WAL directory. "
                     "If not set, the WAL directory will be located inside the regular data directory",
                     new StringParameter(&_walDirectory));

  options->addOption("--rocksdb.transaction-lock-timeout",
                     "If positive, specifies the wait timeout in milliseconds when "
                     " a transaction attempts to lock a document. A negative value "
                     "is not recommended as it can lead to deadlocks (0 = no waiting, < 0 no timeout)",
                     new Int64Parameter(&_transactionLockTimeout));

  options->addOption("--rocksdb.write-buffer-size",
                     "amount of data to build up in memory before converting "
                     "to a sorted on-disk file (0 = disabled)",
                     new UInt64Parameter(&_writeBufferSize));

  options->addOption("--rocksdb.max-write-buffer-number",
                     "maximum number of write buffers that built up in memory",
                     new UInt64Parameter(&_maxWriteBufferNumber));

  options->addOption("--rocksdb.max-total-wal-size",
                     "maximum total size of WAL files that will force flush stale column families",
                     new UInt64Parameter(&_maxTotalWalSize));

  options->addHiddenOption(
      "--rocksdb.delayed_write_rate",
      "limited write rate to DB (in bytes per second) if we are writing to the "
      "last mem-table allowed and we allow more than 3 mem-tables, or if we "
      "have surpassed a certain number of level-0 files and need to slowdown "
      "writes",
      new UInt64Parameter(&_delayedWriteRate));

  options->addOption("--rocksdb.min-write-buffer-number-to-merge",
                     "minimum number of write buffers that will be merged "
                     "together before writing "
                     "to storage",
                     new UInt64Parameter(&_minWriteBufferNumberToMerge));

  options->addOption("--rocksdb.num-levels",
                     "number of levels for the database",
                     new UInt64Parameter(&_numLevels));

  options->addOption("--rocksdb.num-uncompressed-levels",
                     "number of uncompressed levels for the database",
                     new UInt64Parameter(&_numUncompressedLevels));

  options->addOption("--rocksdb.dynamic-level-bytes",
                     "if true, determine the number of bytes for each level "
                     "dynamically to minimize space amplification",
                     new BooleanParameter(&_dynamicLevelBytes));

  options->addOption("--rocksdb.max-bytes-for-level-base",
                     "if not using dynamic level sizes, this controls the "
                     "maximum total data size for level-1",
                     new UInt64Parameter(&_maxBytesForLevelBase));

  options->addOption("--rocksdb.max-bytes-for-level-multiplier",
                     "if not using dynamic level sizes, the maximum number of "
                     "bytes for level L can be calculated as "
                     " max-bytes-for-level-base * "
                     "(max-bytes-for-level-multiplier ^ (L-1))",
                     new DoubleParameter(&_maxBytesForLevelMultiplier));

  options->addOption("--rocksdb.enable-pipelined-write",
                     "if true, use a two stage write queue for WAL writes and memtable writes",
                     new BooleanParameter(&_enablePipelinedWrite));

  options->addOption("--rocksdb.enable-statistics",
                     "whether or not RocksDB statistics should be turned on",
                     new BooleanParameter(&_enableStatistics));

  options->addHiddenOption(
      "--rocksdb.optimize-filters-for-hits",
      "this flag specifies that the implementation should optimize the filters "
      "mainly for cases where keys are found rather than also optimize for "
      "keys missed. This would be used in cases where the application knows "
      "that there are very few misses or the performance in the case of "
      "misses is not important",
      new BooleanParameter(&_optimizeFiltersForHits));

#ifdef __linux__
  options->addHiddenOption("--rocksdb.use-direct-reads",
                           "use O_DIRECT for reading files",
                           new BooleanParameter(&_useDirectReads));

  options->addHiddenOption("--rocksdb.use-direct-io-for-flush-and-compaction",
                           "use O_DIRECT for flush and compaction",
                           new BooleanParameter(&_useDirectIoForFlushAndCompaction));
#endif

  options->addHiddenOption("--rocksdb.use-fsync",
                           "issue an fsync when writing to disk (set to true "
                           "for issuing fdatasync only)",
                           new BooleanParameter(&_useFSync));

  options->addHiddenOption(
      "--rocksdb.max-background-jobs",
      "Maximum number of concurrent background jobs (compactions and flushes)",
      new Int32Parameter(&_maxBackgroundJobs));

  options->addOption("--rocksdb.max-subcompactions",
                     "maximum number of concurrent subjobs for a background "
                     "compaction",
                     new UInt64Parameter(&_maxSubcompactions));

  options->addOption("--rocksdb.level0-compaction-trigger",
                     "number of level-0 files that triggers a compaction",
                     new Int64Parameter(&_level0CompactionTrigger));

  options->addOption("--rocksdb.level0-slowdown-trigger",
                     "number of level-0 files that triggers a write slowdown",
                     new Int64Parameter(&_level0SlowdownTrigger));

  options->addOption("--rocksdb.level0-stop-trigger",
                     "number of level-0 files that triggers a full write stall",
                     new Int64Parameter(&_level0StopTrigger));

  options->addOption(
      "--rocksdb.num-threads-priority-high",
      "number of threads for high priority operations (e.g. flush)",
      new UInt32Parameter(&_numThreadsHigh));

  options->addOption(
      "--rocksdb.num-threads-priority-low",
      "number of threads for low priority operations (e.g. compaction)",
      new UInt32Parameter(&_numThreadsLow));

  options->addOption("--rocksdb.block-cache-size",
                     "size of block cache in bytes",
                     new UInt64Parameter(&_blockCacheSize));

  options->addOption("--rocksdb.block-cache-shard-bits",
                     "number of shard bits to use for block cache (use -1 for default value)",
                     new Int64Parameter(&_blockCacheShardBits));

  options->addOption("--rocksdb.table-block-size",
                     "approximate size (in bytes) of user data packed per block",
                     new UInt64Parameter(&_tableBlockSize));

  options->addHiddenOption("--rocksdb.recycle-log-file-num",
                           "number of log files to keep around for recycling",
                           new UInt64Parameter(&_recycleLogFileNum));

  options->addOption(
      "--rocksdb.compaction-read-ahead-size",
      "if non-zero, we perform bigger reads when doing compaction. If you're "
      "running RocksDB on spinning disks, you should set this to at least 2MB. "
      "that way RocksDB's compaction is doing sequential instead of random "
      "reads.",
      new UInt64Parameter(&_compactionReadaheadSize));

  options->addHiddenOption("--rocksdb.wal-recovery-skip-corrupted",
                           "skip corrupted records in WAL recovery",
                           new BooleanParameter(&_skipCorrupted));
}

void RocksDBOptionFeature::validateOptions(
    std::shared_ptr<ProgramOptions> options) {
  if (_writeBufferSize > 0 && _writeBufferSize < 1024 * 1024) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.write-buffer-size'";
    FATAL_ERROR_EXIT();
  }
  if (_maxBytesForLevelMultiplier <= 0.0) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.max-bytes-for-level-multiplier'";
    FATAL_ERROR_EXIT();
  }
  if (_numLevels < 1 || _numLevels > 20) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-levels'";
    FATAL_ERROR_EXIT();
  }

  if (_maxBackgroundJobs != -1 &&
      (_maxBackgroundJobs < 1 || _maxBackgroundJobs > 128)) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.max-background-jobs'";
    FATAL_ERROR_EXIT();
  }

  if (_numThreadsHigh > 64) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-threads-priority-high'";
    FATAL_ERROR_EXIT();
  }
  if ( _numThreadsLow > 256) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-threads-priority-low'";
    FATAL_ERROR_EXIT();
  }
  if (_maxSubcompactions > _numThreadsLow) {
    _maxSubcompactions = _numThreadsLow;
  }
  if (_blockCacheShardBits >= 20 || _blockCacheShardBits < -1) {
    // -1 is RocksDB default value, but anything less is invalid
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.block-cache-shard-bits'";
    FATAL_ERROR_EXIT();
  }
}

void RocksDBOptionFeature::start() {
  uint32_t max = _maxBackgroundJobs / 2;
  uint32_t clamped = std::max(std::min((uint32_t)TRI_numberProcessors(), max), 1U);
  // lets test this out
  if (_numThreadsHigh == 0) {
    _numThreadsHigh = clamped;
  }
  if (_numThreadsLow == 0) {
    _numThreadsLow = clamped;
  }

  LOG_TOPIC(TRACE, Logger::ROCKSDB) << "using RocksDB options:"
                                    << " wal_dir: " << _walDirectory << "'"
                                    << ", write_buffer_size: " << _writeBufferSize
                                    << ", max_write_buffer_number: " << _maxWriteBufferNumber
                                    << ", max_total_wal_size: " << _maxTotalWalSize
                                    << ", delayed_write_rate: " << _delayedWriteRate
                                    << ", min_write_buffer_number_to_merge: " << _minWriteBufferNumberToMerge
                                    << ", num_levels: " << _numLevels
                                    << ", num_uncompressed_levels: " << _numUncompressedLevels
                                    << ", max_bytes_for_level_base: " << _maxBytesForLevelBase
                                    << ", max_bytes_for_level_multiplier: " << _maxBytesForLevelMultiplier
                                    << ", max_background_jobs: " << _maxBackgroundJobs
                                    << ", max_sub_compactions: " << _maxSubcompactions
                                    << ", num_threads_high: " << _numThreadsHigh
                                    << ", num_threads_low: " << _numThreadsLow
                                    << ", block_cache_size: " << _blockCacheSize
                                    << ", block_cache_shard_bits: " << _blockCacheShardBits
                                    << ", table_block_size: " << _tableBlockSize
                                    << ", recycle_log_file_num: " << _recycleLogFileNum
                                    << ", compaction_read_ahead_size: " << _compactionReadaheadSize
                                    << ", level0_compaction_trigger: " << _level0CompactionTrigger
                                    << ", level0_slowdown_trigger: " << _level0SlowdownTrigger
                                    << ", enable_pipelined_write: " << _enablePipelinedWrite
                                    << ", optimize_filters_for_hits: " << _optimizeFiltersForHits
                                    << ", use_direct_reads: " << _useDirectReads
                                    << ", use_direct_io_for_flush_and_compaction: " << _useDirectIoForFlushAndCompaction
                                    << ", use_fsync: " << _useFSync
                                    << ", dynamic_level_bytes: " << std::boolalpha << _dynamicLevelBytes;
}
<commit_msg>fix small typo<commit_after>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Jan Christoph Uhde
////////////////////////////////////////////////////////////////////////////////

#include "RocksDBOptionFeature.h"
#include "Basics/process-utils.h"
#include "Logger/Logger.h"
#include "ProgramOptions/ProgramOptions.h"
#include "ProgramOptions/Section.h"

#include <rocksdb/options.h>
#include <rocksdb/table.h>
#include <rocksdb/utilities/transaction_db.h>

using namespace arangodb;
using namespace arangodb::application_features;
using namespace arangodb::options;

namespace {
rocksdb::TransactionDBOptions rocksDBTrxDefaults;
rocksdb::Options rocksDBDefaults;
rocksdb::BlockBasedTableOptions rocksDBTableOptionsDefaults;
}

RocksDBOptionFeature::RocksDBOptionFeature(
    application_features::ApplicationServer* server)
    : application_features::ApplicationFeature(server, "RocksDBOption"),
      _transactionLockTimeout(rocksDBTrxDefaults.transaction_lock_timeout),
      _writeBufferSize(rocksDBDefaults.write_buffer_size),
      _maxWriteBufferNumber(rocksDBDefaults.max_write_buffer_number),
      _maxTotalWalSize(80 << 20),
      _delayedWriteRate(rocksDBDefaults.delayed_write_rate),
      _minWriteBufferNumberToMerge(
          rocksDBDefaults.min_write_buffer_number_to_merge),
      _numLevels(rocksDBDefaults.num_levels),
      _numUncompressedLevels(2),
      _maxBytesForLevelBase(rocksDBDefaults.max_bytes_for_level_base),
      _maxBytesForLevelMultiplier(
          rocksDBDefaults.max_bytes_for_level_multiplier),
      _maxBackgroundJobs(rocksDBDefaults.max_background_jobs),
      _maxSubcompactions(rocksDBDefaults.max_subcompactions),
      _numThreadsHigh(0),
      _numThreadsLow(0),
      _blockCacheSize((TRI_PhysicalMemory >= (static_cast<uint64_t>(4) << 30))
        ? static_cast<uint64_t>(((TRI_PhysicalMemory - (static_cast<uint64_t>(2) << 30)) * 0.3))
        : (256 << 20)),
      _blockCacheShardBits(-1),
      _tableBlockSize(std::max(rocksDBTableOptionsDefaults.block_size, static_cast<decltype(rocksDBTableOptionsDefaults.block_size)>(16 * 1024))),
      _recycleLogFileNum(rocksDBDefaults.recycle_log_file_num),
      _compactionReadaheadSize(2 * 1024 * 1024),//rocksDBDefaults.compaction_readahead_size
      _level0CompactionTrigger(2),
      _level0SlowdownTrigger(rocksDBDefaults.level0_slowdown_writes_trigger),
      _level0StopTrigger(rocksDBDefaults.level0_stop_writes_trigger),
      _enablePipelinedWrite(rocksDBDefaults.enable_pipelined_write),
      _optimizeFiltersForHits(rocksDBDefaults.optimize_filters_for_hits),
      _useDirectReads(rocksDBDefaults.use_direct_reads),
      _useDirectIoForFlushAndCompaction(rocksDBDefaults.use_direct_io_for_flush_and_compaction),
      _useFSync(rocksDBDefaults.use_fsync),
      _skipCorrupted(false),
      _dynamicLevelBytes(true),
      _enableStatistics(false) {
  // setting the number of background jobs to
  _maxBackgroundJobs = static_cast<int32_t>(std::max((size_t)2,
                                                     std::min(TRI_numberProcessors(), (size_t)8)));
#ifdef _WIN32
  // Windows code does not (yet) support lowering thread priority of
  //  compactions.  Therefore it is possible for rocksdb to use all
  //  CPU time on compactions.  Essential network communications can be lost.
  //  Save one CPU for ArangoDB network and other activities.
  if (2 < _maxBackgroundJobs) {
    --_maxBackgroundJobs;
  } // if
#endif

  setOptional(true);
  startsAfter("Daemon");
  startsAfter("DatabasePath");
}

void RocksDBOptionFeature::collectOptions(
    std::shared_ptr<ProgramOptions> options) {
  options->addSection("rocksdb", "Configure the RocksDB engine");

  options->addObsoleteOption("--rocksdb.enabled",
                             "obsolete always active - Whether or not the "
                             "RocksDB engine is enabled for the persistent "
                             "index",
                             true);

  options->addOption("--rocksdb.wal-directory",
                     "optional path to the RocksDB WAL directory. "
                     "If not set, the WAL directory will be located inside the regular data directory",
                     new StringParameter(&_walDirectory));

  options->addOption("--rocksdb.transaction-lock-timeout",
                     "If positive, specifies the wait timeout in milliseconds when "
                     " a transaction attempts to lock a document. A negative value "
                     "is not recommended as it can lead to deadlocks (0 = no waiting, < 0 no timeout)",
                     new Int64Parameter(&_transactionLockTimeout));

  options->addOption("--rocksdb.write-buffer-size",
                     "amount of data to build up in memory before converting "
                     "to a sorted on-disk file (0 = disabled)",
                     new UInt64Parameter(&_writeBufferSize));

  options->addOption("--rocksdb.max-write-buffer-number",
                     "maximum number of write buffers that built up in memory",
                     new UInt64Parameter(&_maxWriteBufferNumber));

  options->addOption("--rocksdb.max-total-wal-size",
                     "maximum total size of WAL files that will force flush stale column families",
                     new UInt64Parameter(&_maxTotalWalSize));

  options->addHiddenOption(
      "--rocksdb.delayed_write_rate",
      "limited write rate to DB (in bytes per second) if we are writing to the "
      "last mem-table allowed and we allow more than 3 mem-tables, or if we "
      "have surpassed a certain number of level-0 files and need to slowdown "
      "writes",
      new UInt64Parameter(&_delayedWriteRate));

  options->addOption("--rocksdb.min-write-buffer-number-to-merge",
                     "minimum number of write buffers that will be merged "
                     "together before writing "
                     "to storage",
                     new UInt64Parameter(&_minWriteBufferNumberToMerge));

  options->addOption("--rocksdb.num-levels",
                     "number of levels for the database",
                     new UInt64Parameter(&_numLevels));

  options->addOption("--rocksdb.num-uncompressed-levels",
                     "number of uncompressed levels for the database",
                     new UInt64Parameter(&_numUncompressedLevels));

  options->addOption("--rocksdb.dynamic-level-bytes",
                     "if true, determine the number of bytes for each level "
                     "dynamically to minimize space amplification",
                     new BooleanParameter(&_dynamicLevelBytes));

  options->addOption("--rocksdb.max-bytes-for-level-base",
                     "if not using dynamic level sizes, this controls the "
                     "maximum total data size for level-1",
                     new UInt64Parameter(&_maxBytesForLevelBase));

  options->addOption("--rocksdb.max-bytes-for-level-multiplier",
                     "if not using dynamic level sizes, the maximum number of "
                     "bytes for level L can be calculated as "
                     " max-bytes-for-level-base * "
                     "(max-bytes-for-level-multiplier ^ (L-1))",
                     new DoubleParameter(&_maxBytesForLevelMultiplier));

  options->addOption("--rocksdb.enable-pipelined-write",
                     "if true, use a two stage write queue for WAL writes and memtable writes",
                     new BooleanParameter(&_enablePipelinedWrite));

  options->addOption("--rocksdb.enable-statistics",
                     "whether or not RocksDB statistics should be turned on",
                     new BooleanParameter(&_enableStatistics));

  options->addHiddenOption(
      "--rocksdb.optimize-filters-for-hits",
      "this flag specifies that the implementation should optimize the filters "
      "mainly for cases where keys are found rather than also optimize for "
      "keys missed. This would be used in cases where the application knows "
      "that there are very few misses or the performance in the case of "
      "misses is not important",
      new BooleanParameter(&_optimizeFiltersForHits));

#ifdef __linux__
  options->addHiddenOption("--rocksdb.use-direct-reads",
                           "use O_DIRECT for reading files",
                           new BooleanParameter(&_useDirectReads));

  options->addHiddenOption("--rocksdb.use-direct-io-for-flush-and-compaction",
                           "use O_DIRECT for flush and compaction",
                           new BooleanParameter(&_useDirectIoForFlushAndCompaction));
#endif

  options->addHiddenOption("--rocksdb.use-fsync",
                           "issue an fsync when writing to disk (set to true "
                           "for issuing fdatasync only)",
                           new BooleanParameter(&_useFSync));

  options->addHiddenOption(
      "--rocksdb.max-background-jobs",
      "Maximum number of concurrent background jobs (compactions and flushes)",
      new Int32Parameter(&_maxBackgroundJobs));

  options->addOption("--rocksdb.max-subcompactions",
                     "maximum number of concurrent subjobs for a background "
                     "compaction",
                     new UInt64Parameter(&_maxSubcompactions));

  options->addOption("--rocksdb.level0-compaction-trigger",
                     "number of level-0 files that triggers a compaction",
                     new Int64Parameter(&_level0CompactionTrigger));

  options->addOption("--rocksdb.level0-slowdown-trigger",
                     "number of level-0 files that triggers a write slowdown",
                     new Int64Parameter(&_level0SlowdownTrigger));

  options->addOption("--rocksdb.level0-stop-trigger",
                     "number of level-0 files that triggers a full write stall",
                     new Int64Parameter(&_level0StopTrigger));

  options->addOption(
      "--rocksdb.num-threads-priority-high",
      "number of threads for high priority operations (e.g. flush)",
      new UInt32Parameter(&_numThreadsHigh));

  options->addOption(
      "--rocksdb.num-threads-priority-low",
      "number of threads for low priority operations (e.g. compaction)",
      new UInt32Parameter(&_numThreadsLow));

  options->addOption("--rocksdb.block-cache-size",
                     "size of block cache in bytes",
                     new UInt64Parameter(&_blockCacheSize));

  options->addOption("--rocksdb.block-cache-shard-bits",
                     "number of shard bits to use for block cache (use -1 for default value)",
                     new Int64Parameter(&_blockCacheShardBits));

  options->addOption("--rocksdb.table-block-size",
                     "approximate size (in bytes) of user data packed per block",
                     new UInt64Parameter(&_tableBlockSize));

  options->addHiddenOption("--rocksdb.recycle-log-file-num",
                           "number of log files to keep around for recycling",
                           new UInt64Parameter(&_recycleLogFileNum));

  options->addOption(
      "--rocksdb.compaction-read-ahead-size",
      "if non-zero, we perform bigger reads when doing compaction. If you're "
      "running RocksDB on spinning disks, you should set this to at least 2MB. "
      "that way RocksDB's compaction is doing sequential instead of random "
      "reads.",
      new UInt64Parameter(&_compactionReadaheadSize));

  options->addHiddenOption("--rocksdb.wal-recovery-skip-corrupted",
                           "skip corrupted records in WAL recovery",
                           new BooleanParameter(&_skipCorrupted));
}

void RocksDBOptionFeature::validateOptions(
    std::shared_ptr<ProgramOptions> options) {
  if (_writeBufferSize > 0 && _writeBufferSize < 1024 * 1024) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.write-buffer-size'";
    FATAL_ERROR_EXIT();
  }
  if (_maxBytesForLevelMultiplier <= 0.0) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.max-bytes-for-level-multiplier'";
    FATAL_ERROR_EXIT();
  }
  if (_numLevels < 1 || _numLevels > 20) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-levels'";
    FATAL_ERROR_EXIT();
  }

  if (_maxBackgroundJobs != -1 &&
      (_maxBackgroundJobs < 1 || _maxBackgroundJobs > 128)) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.max-background-jobs'";
    FATAL_ERROR_EXIT();
  }

  if (_numThreadsHigh > 64) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-threads-priority-high'";
    FATAL_ERROR_EXIT();
  }
  if ( _numThreadsLow > 256) {
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.num-threads-priority-low'";
    FATAL_ERROR_EXIT();
  }
  if (_maxSubcompactions > _numThreadsLow) {
    _maxSubcompactions = _numThreadsLow;
  }
  if (_blockCacheShardBits >= 20 || _blockCacheShardBits < -1) {
    // -1 is RocksDB default value, but anything less is invalid
    LOG_TOPIC(FATAL, arangodb::Logger::FIXME)
        << "invalid value for '--rocksdb.block-cache-shard-bits'";
    FATAL_ERROR_EXIT();
  }
}

void RocksDBOptionFeature::start() {
  uint32_t max = _maxBackgroundJobs / 2;
  uint32_t clamped = std::max(std::min((uint32_t)TRI_numberProcessors(), max), 1U);
  // lets test this out
  if (_numThreadsHigh == 0) {
    _numThreadsHigh = clamped;
  }
  if (_numThreadsLow == 0) {
    _numThreadsLow = clamped;
  }

  LOG_TOPIC(TRACE, Logger::ROCKSDB) << "using RocksDB options:"
                                    << " wal_dir: '" << _walDirectory << "'"
                                    << ", write_buffer_size: " << _writeBufferSize
                                    << ", max_write_buffer_number: " << _maxWriteBufferNumber
                                    << ", max_total_wal_size: " << _maxTotalWalSize
                                    << ", delayed_write_rate: " << _delayedWriteRate
                                    << ", min_write_buffer_number_to_merge: " << _minWriteBufferNumberToMerge
                                    << ", num_levels: " << _numLevels
                                    << ", num_uncompressed_levels: " << _numUncompressedLevels
                                    << ", max_bytes_for_level_base: " << _maxBytesForLevelBase
                                    << ", max_bytes_for_level_multiplier: " << _maxBytesForLevelMultiplier
                                    << ", max_background_jobs: " << _maxBackgroundJobs
                                    << ", max_sub_compactions: " << _maxSubcompactions
                                    << ", num_threads_high: " << _numThreadsHigh
                                    << ", num_threads_low: " << _numThreadsLow
                                    << ", block_cache_size: " << _blockCacheSize
                                    << ", block_cache_shard_bits: " << _blockCacheShardBits
                                    << ", table_block_size: " << _tableBlockSize
                                    << ", recycle_log_file_num: " << _recycleLogFileNum
                                    << ", compaction_read_ahead_size: " << _compactionReadaheadSize
                                    << ", level0_compaction_trigger: " << _level0CompactionTrigger
                                    << ", level0_slowdown_trigger: " << _level0SlowdownTrigger
                                    << ", enable_pipelined_write: " << _enablePipelinedWrite
                                    << ", optimize_filters_for_hits: " << _optimizeFiltersForHits
                                    << ", use_direct_reads: " << _useDirectReads
                                    << ", use_direct_io_for_flush_and_compaction: " << _useDirectIoForFlushAndCompaction
                                    << ", use_fsync: " << _useFSync
                                    << ", dynamic_level_bytes: " << std::boolalpha << _dynamicLevelBytes;
}
<|endoftext|>


<commit_before>//===- SourceCoverageViewText.cpp - A text-based code coverage view -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements the text-based coverage renderer.
//
//===----------------------------------------------------------------------===//

#include "SourceCoverageViewText.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/StringExtras.h"

using namespace llvm;

namespace {

constexpr unsigned LineCoverageColumnWidth = 7;
constexpr unsigned LineNumberColumnWidth = 5;

/// \brief Get the width of the leading columns.
unsigned getCombinedColumnWidth(const CoverageViewOptions &Opts) {
  return (Opts.ShowLineStats ? LineCoverageColumnWidth + 1 : 0) +
         (Opts.ShowLineNumbers ? LineNumberColumnWidth + 1 : 0);
}

/// \brief The width of the line that is used to divide between the view and
/// the subviews.
unsigned getDividerWidth(const CoverageViewOptions &Opts) {
  return getCombinedColumnWidth(Opts) + 4;
}

} // anonymous namespace

void SourceCoverageViewText::renderSourceName(raw_ostream &OS) {
  getOptions().colored_ostream(OS, raw_ostream::CYAN) << getSourceName()
                                                      << ":\n";
}

void SourceCoverageViewText::renderLinePrefix(raw_ostream &OS,
                                              unsigned ViewDepth) {
  for (unsigned I = 0; I < ViewDepth; ++I)
    OS << "  |";
}

void SourceCoverageViewText::renderViewDivider(raw_ostream &OS,
                                               unsigned ViewDepth) {
  assert(ViewDepth != 0 && "Cannot render divider at top level");
  renderLinePrefix(OS, ViewDepth - 1);
  OS.indent(2);
  unsigned Length = getDividerWidth(getOptions());
  for (unsigned I = 0; I < Length; ++I)
    OS << '-';
  OS << '\n';
}

void SourceCoverageViewText::renderLine(
    raw_ostream &OS, LineRef L,
    const coverage::CoverageSegment *WrappedSegment,
    CoverageSegmentArray Segments, unsigned ExpansionCol, unsigned ViewDepth) {
  StringRef Line = L.Line;
  unsigned LineNumber = L.LineNo;

  Optional<raw_ostream::Colors> Highlight;
  SmallVector<std::pair<unsigned, unsigned>, 2> HighlightedRanges;

  // The first segment overlaps from a previous line, so we treat it specially.
  if (WrappedSegment && WrappedSegment->HasCount && WrappedSegment->Count == 0)
    Highlight = raw_ostream::RED;

  // Output each segment of the line, possibly highlighted.
  unsigned Col = 1;
  for (const auto *S : Segments) {
    unsigned End = std::min(S->Col, static_cast<unsigned>(Line.size()) + 1);
    colored_ostream(OS, Highlight ? *Highlight : raw_ostream::SAVEDCOLOR,
                    getOptions().Colors && Highlight, /*Bold=*/false,
                    /*BG=*/true)
        << Line.substr(Col - 1, End - Col);
    if (getOptions().Debug && Highlight)
      HighlightedRanges.push_back(std::make_pair(Col, End));
    Col = End;
    if (Col == ExpansionCol)
      Highlight = raw_ostream::CYAN;
    else if (S->HasCount && S->Count == 0)
      Highlight = raw_ostream::RED;
    else
      Highlight = None;
  }

  // Show the rest of the line.
  colored_ostream(OS, Highlight ? *Highlight : raw_ostream::SAVEDCOLOR,
                  getOptions().Colors && Highlight, /*Bold=*/false, /*BG=*/true)
      << Line.substr(Col - 1, Line.size() - Col + 1);
  OS << '\n';

  if (getOptions().Debug) {
    for (const auto &Range : HighlightedRanges)
      errs() << "Highlighted line " << LineNumber << ", " << Range.first
             << " -> " << Range.second << '\n';
    if (Highlight)
      errs() << "Highlighted line " << LineNumber << ", " << Col << " -> ?\n";
  }
}

void SourceCoverageViewText::renderLineCoverageColumn(
    raw_ostream &OS, const LineCoverageStats &Line) {
  if (!Line.isMapped()) {
    OS.indent(LineCoverageColumnWidth) << '|';
    return;
  }
  std::string C = formatCount(Line.ExecutionCount);
  OS.indent(LineCoverageColumnWidth - C.size());
  colored_ostream(OS, raw_ostream::MAGENTA,
                  Line.hasMultipleRegions() && getOptions().Colors)
      << C;
  OS << '|';
}

void SourceCoverageViewText::renderLineNumberColumn(raw_ostream &OS,
                                                    unsigned LineNo) {
  SmallString<32> Buffer;
  raw_svector_ostream BufferOS(Buffer);
  BufferOS << LineNo;
  auto Str = BufferOS.str();
  // Trim and align to the right.
  Str = Str.substr(0, std::min(Str.size(), (size_t)LineNumberColumnWidth));
  OS.indent(LineNumberColumnWidth - Str.size()) << Str << '|';
}

void SourceCoverageViewText::renderRegionMarkers(
    raw_ostream &OS, CoverageSegmentArray Segments, unsigned ViewDepth) {
  renderLinePrefix(OS, ViewDepth);
  OS.indent(getCombinedColumnWidth(getOptions()));

  unsigned PrevColumn = 1;
  for (const auto *S : Segments) {
    if (!S->IsRegionEntry)
      continue;
    // Skip to the new region.
    if (S->Col > PrevColumn)
      OS.indent(S->Col - PrevColumn);
    PrevColumn = S->Col + 1;
    std::string C = formatCount(S->Count);
    PrevColumn += C.size();
    OS << '^' << C;
  }
  OS << '\n';

  if (getOptions().Debug)
    for (const auto *S : Segments)
      errs() << "Marker at " << S->Line << ":" << S->Col << " = "
             << formatCount(S->Count) << (S->IsRegionEntry ? "\n" : " (pop)\n");
}

unsigned SourceCoverageViewText::renderExpansionView(
    raw_ostream &OS, ExpansionView &ESV, Optional<LineRef> FirstLine,
    unsigned ExpansionCol, const coverage::CoverageSegment *WrappedSegment,
    CoverageSegmentArray LineSegments, unsigned ViewDepth) {
  unsigned NextExpansionCol = ExpansionCol;

  if (FirstLine.hasValue()) {
    // Re-render the current line and highlight the expansion range for
    // this subview.
    NextExpansionCol = ESV.getStartCol();
    renderLinePrefix(OS, ViewDepth);
    OS.indent(getCombinedColumnWidth(getOptions()) + (ViewDepth == 0 ? 0 : 1));
    renderLine(OS, *FirstLine, WrappedSegment, LineSegments, ExpansionCol,
               ViewDepth);
    renderViewDivider(OS, ViewDepth + 1);
  }

  // Render the child subview.
  if (getOptions().Debug)
    errs() << "Expansion at line " << ESV.getLine() << ", " << ESV.getStartCol()
           << " -> " << ESV.getEndCol() << '\n';
  ESV.View->print(OS, /*WholeFile=*/false, /*ShowSourceName=*/false,
                  ViewDepth + 1);

  return NextExpansionCol;
}

void SourceCoverageViewText::renderInstantiationView(raw_ostream &OS,
                                                     InstantiationView &ISV,
                                                     unsigned ViewDepth) {
  renderViewDivider(OS, ViewDepth);
  renderLinePrefix(OS, ViewDepth);
  OS << ' ';
  ISV.View->print(OS, /*WholeFile=*/false, /*ShowSourceName=*/true, ViewDepth);
}
<commit_msg>Try to fix the MSVC build<commit_after>//===- SourceCoverageViewText.cpp - A text-based code coverage view -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements the text-based coverage renderer.
//
//===----------------------------------------------------------------------===//

#include "SourceCoverageViewText.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/StringExtras.h"

using namespace llvm;

namespace {

static const unsigned LineCoverageColumnWidth = 7;
static const unsigned LineNumberColumnWidth = 5;

/// \brief Get the width of the leading columns.
unsigned getCombinedColumnWidth(const CoverageViewOptions &Opts) {
  return (Opts.ShowLineStats ? LineCoverageColumnWidth + 1 : 0) +
         (Opts.ShowLineNumbers ? LineNumberColumnWidth + 1 : 0);
}

/// \brief The width of the line that is used to divide between the view and
/// the subviews.
unsigned getDividerWidth(const CoverageViewOptions &Opts) {
  return getCombinedColumnWidth(Opts) + 4;
}

} // anonymous namespace

void SourceCoverageViewText::renderSourceName(raw_ostream &OS) {
  getOptions().colored_ostream(OS, raw_ostream::CYAN) << getSourceName()
                                                      << ":\n";
}

void SourceCoverageViewText::renderLinePrefix(raw_ostream &OS,
                                              unsigned ViewDepth) {
  for (unsigned I = 0; I < ViewDepth; ++I)
    OS << "  |";
}

void SourceCoverageViewText::renderViewDivider(raw_ostream &OS,
                                               unsigned ViewDepth) {
  assert(ViewDepth != 0 && "Cannot render divider at top level");
  renderLinePrefix(OS, ViewDepth - 1);
  OS.indent(2);
  unsigned Length = getDividerWidth(getOptions());
  for (unsigned I = 0; I < Length; ++I)
    OS << '-';
  OS << '\n';
}

void SourceCoverageViewText::renderLine(
    raw_ostream &OS, LineRef L,
    const coverage::CoverageSegment *WrappedSegment,
    CoverageSegmentArray Segments, unsigned ExpansionCol, unsigned ViewDepth) {
  StringRef Line = L.Line;
  unsigned LineNumber = L.LineNo;

  Optional<raw_ostream::Colors> Highlight;
  SmallVector<std::pair<unsigned, unsigned>, 2> HighlightedRanges;

  // The first segment overlaps from a previous line, so we treat it specially.
  if (WrappedSegment && WrappedSegment->HasCount && WrappedSegment->Count == 0)
    Highlight = raw_ostream::RED;

  // Output each segment of the line, possibly highlighted.
  unsigned Col = 1;
  for (const auto *S : Segments) {
    unsigned End = std::min(S->Col, static_cast<unsigned>(Line.size()) + 1);
    colored_ostream(OS, Highlight ? *Highlight : raw_ostream::SAVEDCOLOR,
                    getOptions().Colors && Highlight, /*Bold=*/false,
                    /*BG=*/true)
        << Line.substr(Col - 1, End - Col);
    if (getOptions().Debug && Highlight)
      HighlightedRanges.push_back(std::make_pair(Col, End));
    Col = End;
    if (Col == ExpansionCol)
      Highlight = raw_ostream::CYAN;
    else if (S->HasCount && S->Count == 0)
      Highlight = raw_ostream::RED;
    else
      Highlight = None;
  }

  // Show the rest of the line.
  colored_ostream(OS, Highlight ? *Highlight : raw_ostream::SAVEDCOLOR,
                  getOptions().Colors && Highlight, /*Bold=*/false, /*BG=*/true)
      << Line.substr(Col - 1, Line.size() - Col + 1);
  OS << '\n';

  if (getOptions().Debug) {
    for (const auto &Range : HighlightedRanges)
      errs() << "Highlighted line " << LineNumber << ", " << Range.first
             << " -> " << Range.second << '\n';
    if (Highlight)
      errs() << "Highlighted line " << LineNumber << ", " << Col << " -> ?\n";
  }
}

void SourceCoverageViewText::renderLineCoverageColumn(
    raw_ostream &OS, const LineCoverageStats &Line) {
  if (!Line.isMapped()) {
    OS.indent(LineCoverageColumnWidth) << '|';
    return;
  }
  std::string C = formatCount(Line.ExecutionCount);
  OS.indent(LineCoverageColumnWidth - C.size());
  colored_ostream(OS, raw_ostream::MAGENTA,
                  Line.hasMultipleRegions() && getOptions().Colors)
      << C;
  OS << '|';
}

void SourceCoverageViewText::renderLineNumberColumn(raw_ostream &OS,
                                                    unsigned LineNo) {
  SmallString<32> Buffer;
  raw_svector_ostream BufferOS(Buffer);
  BufferOS << LineNo;
  auto Str = BufferOS.str();
  // Trim and align to the right.
  Str = Str.substr(0, std::min(Str.size(), (size_t)LineNumberColumnWidth));
  OS.indent(LineNumberColumnWidth - Str.size()) << Str << '|';
}

void SourceCoverageViewText::renderRegionMarkers(
    raw_ostream &OS, CoverageSegmentArray Segments, unsigned ViewDepth) {
  renderLinePrefix(OS, ViewDepth);
  OS.indent(getCombinedColumnWidth(getOptions()));

  unsigned PrevColumn = 1;
  for (const auto *S : Segments) {
    if (!S->IsRegionEntry)
      continue;
    // Skip to the new region.
    if (S->Col > PrevColumn)
      OS.indent(S->Col - PrevColumn);
    PrevColumn = S->Col + 1;
    std::string C = formatCount(S->Count);
    PrevColumn += C.size();
    OS << '^' << C;
  }
  OS << '\n';

  if (getOptions().Debug)
    for (const auto *S : Segments)
      errs() << "Marker at " << S->Line << ":" << S->Col << " = "
             << formatCount(S->Count) << (S->IsRegionEntry ? "\n" : " (pop)\n");
}

unsigned SourceCoverageViewText::renderExpansionView(
    raw_ostream &OS, ExpansionView &ESV, Optional<LineRef> FirstLine,
    unsigned ExpansionCol, const coverage::CoverageSegment *WrappedSegment,
    CoverageSegmentArray LineSegments, unsigned ViewDepth) {
  unsigned NextExpansionCol = ExpansionCol;

  if (FirstLine.hasValue()) {
    // Re-render the current line and highlight the expansion range for
    // this subview.
    NextExpansionCol = ESV.getStartCol();
    renderLinePrefix(OS, ViewDepth);
    OS.indent(getCombinedColumnWidth(getOptions()) + (ViewDepth == 0 ? 0 : 1));
    renderLine(OS, *FirstLine, WrappedSegment, LineSegments, ExpansionCol,
               ViewDepth);
    renderViewDivider(OS, ViewDepth + 1);
  }

  // Render the child subview.
  if (getOptions().Debug)
    errs() << "Expansion at line " << ESV.getLine() << ", " << ESV.getStartCol()
           << " -> " << ESV.getEndCol() << '\n';
  ESV.View->print(OS, /*WholeFile=*/false, /*ShowSourceName=*/false,
                  ViewDepth + 1);

  return NextExpansionCol;
}

void SourceCoverageViewText::renderInstantiationView(raw_ostream &OS,
                                                     InstantiationView &ISV,
                                                     unsigned ViewDepth) {
  renderViewDivider(OS, ViewDepth);
  renderLinePrefix(OS, ViewDepth);
  OS << ' ';
  ISV.View->print(OS, /*WholeFile=*/false, /*ShowSourceName=*/true, ViewDepth);
}
<|endoftext|>


<commit_before>#include "controllers.h"
#include "triggers.h"

#include "ast/all.h"

using namespace Mugen;

Content getController(const std::string & type){
    if (match(type, "AfterImage")){
    } else if (match(type, "AfterImageTime")){
    } else if (match(type, "AllPalFX")){
    } else if (match(type, "AngleAdd")){
    } else if (match(type, "AngleDraw")){
    } else if (match(type, "AngleMul")){
    } else if (match(type, "AngleSet")){
    } else if (match(type, "AppendToClipboard")){
    } else if (match(type, "AssertSpecial")){
    } else if (match(type, "AttackDist")){
    } else if (match(type, "AttackMulSet")){
    } else if (match(type, "BGPalFX")){
    } else if (match(type, "BindToParent")){
    } else if (match(type, "BindToRoot")){
    } else if (match(type, "BindToTarget")){
    } else if (match(type, "ChangeAnim")){
    } else if (match(type, "ChangeAnim2")){
    } else if (match(type, "ChangeState")){
    } else if (match(type, "ClearClipboard")){
    } else if (match(type, "CtrlSet")){
    } else if (match(type, "DefenceMulSet")){
    } else if (match(type, "DestroySelf")){
    } else if (match(type, "DisplayToClipboard")){
    } else if (match(type, "EnvColor")){
    } else if (match(type, "EnvShake")){
    } else if (match(type, "Explod")){
    } else if (match(type, "ExplodBindTime")){
    } else if (match(type, "FallEnvShake")){
    } else if (match(type, "ForceFeedback")){
    } else if (match(type, "GameMakeAnim")){
    } else if (match(type, "Gravity")){
    } else if (match(type, "Helper")){
    } else if (match(type, "HitAdd")){
    } else if (match(type, "HitBy")){
    } else if (match(type, "HitDef")){
    } else if (match(type, "HitFallDamage")){
    } else if (match(type, "HitFallSet")){
    } else if (match(type, "HitFallVel")){
    } else if (match(type, "HitOverride")){
    } else if (match(type, "HitVelSet")){
    } else if (match(type, "LifeAdd")){
    } else if (match(type, "LifeSet")){
    } else if (match(type, "MakeDust")){
    } else if (match(type, "ModifyExplod")){
    } else if (match(type, "MoveHitReset")){
    } else if (match(type, "NotHitBy")){
    } else if (match(type, "Null")){
        Content null(1,"# Disabled or invalid controller.");
        null.addLine(1, "pass");
        return null;
    } else if (match(type, "Offset")){
    } else if (match(type, "PalFX")){
    } else if (match(type, "ParentVarAdd")){
    } else if (match(type, "ParentVarSet")){
    } else if (match(type, "Pause")){
    } else if (match(type, "PlaySnd")){
    } else if (match(type, "PlayerPush")){
    } else if (match(type, "PosAdd")){
    } else if (match(type, "PosFreeze")){
    } else if (match(type, "PosSet")){
    } else if (match(type, "PowerAdd")){
    } else if (match(type, "PowerSet")){
    } else if (match(type, "Projectile")){
    } else if (match(type, "RemapPal")){
    } else if (match(type, "RemoveExplod")){
    } else if (match(type, "ReversalDef")){
    } else if (match(type, "ScreenBound")){
    } else if (match(type, "SelfState")){
    } else if (match(type, "SndPan")){
    } else if (match(type, "SprPriority")){
    } else if (match(type, "StateTypeSet")){
    } else if (match(type, "StopSnd")){
    } else if (match(type, "SuperPause")){
    } else if (match(type, "TargetBind")){
    } else if (match(type, "TargetDrop")){
    } else if (match(type, "TargetFacing")){
    } else if (match(type, "TargetLifeAdd")){
    } else if (match(type, "TargetPowerAdd")){
    } else if (match(type, "TargetState")){
    } else if (match(type, "TargetVelAdd")){
    } else if (match(type, "TargetVelSet")){
    } else if (match(type, "Trans")){
    } else if (match(type, "Turn")){
    } else if (match(type, "VarAdd")){
    } else if (match(type, "VarRandom")){
    } else if (match(type, "VarRangeSet")){
    } else if (match(type, "VarSet")){
    } else if (match(type, "VelAdd")){
    } else if (match(type, "VelMul")){
    } else if (match(type, "VelSet")){
    } else if (match(type, "VictoryQuote")){
    } else if (match(type, "Width")){
    } else {
        std::cout << "Unhandled Controller Type: " << type << std::endl;
    }

    // Set to pass if state controller can't be determined
    return Content(1, "pass");
}

StateParameterMap::StateParameterMap(){
}

StateParameterMap::StateParameterMap(const StateParameterMap & copy){
    this->parameters = copy.parameters;
}

StateParameterMap::~StateParameterMap(){
}

const StateParameterMap & StateParameterMap::operator=(const StateParameterMap & copy){
    this->parameters = copy.parameters;
}

void StateParameterMap::addTriggerall(Ast::AttributeSimple * simple){
    triggerall.push_back(simple);
}
void StateParameterMap::addTrigger(int number, Ast::AttributeSimple * simple){
    triggers[number].push_back(simple);
}
void StateParameterMap::addVar(Ast::AttributeSimple *simple){
    vars.push_back(simple);
}

void StateParameterMap::add(const std::string & param, Ast::AttributeSimple * simple){
    parameters[param].push_back(simple);
}

std::vector<Ast::AttributeSimple *> * StateParameterMap::find(const std::string & param){
    std::map<std::string, std::vector<Ast::AttributeSimple *> >::iterator found = parameters.find(param);
    if (found != parameters.end()){
        return &(*found).second;
    }
    return NULL;
}

void StateParameterMap::addToDefinition(PythonDefinition & definition){
    for (std::vector<Ast::AttributeSimple *>::iterator i = triggerall.begin(); i != triggerall.end(); ++i){
        Ast::AttributeSimple * simple = (*i);
        ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
        Trigger trigger = Trigger(builder);
        const std::vector<Content> & functions = trigger.getFunctions();
        for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
            definition.addContent(*j);
        }
    }

    for (int i = 0; ; ++i){
        std::map<int, std::vector<Ast::AttributeSimple *> >::iterator trigger = triggers.find(i);
        if (trigger != triggers.end()){
            for (std::vector<Ast::AttributeSimple *>::iterator i = trigger->second.begin(); i != trigger->second.end(); ++i){
                Ast::AttributeSimple * simple = (*i);
                ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
                Trigger trigger = Trigger(builder);
                const std::vector<Content> & functions = trigger.getFunctions();
                for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
                    definition.addContent(*j);
                }
            }
        } else {
            break;
        }
    }

    for (std::vector<Ast::AttributeSimple *>::iterator i = vars.begin(); i != vars.end(); ++i){
        Ast::AttributeSimple * simple = (*i);
        ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
        Trigger trigger = Trigger(builder);
        const std::vector<Content> & functions = trigger.getFunctions();
        for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
            definition.addContent(*j);
        }
    }
}

StateControllerStore::StateControllerStore(){
}
StateControllerStore::~StateControllerStore(){
}

void StateControllerStore::newController(){
    controllerParameters.push_back(StateParameterMap());
}

StateParameterMap & StateControllerStore::getCurrentController(){
    return controllerParameters.back();
}

void StateControllerStore::addToDefinition(PythonDefinition & definition){

    // TODO Go through current parameters and create new classes based on 'type' and add to definition
    for (std::vector<StateParameterMap>::iterator i = controllerParameters.begin(); i != controllerParameters.end(); ++i){
        std::vector<Ast::AttributeSimple *> * type = (*i).find("type");
        if (type != NULL){
            std::stringstream line;
            const std::string & ctrl = type->back()->valueAsString();
            line << "# Found controller type '" << ctrl << "' on line " << type->back()->getLine();
            definition.addContent(Content(1, line.str()));
            definition.addContent(getController(ctrl));
        }
        (*i).addToDefinition(definition);
    }
}

StateController::StateController(const StateParameterMap & parameters):
parameters(parameters){
}

StateController::~StateController(){
}
<commit_msg>Evaluate triggers in controller.<commit_after>#include "controllers.h"
#include "triggers.h"

#include "ast/all.h"

#include <sstream>

using namespace Mugen;

Content getController(const std::string & type){
    if (match(type, "AfterImage")){
    } else if (match(type, "AfterImageTime")){
    } else if (match(type, "AllPalFX")){
    } else if (match(type, "AngleAdd")){
    } else if (match(type, "AngleDraw")){
    } else if (match(type, "AngleMul")){
    } else if (match(type, "AngleSet")){
    } else if (match(type, "AppendToClipboard")){
    } else if (match(type, "AssertSpecial")){
    } else if (match(type, "AttackDist")){
    } else if (match(type, "AttackMulSet")){
    } else if (match(type, "BGPalFX")){
    } else if (match(type, "BindToParent")){
    } else if (match(type, "BindToRoot")){
    } else if (match(type, "BindToTarget")){
    } else if (match(type, "ChangeAnim")){
    } else if (match(type, "ChangeAnim2")){
    } else if (match(type, "ChangeState")){
    } else if (match(type, "ClearClipboard")){
    } else if (match(type, "CtrlSet")){
    } else if (match(type, "DefenceMulSet")){
    } else if (match(type, "DestroySelf")){
    } else if (match(type, "DisplayToClipboard")){
    } else if (match(type, "EnvColor")){
    } else if (match(type, "EnvShake")){
    } else if (match(type, "Explod")){
    } else if (match(type, "ExplodBindTime")){
    } else if (match(type, "FallEnvShake")){
    } else if (match(type, "ForceFeedback")){
    } else if (match(type, "GameMakeAnim")){
    } else if (match(type, "Gravity")){
    } else if (match(type, "Helper")){
    } else if (match(type, "HitAdd")){
    } else if (match(type, "HitBy")){
    } else if (match(type, "HitDef")){
    } else if (match(type, "HitFallDamage")){
    } else if (match(type, "HitFallSet")){
    } else if (match(type, "HitFallVel")){
    } else if (match(type, "HitOverride")){
    } else if (match(type, "HitVelSet")){
    } else if (match(type, "LifeAdd")){
    } else if (match(type, "LifeSet")){
    } else if (match(type, "MakeDust")){
    } else if (match(type, "ModifyExplod")){
    } else if (match(type, "MoveHitReset")){
    } else if (match(type, "NotHitBy")){
    } else if (match(type, "Null")){
        Content null(1,"# Disabled or invalid controller.");
        null.addLine(1, "pass");
        return null;
    } else if (match(type, "Offset")){
    } else if (match(type, "PalFX")){
    } else if (match(type, "ParentVarAdd")){
    } else if (match(type, "ParentVarSet")){
    } else if (match(type, "Pause")){
    } else if (match(type, "PlaySnd")){
    } else if (match(type, "PlayerPush")){
    } else if (match(type, "PosAdd")){
    } else if (match(type, "PosFreeze")){
    } else if (match(type, "PosSet")){
    } else if (match(type, "PowerAdd")){
    } else if (match(type, "PowerSet")){
    } else if (match(type, "Projectile")){
    } else if (match(type, "RemapPal")){
    } else if (match(type, "RemoveExplod")){
    } else if (match(type, "ReversalDef")){
    } else if (match(type, "ScreenBound")){
    } else if (match(type, "SelfState")){
    } else if (match(type, "SndPan")){
    } else if (match(type, "SprPriority")){
    } else if (match(type, "StateTypeSet")){
    } else if (match(type, "StopSnd")){
    } else if (match(type, "SuperPause")){
    } else if (match(type, "TargetBind")){
    } else if (match(type, "TargetDrop")){
    } else if (match(type, "TargetFacing")){
    } else if (match(type, "TargetLifeAdd")){
    } else if (match(type, "TargetPowerAdd")){
    } else if (match(type, "TargetState")){
    } else if (match(type, "TargetVelAdd")){
    } else if (match(type, "TargetVelSet")){
    } else if (match(type, "Trans")){
    } else if (match(type, "Turn")){
    } else if (match(type, "VarAdd")){
    } else if (match(type, "VarRandom")){
    } else if (match(type, "VarRangeSet")){
    } else if (match(type, "VarSet")){
    } else if (match(type, "VelAdd")){
    } else if (match(type, "VelMul")){
    } else if (match(type, "VelSet")){
    } else if (match(type, "VictoryQuote")){
    } else if (match(type, "Width")){
    } else {
        std::cout << "Unhandled Controller Type: " << type << std::endl;
    }

    // Set to pass if state controller can't be determined
    return Content(1, "pass");
}

StateParameterMap::StateParameterMap(){
}

StateParameterMap::StateParameterMap(const StateParameterMap & copy){
    this->parameters = copy.parameters;
}

StateParameterMap::~StateParameterMap(){
}

const StateParameterMap & StateParameterMap::operator=(const StateParameterMap & copy){
    this->parameters = copy.parameters;
}

void StateParameterMap::addTriggerall(Ast::AttributeSimple * simple){
    triggerall.push_back(simple);
}
void StateParameterMap::addTrigger(int number, Ast::AttributeSimple * simple){
    triggers[number].push_back(simple);
}
void StateParameterMap::addVar(Ast::AttributeSimple *simple){
    vars.push_back(simple);
}

void StateParameterMap::add(const std::string & param, Ast::AttributeSimple * simple){
    parameters[param].push_back(simple);
}

std::vector<Ast::AttributeSimple *> * StateParameterMap::find(const std::string & param){
    std::map<std::string, std::vector<Ast::AttributeSimple *> >::iterator found = parameters.find(param);
    if (found != parameters.end()){
        return &(*found).second;
    }
    return NULL;
}

void StateParameterMap::addToDefinition(PythonDefinition & definition){
    Content triggerallContent(1, "# Triggerall");
    triggerallContent.addLine(1, "def triggerall(self):");
    std::string tempString;
    for (std::vector<Ast::AttributeSimple *>::iterator i = triggerall.begin(); i != triggerall.end(); ++i){
        Ast::AttributeSimple * simple = (*i);
        ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
        Trigger trigger = Trigger(builder);
        const std::vector<Content> & functions = trigger.getFunctions();
        for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
            definition.addContent(*j);
        }
        tempString += trigger.getName() + " and ";
    }
    triggerallContent.addLine(2, (tempString.empty() ? "return 1" : "return (" + tempString.substr(0, tempString.size()-5) + ")"));
    definition.addContent(triggerallContent);

    Content allTriggers(1, "# Check all triggers");
    allTriggers.addLine(1, "def checkAll(self):");
    tempString.clear();
    for (int i = 0; ; ++i){
        std::map<int, std::vector<Ast::AttributeSimple *> >::iterator trigger = triggers.find(i);
        if (trigger != triggers.end()){
            std::ostringstream number;
            number << i;
            Content triggerContent(1, "# Trigger " + number.str());
            triggerContent.addLine(1, "def trigger" + number.str() +"(self):");
            std::string triggerString;
            for (std::vector<Ast::AttributeSimple *>::iterator i = trigger->second.begin(); i != trigger->second.end(); ++i){
                Ast::AttributeSimple * simple = (*i);
                ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
                Trigger trigger = Trigger(builder);
                const std::vector<Content> & functions = trigger.getFunctions();
                for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
                    definition.addContent(*j);
                }
                triggerString += "(triggerall() and " + trigger.getName() + ")" + " and ";
            }
            triggerContent.addLine(2, (triggerString.empty() ? "return 1" : "return (" + triggerString.substr(0, triggerString.size()-5) + ")"));
            definition.addContent(triggerContent);

            tempString += "trigger" + number.str() + "() or ";
        } else {
            // No triggers were added
            break;
        }
    }
    // checkAll
    allTriggers.addLine(2, (tempString.empty() ? "return 1" : "return (" + tempString.substr(0, tempString.size()-3) + ")"));
    definition.addContent(allTriggers);

    Content allContent(1, "# Check all triggers");
    allContent.addLine(1, "if checkAll() != 1:");
    allContent.addLine(2, "return");
    definition.addContent(allContent);

    /* FIXME What to do here? */
    for (std::vector<Ast::AttributeSimple *>::iterator i = vars.begin(); i != vars.end(); ++i){
        Ast::AttributeSimple * simple = (*i);
        ExpressionBuilder builder = TriggerHandler::convert(*simple->getValue());
        Trigger trigger = Trigger(builder);
        const std::vector<Content> & functions = trigger.getFunctions();
        for (std::vector<Content>::const_iterator j = functions.begin(); j != functions.end(); ++j){
            definition.addContent(*j);
        }
    }
}

StateControllerStore::StateControllerStore(){
}
StateControllerStore::~StateControllerStore(){
}

void StateControllerStore::newController(){
    controllerParameters.push_back(StateParameterMap());
}

StateParameterMap & StateControllerStore::getCurrentController(){
    return controllerParameters.back();
}

void StateControllerStore::addToDefinition(PythonDefinition & definition){

    // TODO Go through current parameters and create new classes based on 'type' and add to definition
    for (std::vector<StateParameterMap>::iterator i = controllerParameters.begin(); i != controllerParameters.end(); ++i){
        std::vector<Ast::AttributeSimple *> * type = (*i).find("type");
        if (type != NULL){
            std::stringstream line;
            const std::string & ctrl = type->back()->valueAsString();
            line << "# Found controller type '" << ctrl << "' on line " << type->back()->getLine();
            definition.addContent(Content(1, line.str()));
            definition.addContent(getController(ctrl));
        }
        (*i).addToDefinition(definition);
    }
}

StateController::StateController(const StateParameterMap & parameters):
parameters(parameters){
}

StateController::~StateController(){
}
<|endoftext|>