From 04c38220cac3ea13557fa18d30aaf6a6dd47d52a Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 18 Nov 2024 16:43:05 +0100 Subject: [PATCH] Move MostlySend, ThreadLocal, FullySend to their own commit --- crates/milli/src/update/new/extract/cache.rs | 2 +- .../milli/src/update/new/extract/documents.rs | 3 +- .../new/extract/faceted/extract_facets.rs | 4 +- .../milli/src/update/new/extract/geo/mod.rs | 3 +- crates/milli/src/update/new/extract/mod.rs | 5 +- .../extract/searchable/extract_word_docids.rs | 4 +- .../src/update/new/extract/searchable/mod.rs | 4 +- .../src/update/new/extract/vectors/mod.rs | 3 +- .../update/new/indexer/document_changes.rs | 175 +----------------- .../update/new/indexer/document_deletion.rs | 6 +- .../update/new/indexer/document_operation.rs | 3 +- crates/milli/src/update/new/indexer/mod.rs | 3 +- .../src/update/new/indexer/partial_dump.rs | 3 +- .../update/new/indexer/update_by_function.rs | 3 +- crates/milli/src/update/new/mod.rs | 1 + crates/milli/src/update/new/thread_local.rs | 174 +++++++++++++++++ 16 files changed, 203 insertions(+), 193 deletions(-) create mode 100644 crates/milli/src/update/new/thread_local.rs diff --git a/crates/milli/src/update/new/extract/cache.rs b/crates/milli/src/update/new/extract/cache.rs index dd43feefb..9c864372d 100644 --- a/crates/milli/src/update/new/extract/cache.rs +++ b/crates/milli/src/update/new/extract/cache.rs @@ -79,7 +79,7 @@ use roaring::RoaringBitmap; use rustc_hash::FxBuildHasher; use crate::update::del_add::{DelAdd, KvWriterDelAdd}; -use crate::update::new::indexer::document_changes::MostlySend; +use crate::update::new::thread_local::MostlySend; use crate::update::new::KvReaderDelAdd; use crate::update::MergeDeladdCboRoaringBitmaps; use crate::{CboRoaringBitmapCodec, Result}; diff --git a/crates/milli/src/update/new/extract/documents.rs b/crates/milli/src/update/new/extract/documents.rs index c0a2e3d6a..23d93a2c2 100644 --- a/crates/milli/src/update/new/extract/documents.rs +++ b/crates/milli/src/update/new/extract/documents.rs @@ -6,8 +6,9 @@ use hashbrown::HashMap; use super::DelAddRoaringBitmap; use crate::update::new::channel::DocumentsSender; use crate::update::new::document::{write_to_obkv, Document as _}; -use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor, FullySend}; +use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor}; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::FullySend; use crate::update::new::DocumentChange; use crate::vector::EmbeddingConfigs; use crate::Result; diff --git a/crates/milli/src/update/new/extract/faceted/extract_facets.rs b/crates/milli/src/update/new/extract/faceted/extract_facets.rs index 14b1b1bdd..acf211d63 100644 --- a/crates/milli/src/update/new/extract/faceted/extract_facets.rs +++ b/crates/milli/src/update/new/extract/faceted/extract_facets.rs @@ -15,10 +15,10 @@ use crate::heed_codec::facet::OrderedF64Codec; use crate::update::del_add::DelAdd; use crate::update::new::channel::FieldIdDocidFacetSender; use crate::update::new::indexer::document_changes::{ - extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext, - Progress, ThreadLocal, + extract, DocumentChangeContext, DocumentChanges, Extractor, IndexingContext, Progress, }; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::{FullySend, ThreadLocal}; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{DocumentId, FieldId, Index, Result, MAX_FACET_VALUE_LENGTH}; diff --git a/crates/milli/src/update/new/extract/geo/mod.rs b/crates/milli/src/update/new/extract/geo/mod.rs index e883a04cc..c3ea76c42 100644 --- a/crates/milli/src/update/new/extract/geo/mod.rs +++ b/crates/milli/src/update/new/extract/geo/mod.rs @@ -11,8 +11,9 @@ use serde_json::Value; use crate::error::GeoError; use crate::update::new::document::Document; -use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor, MostlySend}; +use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor}; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::MostlySend; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{lat_lng_to_xyz, DocumentId, GeoPoint, Index, InternalError, Result}; diff --git a/crates/milli/src/update/new/extract/mod.rs b/crates/milli/src/update/new/extract/mod.rs index 14cfa83cb..3b2bd77ce 100644 --- a/crates/milli/src/update/new/extract/mod.rs +++ b/crates/milli/src/update/new/extract/mod.rs @@ -13,9 +13,8 @@ pub use geo::*; pub use searchable::*; pub use vectors::EmbeddingExtractor; -use super::indexer::document_changes::{ - DocumentChanges, FullySend, IndexingContext, Progress, ThreadLocal, -}; +use super::indexer::document_changes::{DocumentChanges, IndexingContext, Progress}; +use super::thread_local::{FullySend, ThreadLocal}; use crate::update::GrenadParameters; use crate::Result; diff --git a/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs b/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs index dfb55853f..9822570d0 100644 --- a/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs +++ b/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs @@ -11,10 +11,10 @@ use super::tokenize_document::{tokenizer_builder, DocumentTokenizer}; use crate::update::new::extract::cache::BalancedCaches; use crate::update::new::extract::perm_json_p::contained_in; use crate::update::new::indexer::document_changes::{ - extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext, - MostlySend, Progress, ThreadLocal, + extract, DocumentChangeContext, DocumentChanges, Extractor, IndexingContext, Progress, }; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::{FullySend, MostlySend, ThreadLocal}; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{bucketed_position, DocumentId, FieldId, Index, Result, MAX_POSITION_PER_ATTRIBUTE}; diff --git a/crates/milli/src/update/new/extract/searchable/mod.rs b/crates/milli/src/update/new/extract/searchable/mod.rs index 46a05be4e..2a9078d6e 100644 --- a/crates/milli/src/update/new/extract/searchable/mod.rs +++ b/crates/milli/src/update/new/extract/searchable/mod.rs @@ -14,9 +14,9 @@ use tokenize_document::{tokenizer_builder, DocumentTokenizer}; use super::cache::BalancedCaches; use super::DocidsExtractor; use crate::update::new::indexer::document_changes::{ - extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext, - Progress, ThreadLocal, + extract, DocumentChangeContext, DocumentChanges, Extractor, IndexingContext, Progress, }; +use crate::update::new::thread_local::{FullySend, ThreadLocal}; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{Index, Result, MAX_POSITION_PER_ATTRIBUTE}; diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index 2fb717c71..8ac73a8d7 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -8,7 +8,8 @@ use super::cache::DelAddRoaringBitmap; use crate::error::FaultSource; use crate::prompt::Prompt; use crate::update::new::channel::EmbeddingSender; -use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor, MostlySend}; +use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor}; +use crate::update::new::thread_local::MostlySend; use crate::update::new::vector_document::VectorDocument; use crate::update::new::DocumentChange; use crate::vector::error::{ diff --git a/crates/milli/src/update/new/indexer/document_changes.rs b/crates/milli/src/update/new/indexer/document_changes.rs index e4b088f31..308582002 100644 --- a/crates/milli/src/update/new/indexer/document_changes.rs +++ b/crates/milli/src/update/new/indexer/document_changes.rs @@ -8,182 +8,9 @@ use rayon::iter::IndexedParallelIterator; use super::super::document_change::DocumentChange; use crate::fields_ids_map::metadata::FieldIdMapWithMetadata; use crate::update::new::parallel_iterator_ext::ParallelIteratorExt as _; +use crate::update::new::thread_local::{FullySend, MostlySend, ThreadLocal}; use crate::{FieldsIdsMap, GlobalFieldsIdsMap, Index, InternalError, Result}; -/// A trait for types that are **not** [`Send`] only because they would then allow concurrent access to a type that is not [`Sync`]. -/// -/// The primary example of such a type is `&T`, with `T: !Sync`. -/// -/// In the authors' understanding, a type can be `!Send` for two distinct reasons: -/// -/// 1. Because it contains data that *genuinely* cannot be moved between threads, such as thread-local data. -/// 2. Because sending the type would allow concurrent access to a `!Sync` type, which is undefined behavior. -/// -/// `MostlySend` exists to be used in bounds where you need a type whose data is **not** *attached* to a thread -/// because you might access it from a different thread, but where you will never access the type **concurrently** from -/// multiple threads. -/// -/// Like [`Send`], `MostlySend` assumes properties on types that cannot be verified by the compiler, which is why implementing -/// this trait is unsafe. -/// -/// # Safety -/// -/// Implementers of this trait promises that the following properties hold on the implementing type: -/// -/// 1. Its data can be accessed from any thread and will be the same regardless of the thread accessing it. -/// 2. Any operation that can be performed on the type does not depend on the thread that executes it. -/// -/// As these properties are subtle and are not generally tracked by the Rust type system, great care should be taken before -/// implementing `MostlySend` on a type, especially a foreign type. -/// -/// - An example of a type that verifies (1) and (2) is [`std::rc::Rc`] (when `T` is `Send` and `Sync`). -/// - An example of a type that doesn't verify (1) is thread-local data. -/// - An example of a type that doesn't verify (2) is [`std::sync::MutexGuard`]: a lot of mutex implementations require that -/// a lock is returned to the operating system on the same thread that initially locked the mutex, failing to uphold this -/// invariant will cause Undefined Behavior -/// (see last § in [the nomicon](https://doc.rust-lang.org/nomicon/send-and-sync.html)). -/// -/// It is **always safe** to implement this trait on a type that is `Send`, but no placeholder impl is provided due to limitations in -/// coherency. Use the [`FullySend`] wrapper in this situation. -pub unsafe trait MostlySend {} - -#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Ord, PartialOrd, Hash)] -pub struct FullySend(pub T); - -// SAFETY: a type **fully** send is always mostly send as well. -unsafe impl MostlySend for FullySend where T: Send {} - -unsafe impl MostlySend for RefCell where T: MostlySend {} - -unsafe impl MostlySend for Option where T: MostlySend {} - -impl FullySend { - pub fn into(self) -> T { - self.0 - } -} - -impl From for FullySend { - fn from(value: T) -> Self { - Self(value) - } -} - -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] -struct MostlySendWrapper(T); - -impl MostlySendWrapper { - /// # Safety - /// - /// - (P1) Users of this type will never access the type concurrently from multiple threads without synchronization - unsafe fn new(t: T) -> Self { - Self(t) - } - - fn as_ref(&self) -> &T { - &self.0 - } - - fn as_mut(&mut self) -> &mut T { - &mut self.0 - } - - fn into_inner(self) -> T { - self.0 - } -} - -/// # Safety -/// -/// 1. `T` is [`MostlySend`], so by its safety contract it can be accessed by any thread and all of its operations are available -/// from any thread. -/// 2. (P1) of `MostlySendWrapper::new` forces the user to never access the value from multiple threads concurrently. -unsafe impl Send for MostlySendWrapper {} - -/// A wrapper around [`thread_local::ThreadLocal`] that accepts [`MostlySend`] `T`s. -#[derive(Default)] -pub struct ThreadLocal { - inner: thread_local::ThreadLocal>, - // FIXME: this should be necessary - //_no_send: PhantomData<*mut ()>, -} - -impl ThreadLocal { - pub fn new() -> Self { - Self { inner: thread_local::ThreadLocal::new() } - } - - pub fn with_capacity(capacity: usize) -> Self { - Self { inner: thread_local::ThreadLocal::with_capacity(capacity) } - } - - pub fn clear(&mut self) { - self.inner.clear() - } - - pub fn get(&self) -> Option<&T> { - self.inner.get().map(|t| t.as_ref()) - } - - pub fn get_or(&self, create: F) -> &T - where - F: FnOnce() -> T, - { - /// TODO: move ThreadLocal, MostlySend, FullySend to a dedicated file - self.inner.get_or(|| unsafe { MostlySendWrapper::new(create()) }).as_ref() - } - - pub fn get_or_try(&self, create: F) -> std::result::Result<&T, E> - where - F: FnOnce() -> std::result::Result, - { - self.inner - .get_or_try(|| unsafe { Ok(MostlySendWrapper::new(create()?)) }) - .map(MostlySendWrapper::as_ref) - } - - pub fn get_or_default(&self) -> &T - where - T: Default, - { - self.inner.get_or_default().as_ref() - } - - pub fn iter_mut(&mut self) -> IterMut { - IterMut(self.inner.iter_mut()) - } -} - -impl IntoIterator for ThreadLocal { - type Item = T; - - type IntoIter = IntoIter; - - fn into_iter(self) -> Self::IntoIter { - IntoIter(self.inner.into_iter()) - } -} - -pub struct IterMut<'a, T: MostlySend>(thread_local::IterMut<'a, MostlySendWrapper>); - -impl<'a, T: MostlySend> Iterator for IterMut<'a, T> { - type Item = &'a mut T; - - fn next(&mut self) -> Option { - self.0.next().map(|t| t.as_mut()) - } -} - -pub struct IntoIter(thread_local::IntoIter>); - -impl Iterator for IntoIter { - type Item = T; - - fn next(&mut self) -> Option { - self.0.next().map(|t| t.into_inner()) - } -} - pub struct DocumentChangeContext< 'doc, // covariant lifetime of a single `process` call 'extractor: 'doc, // invariant lifetime of the extractor_allocs diff --git a/crates/milli/src/update/new/indexer/document_deletion.rs b/crates/milli/src/update/new/indexer/document_deletion.rs index 353995a59..2e46be63d 100644 --- a/crates/milli/src/update/new/indexer/document_deletion.rs +++ b/crates/milli/src/update/new/indexer/document_deletion.rs @@ -4,8 +4,9 @@ use rayon::iter::IndexedParallelIterator; use rayon::slice::ParallelSlice as _; use roaring::RoaringBitmap; -use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend}; +use super::document_changes::{DocumentChangeContext, DocumentChanges}; use crate::documents::PrimaryKey; +use crate::update::new::thread_local::MostlySend; use crate::update::new::{Deletion, DocumentChange}; use crate::{DocumentId, Result}; @@ -92,9 +93,10 @@ mod test { use crate::fields_ids_map::metadata::{FieldIdMapWithMetadata, MetadataBuilder}; use crate::index::tests::TempIndex; use crate::update::new::indexer::document_changes::{ - extract, DocumentChangeContext, Extractor, IndexingContext, MostlySend, ThreadLocal, + extract, DocumentChangeContext, Extractor, IndexingContext, }; use crate::update::new::indexer::DocumentDeletion; + use crate::update::new::thread_local::{MostlySend, ThreadLocal}; use crate::update::new::DocumentChange; use crate::DocumentId; diff --git a/crates/milli/src/update/new/indexer/document_operation.rs b/crates/milli/src/update/new/indexer/document_operation.rs index 604dd1786..71d410ea6 100644 --- a/crates/milli/src/update/new/indexer/document_operation.rs +++ b/crates/milli/src/update/new/indexer/document_operation.rs @@ -9,10 +9,11 @@ use serde_json::value::RawValue; use serde_json::Deserializer; use super::super::document_change::DocumentChange; -use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend}; +use super::document_changes::{DocumentChangeContext, DocumentChanges}; use super::retrieve_or_guess_primary_key; use crate::documents::PrimaryKey; use crate::update::new::document::Versions; +use crate::update::new::thread_local::MostlySend; use crate::update::new::{Deletion, Insertion, Update}; use crate::update::{AvailableIds, IndexDocumentsMethod}; use crate::{DocumentId, Error, FieldsIdsMap, Index, InternalError, Result, UserError}; diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs index 71fcdd204..0511dd6a1 100644 --- a/crates/milli/src/update/new/indexer/mod.rs +++ b/crates/milli/src/update/new/indexer/mod.rs @@ -3,7 +3,7 @@ use std::sync::{OnceLock, RwLock}; use std::thread::{self, Builder}; use big_s::S; -use document_changes::{extract, DocumentChanges, IndexingContext, Progress, ThreadLocal}; +use document_changes::{extract, DocumentChanges, IndexingContext, Progress}; pub use document_deletion::DocumentDeletion; pub use document_operation::{DocumentOperation, PayloadStats}; use hashbrown::HashMap; @@ -20,6 +20,7 @@ use super::channel::*; use super::extract::*; use super::facet_search_builder::FacetSearchBuilder; use super::merger::FacetFieldIdsDelta; +use super::thread_local::ThreadLocal; use super::word_fst_builder::{PrefixData, PrefixDelta, WordFstBuilder}; use super::words_prefix_docids::{ compute_word_prefix_docids, compute_word_prefix_fid_docids, compute_word_prefix_position_docids, diff --git a/crates/milli/src/update/new/indexer/partial_dump.rs b/crates/milli/src/update/new/indexer/partial_dump.rs index 2da047824..8b5a8b650 100644 --- a/crates/milli/src/update/new/indexer/partial_dump.rs +++ b/crates/milli/src/update/new/indexer/partial_dump.rs @@ -3,11 +3,12 @@ use std::ops::DerefMut; use rayon::iter::IndexedParallelIterator; use serde_json::value::RawValue; -use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend}; +use super::document_changes::{DocumentChangeContext, DocumentChanges}; use crate::documents::PrimaryKey; use crate::update::concurrent_available_ids::ConcurrentAvailableIds; use crate::update::new::document::Versions; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::MostlySend; use crate::update::new::{DocumentChange, Insertion}; use crate::{Error, InternalError, Result, UserError}; diff --git a/crates/milli/src/update/new/indexer/update_by_function.rs b/crates/milli/src/update/new/indexer/update_by_function.rs index f6df3981d..a8e3e38a8 100644 --- a/crates/milli/src/update/new/indexer/update_by_function.rs +++ b/crates/milli/src/update/new/indexer/update_by_function.rs @@ -4,13 +4,14 @@ use rayon::slice::ParallelSlice as _; use rhai::{Dynamic, Engine, OptimizationLevel, Scope, AST}; use roaring::RoaringBitmap; -use super::document_changes::{DocumentChangeContext, MostlySend}; +use super::document_changes::DocumentChangeContext; use super::DocumentChanges; use crate::documents::Error::InvalidDocumentFormat; use crate::documents::PrimaryKey; use crate::error::{FieldIdMapMissingEntry, InternalError}; use crate::update::new::document::Versions; use crate::update::new::ref_cell_ext::RefCellExt as _; +use crate::update::new::thread_local::MostlySend; use crate::update::new::{Deletion, DocumentChange, KvReaderFieldId, Update}; use crate::{all_obkv_to_json, Error, FieldsIdsMap, Object, Result, UserError}; diff --git a/crates/milli/src/update/new/mod.rs b/crates/milli/src/update/new/mod.rs index 7a749228e..edbbdf497 100644 --- a/crates/milli/src/update/new/mod.rs +++ b/crates/milli/src/update/new/mod.rs @@ -17,6 +17,7 @@ pub mod indexer; mod merger; mod parallel_iterator_ext; mod ref_cell_ext; +pub(crate) mod thread_local; mod top_level_map; pub mod vector_document; mod word_fst_builder; diff --git a/crates/milli/src/update/new/thread_local.rs b/crates/milli/src/update/new/thread_local.rs new file mode 100644 index 000000000..acdc78c7b --- /dev/null +++ b/crates/milli/src/update/new/thread_local.rs @@ -0,0 +1,174 @@ +use std::cell::RefCell; + +/// A trait for types that are **not** [`Send`] only because they would then allow concurrent access to a type that is not [`Sync`]. +/// +/// The primary example of such a type is `&T`, with `T: !Sync`. +/// +/// In the authors' understanding, a type can be `!Send` for two distinct reasons: +/// +/// 1. Because it contains data that *genuinely* cannot be moved between threads, such as thread-local data. +/// 2. Because sending the type would allow concurrent access to a `!Sync` type, which is undefined behavior. +/// +/// `MostlySend` exists to be used in bounds where you need a type whose data is **not** *attached* to a thread +/// because you might access it from a different thread, but where you will never access the type **concurrently** from +/// multiple threads. +/// +/// Like [`Send`], `MostlySend` assumes properties on types that cannot be verified by the compiler, which is why implementing +/// this trait is unsafe. +/// +/// # Safety +/// +/// Implementers of this trait promises that the following properties hold on the implementing type: +/// +/// 1. Its data can be accessed from any thread and will be the same regardless of the thread accessing it. +/// 2. Any operation that can be performed on the type does not depend on the thread that executes it. +/// +/// As these properties are subtle and are not generally tracked by the Rust type system, great care should be taken before +/// implementing `MostlySend` on a type, especially a foreign type. +/// +/// - An example of a type that verifies (1) and (2) is [`std::rc::Rc`] (when `T` is `Send` and `Sync`). +/// - An example of a type that doesn't verify (1) is thread-local data. +/// - An example of a type that doesn't verify (2) is [`std::sync::MutexGuard`]: a lot of mutex implementations require that +/// a lock is returned to the operating system on the same thread that initially locked the mutex, failing to uphold this +/// invariant will cause Undefined Behavior +/// (see last § in [the nomicon](https://doc.rust-lang.org/nomicon/send-and-sync.html)). +/// +/// It is **always safe** to implement this trait on a type that is `Send`, but no placeholder impl is provided due to limitations in +/// coherency. Use the [`FullySend`] wrapper in this situation. +pub unsafe trait MostlySend {} + +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Ord, PartialOrd, Hash)] +pub struct FullySend(pub T); + +// SAFETY: a type **fully** send is always mostly send as well. +unsafe impl MostlySend for FullySend where T: Send {} + +unsafe impl MostlySend for RefCell where T: MostlySend {} + +unsafe impl MostlySend for Option where T: MostlySend {} + +impl FullySend { + pub fn into(self) -> T { + self.0 + } +} + +impl From for FullySend { + fn from(value: T) -> Self { + Self(value) + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +struct MostlySendWrapper(T); + +impl MostlySendWrapper { + /// # Safety + /// + /// - (P1) Users of this type will never access the type concurrently from multiple threads without synchronization + unsafe fn new(t: T) -> Self { + Self(t) + } + + fn as_ref(&self) -> &T { + &self.0 + } + + fn as_mut(&mut self) -> &mut T { + &mut self.0 + } + + fn into_inner(self) -> T { + self.0 + } +} + +/// # Safety +/// +/// 1. `T` is [`MostlySend`], so by its safety contract it can be accessed by any thread and all of its operations are available +/// from any thread. +/// 2. (P1) of `MostlySendWrapper::new` forces the user to never access the value from multiple threads concurrently. +unsafe impl Send for MostlySendWrapper {} + +/// A wrapper around [`thread_local::ThreadLocal`] that accepts [`MostlySend`] `T`s. +#[derive(Default)] +pub struct ThreadLocal { + inner: thread_local::ThreadLocal>, + // FIXME: this should be necessary + //_no_send: PhantomData<*mut ()>, +} + +impl ThreadLocal { + pub fn new() -> Self { + Self { inner: thread_local::ThreadLocal::new() } + } + + pub fn with_capacity(capacity: usize) -> Self { + Self { inner: thread_local::ThreadLocal::with_capacity(capacity) } + } + + pub fn clear(&mut self) { + self.inner.clear() + } + + pub fn get(&self) -> Option<&T> { + self.inner.get().map(|t| t.as_ref()) + } + + pub fn get_or(&self, create: F) -> &T + where + F: FnOnce() -> T, + { + self.inner.get_or(|| unsafe { MostlySendWrapper::new(create()) }).as_ref() + } + + pub fn get_or_try(&self, create: F) -> std::result::Result<&T, E> + where + F: FnOnce() -> std::result::Result, + { + self.inner + .get_or_try(|| unsafe { Ok(MostlySendWrapper::new(create()?)) }) + .map(MostlySendWrapper::as_ref) + } + + pub fn get_or_default(&self) -> &T + where + T: Default, + { + self.inner.get_or_default().as_ref() + } + + pub fn iter_mut(&mut self) -> IterMut { + IterMut(self.inner.iter_mut()) + } +} + +impl IntoIterator for ThreadLocal { + type Item = T; + + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter(self.inner.into_iter()) + } +} + +pub struct IterMut<'a, T: MostlySend>(thread_local::IterMut<'a, MostlySendWrapper>); + +impl<'a, T: MostlySend> Iterator for IterMut<'a, T> { + type Item = &'a mut T; + + fn next(&mut self) -> Option { + self.0.next().map(|t| t.as_mut()) + } +} + +pub struct IntoIter(thread_local::IntoIter>); + +impl Iterator for IntoIter { + type Item = T; + + fn next(&mut self) -> Option { + self.0.next().map(|t| t.into_inner()) + } +}