use std::{
ffi::c_void,
future::Future,
sync::{
Arc,
atomic::{AtomicPtr, AtomicU32},
},
};
use tokio::{
task::{JoinHandle, futures::TaskLocalFuture},
task_local,
};
#[derive(Debug)]
pub struct CompilerContext {
dependenc_id_generator: AtomicU32,
exports_info_artifact_ptr: AtomicPtr<c_void>,
}
task_local! {
pub static CURRENT_COMPILER_CONTEXT: Arc<CompilerContext>;
}
#[allow(clippy::new_without_default)]
impl CompilerContext {
pub fn new() -> Self {
Self {
dependenc_id_generator: AtomicU32::new(0),
exports_info_artifact_ptr: AtomicPtr::new(std::ptr::null_mut()),
}
}
pub fn fetch_new_dependency_id(&self) -> u32 {
self
.dependenc_id_generator
.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
}
pub fn dependency_id(&self) -> u32 {
self
.dependenc_id_generator
.load(std::sync::atomic::Ordering::SeqCst)
}
pub fn set_dependency_id(&self, id: u32) {
self
.dependenc_id_generator
.store(id, std::sync::atomic::Ordering::SeqCst);
}
pub fn exports_info_artifact_ptr(&self) -> Option<*mut c_void> {
let ptr = self
.exports_info_artifact_ptr
.load(std::sync::atomic::Ordering::SeqCst);
(!ptr.is_null()).then_some(ptr)
}
pub fn set_exports_info_artifact_ptr(&self, ptr: Option<*mut c_void>) {
self.exports_info_artifact_ptr.store(
ptr.unwrap_or(std::ptr::null_mut()),
std::sync::atomic::Ordering::SeqCst,
);
}
}
pub fn fetch_new_dependency_id() -> u32 {
CURRENT_COMPILER_CONTEXT.get().fetch_new_dependency_id()
}
pub fn get_current_dependency_id() -> u32 {
CURRENT_COMPILER_CONTEXT.get().dependency_id()
}
pub fn set_current_dependency_id(id: u32) {
CURRENT_COMPILER_CONTEXT.get().set_dependency_id(id);
}
pub fn within_compiler_context<F>(
compiler_context: Arc<CompilerContext>,
f: F,
) -> TaskLocalFuture<Arc<CompilerContext>, F>
where
F: Future,
{
CURRENT_COMPILER_CONTEXT.scope(compiler_context, f)
}
pub fn within_compiler_context_sync<F, R>(compiler_context: Arc<CompilerContext>, f: F) -> R
where
F: FnOnce() -> R,
{
CURRENT_COMPILER_CONTEXT.sync_scope(compiler_context, f)
}
pub fn within_compiler_context_for_testing_sync<F, R>(f: F) -> R
where
F: FnOnce() -> R,
{
CURRENT_COMPILER_CONTEXT.sync_scope(Arc::new(CompilerContext::new()), f)
}
pub fn within_compiler_context_for_testing<F>(f: F) -> TaskLocalFuture<Arc<CompilerContext>, F>
where
F: Future,
{
CURRENT_COMPILER_CONTEXT.scope(Arc::new(CompilerContext::new()), f)
}
pub fn spawn_in_compiler_context<F>(future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let compiler_context = CURRENT_COMPILER_CONTEXT.get();
tokio::spawn(CURRENT_COMPILER_CONTEXT.scope(compiler_context, future))
}
pub fn spawn_in_context<F>(future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
match CURRENT_COMPILER_CONTEXT.try_get() {
Ok(compiler_context) => tokio::spawn(CURRENT_COMPILER_CONTEXT.scope(compiler_context, future)),
Err(_) => tokio::spawn(future),
}
}