MapRef::get_or_init and DefaultPrelim trait

renewedvision · Jul 4, 2024 · 2b5ab81 · 2b5ab81
1 parent 63ce442
commit 2b5ab81
Show file tree

Hide file tree

Showing 5 changed files with 117 additions and 9 deletions.
diff --git a/yrs/src/types/array.rs b/yrs/src/types/array.rs
@@ -3,8 +3,8 @@ use crate::block_iter::BlockIter;
 use crate::moving::StickyIndex;
 use crate::transaction::TransactionMut;
 use crate::types::{
-    event_change_set, AsPrelim, Branch, BranchPtr, Change, ChangeSet, In, Out, Path, RootRef,
-    SharedRef, ToJson, TypeRef,
+    event_change_set, AsPrelim, Branch, BranchPtr, Change, ChangeSet, DefaultPrelim, In, Out, Path,
+    RootRef, SharedRef, ToJson, TypeRef,
 };
 use crate::{Any, Assoc, DeepObservable, IndexedSequence, Observable, ReadTxn, ID};
 use std::borrow::Borrow;
@@ -156,6 +156,15 @@ impl AsPrelim for ArrayRef {
     }
 }
 
+impl DefaultPrelim for ArrayRef {
+    type Prelim = ArrayPrelim;
+
+    #[inline]
+    fn default_prelim() -> Self::Prelim {
+        ArrayPrelim::default()
+    }
+}
+
 pub trait Array: AsRef<Branch> + Sized {
     /// Returns a number of elements stored in current array.
     fn len<T: ReadTxn>(&self, _txn: &T) -> u32 {

diff --git a/yrs/src/types/map.rs b/yrs/src/types/map.rs
@@ -1,8 +1,8 @@
 use crate::block::{EmbedPrelim, ItemContent, ItemPosition, ItemPtr, Prelim};
 use crate::transaction::TransactionMut;
 use crate::types::{
-    event_keys, AsPrelim, Branch, BranchPtr, Entries, EntryChange, In, Out, Path, RootRef,
-    SharedRef, ToJson, TypeRef,
+    event_keys, AsPrelim, Branch, BranchPtr, DefaultPrelim, Entries, EntryChange, In, Out, Path,
+    RootRef, SharedRef, ToJson, TypeRef,
 };
 use crate::*;
 use std::borrow::Borrow;
@@ -137,6 +137,15 @@ impl AsPrelim for MapRef {
     }
 }
 
+impl DefaultPrelim for MapRef {
+    type Prelim = MapPrelim;
+
+    #[inline]
+    fn default_prelim() -> Self::Prelim {
+        MapPrelim::default()
+    }
+}
+
 pub trait Map: AsRef<Branch> + Sized {
     /// Returns a number of entries stored within current map.
     fn len<T: ReadTxn>(&self, _txn: &T) -> u32 {
@@ -195,6 +204,25 @@ pub trait Map: AsRef<Branch> + Sized {
         }
     }
 
+    /// Returns an existing instance of a type stored under a given `key` within current map.
+    /// If the given entry was not found, has been deleted or its type is different from expected,
+    /// that entry will be reset to a given type and its reference will be returned.
+    fn get_or_init<K, V>(&self, txn: &mut TransactionMut, key: K) -> V
+    where
+        K: Into<Arc<str>>,
+        V: DefaultPrelim + TryFrom<Out>,
+    {
+        let key = key.into();
+        let branch = self.as_ref();
+        if let Some(value) = branch.get(txn, &key) {
+            if let Ok(value) = value.try_into() {
+                return value;
+            }
+        }
+        let value = V::default_prelim();
+        self.insert(txn, key, value)
+    }
+
     /// Removes a stored within current map under a given `key`. Returns that value or `None` if
     /// no entry with a given `key` was present in current map.
     ///
@@ -481,8 +509,9 @@ mod test {
     use crate::updates::decoder::Decode;
     use crate::updates::encoder::{Encoder, EncoderV1};
     use crate::{
-        any, Any, Array, ArrayPrelim, ArrayRef, Doc, Map, MapPrelim, MapRef, Observable,
-        StateVector, Text, Transact, Update,
+        any, Any, Array, ArrayPrelim, ArrayRef, Doc, GetString, Map, MapPrelim, MapRef, Observable,
+        StateVector, Text, TextRef, Transact, Update, WriteTxn, XmlFragment, XmlFragmentRef,
+        XmlTextPrelim, XmlTextRef,
     };
     use arc_swap::ArcSwapOption;
     use fastrand::Rng;
@@ -1014,6 +1043,38 @@ mod test {
         );
     }
 
+    #[test]
+    fn get_or_init() {
+        let doc = Doc::with_client_id(1);
+        let mut txn = doc.transact_mut();
+        let map = txn.get_or_insert_map("map");
+
+        let m: MapRef = map.get_or_init(&mut txn, "nested");
+        m.insert(&mut txn, "key", 1);
+        let m: MapRef = map.get_or_init(&mut txn, "nested");
+        assert_eq!(m.get(&txn, "key"), Some(Out::from(1)));
+
+        let m: ArrayRef = map.get_or_init(&mut txn, "nested");
+        m.insert(&mut txn, 0, 1);
+        let m: ArrayRef = map.get_or_init(&mut txn, "nested");
+        assert_eq!(m.get(&txn, 0), Some(Out::from(1)));
+
+        let m: TextRef = map.get_or_init(&mut txn, "nested");
+        m.insert(&mut txn, 0, "a");
+        let m: TextRef = map.get_or_init(&mut txn, "nested");
+        assert_eq!(m.get_string(&txn), "a".to_string());
+
+        let m: XmlFragmentRef = map.get_or_init(&mut txn, "nested");
+        m.insert(&mut txn, 0, XmlTextPrelim::new("b"));
+        let m: XmlFragmentRef = map.get_or_init(&mut txn, "nested");
+        assert_eq!(m.get_string(&txn), "b".to_string());
+
+        let m: XmlTextRef = map.get_or_init(&mut txn, "nested");
+        m.insert(&mut txn, 0, "c");
+        let m: XmlTextRef = map.get_or_init(&mut txn, "nested");
+        assert_eq!(m.get_string(&txn), "c".to_string());
+    }
+
     #[test]
     fn multi_threading() {
         use std::sync::{Arc, RwLock};

diff --git a/yrs/src/types/mod.rs b/yrs/src/types/mod.rs
@@ -352,6 +352,16 @@ pub trait AsPrelim {
     fn as_prelim<T: ReadTxn>(&self, txn: &T) -> Self::Prelim;
 }
 
+/// Trait which allows to generate a [Prelim]-compatible type that - when integrated - will be
+/// converted into an instance of a current type.
+pub trait DefaultPrelim {
+    type Prelim: Prelim<Return = Self>;
+
+    /// Returns an instance of [Prelim]-compatible type, which will turn into reference of a current
+    /// type after being integrated into the document store.
+    fn default_prelim() -> Self::Prelim;
+}
+
 /// Trait implemented by all Y-types, allowing for observing events which are emitted by
 /// nested types.
 #[cfg(feature = "sync")]

diff --git a/yrs/src/types/text.rs b/yrs/src/types/text.rs
@@ -1,7 +1,8 @@
 use crate::block::{EmbedPrelim, Item, ItemContent, ItemPosition, ItemPtr, Prelim, Unused};
 use crate::transaction::TransactionMut;
 use crate::types::{
-    AsPrelim, Attrs, Branch, BranchPtr, Delta, Out, Path, RootRef, SharedRef, TypePtr, TypeRef,
+    AsPrelim, Attrs, Branch, BranchPtr, DefaultPrelim, Delta, Out, Path, RootRef, SharedRef,
+    TypePtr, TypeRef,
 };
 use crate::utils::OptionExt;
 use crate::*;
@@ -471,6 +472,15 @@ impl AsPrelim for TextRef {
     }
 }
 
+impl DefaultPrelim for TextRef {
+    type Prelim = TextPrelim;
+
+    #[inline]
+    fn default_prelim() -> Self::Prelim {
+        TextPrelim::default()
+    }
+}
+
 impl Eq for TextRef {}
 impl PartialEq for TextRef {
     fn eq(&self, other: &Self) -> bool {

diff --git a/yrs/src/types/xml.rs b/yrs/src/types/xml.rs
@@ -12,8 +12,8 @@ use crate::block_iter::BlockIter;
 use crate::transaction::TransactionMut;
 use crate::types::text::{diff_between, TextEvent, YChange};
 use crate::types::{
-    event_change_set, event_keys, AsPrelim, Branch, BranchPtr, Change, ChangeSet, Delta, Entries,
-    EntryChange, MapRef, Out, Path, RootRef, SharedRef, ToJson, TypePtr, TypeRef,
+    event_change_set, event_keys, AsPrelim, Branch, BranchPtr, Change, ChangeSet, DefaultPrelim,
+    Delta, Entries, EntryChange, MapRef, Out, Path, RootRef, SharedRef, ToJson, TypePtr, TypeRef,
 };
 use crate::{
     Any, ArrayRef, BranchID, DeepObservable, GetString, In, IndexedSequence, Map, Observable,
@@ -658,6 +658,15 @@ impl AsPrelim for XmlTextRef {
     }
 }
 
+impl DefaultPrelim for XmlTextRef {
+    type Prelim = XmlTextPrelim;
+
+    #[inline]
+    fn default_prelim() -> Self::Prelim {
+        XmlTextPrelim::default()
+    }
+}
+
 /// A preliminary type that will be materialized into an [XmlTextRef] once it will be integrated
 /// into Yrs document.
 #[repr(transparent)]
@@ -868,6 +877,15 @@ impl AsPrelim for XmlFragmentRef {
     }
 }
 
+impl DefaultPrelim for XmlFragmentRef {
+    type Prelim = XmlFragmentPrelim;
+
+    #[inline]
+    fn default_prelim() -> Self::Prelim {
+        XmlFragmentPrelim::default()
+    }
+}
+
 /// A preliminary type that will be materialized into an [XmlFragmentRef] once it will be integrated
 /// into Yrs document.
 #[derive(Debug, Clone, PartialEq, Default)]