splinter_rs/
cow.rs

1use std::fmt::Debug;
2use std::ops::{
3    BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Deref, RangeBounds, Sub,
4    SubAssign,
5};
6
7use bytes::{BufMut, Bytes};
8use either::Either;
9
10use crate::{
11    Encodable, PartitionRead, PartitionWrite, Splinter, SplinterRef,
12    codec::{DecodeErr, encoder::Encoder},
13    level::High,
14};
15
16/// A clone-on-write splinter that can hold either a reference or an owned value.
17///
18/// `CowSplinter` is an enum that can contain either a [`SplinterRef<B>`] (for zero-copy
19/// read-only access) or an owned [`Splinter`] (for mutable operations). It automatically
20/// converts from read-only to owned when mutation is needed, providing an efficient
21/// way to work with splinter data that might come from different sources.
22///
23/// This is particularly useful when you want to:
24/// - Start with serialized/borrowed data for read-only operations
25/// - Potentially modify the data later
26/// - Avoid unnecessary copying until mutation is actually needed
27/// - Or work with a collection of Splinters which may or may not be owned
28///
29/// # Examples
30///
31/// Starting with a reference and converting to owned when needed:
32///
33/// ```
34/// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
35///
36/// // Create a CowSplinter from a reference
37/// let original = Splinter::from_iter([100, 200]);
38/// let splinter_ref = original.encode_to_splinter_ref();
39/// let mut cow = CowSplinter::from_ref(splinter_ref);
40///
41/// // Read operations work on the reference
42/// assert_eq!(cow.cardinality(), 2);
43///
44/// // First write operation converts to owned
45/// cow.insert(300); // This triggers clone-on-write
46/// assert_eq!(cow.cardinality(), 3);
47/// ```
48///
49/// Creating from different sources:
50///
51/// ```
52/// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
53/// use bytes::Bytes;
54///
55/// // From owned splinter
56/// let owned = Splinter::EMPTY;
57/// let cow1: CowSplinter<Bytes> = CowSplinter::from_owned(owned);
58///
59/// // From iterator
60/// let cow2 = CowSplinter::<Bytes>::from_iter([1u32, 2, 3]);
61///
62/// // From bytes
63/// let bytes = cow2.encode_to_bytes();
64/// let cow3 = CowSplinter::from_bytes(bytes).unwrap();
65/// ```
66#[derive(Clone)]
67pub enum CowSplinter<B> {
68    /// Contains a zero-copy reference to serialized data
69    Ref(SplinterRef<B>),
70    /// Contains an owned, mutable splinter
71    Owned(Splinter),
72}
73
74impl<B> Default for CowSplinter<B> {
75    fn default() -> Self {
76        Self::Owned(Splinter::EMPTY)
77    }
78}
79
80impl<B: Deref<Target = [u8]>> Debug for CowSplinter<B> {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        match self {
83            CowSplinter::Ref(splinter_ref) => f
84                .debug_tuple("CowSplinter::Ref")
85                .field(splinter_ref)
86                .finish(),
87            CowSplinter::Owned(splinter) => {
88                f.debug_tuple("CowSplinter::Owned").field(splinter).finish()
89            }
90        }
91    }
92}
93
94impl<B, K: Into<u32>> FromIterator<K> for CowSplinter<B>
95where
96    B: Deref<Target = [u8]>,
97{
98    fn from_iter<I: IntoIterator<Item = K>>(iter: I) -> Self {
99        Self::Owned(Splinter::from_iter(iter.into_iter().map(|k| k.into())))
100    }
101}
102
103impl<B> From<Splinter> for CowSplinter<B> {
104    fn from(splinter: Splinter) -> Self {
105        Self::Owned(splinter)
106    }
107}
108
109impl<B> From<SplinterRef<B>> for CowSplinter<B> {
110    fn from(splinter_ref: SplinterRef<B>) -> Self {
111        Self::Ref(splinter_ref)
112    }
113}
114
115impl<B: Deref<Target = [u8]>> From<CowSplinter<B>> for Splinter {
116    fn from(cow_splinter: CowSplinter<B>) -> Self {
117        cow_splinter.into_owned()
118    }
119}
120
121impl From<CowSplinter<Bytes>> for SplinterRef<Bytes> {
122    fn from(cow: CowSplinter<Bytes>) -> Self {
123        match cow {
124            CowSplinter::Ref(splinter_ref) => splinter_ref,
125            CowSplinter::Owned(splinter) => splinter.encode_to_splinter_ref(),
126        }
127    }
128}
129
130impl<B> CowSplinter<B> {
131    /// Creates a `CowSplinter` from an owned [`Splinter`].
132    ///
133    /// # Examples
134    ///
135    /// ```
136    /// use splinter_rs::{Splinter, CowSplinter, PartitionRead, Encodable};
137    ///
138    /// let splinter = Splinter::from_iter([42]);
139    /// let cow: CowSplinter<Vec<u8>> = CowSplinter::from_owned(splinter);
140    /// assert!(cow.contains(42));
141    /// ```
142    pub fn from_owned(splinter: Splinter) -> Self {
143        splinter.into()
144    }
145
146    /// Creates a `CowSplinter` from a [`SplinterRef`].
147    ///
148    /// # Examples
149    ///
150    /// ```
151    /// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
152    ///
153    /// let mut splinter = Splinter::EMPTY;
154    /// splinter.insert(42);
155    /// let splinter_ref = splinter.encode_to_splinter_ref();
156    ///
157    /// let cow = CowSplinter::from_ref(splinter_ref);
158    /// assert!(cow.contains(42));
159    /// ```
160    pub fn from_ref(splinter: SplinterRef<B>) -> Self {
161        splinter.into()
162    }
163}
164
165impl<B: Deref<Target = [u8]>> CowSplinter<B> {
166    /// Creates a `CowSplinter` from raw bytes, validating the format.
167    ///
168    /// This is equivalent to creating a `SplinterRef` from the bytes and wrapping
169    /// it in a `CowSplinter::Ref`. All the same validation rules apply.
170    ///
171    /// # Errors
172    ///
173    /// Returns the same errors as [`SplinterRef::from_bytes`].
174    ///
175    /// # Examples
176    ///
177    /// ```
178    /// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
179    ///
180    /// let mut splinter = Splinter::EMPTY;
181    /// splinter.insert(42);
182    /// let bytes = splinter.encode_to_bytes();
183    ///
184    /// let cow = CowSplinter::from_bytes(bytes).unwrap();
185    /// assert!(cow.contains(42));
186    /// ```
187    pub fn from_bytes(data: B) -> Result<Self, DecodeErr> {
188        Ok(Self::Ref(SplinterRef::from_bytes(data)?))
189    }
190
191    /// Converts this `CowSplinter` into an owned [`Splinter`].
192    ///
193    /// If this is already an owned splinter, it returns it directly.
194    /// If this is a reference, it deserializes the data into a new owned splinter.
195    ///
196    /// # Examples
197    ///
198    /// ```
199    /// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
200    ///
201    /// let mut original = Splinter::EMPTY;
202    /// original.insert(100);
203    /// let splinter_ref = original.encode_to_splinter_ref();
204    /// let cow = CowSplinter::from_ref(splinter_ref);
205    ///
206    /// let owned = cow.into_owned();
207    /// assert_eq!(owned.cardinality(), 1);
208    /// assert!(owned.contains(100));
209    /// ```
210    pub fn into_owned(self) -> Splinter {
211        match self {
212            Self::Ref(splinter_ref) => splinter_ref.decode_to_splinter(),
213            Self::Owned(splinter) => splinter,
214        }
215    }
216
217    /// Returns a mutable reference to the underlying [`Splinter`].
218    ///
219    /// This method implements the "clone-on-write" behavior: if the current value
220    /// is a `Ref`, it will be converted to `Owned` by deserializing the data.
221    /// Subsequent calls will return the same mutable reference without additional
222    /// conversions.
223    ///
224    /// # Examples
225    ///
226    /// ```
227    /// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
228    ///
229    /// let mut original = Splinter::EMPTY;
230    /// original.insert(100);
231    /// let splinter_ref = original.encode_to_splinter_ref();
232    /// let mut cow = CowSplinter::from_ref(splinter_ref);
233    ///
234    /// // This triggers the clone-on-write conversion
235    /// let mutable_ref = cow.to_mut();
236    /// mutable_ref.insert(200);
237    ///
238    /// assert_eq!(cow.cardinality(), 2);
239    /// ```
240    pub fn to_mut(&mut self) -> &mut Splinter {
241        match *self {
242            Self::Ref(ref splinter_ref) => {
243                *self = Self::Owned(splinter_ref.decode_to_splinter());
244                match *self {
245                    Self::Ref(..) => unreachable!(),
246                    Self::Owned(ref mut owned) => owned,
247                }
248            }
249            Self::Owned(ref mut owned) => owned,
250        }
251    }
252}
253
254impl CowSplinter<Bytes> {
255    /// Returns the serialized bytes representation of this splinter.
256    ///
257    /// For the `Ref` variant, this clones the underlying `Bytes` (which is efficient
258    /// due to reference counting). For the `Owned` variant, this encodes the splinter
259    /// to bytes.
260    ///
261    /// # Examples
262    ///
263    /// ```
264    /// use splinter_rs::{Splinter, CowSplinter, PartitionWrite, PartitionRead, Encodable};
265    ///
266    /// let mut splinter = Splinter::EMPTY;
267    /// splinter.insert(42);
268    /// let cow = CowSplinter::from_owned(splinter);
269    ///
270    /// let bytes = cow.encode_to_bytes();
271    /// assert!(!bytes.is_empty());
272    /// ```
273    pub fn encode_to_bytes(&self) -> Bytes {
274        match self {
275            CowSplinter::Ref(splinter_ref) => splinter_ref.encode_to_bytes(),
276            CowSplinter::Owned(splinter) => splinter.encode_to_bytes(),
277        }
278    }
279}
280
281impl<B: Deref<Target = [u8]>> Encodable for CowSplinter<B> {
282    fn encoded_size(&self) -> usize {
283        match self {
284            CowSplinter::Ref(splinter_ref) => splinter_ref.encoded_size(),
285            CowSplinter::Owned(splinter) => splinter.encoded_size(),
286        }
287    }
288
289    fn encode<T: BufMut>(&self, encoder: &mut Encoder<T>) {
290        match self {
291            CowSplinter::Ref(splinter_ref) => splinter_ref.encode(encoder),
292            CowSplinter::Owned(splinter) => splinter.encode(encoder),
293        }
294    }
295}
296
297impl<B: Deref<Target = [u8]>> PartitionRead<High> for CowSplinter<B> {
298    fn cardinality(&self) -> usize {
299        match self {
300            CowSplinter::Ref(splinter_ref) => splinter_ref.cardinality(),
301            CowSplinter::Owned(splinter) => splinter.cardinality(),
302        }
303    }
304
305    fn is_empty(&self) -> bool {
306        match self {
307            CowSplinter::Ref(splinter_ref) => splinter_ref.is_empty(),
308            CowSplinter::Owned(splinter) => splinter.is_empty(),
309        }
310    }
311
312    fn contains(&self, value: u32) -> bool {
313        match self {
314            CowSplinter::Ref(splinter_ref) => splinter_ref.contains(value),
315            CowSplinter::Owned(splinter) => splinter.contains(value),
316        }
317    }
318
319    fn position(&self, value: u32) -> Option<usize> {
320        match self {
321            CowSplinter::Ref(splinter_ref) => splinter_ref.position(value),
322            CowSplinter::Owned(splinter) => splinter.position(value),
323        }
324    }
325
326    fn rank(&self, value: u32) -> usize {
327        match self {
328            CowSplinter::Ref(splinter_ref) => splinter_ref.rank(value),
329            CowSplinter::Owned(splinter) => splinter.rank(value),
330        }
331    }
332
333    fn select(&self, idx: usize) -> Option<u32> {
334        match self {
335            CowSplinter::Ref(splinter_ref) => splinter_ref.select(idx),
336            CowSplinter::Owned(splinter) => splinter.select(idx),
337        }
338    }
339
340    fn last(&self) -> Option<u32> {
341        match self {
342            CowSplinter::Ref(splinter_ref) => splinter_ref.last(),
343            CowSplinter::Owned(splinter) => splinter.last(),
344        }
345    }
346
347    fn iter(&self) -> impl Iterator<Item = u32> {
348        match self {
349            CowSplinter::Ref(splinter_ref) => Either::Left(splinter_ref.iter()),
350            CowSplinter::Owned(splinter) => Either::Right(splinter.iter()),
351        }
352    }
353
354    fn contains_all<R: RangeBounds<u32>>(&self, values: R) -> bool {
355        match self {
356            CowSplinter::Ref(splinter_ref) => splinter_ref.contains_all(values),
357            CowSplinter::Owned(splinter) => splinter.contains_all(values),
358        }
359    }
360
361    fn contains_any<R: RangeBounds<u32>>(&self, values: R) -> bool {
362        match self {
363            CowSplinter::Ref(splinter_ref) => splinter_ref.contains_any(values),
364            CowSplinter::Owned(splinter) => splinter.contains_any(values),
365        }
366    }
367}
368
369impl<B: Deref<Target = [u8]>> PartitionWrite<High> for CowSplinter<B> {
370    #[inline]
371    fn insert(&mut self, value: u32) -> bool {
372        self.to_mut().insert(value)
373    }
374
375    #[inline]
376    fn remove(&mut self, value: u32) -> bool {
377        self.to_mut().remove(value)
378    }
379
380    #[inline]
381    fn remove_range<R: RangeBounds<u32>>(&mut self, values: R) {
382        self.to_mut().remove_range(values)
383    }
384}
385
386impl<B: Deref<Target = [u8]>, B2: Deref<Target = [u8]>> PartialEq<CowSplinter<B2>>
387    for CowSplinter<B>
388{
389    fn eq(&self, other: &CowSplinter<B2>) -> bool {
390        use CowSplinter::*;
391        match (self, other) {
392            (Ref(l), Ref(r)) => l == r,
393            (Ref(l), Owned(r)) => l == r,
394            (Owned(l), Ref(r)) => l == r,
395            (Owned(l), Owned(r)) => l == r,
396        }
397    }
398}
399
400impl<B: Deref<Target = [u8]>> Eq for CowSplinter<B> {}
401
402impl<B: Deref<Target = [u8]>> PartialEq<Splinter> for CowSplinter<B> {
403    fn eq(&self, other: &Splinter) -> bool {
404        other == self
405    }
406}
407
408impl<B: Deref<Target = [u8]>, B2: Deref<Target = [u8]>> PartialEq<SplinterRef<B2>>
409    for CowSplinter<B>
410{
411    fn eq(&self, other: &SplinterRef<B2>) -> bool {
412        match self {
413            CowSplinter::Ref(splinter_ref) => splinter_ref == other,
414            CowSplinter::Owned(splinter) => splinter == other,
415        }
416    }
417}
418
419macro_rules! owned_bitop {
420    ($b:ty $([$($for:tt)*])?) => {
421        owned_bitop!($b $([$($for)*])?, BitAnd, bitand, BitAndAssign::bitand_assign);
422        owned_bitop!($b $([$($for)*])?, BitOr, bitor, BitOrAssign::bitor_assign);
423        owned_bitop!($b $([$($for)*])?, BitXor, bitxor, BitXorAssign::bitxor_assign);
424        owned_bitop!($b $([$($for)*])?, Sub, sub, SubAssign::sub_assign);
425    };
426    ($b:ty $([$($for:tt)*])?, $BitOp:ident, $bitop:ident, $bitassign:path) => {
427        impl<B1: Deref<Target = [u8]> $(, $($for)*)?> $BitOp<$b> for CowSplinter<B1> {
428            type Output = Self;
429            fn $bitop(mut self, rhs: $b) -> Self::Output {
430                $bitassign(self.to_mut(), rhs);
431                self
432            }
433        }
434    };
435}
436
437owned_bitop!(Splinter);
438owned_bitop!(&Splinter);
439owned_bitop!(SplinterRef<B2> [B2: Deref<Target=[u8]>]);
440owned_bitop!(&SplinterRef<B2> [B2: Deref<Target=[u8]>]);
441owned_bitop!(CowSplinter<B2> [B2: Deref<Target=[u8]>]);
442owned_bitop!(&CowSplinter<B2> [B2: Deref<Target=[u8]>]);
443
444macro_rules! ref_bitop {
445    ($b:ty $([$($for:tt)*])?) => {
446        ref_bitop!($b $([$($for)*])?, BitAnd, bitand, BitAndAssign::bitand_assign);
447        ref_bitop!($b $([$($for)*])?, BitOr, bitor, BitOrAssign::bitor_assign);
448        ref_bitop!($b $([$($for)*])?, BitXor, bitxor, BitXorAssign::bitxor_assign);
449        ref_bitop!($b $([$($for)*])?, Sub, sub, SubAssign::sub_assign);
450    };
451    ($b:ty $([$($for:tt)*])?, $BitOp:ident, $bitop:ident, $bitassign:path) => {
452        impl<B1: Deref<Target = [u8]> + Clone $(, $($for)*)?> $BitOp<$b>
453            for &CowSplinter<B1>
454        {
455            type Output = CowSplinter<B1>;
456            fn $bitop(self, rhs: $b) -> Self::Output {
457                let mut out = self.clone();
458                $bitassign(out.to_mut(), rhs);
459                out
460            }
461        }
462    };
463}
464
465ref_bitop!(Splinter);
466ref_bitop!(&Splinter);
467ref_bitop!(SplinterRef<B2> [B2: Deref<Target=[u8]>]);
468ref_bitop!(&SplinterRef<B2> [B2: Deref<Target=[u8]>]);
469ref_bitop!(CowSplinter<B2> [B2: Deref<Target=[u8]>]);
470ref_bitop!(&CowSplinter<B2> [B2: Deref<Target=[u8]>]);