Move the RefCellExt trait in a dedicated module

crates/milli/src/update/new/extract/documents.rs

@@ -6,9 +6,8 @@ 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::{
+use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor, FullySend};
-    DocumentChangeContext, Extractor, FullySend, RefCellExt as _,
+use crate::update::new::ref_cell_ext::RefCellExt as _;
-};
 use crate::update::new::DocumentChange;
 use crate::vector::EmbeddingConfigs;
 use crate::Result;

crates/milli/src/update/new/extract/faceted/extract_facets.rs

@@ -13,8 +13,9 @@ use crate::facet::value_encoding::f64_into_bytes;
 use crate::update::new::extract::DocidsExtractor;
 use crate::update::new::indexer::document_changes::{
     extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext,
-    Progress, RefCellExt, ThreadLocal,
+    Progress, ThreadLocal,
 };
+use crate::update::new::ref_cell_ext::RefCellExt as _;
 use crate::update::new::DocumentChange;
 use crate::update::GrenadParameters;
 use crate::{DocumentId, FieldId, Index, Result, MAX_FACET_VALUE_LENGTH};

crates/milli/src/update/new/extract/searchable/extract_word_docids.rs

@@ -12,8 +12,9 @@ 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, RefCellExt, ThreadLocal,
+    MostlySend, Progress, ThreadLocal,
 };
+use crate::update::new::ref_cell_ext::RefCellExt as _;
 use crate::update::new::DocumentChange;
 use crate::update::GrenadParameters;
 use crate::{bucketed_position, DocumentId, FieldId, Index, Result, MAX_POSITION_PER_ATTRIBUTE};

crates/milli/src/update/new/extract/searchable/extract_word_pair_proximity_docids.rs

@@ -9,7 +9,8 @@ use super::SearchableExtractor;
 use crate::proximity::{index_proximity, MAX_DISTANCE};
 use crate::update::new::document::Document;
 use crate::update::new::extract::cache::BalancedCaches;
-use crate::update::new::indexer::document_changes::{DocumentChangeContext, RefCellExt};
+use crate::update::new::indexer::document_changes::DocumentChangeContext;
+use crate::update::new::ref_cell_ext::RefCellExt as _;
 use crate::update::new::DocumentChange;
 use crate::{FieldId, GlobalFieldsIdsMap, Index, Result};
 

crates/milli/src/update/new/indexer/document_changes.rs

@@ -1,4 +1,4 @@
-use std::cell::{Cell, Ref, RefCell, RefMut};
+use std::cell::{Cell, RefCell};
 use std::sync::{Arc, RwLock};
 
 use bumpalo::Bump;
@@ -10,57 +10,6 @@ use crate::fields_ids_map::metadata::FieldIdMapWithMetadata;
 use crate::update::new::parallel_iterator_ext::ParallelIteratorExt as _;
 use crate::{FieldsIdsMap, GlobalFieldsIdsMap, Index, InternalError, Result};
 
-pub trait RefCellExt<T: ?Sized> {
-    fn try_borrow_or_yield(&self) -> std::result::Result<Ref<'_, T>, std::cell::BorrowError>;
-    fn try_borrow_mut_or_yield(
-        &self,
-    ) -> std::result::Result<RefMut<'_, T>, std::cell::BorrowMutError>;
-
-    fn borrow_or_yield(&self) -> Ref<'_, T> {
-        self.try_borrow_or_yield().unwrap()
-    }
-
-    fn borrow_mut_or_yield(&self) -> RefMut<'_, T> {
-        self.try_borrow_mut_or_yield().unwrap()
-    }
-}
-
-impl<T: ?Sized> RefCellExt<T> for RefCell<T> {
-    fn try_borrow_or_yield(&self) -> std::result::Result<Ref<'_, T>, std::cell::BorrowError> {
-        /// TODO: move this trait and impl elsewhere
-        loop {
-            match self.try_borrow() {
-                Ok(borrow) => break Ok(borrow),
-                Err(error) => {
-                    tracing::warn!("dynamic borrow failed, yielding to local tasks");
-                    match rayon::yield_local() {
-                        Some(rayon::Yield::Executed) => continue,
-                        _ => return Err(error),
-                    }
-                }
-            }
-        }
-    }
-
-    fn try_borrow_mut_or_yield(
-        &self,
-    ) -> std::result::Result<RefMut<'_, T>, std::cell::BorrowMutError> {
-        loop {
-            match self.try_borrow_mut() {
-                Ok(borrow) => break Ok(borrow),
-                Err(error) => {
-                    tracing::warn!("dynamic borrow failed, yielding to local tasks");
-
-                    match rayon::yield_local() {
-                        Some(rayon::Yield::Executed) => continue,
-                        _ => return Err(error),
-                    }
-                }
-            }
-        }
-    }
-}
-
 /// 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`.

