Skip to main content

rspack_cacheable/
context.rs

1use std::{any::Any, ops::Deref, path::Path, ptr::NonNull};
2
3use rkyv::{
4  de::{ErasedPtr, Pooling, PoolingState},
5  ser::{Sharing, sharing::SharingState},
6};
7
8use crate::error::{Error, Result};
9
10const CONTEXT_ADDR: usize = 0;
11unsafe fn default_drop(_: ErasedPtr) {}
12
13pub trait CacheableContext: Any {
14  fn project_root(&self) -> Option<&Path>;
15}
16
17// Implement for unit type for convenience in tests and simple cases
18impl CacheableContext for () {
19  fn project_root(&self) -> Option<&Path> {
20    None
21  }
22}
23
24/// A context wrapper that provides shared context methods
25pub struct ContextGuard<'a> {
26  context: &'a dyn CacheableContext,
27}
28
29impl<'a> Deref for ContextGuard<'a> {
30  type Target = dyn CacheableContext;
31  fn deref(&self) -> &Self::Target {
32    self.context
33  }
34}
35
36impl<'a> ContextGuard<'a> {
37  pub fn new(context: &'a dyn CacheableContext) -> Self {
38    Self { context }
39  }
40
41  pub fn add_to_sharing<S: Sharing<Error> + ?Sized>(&self, sharing: &mut S) -> Result<()> {
42    sharing.start_sharing(CONTEXT_ADDR);
43    sharing.finish_sharing(CONTEXT_ADDR, self as *const _ as usize)
44  }
45
46  pub fn sharing_guard<S: Sharing<Error> + ?Sized>(sharing: &'a mut S) -> Result<&'a Self> {
47    match sharing.start_sharing(CONTEXT_ADDR) {
48      SharingState::Finished(addr) => {
49        let guard: &Self = unsafe { &*(addr as *const Self) };
50        Ok(guard)
51      }
52      _ => Err(Error::NoContext),
53    }
54  }
55
56  pub fn add_to_pooling<P: Pooling<Error> + ?Sized>(&self, pooling: &mut P) -> Result<()> {
57    unsafe {
58      let ctx_ptr = ErasedPtr::new(NonNull::new_unchecked(self as *const _ as *mut ()));
59      pooling.start_pooling(CONTEXT_ADDR);
60      pooling.finish_pooling(CONTEXT_ADDR, ctx_ptr, default_drop)
61    }
62  }
63
64  pub fn pooling_guard<P: Pooling<Error> + ?Sized>(pooling: &'a mut P) -> Result<&'a Self> {
65    match pooling.start_pooling(CONTEXT_ADDR) {
66      PoolingState::Finished(ptr) => {
67        let guard: &Self = unsafe { &*(ptr.data_address() as *const Self) };
68        Ok(guard)
69      }
70      _ => Err(Error::NoContext),
71    }
72  }
73
74  pub fn downcast_context<T: 'static>(&'a self) -> Result<&'a T> {
75    (self.context as &dyn Any)
76      .downcast_ref()
77      .ok_or(Error::NoContext)
78  }
79}