No public description
PiperOrigin-RevId: 744075909
Change-Id: Ibc63343d33da09049362cd4de5eb5a7a3442e1be
diff --git a/google_threads/env.h b/google_threads/env.h
new file mode 100644
index 0000000..ef516cd
--- /dev/null
+++ b/google_threads/env.h
@@ -0,0 +1,93 @@
+#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_
