google_threads/env.h

#ifndef THIRD_PARTY_EIGEN3_GOOGLE_THREADS_ENV_H_
#define THIRD_PARTY_EIGEN3_GOOGLE_THREADS_ENV_H_

#include <functional>
#include <string>
#include <utility>

#include "thread/fiber/fiber.h"
#include "thread/thread.h"
#include "thread/thread_options.h"

namespace Eigen {

inline std::string GetGoogleNamePrefix() { return "GoogleEigen_"; }

// Implements the `Environment` interface needed by Eigen's `ThreadPoolTempl`
// class. See Eigen's `ThreadEnvironment.h` for an implementation that uses the
// banned `std::thread` library.
struct GoogleEnvironment {
  struct Task {
    std::function<void()> f;
  };
  struct EnvThread {};
  virtual ~GoogleEnvironment() = default;
  virtual EnvThread* CreateThread(std::function<void()> f) = 0;
  virtual Task CreateTask(std::function<void()> f) = 0;
  virtual void ExecuteTask(const Task& task) = 0;
};

// Uses Google Threads (go/c++-concurrency#threads_fibers) to handle thread
// management.
struct GoogleThreadEnvironment : public GoogleEnvironment {
  // Constructor must start the thread, while destructor must join it.
  struct EnvThread : public GoogleEnvironment::EnvThread {
    EnvThread(std::function<void()> f) : func_thread_(std::move(f)) {
      func_thread_.Start();
    };
    ~EnvThread() { func_thread_.Join(); };
    void OnCancel() {};

   private:
    // Wraps a function in a thread and calls it when the thread starts.
    // Needed to conform to Eigen's `EnvThread` interface.
    class FuncThread : public ::Thread {
     public:
      static ::thread::Options options() {
        ::thread::Options options;
        options.set_joinable(true);
        return options;
      }

      FuncThread(std::function<void()> f)
          : ::Thread(options(), GetGoogleNamePrefix()), f_(std::move(f)) {};
      ~FuncThread() override = default;
      void Run() override { f_(); };

     private:
      std::function<void()> f_;
    };

    // Managed by `EnvThread` internally.
    FuncThread func_thread_;
  };

  EnvThread* CreateThread(std::function<void()> f) {
    return new EnvThread(std::move(f));
  }
  Task CreateTask(std::function<void()> f) { return Task{std::move(f)}; }
  void ExecuteTask(const Task& task) { task.f(); }
};

// Uses Google Fibers (go/fibers) to handle thread management.
struct GoogleFiberEnvironment : public GoogleEnvironment {
  // Constructor must start the thread, while destructor must join it.
  struct EnvThread : public GoogleEnvironment::EnvThread {
    EnvThread(std::function<void()> f) : fiber_(std::move(f)) {};
    ~EnvThread() { fiber_.Join(); };
    void OnCancel() {};

   private:
    ::thread::Fiber fiber_;
  };

  EnvThread* CreateThread(std::function<void()> f) {
    return new EnvThread(std::move(f));
  }
  Task CreateTask(std::function<void()> f) { return Task{std::move(f)}; }
  void ExecuteTask(const Task& task) { task.f(); }
};

}  // namespace Eigen

#endif  // THIRD_PARTY_EIGEN3_GOOGLE_THREADS_ENV_H_