From 8289ada6e73eb8dd8b20f3b526a344a1070dc35a Mon Sep 17 00:00:00 2001
From: Adam Reichold <adam.reichold@t-online.de>
Date: Sat, 18 May 2024 11:02:24 +0200
Subject: [PATCH] WIP: Expose current thread pool via the interface of
 `ThreadPool` to allow reifying the ambient capability

---
 rayon-core/src/thread_pool/mod.rs | 15 +++++++++++++++
 tests/with_current.rs             | 32 +++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 tests/with_current.rs

diff --git a/rayon-core/src/thread_pool/mod.rs b/rayon-core/src/thread_pool/mod.rs
index 5ae6e0f60..bcdf78856 100644
--- a/rayon-core/src/thread_pool/mod.rs
+++ b/rayon-core/src/thread_pool/mod.rs
@@ -198,6 +198,21 @@ impl ThreadPool {
         unsafe { broadcast::broadcast_in(op, &self.registry) }
     }
 
+    /// TODO
+    pub fn current() -> Self {
+        Self {
+            registry: Registry::current(),
+        }
+    }
+
+    /// TODO
+    pub fn with_current<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        Registry::with_current(Some(&self.registry), f)
+    }
+
     /// Returns the (current) number of threads in the thread pool.
     ///
     /// # Future compatibility note
diff --git a/tests/with_current.rs b/tests/with_current.rs
new file mode 100644
index 000000000..92925dd87
--- /dev/null
+++ b/tests/with_current.rs
@@ -0,0 +1,32 @@
+use rayon::prelude::*;
+use rayon::{ThreadPool, ThreadPoolBuilder};
+use std::thread;
+
+#[test]
+#[cfg_attr(any(target_os = "emscripten", target_family = "wasm"), ignore)]
+fn with_current() {
+    fn high_priority_work() {
+        assert_eq!(thread::current().name(), Some("high-priority-thread"));
+    }
+
+    fn regular_work(_item: &()) {
+        assert_eq!(thread::current().name(), None);
+    }
+
+    let items = vec![(); 128];
+
+    let default_pool = ThreadPool::current();
+
+    let high_priority_pool = ThreadPoolBuilder::new()
+        .thread_name(|_| "high-priority-thread".to_owned())
+        .build()
+        .unwrap();
+
+    high_priority_pool.in_place_scope(|scope| {
+        scope.spawn(|_| high_priority_work());
+
+        default_pool.with_current(|| {
+            items.par_iter().for_each(|item| regular_work(item));
+        })
+    });
+}