<commit_before>#include <taichi/visual/gui.h>
#include "../tlang.h"

TLANG_NAMESPACE_BEGIN

// https://github.com/yuanming-hu/topo_opt_private/blob/master/misc/volume_rendering_kernel.cu

auto voxel_renderer = [](const std::vector<std::string> &params) {
  CoreState::set_trigger_gdb_when_crash(true);

  int n = 512;
  int grid_resolution = 256;

  if (params.size() != 2) {
    TC_INFO("Usage: ti voxel renderer filename.bin resolution");
    exit(-1);
  }

  auto f = fopen(params[0].c_str(), "rb");
  TC_ERROR_UNLESS(f, "File {} not found", params[0]);
  std::vector<float32> density_field(pow<3>(grid_resolution));
  std::fread(density_field.data(), sizeof(float32), density_field.size(), f);
  std::fclose(f);

  Program prog(Arch::gpu);
  prog.config.print_ir = true;

  Vector buffer(DataType::f32, 3);
  Global(density, f32);

  layout([&]() {
    root.dense(Index(0), n * n * 2).place(buffer(0), buffer(1), buffer(2));
    root.dense(Indices(0, 1, 2), grid_resolution).place(density);
  });

  auto query_density_int = [&](Vector p_) {
    auto p = p_.cast_elements<int32>();
    auto inside_box =
        Var(0 <= p(0) && p(0) < grid_resolution && 0 <= p(1) &&
            p(1) < grid_resolution && 0 <= p(2) && p(2) < grid_resolution);
    auto ret = Var(0.0f);
    If(inside_box).Then([&] { ret = density[p]; });
    return ret;
  };

  auto get_next_hit = [&](const Vector &eye_o, const Vector &eye_d,
                          Expr &hit_distance, Vector &hit_pos, Vector &normal) {
    auto d = normalized(eye_d);
    auto pos = Var(eye_o + d * 1e-4f);

    auto rinv = Var(1.0f / d);
    auto rsign = Vector(3);
    for (int i = 0; i < 3; i++) {
      rsign(i) = cast<float32>(d(i) > 0.0f) * 2.0f - 1.0f;  // sign...
    }

    auto o = Var(pos * float32(grid_resolution));
    auto ipos = Var(floor(o));
    auto dis = Var((ipos - o + 0.5f + rsign * 0.5f).element_wise_prod(rinv));

    auto running = Var(1);
    auto i = Var(0);
    hit_distance = -1.0f;
    While(running, [&] {
      auto last_sample = Var(query_density_int(ipos));
      If(last_sample > 0.0f)
          .Then([&] {
            // intersect the cube
            auto mini =
                Var((ipos - o + Vector({0.5f, 0.5f, 0.5f}) - rsign * 0.5f)
                        .element_wise_prod(rinv));
            hit_distance =
                max(max(mini(0), mini(1)), mini(2)) * (1.0f / grid_resolution);
            hit_pos = pos + hit_distance * d;
            running = 0;
          })
          .Else([&] {
            auto mm = Var(Vector({0.0f, 0.0f, 0.0f}));
            If(dis(0) <= dis(1) && dis(0) < dis(2))
                .Then([&] { mm(0) = 1.0f; })
                .Else([&] {
                  If(dis(1) <= dis(0) && dis(1) <= dis(2))
                      .Then([&] { mm(1) = 1.0f; })
                      .Else([&] { mm(2) = 1.0f; });
                });
            dis += mm.element_wise_prod(rsign).element_wise_prod(rinv);
            ipos += mm.element_wise_prod(rsign);
            normal = -mm.element_wise_prod(rsign);
          });
      i += 1;
      If(i > 500).Then([&] { running = 0; });
    });
  };

  float32 fov = 0.7;

  auto background = [](Vector dir) {
    /*
    auto dot = Var(dir.dot(Vector({0.6f, 0.75f, 0.15f})));
    auto coeff1 = Var(clamp(dot * 0.5f + 0.5f, 0.0f, 1.0f));
    auto light = Var(coeff1 * Vector({0.9f, 0.7f, 0.3f}) +
                     (1.0f - coeff1) * Vector({0.4f, 0.5f, 0.9f}));
    return light;
    */
    return Vector({0.7f, 0.7f, 0.8f});
  };

  auto out_dir = [&](Vector n) {
    auto u = Var(Vector({1.0f, 0.0f, 0.0f})), v = Var(Vector(3));
    If(abs(n(1)) < 1 - 1e-3f, [&] {
      u = normalized(cross(n, Vector({0.0f, 1.0f, 0.0f})));
    });
    v = cross(n, u);
    auto phi = Var(2 * pi * Rand<float32>());
    auto r = Var(Rand<float32>());
    auto alpha = Var(0.5f * pi * (r * r));
    return sin(alpha) * (cos(phi) * u + sin(phi) * v) + cos(alpha) * n;
  };

  Kernel(main).def([&]() {
    For(0, n * n * 2, [&](Expr i) {
      auto orig = Var(Vector({0.5f, 0.2f, 1.0f}));

      auto c = Var(Vector(
          {fov * ((Rand<float32>() + cast<float32>(i / n)) / float32(n / 2) -
                  2.01f),
           fov * ((Rand<float32>() + cast<float32>(i % n)) / float32(n / 2) -
                  1.01f),
           -1.0f}));

      c = normalized(c);

      int depth_limit = 4;
      auto depth = Var(0);

      auto color = Var(Vector({1.0f, 1.0f, 1.0f}));

      While(depth < depth_limit, [&] {
        auto hit_dist = Var(0.0f);
        auto hit_pos = Var(Vector({1.0f, 1.0f, 1.0f}));
        auto normal = Var(Vector({1.0f, 1.0f, 1.0f}));
        get_next_hit(orig, c, hit_dist, hit_pos, normal);

        depth += 1;
        If(hit_dist > 0.0f)
            .Then([&] {
              c = normalized(out_dir(normal));
              orig = hit_pos;
              color = color.element_wise_prod(Vector({0.7f, 0.4f, 0.4f}));
            })
            .Else([&] {
              color = color.element_wise_prod(background(c));
              depth = depth_limit;
            });
      });

      buffer[i] += color;
    });
  });

  for (int i = 0; i < grid_resolution; i++) {
    for (int j = 0; j < grid_resolution; j++) {
      for (int k = 0; k < grid_resolution; k++) {
        density.val<float32>(i, j, k) =
            density_field[i * grid_resolution * grid_resolution +
                          j * grid_resolution + k];
      }
    }
  }

  GUI gui("Voxel Renderer", Vector2i(n * 2, n));

  auto tone_map = [](real x) { return std::sqrt(x); };
  constexpr int N = 20;
  for (int frame = 0;; frame++) {
    for (int i = 0; i < N; i++)
      main();
    real scale = 1.0f / ((frame + 1) * N);
    for (int i = 0; i < n * n * 2; i++) {
      gui.buffer[i / n][i % n] =
          Vector4(tone_map(scale * buffer(0).val<float32>(i)),
                  tone_map(scale * buffer(1).val<float32>(i)),
                  tone_map(scale * buffer(2).val<float32>(i)), 1);
    }
    gui.update();
  }
};
TC_REGISTER_TASK(voxel_renderer);

TLANG_NAMESPACE_END
<commit_msg>direct lighting for voxel renderer<commit_after>#include <taichi/visual/gui.h>
#include "../tlang.h"

TLANG_NAMESPACE_BEGIN

// https://github.com/yuanming-hu/topo_opt_private/blob/master/misc/volume_rendering_kernel.cu

auto voxel_renderer = [](const std::vector<std::string> &params) {
  CoreState::set_trigger_gdb_when_crash(true);

  int n = 512;
  int grid_resolution = 256;

  if (params.size() != 2) {
    TC_INFO("Usage: ti voxel renderer filename.bin resolution");
    exit(-1);
  }

  auto f = fopen(params[0].c_str(), "rb");
  TC_ERROR_UNLESS(f, "File {} not found", params[0]);
  std::vector<float32> density_field(pow<3>(grid_resolution));
  std::fread(density_field.data(), sizeof(float32), density_field.size(), f);
  std::fclose(f);

  Program prog(Arch::gpu);
  prog.config.print_ir = true;

  Vector buffer(DataType::f32, 3);
  Global(density, f32);

  layout([&]() {
    root.dense(Index(0), n * n * 2).place(buffer(0), buffer(1), buffer(2));
    root.dense(Indices(0, 1, 2), grid_resolution).place(density);
  });

  auto query_density_int = [&](Vector p_) {
    auto p = p_.cast_elements<int32>();
    auto inside_box =
        Var(0 <= p(0) && p(0) < grid_resolution && 0 <= p(1) &&
            p(1) < grid_resolution && 0 <= p(2) && p(2) < grid_resolution);
    auto ret = Var(0.0f);
    If(inside_box).Then([&] { ret = density[p]; });
    return ret;
  };

  auto get_next_hit = [&](const Vector &eye_o, const Vector &eye_d,
                          Expr &hit_distance, Vector &hit_pos, Vector &normal) {
    auto d = normalized(eye_d);
    auto pos = Var(eye_o + d * 1e-4f);

    auto rinv = Var(1.0f / d);
    auto rsign = Vector(3);
    for (int i = 0; i < 3; i++) {
      rsign(i) = cast<float32>(d(i) > 0.0f) * 2.0f - 1.0f;  // sign...
    }

    auto o = Var(pos * float32(grid_resolution));
    auto ipos = Var(floor(o));
    auto dis = Var((ipos - o + 0.5f + rsign * 0.5f).element_wise_prod(rinv));

    auto running = Var(1);
    auto i = Var(0);
    hit_distance = -1.0f;
    While(running, [&] {
      auto last_sample = Var(query_density_int(ipos));
      If(last_sample > 0.0f)
          .Then([&] {
            // intersect the cube
            auto mini =
                Var((ipos - o + Vector({0.5f, 0.5f, 0.5f}) - rsign * 0.5f)
                        .element_wise_prod(rinv));
            hit_distance =
                max(max(mini(0), mini(1)), mini(2)) * (1.0f / grid_resolution);
            hit_pos = pos + hit_distance * d;
            running = 0;
          })
          .Else([&] {
            auto mm = Var(Vector({0.0f, 0.0f, 0.0f}));
            If(dis(0) <= dis(1) && dis(0) < dis(2))
                .Then([&] { mm(0) = 1.0f; })
                .Else([&] {
                  If(dis(1) <= dis(0) && dis(1) <= dis(2))
                      .Then([&] { mm(1) = 1.0f; })
                      .Else([&] { mm(2) = 1.0f; });
                });
            dis += mm.element_wise_prod(rsign).element_wise_prod(rinv);
            ipos += mm.element_wise_prod(rsign);
            normal = -mm.element_wise_prod(rsign);
          });
      i += 1;
      If(i > 500).Then([&] { running = 0; });
    });
  };

  float32 fov = 0.6;

  auto background = [](Vector dir) {
    /*
    auto dot = Var(dir.dot(Vector({0.6f, 0.75f, 0.15f})));
    auto coeff1 = Var(clamp(dot * 0.5f + 0.5f, 0.0f, 1.0f));
    auto light = Var(coeff1 * Vector({0.9f, 0.7f, 0.3f}) +
                     (1.0f - coeff1) * Vector({0.4f, 0.5f, 0.9f}));
    return light;
    */
    return Vector({0.7f, 0.7f, 0.8f});
  };

  auto out_dir = [&](Vector n) {
    auto u = Var(Vector({1.0f, 0.0f, 0.0f})), v = Var(Vector(3));
    If(abs(n(1)) < 1 - 1e-3f, [&] {
      u = normalized(cross(n, Vector({0.0f, 1.0f, 0.0f})));
    });
    v = cross(n, u);
    auto phi = Var(2 * pi * Rand<float32>());
    auto r = Var(Rand<float32>());
    auto alpha = Var(0.5f * pi * (r * r));
    return sin(alpha) * (cos(phi) * u + sin(phi) * v) + cos(alpha) * n;
  };

  Kernel(main).def([&]() {
    For(0, n * n * 2, [&](Expr i) {
      auto orig = Var(Vector({0.5f, 0.2f, 1.0f}));

      auto c = Var(Vector(
          {fov * ((Rand<float32>() + cast<float32>(i / n)) / float32(n / 2) -
                  2.01f),
           fov * ((Rand<float32>() + cast<float32>(i % n)) / float32(n / 2) -
                  1.01f),
           -1.0f}));

      c = normalized(c);

      int depth_limit = 4;
      auto depth = Var(0);

      auto importon = Var(Vector({1.0f, 1.0f, 1.2f}));

      While(depth < depth_limit, [&] {
        auto hit_dist = Var(0.0f);
        auto hit_pos = Var(Vector({1.0f, 1.0f, 1.0f}));
        auto normal = Var(Vector({1.0f, 1.0f, 1.0f}));
        get_next_hit(orig, c, hit_dist, hit_pos, normal);

        depth += 1;
        If(hit_dist > 0.0f)
            .Then([&] {
              c = normalized(out_dir(normal));
              orig = hit_pos;
              importon = importon.element_wise_prod(Vector({0.3f, 0.3f, 0.4f}));

              If(depth == 1).Then([&] {
                // direct lighting on camera ray...
                get_next_hit(orig, Vector({0.5f, 0.3f, -0.1f}), hit_dist,
                             hit_pos, normal);
                If(hit_dist < 0.0f).Then([&] {
                  buffer[i] +=
                      importon.element_wise_prod(Vector({0.3f, 0.3f, 0.3f}));
                });
              });
            })
            .Else([&] {
              buffer[i] += importon.element_wise_prod(background(c));
              depth = depth_limit;
            });
      });
    });
  });

  for (int i = 0; i < grid_resolution; i++) {
    for (int j = 0; j < grid_resolution; j++) {
      for (int k = 0; k < grid_resolution; k++) {
        density.val<float32>(i, j, k) =
            density_field[i * grid_resolution * grid_resolution +
                          j * grid_resolution + k];
      }
    }
  }

  GUI gui("Voxel Renderer", Vector2i(n * 2, n));

  auto tone_map = [](real x) { return std::sqrt(x * 2); };
  constexpr int N = 20;
  for (int frame = 0;; frame++) {
    for (int i = 0; i < N; i++)
      main();
    real scale = 1.0f / ((frame + 1) * N);
    for (int i = 0; i < n * n * 2; i++) {
      gui.buffer[i / n][i % n] =
          Vector4(tone_map(scale * buffer(0).val<float32>(i)),
                  tone_map(scale * buffer(1).val<float32>(i)),
                  tone_map(scale * buffer(2).val<float32>(i)), 1);
    }
    gui.update();
  }
};
TC_REGISTER_TASK(voxel_renderer);

