causal_hub/types/
cache.rs

1use std::sync::{Arc, RwLock};
2
3use crate::{
4    estimators::{CPDEstimator, ParCPDEstimator},
5    models::HasLabels,
6    types::{Error, Labels, Map, Result, Set},
7};
8
9/// A cache for calling a function with a key and value.
10#[derive(Clone, Debug)]
11pub struct Cache<'a, C, K, V> {
12    call: &'a C,
13    cache: Arc<RwLock<Map<K, V>>>,
14}
15
16impl<'a, E, P> Cache<'a, E, (Vec<usize>, Vec<usize>), P>
17where
18    P: Clone,
19{
20    /// Create a new cache.
21    ///
22    /// # Arguments
23    ///
24    /// * `call` - The function to call.
25    ///
26    /// # Returns
27    ///
28    /// A new cache.
29    ///
30    #[inline]
31    pub fn new(call: &'a E) -> Self {
32        // Create a new cache.
33        let cache = Arc::new(RwLock::new(Map::default()));
34
35        Self { call, cache }
36    }
37}
38
39impl<C, K, V> HasLabels for Cache<'_, C, K, V>
40where
41    C: HasLabels,
42{
43    #[inline]
44    fn labels(&self) -> &Labels {
45        self.call.labels()
46    }
47}
48
49impl<E, P> CPDEstimator<P> for Cache<'_, E, (Vec<usize>, Vec<usize>), P>
50where
51    E: CPDEstimator<P>,
52    P: Clone,
53{
54    fn fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<P> {
55        // Get the key.
56        let key: (Vec<_>, Vec<_>) = (
57            x.into_iter().cloned().collect(),
58            z.into_iter().cloned().collect(),
59        );
60        // Check if the key is in the cache.
61        if let Some(value) = self
62            .cache
63            .read()
64            .map_err(|e| Error::Poison(&e.to_string()))?
65            .get(&key)
66        {
67            // If it is, return the value.
68            return Ok(value.clone());
69        }
70        // If it is not, call the function.
71        let value = self.call.fit(x, z)?;
72        // Insert the value into the cache.
73        self.cache
74            .write()
75            .map_err(|e| Error::Poison(&e.to_string()))?
76            .insert(key, value.clone());
77        // Return the value.
78        Ok(value)
79    }
80}
81
82impl<E, P> ParCPDEstimator<P> for Cache<'_, E, (Vec<usize>, Vec<usize>), P>
83where
84    E: ParCPDEstimator<P>,
85    P: Clone,
86{
87    fn par_fit(&self, x: &Set<usize>, z: &Set<usize>) -> Result<P> {
88        // Get the key.
89        let key: (Vec<_>, Vec<_>) = (
90            x.into_iter().cloned().collect(),
91            z.into_iter().cloned().collect(),
92        );
93        // Check if the key is in the cache.
94        if let Some(value) = self
95            .cache
96            .read()
97            .map_err(|e| Error::Poison(&e.to_string()))?
98            .get(&key)
99        {
100            // If it is, return the value.
101            return Ok(value.clone());
102        }
103        // If it is not, call the function.
104        let value = self.call.par_fit(x, z)?;
105        // Insert the value into the cache.
106        self.cache
107            .write()
108            .map_err(|e| Error::Poison(&e.to_string()))?
109            .insert(key, value.clone());
110        // Return the value.
111        Ok(value)
112    }
113}
causal_hub/types/cache.rs

causal_hub/types/
cache.rs
