mlua-luau-scheduler/lib/result_map.rs

#![allow(clippy::inline_always)]

use std::{cell::RefCell, rc::Rc};

use event_listener::Event;
// NOTE: This is the hash algorithm that mlua also uses, so we
// are not adding any additional dependencies / bloat by using it.
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{thread_id::ThreadId, util::ThreadResult};

#[derive(Clone)]
pub(crate) struct ThreadResultMap {
    tracked: Rc<RefCell<FxHashSet<ThreadId>>>,
    results: Rc<RefCell<FxHashMap<ThreadId, ThreadResult>>>,
    events: Rc<RefCell<FxHashMap<ThreadId, Rc<Event>>>>,
}

impl ThreadResultMap {
    pub fn new() -> Self {
        Self {
            tracked: Rc::new(RefCell::new(FxHashSet::default())),
            results: Rc::new(RefCell::new(FxHashMap::default())),
            events: Rc::new(RefCell::new(FxHashMap::default())),
        }
    }

    #[inline(always)]
    pub fn track(&self, id: ThreadId) {
        self.tracked.borrow_mut().insert(id);
    }

    #[inline(always)]
    pub fn is_tracked(&self, id: ThreadId) -> bool {
        self.tracked.borrow().contains(&id)
    }

    pub fn insert(&self, id: ThreadId, result: ThreadResult) {
        debug_assert!(self.is_tracked(id), "Thread must be tracked");
        self.results.borrow_mut().insert(id, result);
        if let Some(event) = self.events.borrow_mut().remove(&id) {
            event.notify(usize::MAX);
        }
    }

    pub async fn listen(&self, id: ThreadId) {
        debug_assert!(self.is_tracked(id), "Thread must be tracked");
        if !self.results.borrow().contains_key(&id) {
            let listener = {
                let mut events = self.events.borrow_mut();
                let event = events.entry(id).or_insert_with(|| Rc::new(Event::new()));
                event.listen()
            };
            listener.await;
        }
    }

    pub fn remove(&self, id: ThreadId) -> Option<ThreadResult> {
        let res = self.results.borrow_mut().remove(&id)?;
        self.tracked.borrow_mut().remove(&id);
        self.events.borrow_mut().remove(&id);
        Some(res)
    }
}