TLANG_NAMESPACE_END
<|endoftext|>


<commit_before>#include "globals.h"
#include "Thread.h"
#include "filesystem/VFS.h"
#include "filesystem/RawFD.h"
#include <stdarg.h>
#include <pthread.h>
#include <semaphore.h>

typeof(Options) Options;

static const char* const environment[] =
{
	NULL
};

void log(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu): ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
}

void Warning(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu) warning: ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
}

void error(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu) error: ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
	_exit(-1);
}

void Error(const string& message)
{
	writef(2, "lbw error: %s\n", message.c_str());
	_exit(-1);
}

#if 0
static int init_fifo(void)
{
	const char* homedir = getenv("HOME");
	char buffer[PATH_MAX];

	snprintf(buffer, sizeof(buffer),
			"%s/.lbw", homedir);
	mkdir(buffer, S_IRWXU); /* ignore errors */

	snprintf(buffer, sizeof(buffer),
			"%s/.lbw/%d", homedir, slavepid);
	int e = mknod(buffer, S_IFIFO | S_IRWXU, 0);
	if (e)
		return e;

	slavefd = open(buffer, O_RDWR);
	if (slavefd == -1)
		return -1;
	return 0;
}
#endif

class ArgumentParser
{
public:
	ArgumentParser():
		Chroot("/"),
		FakeRoot(false)
	{
		char buffer[PATH_MAX];
		getcwd(buffer, sizeof(buffer));
		CWD = buffer;
	}

	void Parse(int& argc, const char**& argv)
	{
		argc--;
		argv++;

		while (*argv)
		{
			string option = *argv;
			if (option.empty())
				BadOption();

			if (option[0] != '-')
				break;

			string argument;
			if (argv[1])
				argument = argv[1];

			int delta = Option(option, argument);

			argc -= delta;
			argv += delta;
		}
	}

	void BadOption()
	{
		Error("lbw: invalid syntax. Try --help");
	}

	int Option(const string& option, const string& argument)
	{
		if (option == "--help")
			Error("lbw: no help yet. Sorry.");
		else if (option == "--fakeroot")
		{
			FakeRoot = true;
			return 1;
		}
		else if (option == "--chroot")
		{
			Chroot = argument;
			CWD = "/";
			return 2;
		}
		else
			BadOption();
		return 1;
	}

public:
	string Chroot;
	string CWD;
	bool FakeRoot : 1;
};

static void add_std_fd(int fd)
{
	Ref<RawFD> fdo = new RawFD();
	fdo->Reference(); // never allow these objects to be destroyed
	fdo->Open(fd);
	assert(FD::New(fdo) == fd);
}

int main(int argc, const char* argv[], const char* environ[])
{
	{
		ArgumentParser* ap = new ArgumentParser();
		ap->Parse(argc, argv);

		Options.FakeRoot = ap->FakeRoot;
		VFS::SetRoot(ap->Chroot);
		VFS::SetCWD(ap->CWD);

		delete ap;
	}

	if (argc == 0)
		Error("lbw: you must specify a binary to run. Try --help.");

	add_std_fd(0);
	add_std_fd(1);
	add_std_fd(2);

	StartMonitor();
	Exec(argc, argv, environ);
	return 0;
}

<commit_msg>Added a proper CLI front end.<commit_after>#include "globals.h"
#include "Thread.h"
#include "filesystem/VFS.h"
#include "filesystem/RawFD.h"
#include <stdarg.h>
#include <pthread.h>
#include <semaphore.h>

typeof(Options) Options;

static const char* const environment[] =
{
	NULL
};

void log(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu): ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
}

void Warning(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu) warning: ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
}

void error(const char* format, ...)
{
	va_list ap;
	va_start(ap, format);

	writef(2, "lbw(%lu) error: ", getpid());
	vwritef(2, format, ap);
	writef(2, "\n");

	va_end(ap);
	_exit(-1);
}

void Error(const string& message)
{
	writef(2, "lbw error: %s\n", message.c_str());
	_exit(-1);
}

#if 0
static int init_fifo(void)
{
	const char* homedir = getenv("HOME");
	char buffer[PATH_MAX];

	snprintf(buffer, sizeof(buffer),
			"%s/.lbw", homedir);
	mkdir(buffer, S_IRWXU); /* ignore errors */

	snprintf(buffer, sizeof(buffer),
			"%s/.lbw/%d", homedir, slavepid);
	int e = mknod(buffer, S_IFIFO | S_IRWXU, 0);
	if (e)
		return e;

	slavefd = open(buffer, O_RDWR);
	if (slavefd == -1)
		return -1;
	return 0;
}
#endif

class ArgumentParser
{
public:
	ArgumentParser():
		Chroot("/"),
		FakeRoot(false)
	{
		char buffer[PATH_MAX];
		getcwd(buffer, sizeof(buffer));
		CWD = buffer;
	}

	void Parse(int& argc, const char**& argv)
	{
		argc--;
		argv++;

		while (*argv)
		{
			string option = *argv;
			if (option.empty())
				BadOption();

			if (option[0] != '-')
				break;

			string argument;
			if (argv[1])
				argument = argv[1];

			int delta = Option(option, argument);

			argc -= delta;
			argv += delta;
		}
	}

	void BadOption()
	{
		Error("lbw: invalid syntax. Try --help");
	}

	int Option(const string& option, const string& argument)
	{
		if (option == "--help")
			Error("lbw: no help yet. Sorry.");
		else if (option == "--fakeroot")
		{
			FakeRoot = true;
			return 1;
		}
		else if (option == "--chroot")
		{
			Chroot = argument;
			CWD = "/";
			return 2;
		}
		else
			BadOption();
		return 1;
	}

public:
	string Chroot;
	string CWD;
	bool FakeRoot : 1;
};

static void add_std_fd(int fd)
{
	Ref<RawFD> fdo = new RawFD();
	fdo->Reference(); // never allow these objects to be destroyed
	fdo->Open(fd);
	assert(FD::New(fdo) == fd); // ensure they have the right fd
}

int main(int argc, const char* argv[], const char* environ[])
{
	{
		ArgumentParser* ap = new ArgumentParser();
		ap->Parse(argc, argv);

		Options.FakeRoot = ap->FakeRoot;
		VFS::SetRoot(ap->Chroot);
		VFS::SetCWD(ap->CWD);

		delete ap;
	}

	if (argc == 0)
		Error("lbw: you must specify a binary to run. Try --help.");

	add_std_fd(0);
	add_std_fd(1);
	add_std_fd(2);

	StartMonitor();
	Exec(argv[0], argv, environ);
	return 0;
}

<|endoftext|>


<commit_before>#include <sys/resource.h>

#include "server/control.hpp"
#include "server/server.hpp"
#include "fsck/fsck.hpp"
#include "extract/extract.hpp"
#include "clustering/demo.hpp"
#include "utils.hpp"
#include "help.hpp"
#include "migrate/migrate.hpp"

void print_version_message() {
    printf("rethinkdb " RETHINKDB_VERSION
#ifndef NDEBUG
           " (debug)"
#endif
           "\n");
}

void usage() {
    Help_Pager *help = Help_Pager::instance();
    help->pagef("Usage: rethinkdb COMMAND ...\n"
                "Commands:\n"
                "    help        Display help about rethinkdb and rethinkdb commands.\n"
                "\n"
                "Creating and serving databases:\n"
                "    create      Create an empty database.\n"
                "    serve       Serve an existing database.\n"
                "\n"
                "Administrating databases:\n"
                "    extract     Extract as much data as possible from a corrupted database.\n"
                "    import      Import data from raw memcached commands.\n"
                "    fsck        Check a database for corruption.\n"
                "    migrate     Convert between file versions.\n"
                "\n"
                "Use 'rethinkdb help COMMAND' for help on a single command.\n"
                "Use 'rethinkdb --version' for the current version of rethinkdb.\n");
}

int dispatch_on_args(std::vector<char *> args);

int main(int argc, char *argv[]) {
    initialize_precise_time();
    install_generic_crash_handler();

#ifndef NDEBUG
    rlimit core_limit;
    core_limit.rlim_cur = RLIM_INFINITY;
    core_limit.rlim_max = RLIM_INFINITY;
    setrlimit(RLIMIT_CORE, &core_limit);
#endif

    /* Put arguments into a vector so we can modify them if convenient */
    std::vector<char *> args(argv, argv + argc);

    return dispatch_on_args(args);
}

