avoid deadlocks by no longer returning read guard

1 files changed, 31 insertions, 16 deletions

diff --git a/incria/src/lazy_mapping.rs b/incria/src/lazy_mapping.rs
index bef3a3d..27bd76e 100644
--- a/incria/src/lazy_mapping.rs
+++ b/incria/src/lazy_mapping.rs
@@ -1,5 +1,8 @@
-use std::{collections::HashMap, future::Future, hash::Hash, sync::Arc};
+use std::{
+    collections::HashMap, fmt::Debug, future::Future, hash::Hash, mem::transmute, sync::Arc,
+};
 
+use blink_alloc::{Blink, SyncBlinkAlloc};
 use tokio::sync::{Mutex, Notify, RwLock, RwLockReadGuard};
 
 use crate::{
@@ -12,13 +15,15 @@ pub trait LazyMappingBackend<K, V> {
     fn compute(&self, key: K) -> impl Future<Output = V> + Send + '_;
 }
 
-#[derive(Debug, Default)]
-pub struct LazyMapper<K, V, B> {
+#[derive(Default)]
+pub struct LazyMapper<K, V: 'static, B> {
     /// The backend we use to compute values and check dirtiness
     backend: B,
 
+    values: SyncBlinkAlloc,
+
     /// Map of all values we've calculated already
-    calculated: RwLock<HashMap<K, (V, NodeId)>>,
+    calculated: RwLock<HashMap<K, (&'static V, NodeId)>>,
 
     /// Values we're currently calculating
     /// The Notify will be triggered when computation is done.
@@ -34,6 +39,7 @@ impl<K: Clone + Eq + Hash + Send + Sync, V: Send + Sync, B: LazyMappingBackend<K
             backend,
             calculated: RwLock::default(),
             waiting: Mutex::default(),
+            values: SyncBlinkAlloc::new(),
         }
     }
 
@@ -50,24 +56,24 @@ impl<K: Clone + Eq + Hash + Send + Sync, V: Send + Sync, B: LazyMappingBackend<K
 
 impl<K, V, B> Mapper for LazyMapper<K, V, B>
 where
-    K: Clone + Eq + Hash + Send + Sync,
-    V: Send + Sync,
+    K: Clone + Eq + Hash + Send + Sync + Debug,
+    V: Send + Sync + 'static,
     B: LazyMappingBackend<K, V>,
 {
     type Key = K;
     type Value = V;
-    type Wrapper<'a> = RwLockReadGuard<'a, V> where V: 'a;
+    type Wrapper<'a> = &'a V where V: 'a;
 
     async fn get<'a>(&'a self, key: &Self::Key) -> Self::Wrapper<'a> {
         // Attempt to reuse or evict the existing value
         let mut reuse_dep_id = None;
         {
             let finished = self.calculated.read().await;
-            if let Some((_, node)) = finished.get(key) {
+            if let Some((val, node)) = finished.get(key) {
                 let dirty = deps::is_dirty(*node) || self.backend.is_dirty(key).await;
                 if !dirty {
                     deps::add_dep(*node);
-                    return RwLockReadGuard::map(finished, |hm| &hm.get(key).unwrap().0);
+                    return val;
                 } else {
                     reuse_dep_id = Some(*node);
                     drop(finished);
@@ -85,10 +91,12 @@ where
             // Waiting for completion
             barrier.notified().await;
 
-            let val = RwLockReadGuard::map(self.calculated.read().await, |hm| hm.get(key).unwrap());
-            deps::add_dep(val.1);
+            let (value, node_id) =
+                *RwLockReadGuard::map(self.calculated.read().await, |hm| hm.get(key).unwrap());
+
+            deps::add_dep(node_id);
 
-            return RwLockReadGuard::map(val, |inf| &inf.0);
+            value
         } else {
             // Needs calculated
             let notify = Arc::new(Notify::new());
@@ -105,19 +113,26 @@ where
             };
 
             let val = deps::with_node_id(dep, self.backend.compute(key.clone())).await;
+            let val_ref: &'static V = unsafe { transmute(Blink::new_in(&self.values).put(val)) };
             deps::add_dep(dep);
 
             self.calculated
                 .write()
                 .await
-                .insert(key.clone(), (val, dep));
+                .insert(key.clone(), (val_ref, dep));
             self.waiting.lock().await.remove(key);
 
             notify.notify_waiters();
 
-            return RwLockReadGuard::map(self.calculated.read().await, |hm| {
-                &hm.get(key).unwrap().0
-            });
+            val_ref
         }
     }
 }
+
+impl<K: Debug, V: Debug, B: Debug> Debug for LazyMapper<K, V, B> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LazyMapper")
+            .field("backend", &self.backend)
+            .finish_non_exhaustive()
+    }
+}