crates/milli/src/update/new/indexer/partial_dump.rs

@@ -3,10 +3,11 @@ use std::ops::DerefMut;
 use rayon::iter::IndexedParallelIterator;
 use serde_json::value::RawValue;
 
-use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend, RefCellExt};
+use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend};
 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::{DocumentChange, Insertion};
 use crate::{Error, InternalError, Result, UserError};
 

crates/milli/src/update/new/indexer/update_by_function.rs

@@ -4,12 +4,13 @@ use rayon::slice::ParallelSlice as _;
 use rhai::{Dynamic, Engine, OptimizationLevel, Scope, AST};
 use roaring::RoaringBitmap;
 
-use super::document_changes::{DocumentChangeContext, MostlySend, RefCellExt};
+use super::document_changes::{DocumentChangeContext, MostlySend};
 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::{Deletion, DocumentChange, KvReaderFieldId, Update};
 use crate::{all_obkv_to_json, Error, FieldsIdsMap, Object, Result, UserError};
 

crates/milli/src/update/new/mod.rs

@@ -16,6 +16,7 @@ mod fst_merger_builder;
 pub mod indexer;
 mod merger;
 mod parallel_iterator_ext;
+mod ref_cell_ext;
 mod top_level_map;
 pub mod vector_document;
 mod word_fst_builder;

crates/milli/src/update/new/ref_cell_ext.rs

@@ -0,0 +1,52 @@
+use std::cell::{Ref, RefCell, RefMut};
+
+pub trait RefCellExt<T: ?Sized> {
+    fn try_borrow_or_yield(&self) -> std::result::Result<Ref<'_, T>, std::cell::BorrowError>;
+    fn try_borrow_mut_or_yield(
+        &self,
+    ) -> std::result::Result<RefMut<'_, T>, std::cell::BorrowMutError>;
+
+    fn borrow_or_yield(&self) -> Ref<'_, T> {
+        self.try_borrow_or_yield().unwrap()
+    }
+
+    fn borrow_mut_or_yield(&self) -> RefMut<'_, T> {
+        self.try_borrow_mut_or_yield().unwrap()
+    }
+}
+
+impl<T: ?Sized> RefCellExt<T> for RefCell<T> {
+    fn try_borrow_or_yield(&self) -> std::result::Result<Ref<'_, T>, std::cell::BorrowError> {
+        /// TODO: move this trait and impl elsewhere
+        loop {
+            match self.try_borrow() {
+                Ok(borrow) => break Ok(borrow),
+                Err(error) => {
+                    tracing::warn!("dynamic borrow failed, yielding to local tasks");
+                    match rayon::yield_local() {
+                        Some(rayon::Yield::Executed) => continue,
+                        _ => return Err(error),
+                    }
+                }
+            }
+        }
+    }
+
+    fn try_borrow_mut_or_yield(
+        &self,
+    ) -> std::result::Result<RefMut<'_, T>, std::cell::BorrowMutError> {
+        loop {
+            match self.try_borrow_mut() {
+                Ok(borrow) => break Ok(borrow),
+                Err(error) => {
+                    tracing::warn!("dynamic borrow failed, yielding to local tasks");
+
+                    match rayon::yield_local() {
+                        Some(rayon::Yield::Executed) => continue,
+                        _ => return Err(error),
+                    }
+                }
+            }
+        }
+    }
+}

crates/milli/src/update/new/words_prefix_docids.rs

@@ -10,7 +10,7 @@ use roaring::MultiOps;
 use tempfile::tempfile;
 use thread_local::ThreadLocal;
 
-use super::indexer::document_changes::RefCellExt;
+use super::ref_cell_ext::RefCellExt as _;
 use crate::heed_codec::StrBEU16Codec;
 use crate::{CboRoaringBitmapCodec, Index, Prefix, Result};