int dispatch_on_args(std::vector<char *> args) {
    /* Default to "rethinkdb serve" */
    if (args.size() == 1) args.push_back(const_cast<char *>("serve"));

    /* Switch based on subcommand, then dispatch to the appropriate function */
    if (!strcmp(args[1], "serve") || !strcmp(args[1], "create") || !strcmp(args[1], "import")) {
        return run_server(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "extract")) {
        return run_extract(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "fsck")) {
        return run_fsck(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "migrate")) {
        return run_migrate(args.size(), args.data()); //migrate requires the exec path

    } else if (!strcmp(args[1], "cluster")) {
        return run_cluster(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "help") || !strcmp(args[1], "-h") || !strcmp(args[1], "--help")) {
        if (args.size() >= 3) {
            if (!strcmp(args[2], "serve")) {
                usage_serve();
            } else if (!strcmp(args[2], "create")) {
                usage_create();
            } else if (!strcmp(args[2], "import")) {
                usage_import();
            } else if (!strcmp(args[2], "extract")) {
                extract::usage(args[1]);
            } else if (!strcmp(args[2], "fsck")) {
                fsck::usage(args[1]);
            } else if (!strcmp(args[2], "migrate")) {
                migrate::usage(args[1]);
            } else {
                printf("No such command %s.\n", args[2]);
            }
        } else {
            usage();
        }

    } else if (!strcmp(args[1], "--version")) {
        print_version_message();
    } else {
        /* Default to "rethinkdb serve"; we get here if we are run without an explicit subcommand
        but with at least one flag */
        args.insert(args.begin() + 1, const_cast<char *>("serve"));
        return run_server(args.size() - 1, args.data() + 1);
    }

    return 0;
}
<commit_msg>Call ShutdownProtobufLibrary() at exit. Closes #353.<commit_after>#include <sys/resource.h>
#include <google/protobuf/stubs/common.h>   // for `ShutdownProtobufLibrary()`

#include "server/control.hpp"
#include "server/server.hpp"
#include "fsck/fsck.hpp"
#include "extract/extract.hpp"
#include "clustering/demo.hpp"
#include "utils.hpp"
#include "help.hpp"
#include "migrate/migrate.hpp"

void print_version_message() {
    printf("rethinkdb " RETHINKDB_VERSION
#ifndef NDEBUG
           " (debug)"
#endif
           "\n");
}

void usage() {
    Help_Pager *help = Help_Pager::instance();
    help->pagef("Usage: rethinkdb COMMAND ...\n"
                "Commands:\n"
                "    help        Display help about rethinkdb and rethinkdb commands.\n"
                "\n"
                "Creating and serving databases:\n"
                "    create      Create an empty database.\n"
                "    serve       Serve an existing database.\n"
                "\n"
                "Administrating databases:\n"
                "    extract     Extract as much data as possible from a corrupted database.\n"
                "    import      Import data from raw memcached commands.\n"
                "    fsck        Check a database for corruption.\n"
                "    migrate     Convert between file versions.\n"
                "\n"
                "Use 'rethinkdb help COMMAND' for help on a single command.\n"
                "Use 'rethinkdb --version' for the current version of rethinkdb.\n");
}

int dispatch_on_args(std::vector<char *> args);

int main(int argc, char *argv[]) {
    initialize_precise_time();
    install_generic_crash_handler();

    atexit(&google::protobuf::ShutdownProtobufLibrary);

#ifndef NDEBUG
    rlimit core_limit;
    core_limit.rlim_cur = RLIM_INFINITY;
    core_limit.rlim_max = RLIM_INFINITY;
    setrlimit(RLIMIT_CORE, &core_limit);
#endif

    /* Put arguments into a vector so we can modify them if convenient */
    std::vector<char *> args(argv, argv + argc);

    return dispatch_on_args(args);
}

int dispatch_on_args(std::vector<char *> args) {
    /* Default to "rethinkdb serve" */
    if (args.size() == 1) args.push_back(const_cast<char *>("serve"));

    /* Switch based on subcommand, then dispatch to the appropriate function */
    if (!strcmp(args[1], "serve") || !strcmp(args[1], "create") || !strcmp(args[1], "import")) {
        return run_server(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "extract")) {
        return run_extract(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "fsck")) {
        return run_fsck(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "migrate")) {
        return run_migrate(args.size(), args.data()); //migrate requires the exec path

    } else if (!strcmp(args[1], "cluster")) {
        return run_cluster(args.size() - 1, args.data() + 1);

    } else if (!strcmp(args[1], "help") || !strcmp(args[1], "-h") || !strcmp(args[1], "--help")) {
        if (args.size() >= 3) {
            if (!strcmp(args[2], "serve")) {
                usage_serve();
            } else if (!strcmp(args[2], "create")) {
                usage_create();
            } else if (!strcmp(args[2], "import")) {
                usage_import();
            } else if (!strcmp(args[2], "extract")) {
                extract::usage(args[1]);
            } else if (!strcmp(args[2], "fsck")) {
                fsck::usage(args[1]);
            } else if (!strcmp(args[2], "migrate")) {
                migrate::usage(args[1]);
            } else {
                printf("No such command %s.\n", args[2]);
            }
        } else {
            usage();
        }

    } else if (!strcmp(args[1], "--version")) {
        print_version_message();
    } else {
        /* Default to "rethinkdb serve"; we get here if we are run without an explicit subcommand
        but with at least one flag */
        args.insert(args.begin() + 1, const_cast<char *>("serve"));
        return run_server(args.size() - 1, args.data() + 1);
    }

    return 0;
}
<|endoftext|>


<commit_before>#include <unordered_map>
#include <string>
#include <numeric>
#include <algorithm>
#include <Rcpp.h>
using namespace Rcpp;
using namespace std;

// [[Rcpp::export]]
SEXP cpp_stack(SEXP arlist) {
    auto array_list = as<List>(arlist);
    auto dimnames = vector<vector<string>>(); // dim: names along
    auto axmap = vector<unordered_map<string, int>>(); // dim: element name->index
    auto a2r = vector<vector<vector<int>>>(array_list.size()); // array > dim > element

    // create lookup tables for all present dimension names
    for (int ai=0; ai<Rf_xlength(array_list); ai++) { // array in arlist
        auto a = as<NumericVector>(array_list[ai]);
        auto dn = as<List>(a.attr("dimnames"));
        auto da = as<vector<int>>(a.attr("dim"));

        a2r[ai] = vector<vector<int>>(da.size());
        if (dimnames.size() < da.size()) {
            dimnames.resize(da.size());
            axmap.resize(da.size());
        }

        for (int d=0; d<da.size(); d++) { // dimension in array
            auto dni = as<vector<string>>(dn[d]);

            for (int e=0; e<da[d]; e++) { // element in dimension
                if (axmap[d].count(dni[e]) == 0) {
                    cout << "array " << ai << " dim " << d << ": " << dni[e] << " -> " << axmap[d].size() << "\n";
                    axmap[d].emplace(dni[e], axmap[d].size());
                    dimnames[d].push_back(dni[e]);
                }
                a2r[ai][d].push_back(axmap[d][dni[e]]);
            }
        }
    }

    for (auto ai=0; ai<a2r.size(); ai++)
        for (auto di=0; di<a2r[ai].size(); di++) {
            cout << "*** array " << ai << " dim " << di << ": ";
            copy(a2r[ai][di].begin(), a2r[ai][di].end(), ostream_iterator<int>(cout, " "));
            cout << "\n";
        }

    // create result array with attributes
    auto rdim = IntegerVector(dimnames.size());
    auto rdnames = List(dimnames.size());
    for (int i=0; i<dimnames.size(); i++) {
        rdim[i] = dimnames[i].size();
        rdnames[i] = CharacterVector(dimnames[i].begin(), dimnames[i].end());
    }
    auto n = accumulate(rdim.begin(), rdim.end(), 1, multiplies<int>());
    auto result = NumericVector(n);
    result.attr("dim") = rdim;
    result.attr("dimnames") = rdnames;

    // fill the result array
    int maxdim = rdim.size() - 1;
    for (int ai=0; ai<Rf_xlength(array_list); ai++) {
        auto a = as<NumericVector>(array_list[ai]);
        int aidx = 0; // consecutive elements in original array

        auto it = vector<vector<int>::iterator>(a2r[ai].size()); // one for each result dim
        for (int d=0; d<it.size(); d++)
            it[d] = a2r[ai][d].begin();

        int dim_offset;
        bool new_offset = true;
        do {
            if (new_offset) { // calculate new offset if we're jumping dimensions
                dim_offset = 0;
                for (int d=0; d<maxdim; d++)
                    dim_offset += rdim[d] * *it[d+1];
                new_offset = false;
            }

            int cur_offset = *it[0];
            it[0]++;
            for (int d=0; d<maxdim; d++) { // check if we're jumping dimensions
                if (it[d] != a2r[ai][d].end()) // higher-order jump needs lower-order
                    break;

                new_offset = true;
                it[d] = a2r[ai][d].begin();
                it[d+1]++;
            }

            cout << "result[" << cur_offset + dim_offset << "] = a[" << ai << "][" << aidx << "]\n";
            result[cur_offset + dim_offset] = a[aidx++];
        } while(it[maxdim] != a2r[ai][maxdim].end());
    }

    return result;
}
<commit_msg>simplify cur_offset<commit_after>#include <unordered_map>
#include <string>
#include <numeric>
#include <algorithm>
#include <Rcpp.h>
using namespace Rcpp;
using namespace std;

