bitmac/
union.rs

1use crate::{
2    container::{ContainerRead, ContainerWrite},
3    number::Number,
4    with_slots::TryWithSlots,
5    BitAccess, SmallContainerSizeError, UnionError,
6};
7
8/// Union operator (a | b).
9pub trait Union<Rhs, N, B>
10where
11    Rhs: ContainerRead<B, Slot = N>,
12    N: Number,
13    B: BitAccess,
14{
15    /// Calculates union in-place. Result will be stored in `dst`.
16    ///
17    /// ## Panic
18    ///
19    /// Panics if `dst` cannot fit the entire result.
20    /// See non-panic function [`try_union_in`].
21    ///
22    /// [`try_union_in`]: crate::union::Union::try_union_in
23    fn union_in<Dst>(&self, rhs: &Rhs, dst: &mut Dst)
24    where
25        Dst: ContainerWrite<B, Slot = N>;
26
27    /// Calculates union in-place. Result will be stored in `dst`.
28    ///
29    /// Returns `Err(_)` if `dst` cannot fit the entire result.
30    fn try_union_in<Dst>(&self, rhs: &Rhs, dst: &mut Dst) -> Result<(), UnionError>
31    where
32        Dst: ContainerWrite<B, Slot = N>;
33
34    /// Calculates union. Result container will be created with [`try_with_slots`] function.
35    ///
36    /// ## Panic
37    ///
38    /// Panics if `Dst` cannot fit the entire result.
39    /// See non-panic function [`try_union`].
40    ///
41    /// [`try_union`]: crate::union::Union::try_union
42    /// [`try_with_slots`]: crate::with_slots::TryWithSlots::try_with_slots
43    fn union<Dst>(&self, rhs: &Rhs) -> Dst
44    where
45        Dst: ContainerWrite<B, Slot = N> + TryWithSlots;
46
47    /// Calculates union. Result container will be created with [`try_with_slots`] function.
48    ///
49    /// Returns `Err(_)` if `Dst` cannot fit the entire result.
50    ///
51    /// [`try_with_slots`]: crate::with_slots::TryWithSlots::try_with_slots
52    fn try_union<Dst>(&self, rhs: &Rhs) -> Result<Dst, UnionError>
53    where
54        Dst: ContainerWrite<B, Slot = N> + TryWithSlots;
55
56    /// Calculates union length - ones count. It doesn't allocate for storing union result.
57    ///
58    /// Useful if you need to create some storage that relies on the number of bits presented in the bitmap.
59    fn union_len(&self, rhs: &Rhs) -> usize;
60}
61
62pub(crate) fn try_union_in_impl<Lhs, Rhs, Dst, N, B>(
63    lhs: &Lhs,
64    rhs: &Rhs,
65    dst: &mut Dst,
66) -> Result<(), UnionError>
67where
68    Lhs: ContainerRead<B, Slot = N>,
69    Rhs: ContainerRead<B, Slot = N>,
70    Dst: ContainerWrite<B, Slot = N>,
71    N: Number,
72    B: BitAccess,
73{
74    // TODO: shrink size
75    let required_dst_len = usize::max(lhs.slots_count(), rhs.slots_count());
76    if dst.slots_count() < required_dst_len {
77        return Err(SmallContainerSizeError::new(format!(
78            "size of container should be >= {}, but handled {}",
79            required_dst_len,
80            dst.slots_count()
81        ))
82        .into());
83    }
84
85    let head_max_idx = usize::min(lhs.slots_count(), rhs.slots_count());
86    for i in 0..head_max_idx {
87        let dst_slot = dst.get_mut_slot(i);
88        let lhs_slot = lhs.get_slot(i);
89        let rhs_slot = rhs.get_slot(i);
90
91        *dst_slot = lhs_slot | rhs_slot;
92    }
93
94    // Clone rest tail
95    let tail_max_idx = usize::max(lhs.slots_count(), rhs.slots_count());
96    for i in head_max_idx..tail_max_idx {
97        let dst_slot = dst.get_mut_slot(i);
98        let rest_slot = if lhs.slots_count() >= rhs.slots_count() {
99            lhs.get_slot(i)
100        } else {
101            rhs.get_slot(i)
102        };
103
104        *dst_slot = rest_slot
105    }
106
107    Ok(())
108}
109
110pub(crate) fn try_union_impl<Lhs, Rhs, Dst, N, B>(lhs: &Lhs, rhs: &Rhs) -> Result<Dst, UnionError>
111where
112    Lhs: ContainerRead<B, Slot = N>,
113    Rhs: ContainerRead<B, Slot = N>,
114    Dst: ContainerWrite<B, Slot = N> + TryWithSlots,
115    N: Number,
116    B: BitAccess,
117{
118    // TODO: shrink size
119    let slots_count = usize::max(lhs.slots_count(), rhs.slots_count());
120    let mut dst = Dst::try_with_slots(slots_count)?;
121
122    try_union_in_impl(lhs, rhs, &mut dst)?;
123    Ok(dst)
124}
125
126pub(crate) fn union_len_impl<Lhs, Rhs, N, B>(lhs: &Lhs, rhs: &Rhs) -> usize
127where
128    Lhs: ContainerRead<B, Slot = N>,
129    Rhs: ContainerRead<B, Slot = N>,
130    N: Number,
131    B: BitAccess,
132{
133    let head_max_idx = usize::min(lhs.slots_count(), rhs.slots_count());
134
135    let mut len = 0;
136    for i in 0..head_max_idx {
137        let lhs_slot = lhs.get_slot(i);
138        let rhs_slot = rhs.get_slot(i);
139        let intersect = lhs_slot | rhs_slot;
140        len += intersect.count_ones() as usize;
141    }
142
143    // Counting rest tail
144    let tail_max_idx = usize::max(lhs.slots_count(), rhs.slots_count());
145    for i in head_max_idx..tail_max_idx {
146        let rest_slot = if lhs.slots_count() >= rhs.slots_count() {
147            lhs.get_slot(i)
148        } else {
149            rhs.get_slot(i)
150        };
151
152        len += rest_slot.count_ones() as usize;
153    }
154    len
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use crate::LSB;
161
162    #[test]
163    fn union() {
164        let lhs: u8 = 0b0010_1100;
165        let rhs: u8 = 0b0010_0100;
166        let exp: u8 = 0b0010_1100;
167        assert_eq!(try_union_impl::<_, _, u8, _, LSB>(&lhs, &rhs).unwrap(), exp);
168
169        let lhs: u8 = 0b0010_1100;
170        let rhs: u8 = 0b0010_0100;
171        let exp: [u8; 1] = [0b0010_1100];
172        assert_eq!(
173            try_union_impl::<_, _, [u8; 1], _, LSB>(&lhs, &rhs).unwrap(),
174            exp
175        );
176
177        let lhs: u8 = 0b0010_1100;
178        let rhs: u8 = 0b0010_0100;
179        let exp: Vec<u8> = vec![0b0010_1100];
180        assert_eq!(
181            try_union_impl::<_, _, Vec<u8>, _, LSB>(&lhs, &rhs).unwrap(),
182            exp
183        );
184
185        #[cfg(feature = "bytes")]
186        {
187            use bytes::BytesMut;
188            let lhs: u8 = 0b0010_1100;
189            let rhs: u8 = 0b0010_0100;
190            let exp: BytesMut = BytesMut::from(&[0b0010_1100][..]);
191            assert_eq!(
192                try_union_impl::<_, _, BytesMut, _, LSB>(&lhs, &rhs).unwrap(),
193                exp
194            );
195        }
196
197        #[cfg(feature = "smallvec")]
198        {
199            use smallvec::{smallvec, SmallVec};
200            let lhs: u8 = 0b0010_1100;
201            let rhs: u8 = 0b0010_0100;
202            let exp: SmallVec<[u8; 1]> = smallvec![0b0010_1100];
203            assert_eq!(
204                try_union_impl::<_, _, SmallVec<[u8; 1]>, _, LSB>(&lhs, &rhs).unwrap(),
205                exp
206            );
207        }
208
209        /////////
210
211        let lhs: u8 = 0b0010_1100;
212        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
213        let exp: [u8; 2] = [0b0010_1100, 0b0000_0000];
214        assert_eq!(
215            try_union_impl::<_, _, [u8; 2], _, LSB>(&lhs, &rhs).unwrap(),
216            exp
217        );
218
219        let lhs: u8 = 0b0010_1100;
220        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
221        let exp: Vec<u8> = vec![0b0010_1100, 0b0000_0000];
222        assert_eq!(
223            try_union_impl::<_, _, Vec<u8>, _, LSB>(&lhs, &rhs).unwrap(),
224            exp
225        );
226
227        #[cfg(feature = "bytes")]
228        {
229            use bytes::BytesMut;
230            let lhs: u8 = 0b0010_1100;
231            let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
232            let exp: BytesMut = BytesMut::from(&[0b0010_1100, 0b0000_0000][..]);
233            assert_eq!(
234                try_union_impl::<_, _, BytesMut, _, LSB>(&lhs, &rhs).unwrap(),
235                exp
236            );
237        }
238
239        #[cfg(feature = "smallvec")]
240        {
241            use smallvec::{smallvec, SmallVec};
242            let lhs: u8 = 0b0010_1100;
243            let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
244            let exp: SmallVec<[u8; 1]> = smallvec![0b0010_1100, 0b0000_0000];
245            assert_eq!(
246                try_union_impl::<_, _, SmallVec<[u8; 1]>, _, LSB>(&lhs, &rhs).unwrap(),
247                exp
248            );
249        }
250    }
251
252    #[test]
253    fn try_union() {
254        let lhs: u8 = 0b0010_1100;
255        let rhs: u8 = 0b0010_0100;
256        assert!(try_union_impl::<_, _, [u8; 10], _, LSB>(&lhs, &rhs).is_err());
257
258        let lhs: u8 = 0b0010_1100;
259        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
260        assert!(try_union_impl::<_, _, [u8; 3], _, LSB>(&lhs, &rhs).is_err());
261
262        let lhs: u8 = 0b0010_1100;
263        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
264        assert!(try_union_impl::<_, _, u8, _, LSB>(&lhs, &rhs).is_err());
265    }
266
267    #[test]
268    fn union_in() {
269        let lhs: u8 = 0b0010_1100;
270        let rhs: u8 = 0b0010_0100;
271        let mut dst: u8 = 0b0000_0000;
272        let exp: u8 = 0b0010_1100;
273        try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
274        assert_eq!(dst, exp);
275
276        let lhs: u8 = 0b0010_1100;
277        let rhs: u8 = 0b0010_0100;
278        let mut dst: [u8; 1] = [0b0000_0000];
279        let exp: [u8; 1] = [0b0010_1100];
280        try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
281        assert_eq!(dst, exp);
282
283        let lhs: u8 = 0b0010_1100;
284        let rhs: u8 = 0b0010_0100;
285        let mut dst: Vec<u8> = vec![0b0000_0000];
286        let exp: Vec<u8> = vec![0b0010_1100];
287        try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
288        assert_eq!(dst, exp);
289
290        #[cfg(feature = "bytes")]
291        {
292            use bytes::BytesMut;
293            let lhs: u8 = 0b0010_1100;
294            let rhs: u8 = 0b0010_0100;
295            let mut dst: BytesMut = BytesMut::from(&[0b0000_0000][..]);
296            let exp: BytesMut = BytesMut::from(&[0b0010_1100][..]);
297            try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
298            assert_eq!(dst, exp);
299        }
300
301        #[cfg(feature = "smallvec")]
302        {
303            use smallvec::{smallvec, SmallVec};
304            let lhs: u8 = 0b0010_1100;
305            let rhs: u8 = 0b0010_0100;
306            let mut dst: SmallVec<[u8; 1]> = smallvec![0b0000_0000];
307            let exp: SmallVec<[u8; 1]> = smallvec![0b0010_1100];
308            try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
309            assert_eq!(dst, exp);
310        }
311
312        /////////
313
314        let lhs: u8 = 0b0010_1100;
315        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
316        let mut dst: [u8; 2] = [0b0000_0000, 0b0000_0000];
317        let exp: [u8; 2] = [0b0010_1100, 0b0000_0000];
318        try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
319        assert_eq!(dst, exp);
320
321        let lhs: u8 = 0b0010_1100;
322        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
323        let mut dst: Vec<u8> = vec![0b0000_0000, 0b0000_0000];
324        let exp: Vec<u8> = vec![0b0010_1100, 0b0000_0000];
325        try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
326        assert_eq!(dst, exp);
327
328        #[cfg(feature = "bytes")]
329        {
330            use bytes::BytesMut;
331            let lhs: u8 = 0b0010_1100;
332            let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
333            let mut dst: BytesMut = BytesMut::from(&[0b0000_0000, 0b0000_0000][..]);
334            let exp: BytesMut = BytesMut::from(&[0b0010_1100, 0b0000_0000][..]);
335            try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
336            assert_eq!(dst, exp);
337        }
338
339        #[cfg(feature = "smallvec")]
340        {
341            use smallvec::{smallvec, SmallVec};
342            let lhs: u8 = 0b0010_1100;
343            let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
344            let mut dst: SmallVec<[u8; 1]> = smallvec![0b0000_0000, 0b0000_0000];
345            let exp: SmallVec<[u8; 1]> = smallvec![0b0010_1100, 0b0000_0000];
346            try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).unwrap();
347            assert_eq!(dst, exp);
348        }
349    }
350
351    #[test]
352    fn try_union_in() {
353        let lhs: [u8; 2] = [0b0010_1100, 0b0000_0000];
354        let rhs: [u8; 3] = [0b0010_0100, 0b0000_0000, 0b0000_0000];
355        let mut dst: [u8; 2] = [0b0000_0000; 2];
356        assert!(try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).is_err());
357
358        let lhs: [u8; 2] = [0b0010_1100, 0b0000_0000];
359        let rhs: [u8; 3] = [0b0010_0100, 0b0000_0000, 0b0000_0000];
360        let mut dst: Vec<u8> = vec![0b0000_0000; 2];
361        assert!(try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).is_err());
362
363        let lhs: [u8; 2] = [0b0010_1100, 0b0000_0000];
364        let rhs: [u8; 3] = [0b0010_0100, 0b0000_0000, 0b0000_0000];
365        let v = &mut [0b0000_0000, 0b0000_0000][..];
366        let mut dst: &mut [u8] = v;
367        assert!(try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).is_err());
368
369        #[cfg(feature = "bytes")]
370        {
371            use bytes::BytesMut;
372            let lhs: [u8; 2] = [0b0010_1100, 0b0000_0000];
373            let rhs: [u8; 3] = [0b0010_0100, 0b0000_0000, 0b0000_0000];
374            let mut dst: BytesMut = BytesMut::from(&[0b0000_0000, 0b0000_0000][..]);
375            assert!(try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).is_err());
376        }
377
378        #[cfg(feature = "smallvec")]
379        {
380            use smallvec::{smallvec, SmallVec};
381            let lhs: [u8; 2] = [0b0010_1100, 0b0000_0000];
382            let rhs: [u8; 3] = [0b0010_0100, 0b0000_0000, 0b0000_0000];
383            let mut dst: SmallVec<[u8; 1]> = smallvec![0b0000_0000; 2];
384            assert!(try_union_in_impl::<_, _, _, _, LSB>(&lhs, &rhs, &mut dst).is_err());
385        }
386    }
387
388    #[test]
389    fn union_len() {
390        let lhs: u8 = 0b0010_1100;
391        let rhs: u8 = 0b0010_0100;
392        assert_eq!(union_len_impl::<_, _, _, LSB>(&lhs, &rhs), 3);
393
394        let lhs: u8 = 0b0010_1100;
395        let rhs: u8 = 0b0010_0110;
396        assert_eq!(union_len_impl::<_, _, _, LSB>(&lhs, &rhs), 4);
397
398        /////////
399
400        let lhs: u8 = 0b0010_1100;
401        let rhs: [u8; 2] = [0b0010_0100, 0b0000_0000];
402        assert_eq!(union_len_impl::<_, _, _, LSB>(&lhs, &rhs), 3);
403
404        let lhs: u8 = 0b0010_1100;
405        let rhs: [u8; 2] = [0b0010_0100, 0b0101_0000];
406        assert_eq!(union_len_impl::<_, _, _, LSB>(&lhs, &rhs), 5);
407    }
408}