diff options
author | tcmal <tcmal> | 2023-05-24 18:05:53 +0000 |
---|---|---|
committer | Aria <me@aria.rip> | 2023-10-01 17:31:29 +0100 |
commit | f31fcefd85cb69f0c7aa8922f3a579a86eac415d (patch) | |
tree | e0019fbf3383a2cc638d7d0365263c0ae9047347 /incria | |
parent | 3903363da90e127454365695364b778f45c59199 (diff) |
initial commit
Diffstat (limited to 'incria')
-rw-r--r-- | incria/Cargo.lock | 392 | ||||
-rw-r--r-- | incria/Cargo.toml | 9 | ||||
-rw-r--r-- | incria/examples/spreadsheet.rs | 52 | ||||
-rw-r--r-- | incria/src/deps.rs | 182 | ||||
-rw-r--r-- | incria/src/lib.rs | 7 | ||||
-rw-r--r-- | incria/src/mapping.rs | 25 | ||||
-rw-r--r-- | incria/src/thunk.rs | 104 |
7 files changed, 771 insertions, 0 deletions
diff --git a/incria/Cargo.lock b/incria/Cargo.lock new file mode 100644 index 0000000..d815c99 --- /dev/null +++ b/incria/Cargo.lock @@ -0,0 +1,392 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "autocfg" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bytes" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "hermit-abi" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee512640fe35acbfb4bb779db6f0d80704c2cacfa2e39b601ef3e3f47d1ae4c7" +dependencies = [ + "libc", +] + +[[package]] +name = "inc" +version = "0.1.0" +dependencies = [ + "once_cell", + "parking_lot", + "tokio", +] + +[[package]] +name = "libc" +version = "0.2.144" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b00cc1c228a6782d0f076e7b232802e0c5689d41bb5df366f2a6b6621cfdfe1" + +[[package]] +name = "lock_api" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "435011366fe56583b16cf956f9df0095b405b82d76425bc8981c0e22e60ec4df" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "mio" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b9d9a46eff5b4ff64b45a9e316a6d1e0bc719ef429cbec4dc630684212bfdf9" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys 0.45.0", +] + +[[package]] +name = "num_cpus" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "once_cell" +version = "1.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7e5500299e16ebb147ae15a00a942af264cf3688f47923b8fc2cd5858f23ad3" + +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9069cbb9f99e3a5083476ccb29ceb1de18b9118cafa53e90c9551235de2b9521" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-sys 0.45.0", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" + +[[package]] +name = "proc-macro2" +version = "1.0.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa1fb82fc0c281dd9671101b66b771ebbe1eaf967b96ac8740dcba4b70005ca8" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f4f29d145265ec1c483c7c654450edde0bfe043d3938d6972630663356d9500" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + +[[package]] +name = "smallvec" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" + +[[package]] +name = "socket2" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64a4a911eed85daf18834cfaa86a79b7d266ff93ff5ba14005426219480ed662" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "syn" +version = "2.0.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6f671d4b5ffdb8eadec19c0ae67fe2639df8684bd7bc4b83d986b8db549cf01" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tokio" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa32867d44e6f2ce3385e89dceb990188b8bb0fb25b0cf576647a6f98ac5105" +dependencies = [ + "autocfg", + "bytes", + "libc", + "mio", + "num_cpus", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.48.0", +] + +[[package]] +name = "tokio-macros" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5464a87b239f13a63a501f2701565754bae92d243d4bb7eb12f6d57d2269bf4" + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", +] + +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets 0.48.0", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + +[[package]] +name = "windows-targets" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b1eb6f0cd7c80c79759c929114ef071b87354ce476d9d94271031c0497adfd5" +dependencies = [ + "windows_aarch64_gnullvm 0.48.0", + "windows_aarch64_msvc 0.48.0", + "windows_i686_gnu 0.48.0", + "windows_i686_msvc 0.48.0", + "windows_x86_64_gnu 0.48.0", + "windows_x86_64_gnullvm 0.48.0", + "windows_x86_64_msvc 0.48.0", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91ae572e1b79dba883e0d315474df7305d12f569b400fcf90581b06062f7e1bc" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2ef27e0d7bdfcfc7b868b317c1d32c641a6fe4629c171b8928c7b08d98d7cf3" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_gnu" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622a1962a7db830d6fd0a69683c80a18fda201879f0f447f065a3b7467daa241" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_i686_msvc" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4542c6e364ce21bf45d69fdd2a8e455fa38d316158cfd43b3ac1c5b1b19f8e00" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2b8a661f7628cbd23440e50b05d705db3686f894fc9580820623656af974b1" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7896dbc1f41e08872e9d5e8f8baa8fdd2677f29468c4e156210174edc7f7b953" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a515f5799fe4961cb532f983ce2b23082366b898e52ffbce459c86f67c8378a" diff --git a/incria/Cargo.toml b/incria/Cargo.toml new file mode 100644 index 0000000..64fe414 --- /dev/null +++ b/incria/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "inc" +version = "0.1.0" +edition = "2021" + +[dependencies] +once_cell = "1.17.1" +parking_lot = "0.12.1" +tokio = { version = "1.28.1", features = ["full"] } diff --git a/incria/examples/spreadsheet.rs b/incria/examples/spreadsheet.rs new file mode 100644 index 0000000..40b1d5c --- /dev/null +++ b/incria/examples/spreadsheet.rs @@ -0,0 +1,52 @@ +use std::{future::Future, pin::Pin}; + +use inc::{ + deps, + thunk::{Thunk, ThunkMapper}, + Mapper, +}; +use once_cell::sync::OnceCell; + +type Key = (usize, usize); +type Value = usize; + +#[derive(Debug, Default)] +struct CellComputer; + +impl Thunk<Key, Value> for CellComputer { + fn compute(&self, key: Key) -> Pin<Box<dyn Future<Output = Value> + Send + '_>> { + Box::pin(async move { + if key.0 > 0 { + let prev1 = *cell_store().get(&(key.0 - 1, key.1 - 1)).await; + let prev2 = *cell_store().get(&(key.0 - 1, key.1)).await; + let prev3 = *cell_store().get(&(key.0 - 1, key.1 + 1)).await; + prev1 + prev2 + prev3 + } else { + key.1 + } + }) + } +} + +type CellStore = ThunkMapper<Key, Value, CellComputer>; +static CELL_STORE: OnceCell<CellStore> = OnceCell::new(); +fn cell_store() -> &'static CellStore { + CELL_STORE.get_or_init(CellStore::default) +} + +const N: usize = 5; + +#[tokio::main] +async fn main() { + deps::with_node_id(deps::next_node_id(), async { + let val = *dbg!(cell_store().get(&(N, N)).await); + // println!("{}", dep_graphviz()); + + deps::mark_dirty(21); + // println!("{}", dep_graphviz()); + + assert_eq!(val, *cell_store().get(&(N, N)).await); + println!("{}", deps::dep_graphviz()); + }) + .await; +} diff --git a/incria/src/deps.rs b/incria/src/deps.rs new file mode 100644 index 0000000..84411c0 --- /dev/null +++ b/incria/src/deps.rs @@ -0,0 +1,182 @@ +//! The main code used for tracking dependencies (mostly) invisibly to the user. +//! +//! Each mapping that could be memoised is represented by a Node. Each Node can have one or more dependencies, forming a directed acyclic graph. +//! +//! [`self::add_dep`] should be called whenever any mapping is retrieved, to record the dependency. +//! +//! [`self::next_node_id`] and [`self::with_node_id`] should be used to create new nodes in the graph when a new mapping is being evaluated. +//! +//! [`self::mark_dirty`] should be called from outside the evaluation of any mapping. In many cases, this requires the dependency ID to be stored externally alongside the key, which can be accomplished with [`self::current_node_id`] +//! +//! [`self::is_dirty`] should be used when a mapping is retrieved to determine if it needs to be re-evaluated. If it does, [`self::clear`] should be used to reset the mapping's dependencies and dirty state. +use std::{ + collections::{HashMap, VecDeque}, + fmt::Write, + future::Future, +}; + +use once_cell::sync::OnceCell; +use parking_lot::Mutex; + +/// Identifier of a node, unique across a program run. +pub type NodeId = usize; + +tokio::task_local! { + /// The ID of the node in the dependency tree for the currently running computation + static NODE_ID: NodeId; +} + +/// Responsible for tracking dependencies. +/// There is one global instance of this, [`DEP_TRACKER`], with convenience functions for most of its functionality. +#[derive(Default)] +struct DepTracker { + /// Neighbour-List representation of dependency tree + deps: Mutex<HashMap<NodeId, NodeInfo>>, + + /// Next node ID to issue + next_node_id: Mutex<NodeId>, +} + +impl DepTracker { + /// See [`self::add_dep`] + fn add_dep(&self, dep: NodeId) { + self.deps + .lock() + .entry(NODE_ID.get()) + .and_modify(|v| v.deps.push(dep)) + .or_insert(NodeInfo { + deps: vec![dep], + dirty: false, + }); + } + + /// See [`self::next_node_id`] + fn next_node_id(&self) -> usize { + let mut lock = self.next_node_id.lock(); + *lock += 1; + *lock - 1 + } + + /// See [`self::mark_dirty`] + fn mark_dirty(&self, node: NodeId) { + let mut lock = self.deps.lock(); + let mut frontier = VecDeque::new(); + frontier.push_back(node); + while let Some(node_id) = frontier.pop_front() { + let node = match lock.get_mut(&node_id) { + Some(x) => x, + None => continue, + }; + + if node.dirty { + continue; + } + node.dirty = true; + + for (dependent, dependent_info) in lock.iter() { + if dependent_info.deps.contains(&node_id) { + frontier.push_back(*dependent); + } + } + } + } + + /// See [`self::is_dirty`] + fn is_dirty(&self, node: NodeId) -> bool { + self.deps + .lock() + .get(&node) + .map(|v| v.dirty) + .unwrap_or(false) + } + + /// See [`self::clear`] + fn clear(&self, node: NodeId) { + let mut lock = self.deps.lock(); + let node = match lock.get_mut(&node) { + Some(x) => x, + None => return, + }; + + node.dirty = false; + node.deps = vec![]; + } +} + +/// Info about a single dependency node +struct NodeInfo { + deps: Vec<NodeId>, + dirty: bool, +} + +/// The global instance of the dependency tracker +static DEP_TRACKER: OnceCell<DepTracker> = OnceCell::new(); +fn dep_tracker() -> &'static DepTracker { + DEP_TRACKER.get_or_init(DepTracker::default) +} + +/// Register a dependency for the current node +pub fn add_dep(dep: NodeId) { + dep_tracker().add_dep(dep); +} + +/// Get the next node id to use +pub fn next_node_id() -> NodeId { + dep_tracker().next_node_id() +} + +/// Get the ID of the current node +pub fn current_node_id() -> NodeId { + NODE_ID.get() +} + +/// Mark the given node and all nodes that depend on it 'dirty'. +/// +/// This should be used in conjunction with [`self::is_dirty`] to recompute values when necessary. +pub fn mark_dirty(dep: NodeId) { + dep_tracker().mark_dirty(dep) +} + +/// Check if the given node is dirty, indicating it should be recomputed. +pub fn is_dirty(dep: NodeId) -> bool { + dep_tracker().is_dirty(dep) +} + +/// Clear all information about a given node. +/// +/// This should be used when a node is recomputed, to ensure stale dependencies are removed. +pub fn clear(dep: NodeId) { + dep_tracker().clear(dep) +} + +/// Generate a graphviz representation of the current dependency graph, for debugging. +/// +/// Dirty nodes will be coloured grey. +pub fn dep_graphviz() -> String { + let mut out = String::new(); + let deps = dep_tracker().deps.lock(); + writeln!(&mut out, "digraph G {{").unwrap(); + for (id, info) in deps.iter() { + if *id == 0 { + writeln!(&mut out, "\t{} [label=Root]", id).unwrap(); + } else if info.dirty { + writeln!(&mut out, "\t{} [style=filled,color=lightgrey]", id).unwrap(); + } + for dep in info.deps.iter() { + writeln!(&mut out, "\t{} -> {}", id, dep).unwrap(); + } + } + + writeln!(&mut out, "}}").unwrap(); + + out +} + +/// Decorate the given future to have the given dependency ID. +pub async fn with_node_id<F>(dep: NodeId, f: F) -> F::Output +where + F: Future + Send, + F::Output: Send, +{ + NODE_ID.scope(dep, f).await +} diff --git a/incria/src/lib.rs b/incria/src/lib.rs new file mode 100644 index 0000000..2f52fbb --- /dev/null +++ b/incria/src/lib.rs @@ -0,0 +1,7 @@ +#![feature(async_fn_in_trait)] + +pub mod deps; +mod mapping; +pub mod thunk; + +pub use mapping::Mapper; diff --git a/incria/src/mapping.rs b/incria/src/mapping.rs new file mode 100644 index 0000000..6cfac67 --- /dev/null +++ b/incria/src/mapping.rs @@ -0,0 +1,25 @@ +use std::ops::Deref; + +/// A mapper of keys to values, that memoises the results and re-evaluates when necessary. +/// +/// This mapping can be: +/// * Storage or access of externally-modifiable values like files or resources shared with other threads (cells) +/// * Computations that are pure, except for their use of cells above (thunks) +/// +/// Mappers should be mainly responsible for tracking and re-performing new computations as appropriate, using the API in [`crate::deps`]. +pub trait Mapper { + /// The key or input type + type Key; + + /// The result + type Value; + + /// A wrapper around the result. + /// This allows things like read guards or Arcs to be returned, so long as they dereference to the correct type. + type Wrapper<'a>: Deref<Target = Self::Value> + 'a + where + Self::Value: 'a; + + /// Get a value by the given key, calculating it if necessary. + async fn get<'a>(&'a self, key: &Self::Key) -> Self::Wrapper<'a>; +} diff --git a/incria/src/thunk.rs b/incria/src/thunk.rs new file mode 100644 index 0000000..6aa44dd --- /dev/null +++ b/incria/src/thunk.rs @@ -0,0 +1,104 @@ +//! A mapper based on a function from keys to values. +use std::{collections::HashMap, hash::Hash, pin::Pin, sync::Arc}; + +use parking_lot::{MappedRwLockReadGuard, Mutex, RwLock, RwLockReadGuard}; +use tokio::sync::Notify; + +use crate::{ + deps::{self, NodeId}, + Mapper, +}; + +/// A mapping that lazily computes values with a given async thunk and memoizes them. +#[derive(Debug, Default)] +pub struct ThunkMapper<K, V, T> { + /// The thunk we use to compute values + thunk: T, + + /// Map of all values we've calculated already + calculated: RwLock<HashMap<K, (V, NodeId)>>, + + /// Values we're currently calculating + /// The Notify will be triggered when computation is done. + waiting: Mutex<HashMap<K, Arc<Notify>>>, +} + +impl<K: Clone + Eq + Hash + Send + Sync, V: Send + Sync, T: Thunk<K, V>> ThunkMapper<K, V, T> { + /// Create a new instance of the computing store. + pub fn new(thunk: T) -> Self { + Self { + thunk, + calculated: RwLock::default(), + waiting: Mutex::default(), + } + } +} + +impl<K: Clone + Eq + Hash + Send + Sync, V: Send + Sync, C: Thunk<K, V>> Mapper + for ThunkMapper<K, V, C> +{ + type Key = K; + type Value = V; + type Wrapper<'a> = MappedRwLockReadGuard<'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(); + if let Some((_, dep)) = finished.get(key) { + if !deps::is_dirty(*dep) { + deps::add_dep(*dep); + return RwLockReadGuard::map(finished, |hm| &hm.get(key).unwrap().0); + } else { + reuse_dep_id = Some(*dep); + drop(finished); + // Dirty, so we'll recompute below but we should remove it now + if self.calculated.write().remove(key).is_none() { + // Someone else already noticed it was dirty and removed it before us, so we need to deal with that + todo!("dirty value removed between us noticing and us doing something") + } + } + } + } + + let barrier = self.waiting.lock().get(key).cloned(); + if let Some(barrier) = barrier { + // Waiting for completion + barrier.notified().await; + + let val = RwLockReadGuard::map(self.calculated.read(), |hm| hm.get(key).unwrap()); + deps::add_dep(val.1); + + return MappedRwLockReadGuard::map(val, |inf| &inf.0); + } else { + // Needs calculated + let notify = Arc::new(Notify::new()); + self.waiting.lock().insert(key.clone(), notify.clone()); + + let dep = if let Some(x) = reuse_dep_id { + deps::clear(x); + x + } else { + deps::next_node_id() + }; + + let val = deps::with_node_id(dep, self.thunk.compute(key.clone())).await; + deps::add_dep(dep); + + self.calculated.write().insert(key.clone(), (val, dep)); + self.waiting.lock().remove(key); + + notify.notify_waiters(); + + return RwLockReadGuard::map(self.calculated.read(), |hm| &hm.get(key).unwrap().0); + } + } +} + +/// A function from keys to values. +/// +/// Should be pure, except for use of other mappings. This ensures recomputation is done when needed. +pub trait Thunk<K, V>: Send + 'static { + fn compute(&self, key: K) -> Pin<Box<dyn std::future::Future<Output = V> + Send + '_>>; +} |