// [[Rcpp::export]]
SEXP cpp_stack(SEXP arlist) {
    auto array_list = as<List>(arlist);
    auto dimnames = vector<vector<string>>(); // dim: names along
    auto axmap = vector<unordered_map<string, int>>(); // dim: element name->index
    auto a2r = vector<vector<vector<int>>>(array_list.size()); // array > dim > element

    // create lookup tables for all present dimension names
    for (int ai=0; ai<Rf_xlength(array_list); ai++) { // array in arlist
        auto a = as<NumericVector>(array_list[ai]);
        auto dn = as<List>(a.attr("dimnames"));
        auto da = as<vector<int>>(a.attr("dim"));

        a2r[ai] = vector<vector<int>>(da.size());
        if (dimnames.size() < da.size()) {
            dimnames.resize(da.size());
            axmap.resize(da.size());
        }

        for (int d=0; d<da.size(); d++) { // dimension in array
            auto dni = as<vector<string>>(dn[d]);

            for (int e=0; e<da[d]; e++) { // element in dimension
                if (axmap[d].count(dni[e]) == 0) {
                    cout << "array " << ai << " dim " << d << ": " << dni[e] << " -> " << axmap[d].size() << "\n";
                    axmap[d].emplace(dni[e], axmap[d].size());
                    dimnames[d].push_back(dni[e]);
                }
                a2r[ai][d].push_back(axmap[d][dni[e]]);
            }
        }
    }

    for (auto ai=0; ai<a2r.size(); ai++)
        for (auto di=0; di<a2r[ai].size(); di++) {
            cout << "*** array " << ai << " dim " << di << ": ";
            copy(a2r[ai][di].begin(), a2r[ai][di].end(), ostream_iterator<int>(cout, " "));
            cout << "\n";
        }

    // create result array with attributes
    auto rdim = IntegerVector(dimnames.size());
    auto rdnames = List(dimnames.size());
    for (int i=0; i<dimnames.size(); i++) {
        rdim[i] = dimnames[i].size();
        rdnames[i] = CharacterVector(dimnames[i].begin(), dimnames[i].end());
    }
    auto n = accumulate(rdim.begin(), rdim.end(), 1, multiplies<int>());
    auto result = NumericVector(n);
    result.attr("dim") = rdim;
    result.attr("dimnames") = rdnames;

    // fill the result array
    int maxdim = rdim.size() - 1;
    for (int ai=0; ai<Rf_xlength(array_list); ai++) {
        auto a = as<NumericVector>(array_list[ai]);
        int aidx = 0; // consecutive elements in original array

        auto it = vector<vector<int>::iterator>(a2r[ai].size()); // one for each result dim
        for (int d=0; d<it.size(); d++)
            it[d] = a2r[ai][d].begin();

        int dim_offset;
        bool new_offset = true;
        do {
            if (new_offset) { // calculate new offset if we're jumping dimensions
                dim_offset = 0;
                for (int d=0; d<maxdim; d++)
                    dim_offset += rdim[d] * *it[d+1];
                new_offset = false;
            }

            cout << "result[" << *it[0] + dim_offset << "] = a[" << ai << "][" << aidx << "]\n";
            result[*it[0] + dim_offset] = a[aidx++];

            it[0]++;
            for (int d=0; d<maxdim; d++) { // check if we're jumping dimensions
                if (it[d] != a2r[ai][d].end()) // higher-order jump needs lower-order
                    break;

                new_offset = true;
                it[d] = a2r[ai][d].begin();
                it[d+1]++;
            }
        } while(it[maxdim] != a2r[ai][maxdim].end());
    }

    return result;
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2003-2014, John Wiegley.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 * - Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * - Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 *
 * - Neither the name of New Artisans LLC nor the names of its
 *   contributors may be used to endorse or promote products derived from
 *   this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <system.hh>

#include "stream.h"

namespace ledger {

namespace {
  /**
   * @brief Forks a child process so that Ledger may handle running a
   * pager
   *
   * In order for the pager option to work, Ledger has to run the pager
   * itself, which requires Ledger to fork a new process in order to run
   * the pager.  This function does the necessary fork.  After the fork,
   * two processes exist.  One of them is exec'd to create the pager;
   * the other is still Ledger.
   *
   * This function returns only for the process that is still Ledger.
   *
   * @param pager_path Path to the pager command.
   *
   * @return The file descriptor of the pipe to the pager.  The caller
   * is responsible for cleaning this up.
   *
   * @exception std::logic_error Some problem was encountered, such as
   * failure to create a pipe or failure to fork a child process.
   */
  int do_fork(std::ostream ** os, const path& pager_path)
  {
#ifndef WIN32
    int pfd[2];

    int status = pipe(pfd);
    if (status == -1)
      throw std::logic_error(_("Failed to create pipe"));

    status = fork();
    if (status < 0) {
      throw std::logic_error(_("Failed to fork child process"));
    }
    else if (status == 0) {     // child
      // Duplicate pipe's reading end into stdin
      status = dup2(pfd[0], STDIN_FILENO);
      if (status == -1)
        perror("dup2");

      // Close unuseful file descriptors: the pipe's writing and reading
      // ends (the latter is not needed anymore, after the duplication).
      close(pfd[1]);
      close(pfd[0]);

      // Find command name: its the substring starting right of the
      // rightmost '/' character in the pager pathname.  See manpage for
      // strrchr.
#if BOOST_VERSION >= 103700
      path basename(pager_path.filename());
#else
      path basename(pager_path.leaf());
#endif
      execlp(pager_path.string().c_str(), basename.string().c_str(), NULL);

      // We should never, ever reach here
      perror((std::string("execlp: ") + pager_path.string()).c_str());
      exit(1);
    }
    else {                      // parent
      close(pfd[0]);
      typedef iostreams::stream<iostreams::file_descriptor_sink> fdstream;
#if BOOST_VERSION >= 104400
      *os = new fdstream(pfd[1], iostreams::never_close_handle);
#else // BOOST_VERSION >= 104400
      *os = new fdstream(pfd[1]);
#endif // BOOST_VERSION >= 104400
    }
    return pfd[1];
#else
    return 0;
#endif
  }
}

void output_stream_t::initialize(const optional<path>& output_file,
                                 const optional<path>& pager_path)
{
  if (output_file && *output_file != "-")
    os = new ofstream(*output_file);
  else if (pager_path)
    pipe_to_pager_fd = do_fork(&os, *pager_path);
  else
    os = &std::cout;
}

void output_stream_t::close()
{
#ifndef WIN32
  if (os != &std::cout) {
    checked_delete(os);
    os = &std::cout;
  }

  if (pipe_to_pager_fd != -1) {
    ::close(pipe_to_pager_fd);
    pipe_to_pager_fd = -1;

    int status;
    wait(&status);
    if (! WIFEXITED(status) || WEXITSTATUS(status) != 0)
      throw std::logic_error(_("Error in the pager"));
  }
#endif
}

} // namespace ledger
<commit_msg>use /bin/sh to run the pager<commit_after>/*
 * Copyright (c) 2003-2014, John Wiegley.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 * - Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * - Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 *
 * - Neither the name of New Artisans LLC nor the names of its
 *   contributors may be used to endorse or promote products derived from
 *   this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <system.hh>

#include "stream.h"

namespace ledger {

namespace {
  /**
   * @brief Forks a child process so that Ledger may handle running a
   * pager
   *
   * In order for the pager option to work, Ledger has to run the pager
   * itself, which requires Ledger to fork a new process in order to run
   * the pager.  This function does the necessary fork.  After the fork,
   * two processes exist.  One of them is exec'd to create the pager;
   * the other is still Ledger.
   *
   * This function returns only for the process that is still Ledger.
   *
   * @param pager_path Path to the pager command.
   *
   * @return The file descriptor of the pipe to the pager.  The caller
   * is responsible for cleaning this up.
   *
   * @exception std::logic_error Some problem was encountered, such as
   * failure to create a pipe or failure to fork a child process.
   */
  int do_fork(std::ostream ** os, const path& pager_path)
  {
#ifndef WIN32
    int pfd[2];

    int status = pipe(pfd);
    if (status == -1)
      throw std::logic_error(_("Failed to create pipe"));

    status = fork();
    if (status < 0) {
      throw std::logic_error(_("Failed to fork child process"));
    }
    else if (status == 0) {     // child
      // Duplicate pipe's reading end into stdin
      status = dup2(pfd[0], STDIN_FILENO);
      if (status == -1)
        perror("dup2");

      // Close unuseful file descriptors: the pipe's writing and reading
      // ends (the latter is not needed anymore, after the duplication).
      close(pfd[1]);
      close(pfd[0]);

      execlp("/bin/sh", "/bin/sh", "-c", pager_path.string().c_str(), NULL);

      // We should never, ever reach here
      perror("execlp: /bin/sh");
      exit(1);
    }
    else {                      // parent
      close(pfd[0]);
      typedef iostreams::stream<iostreams::file_descriptor_sink> fdstream;
#if BOOST_VERSION >= 104400
      *os = new fdstream(pfd[1], iostreams::never_close_handle);
#else // BOOST_VERSION >= 104400
      *os = new fdstream(pfd[1]);
#endif // BOOST_VERSION >= 104400
    }
    return pfd[1];
#else
    return 0;
#endif
  }
}

void output_stream_t::initialize(const optional<path>& output_file,
                                 const optional<path>& pager_path)
{
  if (output_file && *output_file != "-")
    os = new ofstream(*output_file);
  else if (pager_path)
    pipe_to_pager_fd = do_fork(&os, *pager_path);
  else
    os = &std::cout;
}

void output_stream_t::close()
{
#ifndef WIN32
  if (os != &std::cout) {
    checked_delete(os);
    os = &std::cout;
  }

  if (pipe_to_pager_fd != -1) {
    ::close(pipe_to_pager_fd);
    pipe_to_pager_fd = -1;

    int status;
    wait(&status);
    if (! WIFEXITED(status) || WEXITSTATUS(status) != 0)
      throw std::logic_error(_("Error in the pager"));
  }
#endif
}

} // namespace ledger
<|endoftext|>


<commit_before>/*
    KNode, the KDE newsreader
    Copyright (c) 2004-2005 Volker Krause <vkrause@kde.org>
    Copyright (c) 2009 Olivier Trichet <nive@nivalis.org>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.
    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software Foundation,
    Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, US
*/

#include "kncollectionview.h"

#include "knglobals.h"
#include "knconfig.h"
#include "knnntpaccount.h"
#include "knaccountmanager.h"
#include "kngroup.h"
#include "kngroupmanager.h"
#include "knfolder.h"
#include "knfoldermanager.h"
#include "kncollectionviewitem.h"
#include "settings.h"

#include <kiconloader.h>
#include <klocale.h>

#include <QDropEvent>
#include <QPainter>

KNCollectionView::KNCollectionView( QWidget *parent ) :
  FolderTreeWidget( parent ),
  mActiveItem( 0 ),
  mDragTargetItem( 0 )
{
  setDragEnabled(true);
  setDropIndicatorShown( false ); // We draw our own (see paintEvent)

  // add unread and total columns if necessary
  loadLayout();

  readConfig();

  // load accounts and folders
  reloadAccounts();
  reloadFolders();

  // connect to the account manager
  KNAccountManager* am = knGlobals.accountManager();
  connect(am, SIGNAL(accountAdded(KNNntpAccount*)), SLOT(addAccount(KNNntpAccount*)));
  connect(am, SIGNAL(accountRemoved(KNNntpAccount*)), SLOT(removeAccount(KNNntpAccount*)));
  connect(am, SIGNAL(accountModified(KNNntpAccount*)), SLOT(updateAccount(KNNntpAccount*)));

  // connect to the group manager
  KNGroupManager* gm = knGlobals.groupManager();
  connect(gm, SIGNAL(groupAdded(KNGroup*)), SLOT(addGroup(KNGroup*)));
  connect(gm, SIGNAL(groupRemoved(KNGroup*)), SLOT(removeGroup(KNGroup*)));
  connect(gm, SIGNAL(groupUpdated(KNGroup*)), SLOT(updateGroup(KNGroup*)));

  // connect to the folder manager
  KNFolderManager* fm = knGlobals.folderManager();
  connect(fm, SIGNAL(folderAdded(KNFolder*)), SLOT(addPendingFolders()));
  connect(fm, SIGNAL(folderRemoved(KNFolder*)), SLOT(removeFolder(KNFolder*)));
  connect(fm, SIGNAL(folderActivated(KNFolder*)), SLOT(activateFolder(KNFolder*)));

  // Edition of label
  setEditTriggers( QAbstractItemView::NoEditTriggers );
}

KNCollectionView::~KNCollectionView()
{
  writeConfig();
}

void KNCollectionView::loadLayout()
{
  addLabelColumn( i18n("Name") );
  addUnreadColumn( i18n("Unread") );
  addTotalColumn( i18n("Total") );

  restoreLayout( knGlobals.config(), "GroupView" );

  if ( sortColumn() == -1 ) {
    sortByColumn( labelColumnIndex(), Qt::AscendingOrder );
  }

}

void KNCollectionView::readConfig()
{
  // font
  setFont( knGlobals.settings()->groupListFont() );
}

void KNCollectionView::writeConfig()
{
  saveLayout( knGlobals.config(), "GroupView" );
}

void KNCollectionView::addAccount(KNNntpAccount *a)
{
  // add account item
  KNCollectionViewItem* item = new KNCollectionViewItem( this, FolderTreeWidgetItem::News );
  a->setListItem( item );
  item->setExpanded( a->wasOpen() );

  // add groups for this account
  KNGroup::List groups = knGlobals.groupManager()->groupsOfAccount( a );
  for ( KNGroup::List::Iterator it = groups.begin(); it != groups.end(); ++it ) {
    KNCollectionViewItem *gitem = new KNCollectionViewItem( item, FolderTreeWidgetItem::News );
    (*it)->setListItem( gitem );
    (*it)->updateListItem();
  }
}

void KNCollectionView::removeAccount(KNNntpAccount *a)
{
  if(!a->listItem())
    return;
  KNCollectionViewItem *child = 0;
  KNCollectionViewItem *aitem = a->listItem();
  while ( ( child = static_cast<KNCollectionViewItem*>( aitem->takeChild( 0 ) ) ) ) {
    removeGroup( static_cast<KNGroup*>( child->collection() ) );
  }
  delete aitem;
  a->setListItem(0);
}

void KNCollectionView::updateAccount(KNNntpAccount *a)
{
  a->updateListItem();
}

void KNCollectionView::reloadAccounts()
{
  KNAccountManager::List list = knGlobals.accountManager()->accounts();
  for ( KNAccountManager::List::Iterator it = list.begin(); it != list.end(); ++it ) {
    removeAccount( *it );
    addAccount( *it );
  }
}

void KNCollectionView::addGroup(KNGroup *g)
{
  if (!g->account()->listItem())
    return;

  KNCollectionViewItem *gitem =
      new KNCollectionViewItem( g->account()->listItem(), FolderTreeWidgetItem::News );
  g->setListItem(gitem);
  updateGroup(g);
}

void KNCollectionView::removeGroup(KNGroup *g)
{
  if (!g->listItem())
    return;

  delete g->listItem();
  g->setListItem(0);
}

void KNCollectionView::updateGroup(KNGroup *g)
{
  g->updateListItem();
}

