baa/array/
owned.rs

1// Copyright 2023-2024 The Regents of the University of California
2// Copyright 2024 Cornell University
3// released under BSD 3-Clause License
4// author: Kevin Laeufer <laeufer@cornell.edu>
5
6use crate::array::ops::{ArrayMutOps, ArrayOps};
7use crate::{
8    BitVecMutOps, BitVecOps, BitVecValue, BitVecValueMutRef, BitVecValueRef, WidthInt, Word,
9};
10#[cfg(feature = "rand1")]
11use rand::Rng;
12#[cfg(feature = "rand1")]
13use rand::distr::Uniform;
14use std::collections::HashMap;
15use std::fmt::{Debug, Formatter};
16
17/// Owned Array Container.
18#[derive(Clone)]
19#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
20pub struct ArrayValue {
21    data: ArrayImpl,
22}
23
24#[derive(Clone)]
25#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
26enum ArrayImpl {
27    Sparse(SparseArrayValue),
28    Dense(DenseArrayValue),
29}
30
31impl Debug for ArrayValue {
32    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
33        match &self.data {
34            ArrayImpl::Sparse(s) => s.fmt(f),
35            ArrayImpl::Dense(d) => d.fmt(f),
36        }
37    }
38}
39
40impl ArrayValue {
41    pub fn new_dense<'a>(index_width: WidthInt, default: impl Into<BitVecValueRef<'a>>) -> Self {
42        let data = ArrayImpl::Dense(DenseArrayValue::new(index_width, default));
43        Self { data }
44    }
45    pub fn new_sparse<'a>(index_width: WidthInt, default: impl Into<BitVecValueRef<'a>>) -> Self {
46        let data = ArrayImpl::Sparse(SparseArrayValue::new(index_width, default));
47        Self { data }
48    }
49
50    pub fn is_dense(&self) -> bool {
51        matches!(self.data, ArrayImpl::Dense(_))
52    }
53
54    pub fn is_sparse(&self) -> bool {
55        matches!(self.data, ArrayImpl::Sparse(_))
56    }
57
58    /// Turns the underlying storage into a dense representation if it is not already.
59    pub fn make_dense(&mut self) {
60        if self.is_dense() {
61            // nothing to do
62            return;
63        }
64
65        let value: DenseArrayValue = match &self.data {
66            ArrayImpl::Dense(_) => {
67                unreachable!("should have passed the self.is_dense check earlier!");
68            }
69            ArrayImpl::Sparse(data) => data.into(),
70        };
71        self.data = ArrayImpl::Dense(value);
72    }
73
74    /// Checks equivalence for two arrays of the same type (index_width x data_width).
75    pub fn is_equal(&self, other: &Self) -> Option<bool> {
76        match (&self.data, &other.data) {
77            (ArrayImpl::Dense(a), ArrayImpl::Dense(b)) => a.is_equal(b),
78            (ArrayImpl::Sparse(a), ArrayImpl::Sparse(b)) => a.is_equal(b),
79            (ArrayImpl::Sparse(a), ArrayImpl::Dense(b)) => is_equal_mixed(b, a),
80            (ArrayImpl::Dense(a), ArrayImpl::Sparse(b)) => is_equal_mixed(a, b),
81        }
82    }
83
84    #[cfg(feature = "rand1")]
85    pub fn random(rng: &mut impl rand::Rng, index_width: WidthInt, data_width: WidthInt) -> Self {
86        let dense = DenseArrayValue::random(rng, index_width, data_width);
87        Self {
88            data: ArrayImpl::Dense(dense),
89        }
90    }
91}
92
93impl From<&[bool]> for ArrayValue {
94    fn from(value: &[bool]) -> Self {
95        Self {
96            data: ArrayImpl::Dense(value.into()),
97        }
98    }
99}
100
101impl From<&[u8]> for ArrayValue {
102    fn from(value: &[u8]) -> Self {
103        Self {
104            data: ArrayImpl::Dense(value.into()),
105        }
106    }
107}
108
109impl From<&[u64]> for ArrayValue {
110    fn from(value: &[u64]) -> Self {
111        Self {
112            data: ArrayImpl::Dense(value.into()),
113        }
114    }
115}
116
117impl From<ArrayValue> for SparseArrayValue {
118    fn from(value: ArrayValue) -> Self {
119        match value.data {
120            ArrayImpl::Sparse(s) => s,
121            ArrayImpl::Dense(d) => (&d).into(),
122        }
123    }
124}
125
126impl From<&ArrayValue> for SparseArrayValue {
127    fn from(value: &ArrayValue) -> Self {
128        match &value.data {
129            ArrayImpl::Sparse(s) => s.clone(),
130            ArrayImpl::Dense(d) => d.into(),
131        }
132    }
133}
134
135impl PartialEq<ArrayValue> for ArrayValue {
136    fn eq(&self, other: &ArrayValue) -> bool {
137        self.is_equal(other).unwrap_or_default()
138    }
139}
140
141fn is_equal_mixed(dense: &DenseArrayValue, sparse: &SparseArrayValue) -> Option<bool> {
142    // check early to avoid costly computations
143    if dense.index_width != sparse.index_width || dense.data_width != sparse.data_width {
144        return None;
145    }
146
147    // TODO: implement more efficient version which does not densify
148    let other_dense = DenseArrayValue::from(sparse);
149    other_dense.is_equal(dense)
150}
151
152impl ArrayOps for ArrayValue {
153    fn index_width(&self) -> WidthInt {
154        match &self.data {
155            ArrayImpl::Sparse(v) => v.index_width(),
156            ArrayImpl::Dense(v) => v.index_width(),
157        }
158    }
159
160    fn data_width(&self) -> WidthInt {
161        match &self.data {
162            ArrayImpl::Sparse(v) => v.data_width(),
163            ArrayImpl::Dense(v) => v.data_width(),
164        }
165    }
166
167    fn select<'a>(&self, index: impl Into<BitVecValueRef<'a>>) -> BitVecValue {
168        match &self.data {
169            ArrayImpl::Sparse(v) => v.select(index),
170            ArrayImpl::Dense(v) => v.select(index),
171        }
172    }
173}
174
175impl ArrayMutOps for ArrayValue {
176    fn store<'a, 'b>(
177        &mut self,
178        index: impl Into<BitVecValueRef<'a>>,
179        data: impl Into<BitVecValueRef<'b>>,
180    ) {
181        match &mut self.data {
182            ArrayImpl::Sparse(v) => v.store(index, data),
183            ArrayImpl::Dense(v) => v.store(index, data),
184        }
185    }
186
187    fn clear(&mut self) {
188        match &mut self.data {
189            ArrayImpl::Sparse(v) => v.clear(),
190            ArrayImpl::Dense(v) => v.clear(),
191        }
192    }
193}
194
195#[derive(Clone)]
196#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
197struct DenseArrayValue {
198    index_width: WidthInt,
199    data_width: WidthInt,
200    data: DenseArrayImpl,
201}
202
203impl Debug for DenseArrayValue {
204    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
205        write!(f, "DenseArrayValue(..)")
206    }
207}
208
209#[inline]
210fn index_width_from_len(len: usize) -> WidthInt {
211    if len == 0 {
212        0
213    } else {
214        usize::BITS - (len - 1).leading_zeros()
215    }
216}
217
218#[inline]
219fn num_elements_from_index_width(index_width: WidthInt) -> usize {
220    1usize << index_width
221}
222
223impl From<&[bool]> for DenseArrayValue {
224    fn from(values: &[bool]) -> Self {
225        let index_width = index_width_from_len(values.len());
226        let num_elements = num_elements_from_index_width(index_width);
227        let mut data = BitVecValue::zero(num_elements as WidthInt);
228        for (ii, v) in values.iter().enumerate() {
229            if *v {
230                data.set_bit(ii as WidthInt);
231            }
232        }
233        Self {
234            index_width,
235            data_width: 1,
236            data: DenseArrayImpl::Bit(data),
237        }
238    }
239}
240
241impl From<&[u8]> for DenseArrayValue {
242    fn from(values: &[u8]) -> Self {
243        let index_width = index_width_from_len(values.len());
244        let num_elements = num_elements_from_index_width(index_width);
245        let mut data: Vec<u8> = values.into();
246        // pad with zeros if possible
247        data.resize(num_elements, 0);
248        data.shrink_to_fit();
249        Self {
250            index_width,
251            data_width: 8,
252            data: DenseArrayImpl::U8(data),
253        }
254    }
255}
256
257impl From<&[u64]> for DenseArrayValue {
258    fn from(values: &[u64]) -> Self {
259        let index_width = index_width_from_len(values.len());
260        let num_elements = num_elements_from_index_width(index_width);
261        let mut data: Vec<u64> = values.into();
262        // pad with zeros if possible
263        data.resize(num_elements, 0);
264        data.shrink_to_fit();
265        Self {
266            index_width,
267            data_width: 64,
268            data: DenseArrayImpl::U64(data),
269        }
270    }
271}
272
273#[derive(Clone)]
274#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
275enum DenseArrayImpl {
276    Bit(BitVecValue),
277    U8(Vec<u8>),
278    U64(Vec<u64>),
279    Big(Vec<Word>),
280}
281
282/// 1 GiB max for a dense array, otherwise something is fishy!
283const DENSE_ARRAY_MAX_BYTES: usize = 1024 * 1024 * 1024;
284
285fn approx_dense_storage_size(index_width: WidthInt, data_width: WidthInt) -> usize {
286    let elements = 1usize << index_width;
287    if data_width == 1 {
288        // using the Bit representation
289        elements.div_ceil(u8::BITS as usize)
290    } else if data_width <= u8::BITS {
291        // using the U8 representation
292        elements
293    } else if data_width <= u64::BITS {
294        // using the U64 representation
295        elements * (u64::BITS / u8::BITS) as usize
296    } else {
297        // using the Big representation
298        elements * data_width.div_ceil(Word::BITS) as usize
299    }
300}
301
302impl DenseArrayValue {
303    fn new<'a>(index_width: WidthInt, default: impl Into<BitVecValueRef<'a>>) -> Self {
304        let default = default.into();
305        let data_width = default.width();
306        let storage_size = approx_dense_storage_size(index_width, data_width);
307        debug_assert!(
308            storage_size <= DENSE_ARRAY_MAX_BYTES,
309            "array would take up too much space: {storage_size} bytes"
310        );
311
312        let elements = 1usize << index_width;
313        let data = if data_width == 1 {
314            if default.is_true() {
315                DenseArrayImpl::Bit(BitVecValue::ones(elements as WidthInt))
316            } else {
317                debug_assert!(default.is_false());
318                DenseArrayImpl::Bit(BitVecValue::zero(elements as WidthInt))
319            }
320        } else if data_width <= u8::BITS {
321            let default = default.to_u64().unwrap() as u8;
322            DenseArrayImpl::U8(vec![default; elements])
323        } else if data_width <= u64::BITS {
324            let default = default.to_u64().unwrap();
325            DenseArrayImpl::U64(vec![default; elements])
326        } else {
327            let mut v = Vec::with_capacity(elements * default.words().len());
328            for _ in 0..elements {
329                v.extend_from_slice(default.words());
330            }
331            DenseArrayImpl::Big(v)
332        };
333        Self {
334            index_width,
335            data_width,
336            data,
337        }
338    }
339
340    #[cfg(feature = "rand1")]
341    pub fn random(rng: &mut impl rand::Rng, index_width: WidthInt, data_width: WidthInt) -> Self {
342        let storage_size = approx_dense_storage_size(index_width, data_width);
343        debug_assert!(
344            storage_size <= DENSE_ARRAY_MAX_BYTES,
345            "array would take up too much space: {storage_size} bytes"
346        );
347
348        let elements = 1usize << index_width;
349        let data = if data_width == 1 {
350            DenseArrayImpl::Bit(BitVecValue::random(rng, elements as WidthInt))
351        } else if data_width <= u8::BITS {
352            let range = Uniform::try_from(0..(u8::MAX >> (u8::BITS - data_width))).unwrap();
353            DenseArrayImpl::U8(rng.sample_iter(&range).take(elements).collect())
354        } else if data_width <= u64::BITS {
355            let range = Uniform::try_from(0..(u64::MAX >> (u64::BITS - data_width))).unwrap();
356            DenseArrayImpl::U64(rng.sample_iter(&range).take(elements).collect())
357        } else {
358            let words_per_element = data_width.div_ceil(Word::BITS) as usize;
359            let mut v = vec![0 as Word; elements * words_per_element];
360            for ii in 0..elements {
361                let offset = ii * words_per_element;
362                let mut value =
363                    BitVecValueMutRef::new(data_width, &mut v[offset..offset + words_per_element]);
364                value.randomize(rng);
365            }
366            DenseArrayImpl::Big(v)
367        };
368        Self {
369            index_width,
370            data_width,
371            data,
372        }
373    }
374
375    /// Checks equivalence for two arrays of the same type (index_width x data_width).
376    fn is_equal(&self, other: &Self) -> Option<bool> {
377        if other.index_width != self.index_width || other.data_width != self.data_width {
378            return None;
379        }
380        match (&self.data, &other.data) {
381            (DenseArrayImpl::Bit(value_a), DenseArrayImpl::Bit(value_b)) => {
382                Some(value_a.is_equal(value_b))
383            }
384            (DenseArrayImpl::U8(values_a), DenseArrayImpl::U8(values_b)) => {
385                Some(values_a == values_b)
386            }
387            (DenseArrayImpl::U64(values_a), DenseArrayImpl::U64(values_b)) => {
388                Some(values_a == values_b)
389            }
390            (DenseArrayImpl::Big(values_a), DenseArrayImpl::Big(values_b)) => {
391                Some(values_a == values_b)
392            }
393            _ => unreachable!(
394                "the representation for two arrays of the same type should always be the same!"
395            ),
396        }
397    }
398
399    fn select_usize(&self, index: usize) -> BitVecValue {
400        match &self.data {
401            DenseArrayImpl::Bit(value) => value.is_bit_set(index as WidthInt).into(),
402            DenseArrayImpl::U8(values) => {
403                BitVecValue::from_u64(values[index] as u64, self.data_width)
404            }
405            DenseArrayImpl::U64(values) => BitVecValue::from_u64(values[index], self.data_width),
406            DenseArrayImpl::Big(values) => {
407                let start = self.words_per_element() * index;
408                let end = start + self.words_per_element();
409                let value_ref = BitVecValueRef::new(&values[start..end], self.data_width());
410                value_ref.into()
411            }
412        }
413    }
414}
415
416impl ArrayOps for DenseArrayValue {
417    fn index_width(&self) -> WidthInt {
418        self.index_width
419    }
420
421    fn data_width(&self) -> WidthInt {
422        self.data_width
423    }
424
425    fn select<'a>(&self, index: impl Into<BitVecValueRef<'a>>) -> BitVecValue {
426        let index = index.into();
427        debug_assert_eq!(index.width(), self.index_width);
428        self.select_usize(index.to_u64().unwrap() as usize)
429    }
430}
431
432impl ArrayMutOps for DenseArrayValue {
433    fn store<'a, 'b>(
434        &mut self,
435        index: impl Into<BitVecValueRef<'a>>,
436        data: impl Into<BitVecValueRef<'b>>,
437    ) {
438        let index = index.into();
439        debug_assert_eq!(index.width(), self.index_width);
440        let index = index.to_u64().unwrap() as usize;
441        let data = data.into();
442        debug_assert_eq!(data.width(), self.data_width);
443        let data_width = self.data_width();
444        let words_per_element = self.words_per_element();
445        match &mut self.data {
446            DenseArrayImpl::Bit(value) => {
447                if data.is_true() {
448                    value.set_bit(index as WidthInt);
449                } else {
450                    value.clear_bit(index as WidthInt);
451                }
452            }
453            DenseArrayImpl::U8(values) => {
454                values[index] = data.to_u64().unwrap() as u8;
455            }
456            DenseArrayImpl::U64(values) => {
457                values[index] = data.to_u64().unwrap();
458            }
459            DenseArrayImpl::Big(values) => {
460                let start = words_per_element * index;
461                let end = start + words_per_element;
462                let mut value_ref = BitVecValueMutRef::new(data_width, &mut values[start..end]);
463                value_ref.assign(data);
464            }
465        }
466    }
467
468    fn clear(&mut self) {
469        match &mut self.data {
470            DenseArrayImpl::Bit(value) => value.clear(),
471            DenseArrayImpl::U8(values) => values.iter_mut().for_each(|v| *v = 0),
472            DenseArrayImpl::U64(values) => values.iter_mut().for_each(|v| *v = 0),
473            DenseArrayImpl::Big(values) => values.iter_mut().for_each(|v| *v = 0),
474        }
475    }
476}
477
478impl From<&SparseArrayValue> for DenseArrayValue {
479    fn from(value: &SparseArrayValue) -> Self {
480        let mut dense = Self::new(value.index_width, &value.default());
481        debug_assert_eq!(value.data_width, dense.data_width);
482        debug_assert_eq!(value.index_width, dense.index_width);
483        match &value.data {
484            SparseArrayImpl::U64U64(_, m) => {
485                for (index, v) in m.iter() {
486                    match &mut dense.data {
487                        DenseArrayImpl::Bit(value) => {
488                            if *v == 1 {
489                                value.set_bit(*index as WidthInt);
490                            } else {
491                                debug_assert_eq!(*v, 0);
492                                value.clear_bit(*index as WidthInt);
493                            }
494                        }
495                        DenseArrayImpl::U8(values) => {
496                            values[*index as usize] = *v as u8;
497                        }
498                        DenseArrayImpl::U64(values) => {
499                            values[*index as usize] = *v;
500                        }
501                        DenseArrayImpl::Big(_) => {
502                            unreachable!("data_width does not match")
503                        }
504                    }
505                }
506            }
507            SparseArrayImpl::U64Big(_, m) => {
508                let mut index_bv = BitVecValue::zero(value.data_width);
509                for (k, v) in m.iter() {
510                    index_bv.assign_from_u64(*k);
511                    dense.store(&index_bv, v);
512                }
513            }
514            SparseArrayImpl::BigBig(_, _) => {
515                unreachable!(
516                    "index size is too large, creating the dense array should have failed!"
517                )
518            }
519        }
520        dense
521    }
522}
523
524const MAX_HASH_TABLE_SIZE_FOR_DEFAULT_SEARCH: usize = 1024 * 1024;
525
526impl From<&DenseArrayValue> for SparseArrayValue {
527    fn from(value: &DenseArrayValue) -> Self {
528        match &value.data {
529            DenseArrayImpl::Bit(data) => {
530                // find default value
531                let mut count = [0usize; 2];
532                for bit in 0..value.num_elements() {
533                    count[data.is_bit_set(bit as WidthInt) as usize] += 1;
534                }
535                let default = count[1] > count[0];
536                let bv_default = BitVecValue::from(default);
537                let mut out = SparseArrayValue::new(value.index_width(), &bv_default);
538
539                // find all indices that do not carry the default value
540                for bit in 0..value.num_elements() {
541                    if data.is_bit_set(bit as WidthInt) != default {
542                        out.store_u64_u64(bit as u64, !default as u64);
543                    }
544                }
545                out
546            }
547            DenseArrayImpl::U8(data) => {
548                // find default value
549                let mut count = [0usize; u8::BITS as usize];
550                for bit in 0..value.num_elements() {
551                    count[data[bit] as usize] += 1;
552                }
553                let default = count
554                    .into_iter()
555                    .enumerate()
556                    .max_by_key(|(_, count)| *count)
557                    .map(|(ii, _)| ii as u8)
558                    .unwrap();
559                let bv_default = BitVecValue::from_u64(default as u64, value.data_width());
560                let mut out = SparseArrayValue::new(value.index_width(), &bv_default);
561
562                // find all indices that do not carry the default value
563                for (bit, &data) in data.iter().enumerate() {
564                    if data != default {
565                        out.store_u64_u64(bit as u64, data as u64);
566                    }
567                }
568                out
569            }
570            DenseArrayImpl::U64(data) => {
571                // find default value
572                let mut count = HashMap::<u64, usize>::new();
573                for &data in data.iter() {
574                    *count.entry(data).or_insert(0) += 1;
575                    if count.len() > MAX_HASH_TABLE_SIZE_FOR_DEFAULT_SEARCH {
576                        // if our hash table is exploding in size
577                        // => give up and just use the data we have collected so far
578                        break;
579                    }
580                }
581                let default = count
582                    .into_iter()
583                    .max_by_key(|(_, count)| *count)
584                    .map(|(v, _)| v)
585                    .unwrap();
586                let bv_default = BitVecValue::from_u64(default, value.data_width());
587                let mut out = SparseArrayValue::new(value.index_width(), &bv_default);
588
589                // find all indices that do not carry the default value
590                for (bit, &data) in data.iter().enumerate() {
591                    if data != default {
592                        out.store_u64_u64(bit as u64, data);
593                    }
594                }
595                out
596            }
597            DenseArrayImpl::Big(_) => {
598                // find default value
599                let mut count = HashMap::<BitVecValue, usize>::new();
600                for bit in 0..value.num_elements() {
601                    let value = value.select_usize(bit);
602                    *count.entry(value).or_insert(0) += 1;
603                    if count.len() > MAX_HASH_TABLE_SIZE_FOR_DEFAULT_SEARCH {
604                        // if our hash table is exploding in size
605                        // => give up and just use the data we have collected so far
606                        break;
607                    }
608                }
609                let default = count
610                    .into_iter()
611                    .max_by_key(|(_, count)| *count)
612                    .map(|(v, _)| v)
613                    .unwrap();
614                let mut out = SparseArrayValue::new(value.index_width(), &default);
615
616                // find all indices that do not carry the default value
617                for bit in 0..value.num_elements() {
618                    let value = value.select_usize(bit);
619                    if value != default {
620                        out.store_u64_big(bit as u64, &value);
621                    }
622                }
623                out
624            }
625        }
626    }
627}
628
629#[derive(Clone)]
630#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
631pub struct SparseArrayValue {
632    index_width: WidthInt,
633    data_width: WidthInt,
634    data: SparseArrayImpl,
635}
636
637impl Debug for SparseArrayValue {
638    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
639        write!(f, "SparseArrayValue(..)")
640    }
641}
642
643#[derive(Clone)]
644#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
645enum SparseArrayImpl {
646    U64U64(u64, HashMap<u64, u64>),
647    U64Big(BitVecValue, HashMap<u64, BitVecValue>),
648    BigBig(BitVecValue, HashMap<BitVecValue, BitVecValue>),
649}
650
651impl SparseArrayValue {
652    pub fn new<'a>(index_width: WidthInt, default: impl Into<BitVecValueRef<'a>>) -> Self {
653        let default = default.into();
654        let data_width = default.width();
655        let data = if index_width > Word::BITS {
656            SparseArrayImpl::BigBig(default.into(), HashMap::new())
657        } else if data_width > Word::BITS {
658            SparseArrayImpl::U64Big(default.into(), HashMap::new())
659        } else {
660            SparseArrayImpl::U64U64(default.to_u64().unwrap(), HashMap::new())
661        };
662        Self {
663            index_width,
664            data_width,
665            data,
666        }
667    }
668
669    /// Checks equivalence for two arrays of the same type (index_width x data_width).
670    fn is_equal(&self, other: &Self) -> Option<bool> {
671        if other.index_width != self.index_width || other.data_width != self.data_width {
672            return None;
673        }
674        match (&self.data, &other.data) {
675            (
676                SparseArrayImpl::U64U64(default_a, map_a),
677                SparseArrayImpl::U64U64(default_b, map_b),
678            ) => {
679                if default_a == default_b {
680                    // here we rely on the fact that the default value may never appear in the map
681                    if map_a.len() == map_b.len() {
682                        return Some(map_a == map_b);
683                    }
684                }
685                Some(false)
686            }
687            (
688                SparseArrayImpl::U64Big(default_a, map_a),
689                SparseArrayImpl::U64Big(default_b, map_b),
690            ) => {
691                if default_a == default_b {
692                    // here we rely on the fact that the default value may never appear in the map
693                    if map_a.len() == map_b.len() {
694                        return Some(map_a == map_b);
695                    }
696                }
697                Some(false)
698            }
699            (
700                SparseArrayImpl::BigBig(default_a, map_a),
701                SparseArrayImpl::BigBig(default_b, map_b),
702            ) => {
703                if default_a == default_b {
704                    // here we rely on the fact that the default value may never appear in the map
705                    if map_a.len() == map_b.len() {
706                        return Some(map_a == map_b);
707                    }
708                }
709                Some(false)
710            }
711            _ => unreachable!(
712                "the representation for two arrays of the same type should always be the same!"
713            ),
714        }
715    }
716
717    pub fn default(&self) -> BitVecValue {
718        match &self.data {
719            SparseArrayImpl::U64U64(default, _) => BitVecValue::from_u64(*default, self.data_width),
720            SparseArrayImpl::U64Big(default, _) => default.clone(),
721            SparseArrayImpl::BigBig(default, _) => default.clone(),
722        }
723    }
724
725    fn store_u64_u64(&mut self, index: u64, data: u64) {
726        match &mut self.data {
727            SparseArrayImpl::U64U64(default, map) => {
728                if data == *default {
729                    // ensures that the default value is used for the given index
730                    map.remove(&index);
731                } else {
732                    map.insert(index, data);
733                }
734            }
735            _ => panic!(
736                "store_u64_u64 can only be used when index and data bit width is at or under 64 bit."
737            ),
738        }
739    }
740
741    fn store_u64_big<'b>(&mut self, index: u64, data: impl Into<BitVecValueRef<'b>>) {
742        let data = data.into();
743        match &mut self.data {
744            SparseArrayImpl::U64U64(default, map) => {
745                let data = data.to_u64().unwrap();
746                if data == *default {
747                    // ensures that the default value is used for the given index
748                    map.remove(&index);
749                } else {
750                    map.insert(index, data);
751                }
752            }
753            SparseArrayImpl::U64Big(default, map) => {
754                if data.is_equal(default) {
755                    // ensures that the default value is used for the given index
756                    map.remove(&index);
757                } else {
758                    map.insert(index, data.into());
759                }
760            }
761            _ => panic!(
762                "store_u64_big can only be used when the index bit width is at or under 64 bit."
763            ),
764        }
765    }
766
767    /// Returns an iterator to all index/data pairs that do not evaluate to the default value.
768    /// Note that the order of entries is non-deterministic. Convert to a Vec and sort of you
769    /// want to ensure a predictable outcome.
770    pub fn non_default_entries(&self) -> SparseArrayEntryIter<'_> {
771        let underlying = match &self.data {
772            SparseArrayImpl::U64U64(_, m) => SparseArrayIterImpl::U64U64(m.iter()),
773            SparseArrayImpl::U64Big(_, m) => SparseArrayIterImpl::U64Big(m.iter()),
774            SparseArrayImpl::BigBig(_, m) => SparseArrayIterImpl::BigBig(m.iter()),
775        };
776        let (index_width, data_width) = (self.index_width, self.data_width);
777        SparseArrayEntryIter {
778            index_width,
779            data_width,
780            underlying,
781        }
782    }
783}
784
785pub struct SparseArrayEntryIter<'a> {
786    index_width: WidthInt,
787    data_width: WidthInt,
788    underlying: SparseArrayIterImpl<'a>,
789}
790
791enum SparseArrayIterImpl<'a> {
792    U64U64(std::collections::hash_map::Iter<'a, u64, u64>),
793    U64Big(std::collections::hash_map::Iter<'a, u64, BitVecValue>),
794    BigBig(std::collections::hash_map::Iter<'a, BitVecValue, BitVecValue>),
795}
796
797impl Iterator for SparseArrayEntryIter<'_> {
798    type Item = (BitVecValue, BitVecValue);
799
800    fn next(&mut self) -> Option<Self::Item> {
801        match &mut self.underlying {
802            SparseArrayIterImpl::U64U64(iter) => {
803                let (index, value) = iter.next()?;
804                Some((
805                    BitVecValue::from_u64(*index, self.index_width),
806                    BitVecValue::from_u64(*value, self.data_width),
807                ))
808            }
809            SparseArrayIterImpl::U64Big(iter) => {
810                let (index, value) = iter.next()?;
811                Some((
812                    BitVecValue::from_u64(*index, self.index_width),
813                    value.clone(),
814                ))
815            }
816            SparseArrayIterImpl::BigBig(iter) => {
817                let (index, value) = iter.next()?;
818                Some((index.clone(), value.clone()))
819            }
820        }
821    }
822}
823
824impl ArrayOps for SparseArrayValue {
825    fn index_width(&self) -> WidthInt {
826        self.index_width
827    }
828    fn data_width(&self) -> WidthInt {
829        self.data_width
830    }
831    fn select<'a>(&self, index: impl Into<BitVecValueRef<'a>>) -> BitVecValue {
832        let index = index.into();
833        debug_assert_eq!(index.width(), self.index_width);
834        match &self.data {
835            SparseArrayImpl::U64U64(default, map) => {
836                let index = index.to_u64().unwrap();
837                let value = map.get(&index).copied().unwrap_or(*default);
838                BitVecValue::from_u64(value, self.data_width)
839            }
840            SparseArrayImpl::U64Big(default, map) => {
841                let index = index.to_u64().unwrap();
842
843                map.get(&index).cloned().unwrap_or_else(|| default.clone())
844            }
845            SparseArrayImpl::BigBig(default, map) => {
846                map.get(&index).cloned().unwrap_or_else(|| default.clone())
847            }
848        }
849    }
850}
851
852impl ArrayMutOps for SparseArrayValue {
853    fn store<'a, 'b>(
854        &mut self,
855        index: impl Into<BitVecValueRef<'a>>,
856        data: impl Into<BitVecValueRef<'b>>,
857    ) {
858        let index = index.into();
859        debug_assert_eq!(index.width(), self.index_width);
860        let data = data.into();
861        debug_assert_eq!(data.width(), self.data_width);
862        match &mut self.data {
863            SparseArrayImpl::U64U64(default, map) => {
864                let index = index.to_u64().unwrap();
865                let data = data.to_u64().unwrap();
866                if data == *default {
867                    // ensures that the default value is used for the given index
868                    map.remove(&index);
869                } else {
870                    map.insert(index, data);
871                }
872            }
873            SparseArrayImpl::U64Big(default, map) => {
874                let index = index.to_u64().unwrap();
875                if data.is_equal(default) {
876                    // ensures that the default value is used for the given index
877                    map.remove(&index);
878                } else {
879                    map.insert(index, data.into());
880                }
881            }
882            SparseArrayImpl::BigBig(default, map) => {
883                if data.is_equal(default) {
884                    // ensures that the default value is used for the given index
885                    map.remove(&index);
886                } else {
887                    map.insert(index.into(), data.into());
888                }
889            }
890        }
891    }
892
893    fn clear(&mut self) {
894        match &mut self.data {
895            SparseArrayImpl::U64U64(default, map) => {
896                *default = 0;
897                map.clear();
898            }
899            SparseArrayImpl::U64Big(default, map) => {
900                *default = BitVecValue::zero(self.data_width);
901                map.clear();
902            }
903            SparseArrayImpl::BigBig(default, map) => {
904                *default = BitVecValue::zero(self.data_width);
905                map.clear();
906            }
907        }
908    }
909}
910
911#[cfg(test)]
912mod tests {
913    use super::*;
914
915    #[test]
916    fn type_size() {
917        // Dense Array Size
918        assert_eq!(std::mem::size_of::<Vec<u8>>(), 24);
919        assert_eq!(std::mem::size_of::<BitVecValue>(), 3 * 8);
920        assert_eq!(std::mem::size_of::<DenseArrayImpl>(), 4 * 8); // BitVecValue size + tag + padding
921        assert_eq!(std::mem::size_of::<DenseArrayValue>(), (4 + 1) * 8); // Impl + size + padding
922
923        // Sparse Array Size
924
925        // the hash table size is independent of the key/value types
926        assert_eq!(std::mem::size_of::<HashMap<u64, u64>>(), 48);
927        assert_eq!(std::mem::size_of::<HashMap<u64, BitVecValue>>(), 48);
928        assert_eq!(std::mem::size_of::<HashMap<BitVecValue, BitVecValue>>(), 48);
929
930        // HashMap + BitVecValue + tag + padding
931        assert_eq!(std::mem::size_of::<SparseArrayImpl>(), 48 + 24 + 8);
932        assert_eq!(std::mem::size_of::<SparseArrayValue>(), 80 + 8); // Impl + size + padding
933    }
934
935    #[test]
936    fn test_index_width_from_len() {
937        assert_eq!(index_width_from_len(0), 0);
938        assert_eq!(index_width_from_len(1), 0);
939        assert_eq!(index_width_from_len(2), 1);
940        assert_eq!(index_width_from_len(3), 2);
941        assert_eq!(index_width_from_len(8), 3);
942        assert_eq!(index_width_from_len(9), 4);
943    }
944
945    #[test]
946    fn test_conversions_bool() {
947        let dense0: DenseArrayValue = [true, false, false, false, true].as_slice().into();
948        let sparse0: SparseArrayValue = (&dense0).into();
949        let dense1: DenseArrayValue = (&sparse0).into();
950        assert!(is_equal_mixed(&dense0, &sparse0).unwrap());
951        assert!(dense0.is_equal(&dense1).unwrap());
952    }
953
954    #[test]
955    fn test_conversions_u8() {
956        let dense0: DenseArrayValue = [1u8, 2, 3, 4, 5, 6, 1, 1, 1].as_slice().into();
957        let sparse0: SparseArrayValue = (&dense0).into();
958        let dense1: DenseArrayValue = (&sparse0).into();
959        assert!(is_equal_mixed(&dense0, &sparse0).unwrap());
960        assert!(dense0.is_equal(&dense1).unwrap());
961    }
962
963    #[test]
964    fn test_conversions_u64() {
965        let dense0: DenseArrayValue = [1u64, 2, 3, 4, 5, 6, 1, 1, 1].as_slice().into();
966        let sparse0: SparseArrayValue = (&dense0).into();
967        let dense1: DenseArrayValue = (&sparse0).into();
968        assert!(is_equal_mixed(&dense0, &sparse0).unwrap());
969        assert!(dense0.is_equal(&dense1).unwrap());
970    }
971}