void KNCollectionView::addFolder(KNFolder *f)
{
  KNCollectionViewItem *it;

  if (!f->parent()) {
    // root folder
    it = new KNCollectionViewItem(this, FolderTreeWidgetItem::Local);
  } else {
    // make sure the parent folder has already been added
    if (!f->parent()->listItem())
      addFolder( static_cast<KNFolder*>(f->parent()) );
    // handle special folders
    FolderTreeWidgetItem::FolderType type = FolderTreeWidgetItem::Other;
    switch ( f->id() ) {
      case 1:
        type = FolderTreeWidgetItem::Drafts; break;
      case 2:
        type = FolderTreeWidgetItem::Outbox; break;
      case 3:
        type = FolderTreeWidgetItem::SentMail; break;
    }
    it = new KNCollectionViewItem( f->parent()->listItem(), FolderTreeWidgetItem::Local, type );
  }
  f->setListItem( it );
  updateFolder( f );
}

void KNCollectionView::removeFolder(KNFolder* f)
{
  if(!f->listItem())
    return;
  KNCollectionViewItem *child = 0;
  KNCollectionViewItem *it = f->listItem();
  while ( ( child = static_cast<KNCollectionViewItem*>( it->takeChild( 0 ) ) ) ) {
    removeFolder( static_cast<KNFolder*>( child->collection() ) );
  }
  delete f->listItem();
  f->setListItem(0);
}

void KNCollectionView::reloadFolders()
{
  // remove existing folder items
  removeFolder(knGlobals.folderManager()->root());

  // add folder items
  addPendingFolders();
}

void KNCollectionView::addPendingFolders()
{
  KNFolderManager::List folders = knGlobals.folderManager()->folders();
  for ( KNFolderManager::List::Iterator it = folders.begin(); it != folders.end(); ++it )
    if ( !(*it)->listItem() )
      addFolder( (*it) );
  // now open the folders if they were open in the last session
  for ( KNFolderManager::List::Iterator it = folders.begin(); it != folders.end(); ++it ) {
    if ( (*it)->listItem()) {
      (*it)->listItem()->setExpanded( (*it)->wasOpen() );
    }
  }
}

void KNCollectionView::activateFolder(KNFolder* f)
{
  if(f->listItem())
    setActive( f->listItem() );
}

void KNCollectionView::updateFolder(KNFolder* f)
{
  f->updateListItem();
}

void KNCollectionView::reload()
{
  reloadAccounts();
  reloadFolders();
}

void KNCollectionView::setActive( QTreeWidgetItem *i )
{
  if (!i || mActiveItem == i)
    return;

  i->setSelected( true );
  setCurrentItem( i );
  mActiveItem = i;
}

void KNCollectionView::nextGroup()
{
  incCurrentFolder();
  setActive( currentItem() );
}

void KNCollectionView::prevGroup()
{
  decCurrentFolder();
  setActive( currentItem() );
}

void KNCollectionView::decCurrentFolder()
{
  QTreeWidgetItemIterator it( currentItem() );
  --it;
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( *it );
  if (fti) {
    setFocus();
    setCurrentItem( fti );
  }
}

void KNCollectionView::incCurrentFolder()
{
  QTreeWidgetItemIterator it( currentItem() );
  ++it;
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( *it );
  if (fti) {
    setFocus();
    setCurrentItem( fti );
  }
}

void KNCollectionView::selectCurrentFolder()
{
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( currentItem() );
  if (fti) {
    setActive( fti );
  }
}

void KNCollectionView::contextMenuEvent( QContextMenuEvent *event )
{
  QTreeWidgetItem *item = itemAt( event->pos() );
  if(item) {
    emit contextMenu( item, event->globalPos() );
  }
}

void KNCollectionView::paintEvent( QPaintEvent *event )
{
  FolderTreeWidget::paintEvent( event );

  if ( mDragTargetItem ) {
    QRect rect = visualItemRect( mDragTargetItem );
    if ( rect.isValid() ) {
      QPainter p( viewport() );
      QBrush brush = KColorScheme( QPalette::Active, KColorScheme::Selection ).decoration( KColorScheme::HoverColor ).color();
      p.setPen( QPen( brush, 2 ) );
      p.drawRect( rect );
    }
  }
}

void KNCollectionView::dragEnterEvent( QDragEnterEvent *event )
{
  if ( event->mimeData() && event->mimeData()->hasFormat( "x-knode-drag/folder" ) ) {
    event->accept();
  } else {
    event->ignore();
  }
}

void KNCollectionView::dragLeaveEvent( QDragLeaveEvent *event )
{
  FolderTreeWidget::dragLeaveEvent( event );

  mDragTargetItem = 0;
  viewport()->update(); // repaint
}

void KNCollectionView::dragMoveEvent( QDragMoveEvent *event )
{
  // Needed for auto-scrolling
  FolderTreeWidget::dragMoveEvent( event );

  handleDragNDropEvent( event, false );
}

void KNCollectionView::dropEvent( QDropEvent *event )
{
  handleDragNDropEvent( event, true );

  mDragTargetItem = 0;
  viewport()->update(); // repaint
}

void KNCollectionView::handleDragNDropEvent( QDropEvent *event, bool enforceDrop )
{
  QTreeWidgetItem *item = itemAt( event->pos() );
  KNCollectionViewItem *fti = static_cast<KNCollectionViewItem*>( item );
  bool accepted = false;
  if ( fti && fti->collection() && fti->collection()->type()==KNCollection::CTfolder ) {
    const QMimeData *md = event->mimeData();
    if( md && md->hasFormat( "x-knode-drag/folder" ) ) {
      KNFolder *dest = static_cast<KNFolder*>( fti->collection() );
      KNFolderManager *folderManager = KNGlobals::self()->folderManager();
      if ( !enforceDrop ) {
        // Notify that the move is possible.
        accepted = folderManager->canMoveFolder( folderManager->currentFolder(), dest );
      } else {
        // Informs whether the move succeeded.
        accepted = folderManager->moveFolder( folderManager->currentFolder(), dest );
      }
    }
    /* TODO: redo drag and drop of articles.
    else if ( md && md->hasFormat( "message/news" ) ) {
      if ( !static_cast<KNFolder*>( fti->collection() )->isRootFolder() ) {   // don't drop articles on the root folder
        if ( !enforceDrop ) {
          accepted = true;
        } else if(f_olManager->currentFolder()) {
            if (e->dropAction() == Qt::MoveAction) {
              KNLocalArticle::List l;
              getSelectedArticles(l);
              a_rtManager->moveIntoFolder(l, dest);
              accepted = true;
            } else {
              KNArticle::List l;
              getSelectedArticles(l);
              a_rtManager->copyIntoFolder(l, dest);
              accepted = true;
            }
          }
          else if(g_rpManager->currentGroup()) {
            KNArticle::List l;
            getSelectedArticles(l);
            a_rtManager->copyIntoFolder(l, dest);
            accepted = true;
          }
      }
    }
    */
  }

  // Trigger an update of the target item (drawing of target's borders)
  bool repaintTarget =    ( mDragTargetItem == item )       /* target changes */
                       || ( !accepted && mDragTargetItem ); /* just entered an invalid target */
  mDragTargetItem = accepted ? item : 0;
  if ( repaintTarget ) {
    viewport()->update();
  }

  event->setAccepted( accepted );
}

QStringList KNCollectionView::mimeTypes() const
{
  QStringList l;
  l << "x-knode-drag/folder";
  return l;
}

void KNCollectionView::startDrag( Qt::DropActions supportedActions )
{
  Q_UNUSED( supportedActions );

  KNCollectionViewItem *item = static_cast<KNCollectionViewItem*>( currentItem() );
  // Only folders can be dragged
  if ( !item || !item->collection() || item->collection()->type()!=KNCollection::CTfolder ) {
    return;
  }

  KNFolder *folder = static_cast<KNFolder*>( item->collection() );

  // Can not drag special folders
  if ( folder->isRootFolder() || folder->isStandardFolder() ) {
    return;
  }

  QMimeData *mimeData = new QMimeData();
  const QByteArray id( QString::number( folder->id() ).toAscii() );
  mimeData->setData( "x-knode-drag/folder", id );

  QDrag *drag = new QDrag( this );
  drag->setMimeData( mimeData );
  drag->setPixmap( SmallIcon( "folder" ) );

  drag->exec( Qt::MoveAction );
}

#include "kncollectionview.moc"
<commit_msg>Backport 4cb6e7c by Olivier Trichet (made to the master branch) to 4.4.<commit_after>/*
    KNode, the KDE newsreader
    Copyright (c) 2004-2005 Volker Krause <vkrause@kde.org>
    Copyright (c) 2009 Olivier Trichet <nive@nivalis.org>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.
    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software Foundation,
    Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, US
*/

#include "kncollectionview.h"

#include "knglobals.h"
#include "knconfig.h"
#include "knnntpaccount.h"
#include "knaccountmanager.h"
#include "kngroup.h"
#include "kngroupmanager.h"
#include "knfolder.h"
#include "knfoldermanager.h"
#include "kncollectionviewitem.h"
#include "settings.h"

#include <kiconloader.h>
#include <klocale.h>

#include <QDropEvent>
#include <QPainter>

KNCollectionView::KNCollectionView( QWidget *parent ) :
  FolderTreeWidget( parent ),
  mActiveItem( 0 ),
  mDragTargetItem( 0 )
{
  setDragEnabled(true);
  setDropIndicatorShown( false ); // We draw our own (see paintEvent)

  // add unread and total columns if necessary
  loadLayout();

  readConfig();

  // load accounts and folders
  reloadAccounts();
  reloadFolders();

  // connect to the account manager
  KNAccountManager* am = knGlobals.accountManager();
  connect(am, SIGNAL(accountAdded(KNNntpAccount*)), SLOT(addAccount(KNNntpAccount*)));
  connect(am, SIGNAL(accountRemoved(KNNntpAccount*)), SLOT(removeAccount(KNNntpAccount*)));
  connect(am, SIGNAL(accountModified(KNNntpAccount*)), SLOT(updateAccount(KNNntpAccount*)));

  // connect to the group manager
  KNGroupManager* gm = knGlobals.groupManager();
  connect(gm, SIGNAL(groupAdded(KNGroup*)), SLOT(addGroup(KNGroup*)));
  connect(gm, SIGNAL(groupRemoved(KNGroup*)), SLOT(removeGroup(KNGroup*)));
  connect(gm, SIGNAL(groupUpdated(KNGroup*)), SLOT(updateGroup(KNGroup*)));

  // connect to the folder manager
  KNFolderManager* fm = knGlobals.folderManager();
  connect(fm, SIGNAL(folderAdded(KNFolder*)), SLOT(addPendingFolders()));
  connect(fm, SIGNAL(folderRemoved(KNFolder*)), SLOT(removeFolder(KNFolder*)));
  connect(fm, SIGNAL(folderActivated(KNFolder*)), SLOT(activateFolder(KNFolder*)));

  // Edition of label
  setEditTriggers( QAbstractItemView::NoEditTriggers );
}

KNCollectionView::~KNCollectionView()
{
  writeConfig();
}

void KNCollectionView::loadLayout()
{
  addLabelColumn( i18n("Name") );
  addUnreadColumn( i18n("Unread") );
  addTotalColumn( i18n("Total") );

  restoreLayout( knGlobals.config(), "GroupView" );

  if ( sortColumn() == -1 ) {
    sortByColumn( labelColumnIndex(), Qt::AscendingOrder );
  }

}

void KNCollectionView::readConfig()
{
  // font
  setFont( knGlobals.settings()->groupListFont() );
}

void KNCollectionView::writeConfig()
{
  saveLayout( knGlobals.config(), "GroupView" );
}

void KNCollectionView::addAccount(KNNntpAccount *a)
{
  // add account item
  KNCollectionViewItem* item = new KNCollectionViewItem( this, FolderTreeWidgetItem::News );
  a->setListItem( item );
  item->setExpanded( a->wasOpen() );

  // add groups for this account
  KNGroup::List groups = knGlobals.groupManager()->groupsOfAccount( a );
  for ( KNGroup::List::Iterator it = groups.begin(); it != groups.end(); ++it ) {
    KNCollectionViewItem *gitem = new KNCollectionViewItem( item, FolderTreeWidgetItem::News );
    (*it)->setListItem( gitem );
    (*it)->updateListItem();
  }
}

void KNCollectionView::removeAccount(KNNntpAccount *a)
{
  if(!a->listItem())
    return;
  KNCollectionViewItem *child = 0;
  KNCollectionViewItem *aitem = a->listItem();
  while ( ( child = static_cast<KNCollectionViewItem*>( aitem->takeChild( 0 ) ) ) ) {
    removeGroup( static_cast<KNGroup*>( child->collection() ) );
  }
  delete aitem;
  a->setListItem(0);
}

void KNCollectionView::updateAccount(KNNntpAccount *a)
{
  a->updateListItem();
}

void KNCollectionView::reloadAccounts()
{
  KNAccountManager::List list = knGlobals.accountManager()->accounts();
  for ( KNAccountManager::List::Iterator it = list.begin(); it != list.end(); ++it ) {
    removeAccount( *it );
    addAccount( *it );
  }
}

void KNCollectionView::addGroup(KNGroup *g)
{
  if (!g->account()->listItem())
    return;

  KNCollectionViewItem *gitem =
      new KNCollectionViewItem( g->account()->listItem(), FolderTreeWidgetItem::News );
  g->setListItem(gitem);
  updateGroup(g);
}

void KNCollectionView::removeGroup(KNGroup *g)
{
  if (!g->listItem())
    return;

  g->listItem()->setHidden( true ); // work around bug 248256
  delete g->listItem();
  g->setListItem(0);
}

void KNCollectionView::updateGroup(KNGroup *g)
{
  g->updateListItem();
}

void KNCollectionView::addFolder(KNFolder *f)
{
  KNCollectionViewItem *it;

  if (!f->parent()) {
    // root folder
    it = new KNCollectionViewItem(this, FolderTreeWidgetItem::Local);
  } else {
    // make sure the parent folder has already been added
    if (!f->parent()->listItem())
      addFolder( static_cast<KNFolder*>(f->parent()) );
    // handle special folders
    FolderTreeWidgetItem::FolderType type = FolderTreeWidgetItem::Other;
    switch ( f->id() ) {
      case 1:
        type = FolderTreeWidgetItem::Drafts; break;
      case 2:
        type = FolderTreeWidgetItem::Outbox; break;
      case 3:
        type = FolderTreeWidgetItem::SentMail; break;
    }
    it = new KNCollectionViewItem( f->parent()->listItem(), FolderTreeWidgetItem::Local, type );
  }
  f->setListItem( it );
  updateFolder( f );
}

void KNCollectionView::removeFolder(KNFolder* f)
{
  if(!f->listItem())
    return;
  KNCollectionViewItem *child = 0;
  KNCollectionViewItem *it = f->listItem();
  while ( ( child = static_cast<KNCollectionViewItem*>( it->takeChild( 0 ) ) ) ) {
    removeFolder( static_cast<KNFolder*>( child->collection() ) );
  }
  delete f->listItem();
  f->setListItem(0);
}

void KNCollectionView::reloadFolders()
{
  // remove existing folder items
  removeFolder(knGlobals.folderManager()->root());

  // add folder items
  addPendingFolders();
}

void KNCollectionView::addPendingFolders()
{
  KNFolderManager::List folders = knGlobals.folderManager()->folders();
  for ( KNFolderManager::List::Iterator it = folders.begin(); it != folders.end(); ++it )
    if ( !(*it)->listItem() )
      addFolder( (*it) );
  // now open the folders if they were open in the last session
  for ( KNFolderManager::List::Iterator it = folders.begin(); it != folders.end(); ++it ) {
    if ( (*it)->listItem()) {
      (*it)->listItem()->setExpanded( (*it)->wasOpen() );
    }
  }
}

void KNCollectionView::activateFolder(KNFolder* f)
{
  if(f->listItem())
    setActive( f->listItem() );
}

void KNCollectionView::updateFolder(KNFolder* f)
{
  f->updateListItem();
}

void KNCollectionView::reload()
{
  reloadAccounts();
  reloadFolders();
}

void KNCollectionView::setActive( QTreeWidgetItem *i )
{
  if (!i || mActiveItem == i)
    return;

  i->setSelected( true );
  setCurrentItem( i );
  mActiveItem = i;
}

void KNCollectionView::nextGroup()
{
  incCurrentFolder();
  setActive( currentItem() );
}

void KNCollectionView::prevGroup()
{
  decCurrentFolder();
  setActive( currentItem() );
}

void KNCollectionView::decCurrentFolder()
{
  QTreeWidgetItemIterator it( currentItem() );
  --it;
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( *it );
  if (fti) {
    setFocus();
    setCurrentItem( fti );
  }
}

void KNCollectionView::incCurrentFolder()
{
  QTreeWidgetItemIterator it( currentItem() );
  ++it;
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( *it );
  if (fti) {
    setFocus();
    setCurrentItem( fti );
  }
}

void KNCollectionView::selectCurrentFolder()
{
  FolderTreeWidgetItem* fti = static_cast<FolderTreeWidgetItem*>( currentItem() );
  if (fti) {
    setActive( fti );
  }
}

void KNCollectionView::contextMenuEvent( QContextMenuEvent *event )
{
  QTreeWidgetItem *item = itemAt( event->pos() );
  if(item) {
    emit contextMenu( item, event->globalPos() );
  }
}

void KNCollectionView::paintEvent( QPaintEvent *event )
{
  FolderTreeWidget::paintEvent( event );

  if ( mDragTargetItem ) {
    QRect rect = visualItemRect( mDragTargetItem );
    if ( rect.isValid() ) {
      QPainter p( viewport() );
      QBrush brush = KColorScheme( QPalette::Active, KColorScheme::Selection ).decoration( KColorScheme::HoverColor ).color();
      p.setPen( QPen( brush, 2 ) );
      p.drawRect( rect );
    }
  }
}

void KNCollectionView::dragEnterEvent( QDragEnterEvent *event )
{
  if ( event->mimeData() && event->mimeData()->hasFormat( "x-knode-drag/folder" ) ) {
    event->accept();
  } else {
    event->ignore();
  }
}

void KNCollectionView::dragLeaveEvent( QDragLeaveEvent *event )
{
  FolderTreeWidget::dragLeaveEvent( event );

  mDragTargetItem = 0;
  viewport()->update(); // repaint
}

void KNCollectionView::dragMoveEvent( QDragMoveEvent *event )
{
  // Needed for auto-scrolling
  FolderTreeWidget::dragMoveEvent( event );

  handleDragNDropEvent( event, false );
}

void KNCollectionView::dropEvent( QDropEvent *event )
{
  handleDragNDropEvent( event, true );

  mDragTargetItem = 0;
  viewport()->update(); // repaint
}

void KNCollectionView::handleDragNDropEvent( QDropEvent *event, bool enforceDrop )
{
  QTreeWidgetItem *item = itemAt( event->pos() );
  KNCollectionViewItem *fti = static_cast<KNCollectionViewItem*>( item );
  bool accepted = false;
  if ( fti && fti->collection() && fti->collection()->type()==KNCollection::CTfolder ) {
    const QMimeData *md = event->mimeData();
    if( md && md->hasFormat( "x-knode-drag/folder" ) ) {
      KNFolder *dest = static_cast<KNFolder*>( fti->collection() );
      KNFolderManager *folderManager = KNGlobals::self()->folderManager();
      if ( !enforceDrop ) {
        // Notify that the move is possible.
        accepted = folderManager->canMoveFolder( folderManager->currentFolder(), dest );
      } else {
        // Informs whether the move succeeded.
        accepted = folderManager->moveFolder( folderManager->currentFolder(), dest );
      }
    }
    /* TODO: redo drag and drop of articles.
    else if ( md && md->hasFormat( "message/news" ) ) {
      if ( !static_cast<KNFolder*>( fti->collection() )->isRootFolder() ) {   // don't drop articles on the root folder
        if ( !enforceDrop ) {
          accepted = true;
        } else if(f_olManager->currentFolder()) {
            if (e->dropAction() == Qt::MoveAction) {
              KNLocalArticle::List l;
              getSelectedArticles(l);
              a_rtManager->moveIntoFolder(l, dest);
              accepted = true;
            } else {
              KNArticle::List l;
              getSelectedArticles(l);
              a_rtManager->copyIntoFolder(l, dest);
              accepted = true;
            }
          }
          else if(g_rpManager->currentGroup()) {
            KNArticle::List l;
            getSelectedArticles(l);
            a_rtManager->copyIntoFolder(l, dest);
            accepted = true;
          }
      }
    }
    */
  }

  // Trigger an update of the target item (drawing of target's borders)
  bool repaintTarget =    ( mDragTargetItem == item )       /* target changes */
                       || ( !accepted && mDragTargetItem ); /* just entered an invalid target */
  mDragTargetItem = accepted ? item : 0;
  if ( repaintTarget ) {
    viewport()->update();
  }

  event->setAccepted( accepted );
}

QStringList KNCollectionView::mimeTypes() const
{
  QStringList l;
  l << "x-knode-drag/folder";
  return l;
}

void KNCollectionView::startDrag( Qt::DropActions supportedActions )
{
  Q_UNUSED( supportedActions );

  KNCollectionViewItem *item = static_cast<KNCollectionViewItem*>( currentItem() );
  // Only folders can be dragged
  if ( !item || !item->collection() || item->collection()->type()!=KNCollection::CTfolder ) {
    return;
  }

  KNFolder *folder = static_cast<KNFolder*>( item->collection() );

  // Can not drag special folders
  if ( folder->isRootFolder() || folder->isStandardFolder() ) {
    return;
  }

  QMimeData *mimeData = new QMimeData();
  const QByteArray id( QString::number( folder->id() ).toAscii() );
  mimeData->setData( "x-knode-drag/folder", id );

  QDrag *drag = new QDrag( this );
  drag->setMimeData( mimeData );
  drag->setPixmap( SmallIcon( "folder" ) );

  drag->exec( Qt::MoveAction );
}

#include "kncollectionview.moc"
<|endoftext|>


<commit_before>/*******************************************************************
 This file is part of KNotes.

 Copyright (c) 2004, Bo Thorsen <bo@sonofthor.dk>
               2004, Michael Brade <brade@kde.org>

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 MA  02110-1301, USA.

 In addition, as a special exception, the copyright holders give
 permission to link the code of this program with any edition of
 the Qt library by Trolltech AS, Norway (or with modified versions
 of Qt that use the same license as Qt), and distribute linked
 combinations including the two.  You must obey the GNU General
 Public License in all respects for all of the code used other than
 Qt.  If you modify this file, you may extend this exception to
 your version of the file, but you are not obligated to do so.  If
 you do not wish to do so, delete this exception statement from
 your version.
*******************************************************************/

#include "knotes/resourcelocal.h"
#include "knotes/resourcemanager.h"

#include <kcal/journal.h>

KNotesResourceManager::KNotesResourceManager()
  : QObject( 0 )
{
  setObjectName( "KNotes Resource Manager" );
  m_manager = new KRES::Manager<ResourceNotes>( "notes" );
  m_manager->addObserver( this );
  m_manager->readConfig();
}

KNotesResourceManager::~KNotesResourceManager()
{
  delete m_manager;
}

void KNotesResourceManager::load()
{
  if ( !m_manager->standardResource() ) {
    kWarning( 5500 ) << "No standard resource yet.";
    ResourceNotes *resource = new ResourceLocal();
    m_manager->add( resource );
    m_manager->setStandardResource( resource );
  }

  // Open all active resources
  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    kDebug( 5500 ) << "Opening resource " + ( *it )->resourceName();
    ( *it )->setManager( this );
    if ( ( *it )->open() ) {
      ( *it )->load();
    }
  }
}

void KNotesResourceManager::save()
{
  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    ( *it )->save();
  }
}

// when adding a new note, make sure a config file exists!!

void KNotesResourceManager::addNewNote( KCal::Journal *journal )
{
  // TODO: Make this configurable
  ResourceNotes *resource = m_manager->standardResource();
  if ( resource ) {
    resource->addNote( journal );
    registerNote( resource, journal );
  } else {
    kWarning( 5500 ) << "no resource!";
  }
}

void KNotesResourceManager::registerNote( ResourceNotes *resource,
                                          KCal::Journal *journal )
{
  // TODO: only emit the signal if the journal is new?
  m_resourceMap.insert( journal->uid(), resource );
  emit sigRegisteredNote( journal );
}

void KNotesResourceManager::deleteNote( KCal::Journal *journal )
{
  QString uid = journal->uid();

  // Remove the journal from the resource it came from
  m_resourceMap[ uid ]->deleteNote( journal );
  m_resourceMap.remove( uid );

  // libkcal does not delete the journal immediately, therefore it is ok to
  // emit the journal here
  emit sigDeregisteredNote( journal );
}

KCal::Alarm::List KNotesResourceManager::alarms( const KDateTime &from,
                                                 const KDateTime &to )
{
  KCal::Alarm::List result;

  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    KCal::Alarm::List list = ( *it )->alarms( from, to );
    KCal::Alarm::List::Iterator it;
    for ( it = list.begin(); it != list.end(); ++it ) {
      result.append( *it );
    }
  }

  return result;
}

void KNotesResourceManager::resourceAdded( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource added:" << resource->resourceName();

  if ( !resource->isActive() ) {
    return;
  }

  resource->setManager( this );
  if ( resource->open() ) {
    resource->load();
  }
}

void KNotesResourceManager::resourceModified( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource modified:" << resource->resourceName();
}

void KNotesResourceManager::resourceDeleted( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource deleted:" << resource->resourceName();
}

#include "resourcemanager.moc"
<commit_msg>const'ify<commit_after>/*******************************************************************
 This file is part of KNotes.

 Copyright (c) 2004, Bo Thorsen <bo@sonofthor.dk>
               2004, Michael Brade <brade@kde.org>

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 MA  02110-1301, USA.

 In addition, as a special exception, the copyright holders give
 permission to link the code of this program with any edition of
 the Qt library by Trolltech AS, Norway (or with modified versions
 of Qt that use the same license as Qt), and distribute linked
 combinations including the two.  You must obey the GNU General
 Public License in all respects for all of the code used other than
 Qt.  If you modify this file, you may extend this exception to
 your version of the file, but you are not obligated to do so.  If
 you do not wish to do so, delete this exception statement from
 your version.
*******************************************************************/

#include "knotes/resourcelocal.h"
#include "knotes/resourcemanager.h"

#include <kcal/journal.h>

KNotesResourceManager::KNotesResourceManager()
  : QObject( 0 )
{
  setObjectName( "KNotes Resource Manager" );
  m_manager = new KRES::Manager<ResourceNotes>( "notes" );
  m_manager->addObserver( this );
  m_manager->readConfig();
}

KNotesResourceManager::~KNotesResourceManager()
{
  delete m_manager;
}

void KNotesResourceManager::load()
{
  if ( !m_manager->standardResource() ) {
    kWarning( 5500 ) << "No standard resource yet.";
    ResourceNotes *resource = new ResourceLocal();
    m_manager->add( resource );
    m_manager->setStandardResource( resource );
  }

  // Open all active resources
  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    kDebug( 5500 ) << "Opening resource " + ( *it )->resourceName();
    ( *it )->setManager( this );
    if ( ( *it )->open() ) {
      ( *it )->load();
    }
  }
}

void KNotesResourceManager::save()
{
  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    ( *it )->save();
  }
}

// when adding a new note, make sure a config file exists!!

void KNotesResourceManager::addNewNote( KCal::Journal *journal )
{
  // TODO: Make this configurable
  ResourceNotes *resource = m_manager->standardResource();
  if ( resource ) {
    resource->addNote( journal );
    registerNote( resource, journal );
  } else {
    kWarning( 5500 ) << "no resource!";
  }
}

void KNotesResourceManager::registerNote( ResourceNotes *resource,
                                          KCal::Journal *journal )
{
  // TODO: only emit the signal if the journal is new?
  m_resourceMap.insert( journal->uid(), resource );
  emit sigRegisteredNote( journal );
}

void KNotesResourceManager::deleteNote( KCal::Journal *journal )
{
  QString uid = journal->uid();

  // Remove the journal from the resource it came from
  m_resourceMap[ uid ]->deleteNote( journal );
  m_resourceMap.remove( uid );

  // libkcal does not delete the journal immediately, therefore it is ok to
  // emit the journal here
  emit sigDeregisteredNote( journal );
}

KCal::Alarm::List KNotesResourceManager::alarms( const KDateTime &from,
                                                 const KDateTime &to )
{
  KCal::Alarm::List result;

  KRES::Manager<ResourceNotes>::ActiveIterator it;
  for ( it = m_manager->activeBegin(); it != m_manager->activeEnd(); ++it ) {
    KCal::Alarm::List list = ( *it )->alarms( from, to );
    KCal::Alarm::List::ConstIterator it;
    for ( it = list.constBegin(); it != list.constEnd(); ++it ) {
      result.append( *it );
    }
  }

  return result;
}

void KNotesResourceManager::resourceAdded( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource added:" << resource->resourceName();

  if ( !resource->isActive() ) {
    return;
  }

  resource->setManager( this );
  if ( resource->open() ) {
    resource->load();
  }
}

void KNotesResourceManager::resourceModified( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource modified:" << resource->resourceName();
}

void KNotesResourceManager::resourceDeleted( ResourceNotes *resource )
{
  kDebug( 5500 ) << "Resource deleted:" << resource->resourceName();
}

#include "resourcemanager.moc"
<|endoftext|>


<commit_before>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.h"

#include "grpcpp/generic/generic_stub.h"
#include "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_client_cq_tag.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_state.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/protobuf/eager_service.pb.h"
#include "tensorflow/core/util/env_var.h"

namespace tensorflow {
namespace eager {
namespace {

/*
 * Setting environment variable "TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE" to
 * true will turn on asynchronous execution of remote op. It means that when
 * executing an op on a remote worker, client will not block on waiting
 * for the response anymore. Using follow code as example:
 *
 * with tf.device('worker:0'):
 *   a = tf.matmul(...)
 *   b = tf.matmul(...)
 * logging.into('Requests sent')    # Probably not executed yet
 * logging.info('b: %s', b.numpy()) # Block until 'b' finished.
 *
 * Streaming RPC will preserve order as well. So 'a' must be executed before
 * 'b' on 'worker:0'.
 *
 * When turning on this feature, you should explicitly wait for some result
 * from remote workers at the end of you python program. Otherwise, client may
 * shutdown remote workers without waiting all pending ops.
 *
 * TODO(fishx): When exiting client, make sure all pending ops on remote workers
 * are finished.
 *
 * TODO(b/139210648): Move this comment to eager/execute.py when this feature is
 * on by default.
 */
bool EnableStreaming() {
  bool result;
  TF_CHECK_OK(ReadBoolFromEnvVar("TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE",
                                 true, &result));
  return result;
}

// Ref-counted thread to handle callbacks for completed requests a GRPC
// completion queue. The thread might be shared by multiple eager clients, and
// each one of them should hold a reference count to ensure that the thread
// outlives the clients.
// To ensure that every tag in completion queue is processed, this thread also
// holds a reference to itself and always wait until ref count is one to exit.
class GrpcEagerClientThread : public core::RefCounted {
 public:
  GrpcEagerClientThread() {
    // Hold a reference to ensure every completion tag gets processed.
    Ref();
    thread_.reset(Env::Default()->StartThread(
        ThreadOptions(), "eager_client_thread", [this]() {
          void* tag;
          bool ok;
          while (completion_queue_.Next(&tag, &ok)) {
            VLOG(4) << "GrpcEagerClientThread got next tag";
            GrpcClientCQTag* callback_tag = static_cast<GrpcClientCQTag*>(tag);
            callback_tag->OnCompleted(ok);
            VLOG(4) << "GrpcEagerClientThread blocking for next tag";
            if (RefCountIsOne()) {
              break;
            }
          }
          VLOG(4) << "GrpcEagerClientThread exiting";
          completion_queue_.Shutdown();
          // `this` holds the final reference so cannot directly Unref here.
          // Instead, schedule a separate thread to clean it up.
          Env::Default()->SchedClosure([this]() { this->Unref(); });
        }));
  }

  ~GrpcEagerClientThread() override {}

  ::grpc::CompletionQueue* completion_queue() { return &completion_queue_; }

 private:
  ::grpc::CompletionQueue completion_queue_;
  std::unique_ptr<Thread> thread_;
};

class GrpcEagerClient : public EagerClient {
 public:
  GrpcEagerClient(const tensorflow::SharedGrpcChannelPtr& channel,
                  GrpcEagerClientThread* thread)
      : stub_(channel), thread_(thread) {
    // Hold a reference to make sure the corresponding EagerClientThread
    // outlives the client.
    thread_->Ref();
    cq_ = thread->completion_queue();
  }
  ~GrpcEagerClient() override { thread_->Unref(); }

#define CLIENT_METHOD(method)                                             \
  void method##Async(const method##Request* request,                      \
                     method##Response* response, StatusCallback done)     \
      override {                                                          \
    StatusCallback done_wrapped = callback_wrapper(std::move(done));      \
    new RPCState<protobuf::Message>(                                      \
        &stub_, cq_, "/tensorflow.eager.EagerService/" #method, *request, \
        response, std::move(done_wrapped), /*call_opts=*/nullptr,         \
        /*threadpool=*/nullptr, /*max_retries=*/0, /*fail_fast=*/true);   \
  }

  CLIENT_METHOD(CreateContext);
  CLIENT_METHOD(UpdateContext);
  CLIENT_METHOD(Enqueue);
  CLIENT_METHOD(WaitQueueDone);
  CLIENT_METHOD(KeepAlive);

#undef CLIENT_METHOD

  void CloseContextAsync(const CloseContextRequest* request,
                         CloseContextResponse* response,
                         StatusCallback done) override {
    StatusCallback done_wrapped = callback_wrapper(std::move(done));
    new RPCState<protobuf::Message>(
        &stub_, cq_, "/tensorflow.eager.EagerService/CloseContext", *request,
        response, std::move(done_wrapped), /*call_opts=*/nullptr,
        /*threadpool=*/nullptr);

    VLOG(1) << "Sending RPC to close remote eager context "
            << request->DebugString();

    mutex_lock l(mu_);
    const auto& it = enqueue_dispatchers_.find(request->context_id());
    if (it != enqueue_dispatchers_.end()) {
      it->second.CancelCall();
      enqueue_dispatchers_.erase(it);
    } else if (EnableStreaming()) {
      LOG(ERROR) << "Remote EagerContext with id " << request->context_id()
                 << " does not seem to exist.";
    }
  }

  void StreamingEnqueueAsync(const EnqueueRequest* request,
                             EnqueueResponse* response,
                             StatusCallback done) override {
    StatusCallback done_wrapped = callback_wrapper(std::move(done));
    if (EnableStreaming()) {
      tf_shared_lock l(mu_);
      auto it = enqueue_dispatchers_.find(request->context_id());
      if (it == enqueue_dispatchers_.end()) {
        auto it_and_bool = enqueue_dispatchers_.emplace(
            std::piecewise_construct,
            std::forward_as_tuple(request->context_id()),
            std::forward_as_tuple(
                &stub_, cq_,
                "/tensorflow.eager.EagerService/StreamingEnqueue"));
        it = it_and_bool.first;
      }
      it->second.SendNextRequest(*request, response, std::move(done_wrapped));
    } else {
      Notification n;
      Status status;
      EnqueueAsync(request, response, [&n, &status](const Status& s) {
        status.Update(s);
        n.Notify();
      });
      n.WaitForNotification();
      done_wrapped(status);
    }
  }

 private:
  ::grpc::GenericStub stub_;
  const GrpcEagerClientThread* thread_;

  ::grpc::CompletionQueue* cq_;

  mutable mutex mu_;

  std::unordered_map<uint64, StreamingRPCDispatcher<EnqueueResponse>>
      enqueue_dispatchers_ GUARDED_BY(mu_);

  StatusCallback callback_wrapper(StatusCallback done) {
    Ref();
    return [this, done = std::move(done)](const Status& status) {
      done(status);
      this->Unref();
    };
  }
};

class GrpcEagerClientCache : public EagerClientCache {
 public:
  explicit GrpcEagerClientCache(
      std::shared_ptr<tensorflow::GrpcChannelCache> cache)
      : next_round_robin_assignment_(0), cache_(cache), threads_(4) {
    for (int i = 0; i < threads_.size(); i++) {
      threads_[i].reset(new GrpcEagerClientThread());
    }
  }

  ~GrpcEagerClientCache() override { threads_.clear(); }

  Status GetClient(const string& target,
                   core::RefCountPtr<EagerClient>* client) override {
    auto it = clients_.find(target);
    if (it == clients_.end()) {
      tensorflow::SharedGrpcChannelPtr shared =
          cache_->FindWorkerChannel(target);
      if (shared == nullptr) {
        return errors::InvalidArgument("Client for target ", target,
                                       " not found.");
      }
      int assigned_index = AssignClientToThread(target);
      GrpcEagerClientThread* thread = threads_[assigned_index].get();
      auto worker = new GrpcEagerClient(shared, thread);
      it = clients_.emplace(target, worker).first;
    }

    it->second->Ref();
    client->reset(it->second.get());
    return Status::OK();
  }

 private:
  mutex assignment_mu_;
  std::unordered_map<std::string, size_t> target_assignments_
      GUARDED_BY(assignment_mu_);
  size_t next_round_robin_assignment_ GUARDED_BY(assignment_mu_);

  size_t AssignClientToThread(const string& target) {
    // Round-robin target assignment, but keeps the same target on the same
    // polling thread always, as this is important for gRPC performance
    mutex_lock lock(assignment_mu_);
    auto it = target_assignments_.find(target);
    if (it == target_assignments_.end()) {
      it = target_assignments_
               .insert(std::make_pair(
                   target, (next_round_robin_assignment_++) % threads_.size()))
               .first;
    }
    return it->second;
  }

  std::shared_ptr<tensorflow::GrpcChannelCache> cache_;
  std::unordered_map<string, core::RefCountPtr<EagerClient>> clients_;
  std::vector<core::RefCountPtr<GrpcEagerClientThread>> threads_;
};

}  // namespace

EagerClientCache* NewGrpcEagerClientCache(
    std::shared_ptr<tensorflow::GrpcChannelCache> channel) {
  return new GrpcEagerClientCache(channel);
}

}  // namespace eager
}  // namespace tensorflow
<commit_msg>Fix Kokoro build.<commit_after>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.h"

#include "grpcpp/generic/generic_stub.h"
#include "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_client_cq_tag.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_state.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/protobuf/eager_service.pb.h"
#include "tensorflow/core/util/env_var.h"

namespace tensorflow {
namespace eager {
namespace {

/*
 * Setting environment variable "TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE" to
 * true will turn on asynchronous execution of remote op. It means that when
 * executing an op on a remote worker, client will not block on waiting
 * for the response anymore. Using follow code as example:
 *
 * with tf.device('worker:0'):
 *   a = tf.matmul(...)
 *   b = tf.matmul(...)
 * logging.into('Requests sent')    # Probably not executed yet
 * logging.info('b: %s', b.numpy()) # Block until 'b' finished.
 *
 * Streaming RPC will preserve order as well. So 'a' must be executed before
 * 'b' on 'worker:0'.
 *
 * When turning on this feature, you should explicitly wait for some result
 * from remote workers at the end of you python program. Otherwise, client may
 * shutdown remote workers without waiting all pending ops.
 *
 * TODO(fishx): When exiting client, make sure all pending ops on remote workers
 * are finished.
 *
 * TODO(b/139210648): Move this comment to eager/execute.py when this feature is
 * on by default.
 */
bool EnableStreaming() {
  bool result;
  TF_CHECK_OK(ReadBoolFromEnvVar("TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE",
                                 true, &result));
  return result;
}

// Ref-counted thread to handle callbacks for completed requests a GRPC
// completion queue. The thread might be shared by multiple eager clients, and
// each one of them should hold a reference count to ensure that the thread
// outlives the clients.
// To ensure that every tag in completion queue is processed, this thread also
// holds a reference to itself and always wait until ref count is one to exit.
class GrpcEagerClientThread : public core::RefCounted {
 public:
  GrpcEagerClientThread() {
    // Hold a reference to ensure every completion tag gets processed.
    Ref();
    thread_.reset(Env::Default()->StartThread(
        ThreadOptions(), "eager_client_thread", [this]() {
          void* tag;
          bool ok;
          while (completion_queue_.Next(&tag, &ok)) {
            VLOG(4) << "GrpcEagerClientThread got next tag";
            GrpcClientCQTag* callback_tag = static_cast<GrpcClientCQTag*>(tag);
            callback_tag->OnCompleted(ok);
            VLOG(4) << "GrpcEagerClientThread blocking for next tag";
            if (RefCountIsOne()) {
              break;
            }
          }
          VLOG(4) << "GrpcEagerClientThread exiting";
          completion_queue_.Shutdown();
          // `this` holds the final reference so cannot directly Unref here.
          // Instead, schedule a separate thread to clean it up.
          Env::Default()->SchedClosure([this]() { this->Unref(); });
        }));
  }

  ~GrpcEagerClientThread() override {}

  ::grpc::CompletionQueue* completion_queue() { return &completion_queue_; }

 private:
  ::grpc::CompletionQueue completion_queue_;
  std::unique_ptr<Thread> thread_;
};

class GrpcEagerClient : public EagerClient {
 public:
  GrpcEagerClient(const tensorflow::SharedGrpcChannelPtr& channel,
                  GrpcEagerClientThread* thread)
      : stub_(channel), thread_(thread) {
    // Hold a reference to make sure the corresponding EagerClientThread
    // outlives the client.
    thread_->Ref();
    cq_ = thread->completion_queue();
  }
  ~GrpcEagerClient() override { thread_->Unref(); }

#define CLIENT_METHOD(method)                                             \
  void method##Async(const method##Request* request,                      \
                     method##Response* response, StatusCallback done)     \
      override {                                                          \
    StatusCallback done_wrapped = callback_wrapper(std::move(done));      \
    new RPCState<protobuf::Message>(                                      \
        &stub_, cq_, "/tensorflow.eager.EagerService/" #method, *request, \
        response, std::move(done_wrapped), /*call_opts=*/nullptr,         \
        /*threadpool=*/nullptr, /*max_retries=*/0, /*fail_fast=*/true);   \
  }

  CLIENT_METHOD(CreateContext);
  CLIENT_METHOD(UpdateContext);
  CLIENT_METHOD(Enqueue);
  CLIENT_METHOD(WaitQueueDone);
  CLIENT_METHOD(KeepAlive);

#undef CLIENT_METHOD

  void CloseContextAsync(const CloseContextRequest* request,
                         CloseContextResponse* response,
                         StatusCallback done) override {
    StatusCallback done_wrapped = callback_wrapper(std::move(done));
    new RPCState<protobuf::Message>(
        &stub_, cq_, "/tensorflow.eager.EagerService/CloseContext", *request,
        response, std::move(done_wrapped), /*call_opts=*/nullptr,
        /*threadpool=*/nullptr);

    VLOG(1) << "Sending RPC to close remote eager context "
            << request->DebugString();

    mutex_lock l(mu_);
    const auto& it = enqueue_dispatchers_.find(request->context_id());
    if (it != enqueue_dispatchers_.end()) {
      it->second.CancelCall();
      enqueue_dispatchers_.erase(it);
    } else if (EnableStreaming()) {
      LOG(ERROR) << "Remote EagerContext with id " << request->context_id()
                 << " does not seem to exist.";
    }
  }

  void StreamingEnqueueAsync(const EnqueueRequest* request,
                             EnqueueResponse* response,
                             StatusCallback done) override {
    StatusCallback done_wrapped = callback_wrapper(std::move(done));
    if (EnableStreaming()) {
      tf_shared_lock l(mu_);
      auto it = enqueue_dispatchers_.find(request->context_id());
      if (it == enqueue_dispatchers_.end()) {
        auto it_and_bool = enqueue_dispatchers_.emplace(
            std::piecewise_construct,
            std::forward_as_tuple(request->context_id()),
            std::forward_as_tuple(
                &stub_, cq_,
                "/tensorflow.eager.EagerService/StreamingEnqueue"));
        it = it_and_bool.first;
      }
      it->second.SendNextRequest(*request, response, std::move(done_wrapped));
    } else {
      Notification n;
      Status status;
      EnqueueAsync(request, response, [&n, &status](const Status& s) {
        status.Update(s);
        n.Notify();
      });
      n.WaitForNotification();
      done_wrapped(status);
    }
  }

 private:
  ::grpc::GenericStub stub_;
  const GrpcEagerClientThread* thread_;

  ::grpc::CompletionQueue* cq_;

  mutable mutex mu_;

  std::unordered_map<uint64, StreamingRPCDispatcher<EnqueueResponse>>
      enqueue_dispatchers_ GUARDED_BY(mu_);

  StatusCallback callback_wrapper(StatusCallback done) {
    Ref();
    return [this, done = std::move(done)](const Status& status) {
      done(status);
      this->Unref();
    };
  }
};

class GrpcEagerClientCache : public EagerClientCache {
 public:
  explicit GrpcEagerClientCache(
      std::shared_ptr<tensorflow::GrpcChannelCache> cache)
      : next_round_robin_assignment_(0), cache_(cache), threads_(4) {
    for (int i = 0; i < threads_.size(); i++) {
      threads_[i].reset(new GrpcEagerClientThread());
    }
  }

  ~GrpcEagerClientCache() override { threads_.clear(); }

  Status GetClient(const string& target,
                   core::RefCountPtr<EagerClient>* client) override {
    auto it = clients_.find(target);
    if (it == clients_.end()) {
      tensorflow::SharedGrpcChannelPtr shared =
          cache_->FindWorkerChannel(target);
      if (shared == nullptr) {
        return errors::InvalidArgument("Client for target ", target,
                                       " not found.");
      }
      int assigned_index = AssignClientToThread(target);
      GrpcEagerClientThread* thread = threads_[assigned_index].get();
      core::RefCountPtr<EagerClient> worker(
          new GrpcEagerClient(shared, thread));
      it = clients_.emplace(target, std::move(worker)).first;
    }

    it->second->Ref();
    client->reset(it->second.get());
    return Status::OK();
  }

 private:
  mutex assignment_mu_;
  std::unordered_map<std::string, size_t> target_assignments_
      GUARDED_BY(assignment_mu_);
  size_t next_round_robin_assignment_ GUARDED_BY(assignment_mu_);

  size_t AssignClientToThread(const string& target) {
    // Round-robin target assignment, but keeps the same target on the same
    // polling thread always, as this is important for gRPC performance
    mutex_lock lock(assignment_mu_);
    auto it = target_assignments_.find(target);
    if (it == target_assignments_.end()) {
      it = target_assignments_
               .insert(std::make_pair(
                   target, (next_round_robin_assignment_++) % threads_.size()))
               .first;
    }
    return it->second;
  }

  std::shared_ptr<tensorflow::GrpcChannelCache> cache_;
  std::unordered_map<string, core::RefCountPtr<EagerClient>> clients_;
  std::vector<core::RefCountPtr<GrpcEagerClientThread>> threads_;
};

}  // namespace

EagerClientCache* NewGrpcEagerClientCache(
    std::shared_ptr<tensorflow::GrpcChannelCache> channel) {
  return new GrpcEagerClientCache(channel);
}

}  // namespace eager
}  // namespace tensorflow
<|endoftext|>