1#![no_std]
32
33extern crate alloc;
34
35use alloc::{boxed::Box, vec, vec::Vec};
36
37use core::cmp::max;
38use core::fmt;
39use core::hash;
40use core::mem::{forget, replace, size_of};
41use core::ops::{Index, Range};
42use core::slice;
43
44#[macro_export]
74macro_rules! sbvec {
75 ($elem:expr; $n:expr) => (
76 $crate::SmallBitVec::from_elem($n, $elem)
77 );
78 ($($x:expr),*) => (
79 [$($x),*].iter().cloned().collect::<$crate::SmallBitVec>()
80 );
81 ($($x:expr,)*) => (
82 sbvec![$($x),*]
83 );
84}
85
86macro_rules! const_debug_assert_le {
90 ($left: ident <= $right: expr) => {
91 #[cfg(debug_assertions)]
92 [(); $right + 1][$left];
94 };
95}
96
97#[cfg(test)]
98mod tests;
99
100pub struct SmallBitVec {
105 data: usize,
106}
107
108#[inline(always)]
110const fn inline_bits() -> usize {
111 size_of::<usize>() * 8
112}
113
114#[inline(always)]
119const fn inline_capacity() -> usize {
120 inline_bits() - 2
121}
122
123#[inline(always)]
125const fn inline_shift(n: usize) -> usize {
126 const_debug_assert_le!(n <= inline_capacity());
127 inline_bits() - 1 - n
129}
130
131#[inline(always)]
133const fn inline_index(n: usize) -> usize {
134 1 << inline_shift(n)
135}
136
137#[inline(always)]
139fn inline_ones(n: usize) -> usize {
140 if n == 0 {
141 0
142 } else {
143 !0 << (inline_bits() - n)
144 }
145}
146
147const HEAP_FLAG: usize = 1;
150
151type Storage = usize;
153
154#[inline(always)]
156fn bits_per_storage() -> usize {
157 size_of::<Storage>() * 8
158}
159
160struct Header {
164 len: Storage,
166
167 buffer_len: Storage,
169}
170
171impl Header {
172 fn new(cap: usize, len: usize, val: bool) -> *mut Header {
175 let alloc_len = header_len() + buffer_len(cap);
176 let init = if val { !0 } else { 0 };
177
178 let v: Vec<Storage> = vec![init; alloc_len];
179
180 let buffer_len = v.capacity() - header_len();
181 let header_ptr = v.as_ptr() as *mut Header;
182
183 forget(v);
184
185 unsafe {
186 (*header_ptr).len = len;
187 (*header_ptr).buffer_len = buffer_len;
188 }
189 header_ptr
190 }
191}
192
193#[inline(always)]
195fn header_len() -> usize {
196 size_of::<Header>() / size_of::<Storage>()
197}
198
199#[inline(always)]
201fn buffer_len(cap: usize) -> usize {
202 let bits = bits_per_storage();
203 cap.checked_add(bits - 1)
204 .expect("capacity overflow")
205 / bits
206}
207
208pub enum InternalStorage {
212 Inline(usize),
215
216 Spilled(Box<[usize]>),
218}
219
220impl SmallBitVec {
221 #[inline]
223 pub const fn new() -> SmallBitVec {
224 SmallBitVec {
225 data: inline_index(0),
226 }
227 }
228
229 #[inline]
231 pub fn from_elem(len: usize, val: bool) -> SmallBitVec {
232 if len <= inline_capacity() {
233 return SmallBitVec {
234 data: if val {
235 inline_ones(len + 1)
236 } else {
237 inline_index(len)
238 },
239 };
240 }
241 let header_ptr = Header::new(len, len, val);
242 SmallBitVec {
243 data: (header_ptr as usize) | HEAP_FLAG,
244 }
245 }
246
247 #[inline]
250 pub fn with_capacity(cap: usize) -> SmallBitVec {
251 if cap <= inline_capacity() {
253 return SmallBitVec::new();
254 }
255
256 let header_ptr = Header::new(cap, 0, false);
258 SmallBitVec {
259 data: (header_ptr as usize) | HEAP_FLAG,
260 }
261 }
262
263 #[inline]
265 pub fn len(&self) -> usize {
266 if self.is_inline() {
267 inline_bits() - self.data.trailing_zeros() as usize - 1
270 } else {
271 self.header().len
272 }
273 }
274
275 #[inline]
277 pub fn is_empty(&self) -> bool {
278 self.len() == 0
279 }
280
281 #[inline]
283 pub fn capacity(&self) -> usize {
284 if self.is_inline() {
285 inline_capacity()
286 } else {
287 self.header().buffer_len * bits_per_storage()
288 }
289 }
290
291 #[inline]
293 pub fn get(&self, n: usize) -> Option<bool> {
294 if n < self.len() {
295 Some(unsafe { self.get_unchecked(n) })
296 } else {
297 None
298 }
299 }
300
301 #[inline]
303 pub fn last(&self) -> Option<bool> {
304 self.len()
305 .checked_sub(1)
306 .map(|n| unsafe { self.get_unchecked(n) })
307 }
308
309 #[inline]
311 pub unsafe fn get_unchecked(&self, n: usize) -> bool {
312 if self.is_inline() {
313 self.data & inline_index(n) != 0
314 } else {
315 let buffer = self.buffer();
316 let i = n / bits_per_storage();
317 let offset = n % bits_per_storage();
318 *buffer.get_unchecked(i) & (1 << offset) != 0
319 }
320 }
321
322 #[inline]
324 pub fn set(&mut self, n: usize, val: bool) {
325 assert!(n < self.len(), "Index {} out of bounds", n);
326 unsafe {
327 self.set_unchecked(n, val);
328 }
329 }
330
331 #[inline]
333 pub unsafe fn set_unchecked(&mut self, n: usize, val: bool) {
334 if self.is_inline() {
335 if val {
336 self.data |= inline_index(n);
337 } else {
338 self.data &= !inline_index(n);
339 }
340 } else {
341 let buffer = self.buffer_mut();
342 let i = n / bits_per_storage();
343 let offset = n % bits_per_storage();
344 if val {
345 *buffer.get_unchecked_mut(i) |= 1 << offset;
346 } else {
347 *buffer.get_unchecked_mut(i) &= !(1 << offset);
348 }
349 }
350 }
351
352 #[inline]
363 pub fn push(&mut self, val: bool) {
364 let idx = self.len();
365 if idx == self.capacity() {
366 self.reserve(1);
367 }
368 unsafe {
369 self.set_len(idx + 1);
370 self.set_unchecked(idx, val);
371 }
372 }
373
374 #[inline]
386 pub fn pop(&mut self) -> Option<bool> {
387 self.len().checked_sub(1).map(|last| unsafe {
388 let val = self.get_unchecked(last);
389 self.set_len(last);
390 val
391 })
392 }
393
394 #[inline]
398 pub fn remove(&mut self, idx: usize) -> bool {
399 let len = self.len();
400 let val = self[idx];
401
402 if self.is_inline() {
403 let mask = !inline_ones(idx);
405 let new_vals = (self.data & mask) << 1;
406 self.data = (self.data & !mask) | (new_vals & mask);
407 } else {
408 let first = idx / bits_per_storage();
409 let offset = idx % bits_per_storage();
410 let count = buffer_len(len);
411 {
412 let buf = self.buffer_mut();
414 let mask = !0 << offset;
415 let new_vals = (buf[first] & mask) >> 1;
416 buf[first] = (buf[first] & !mask) | (new_vals & mask);
417 }
418 for i in (first + 1)..count {
420 let bit_idx = i * bits_per_storage();
422 unsafe {
423 let first_bit = self.get_unchecked(bit_idx);
424 self.set_unchecked(bit_idx - 1, first_bit);
425 }
426 self.buffer_mut()[i] >>= 1;
428 }
429 unsafe {
431 self.set_len(len - 1);
432 }
433 }
434 val
435 }
436
437 #[inline]
439 pub fn clear(&mut self) {
440 unsafe {
441 self.set_len(0);
442 }
443 }
444
445 #[inline]
453 pub fn reserve(&mut self, additional: usize) {
454 let old_cap = self.capacity();
455 let new_cap = self
456 .len()
457 .checked_add(additional)
458 .expect("capacity overflow");
459 if new_cap <= old_cap {
460 return;
461 }
462 let double_cap = old_cap.saturating_mul(2);
464 self.reallocate(max(new_cap, double_cap));
465 }
466
467 #[inline]
472 unsafe fn set_len(&mut self, len: usize) {
473 debug_assert!(len <= self.capacity());
474 if self.is_inline() {
475 let sentinel = inline_index(len);
476 let mask = !(sentinel - 1);
477 self.data |= sentinel;
478 self.data &= mask;
479 } else {
480 self.header_mut().len = len;
481 }
482 }
483
484 #[inline]
486 pub fn iter(&self) -> Iter {
487 Iter {
488 vec: self,
489 range: 0..self.len(),
490 }
491 }
492
493 #[inline]
501 pub fn range(&self, range: Range<usize>) -> VecRange {
502 assert!(range.end <= self.len(), "range out of bounds");
503 VecRange { vec: &self, range }
504 }
505
506 #[inline]
510 pub fn all_false(&self) -> bool {
511 let mut len = self.len();
512 if len == 0 {
513 return true;
514 }
515
516 if self.is_inline() {
517 let mask = inline_ones(len);
518 self.data & mask == 0
519 } else {
520 for &storage in self.buffer() {
521 if len >= bits_per_storage() {
522 if storage != 0 {
523 return false;
524 }
525 len -= bits_per_storage();
526 } else {
527 let mask = (1 << len) - 1;
528 if storage & mask != 0 {
529 return false;
530 }
531 break;
532 }
533 }
534 true
535 }
536 }
537
538 #[inline]
542 pub fn all_true(&self) -> bool {
543 let mut len = self.len();
544 if len == 0 {
545 return true;
546 }
547
548 if self.is_inline() {
549 let mask = inline_ones(len);
550 self.data & mask == mask
551 } else {
552 for &storage in self.buffer() {
553 if len >= bits_per_storage() {
554 if storage != !0 {
555 return false;
556 }
557 len -= bits_per_storage();
558 } else {
559 let mask = (1 << len) - 1;
560 if storage & mask != mask {
561 return false;
562 }
563 break;
564 }
565 }
566 true
567 }
568 }
569
570 pub fn truncate(&mut self, len: usize) {
577 unsafe {
578 if len < self.len() {
579 self.set_len(len);
580 }
581 }
582 }
583
584 pub fn resize(&mut self, len: usize, value: bool) {
591 let old_len = self.len();
592
593 if len > old_len {
594 unsafe {
595 self.reallocate(len);
596 self.set_len(len);
597 for i in old_len..len {
598 self.set(i, value);
599 }
600 }
601 } else {
602 self.truncate(len);
603 }
604 }
605
606 fn reallocate(&mut self, cap: usize) {
610 let old_cap = self.capacity();
611 if cap <= old_cap {
612 return;
613 }
614 assert!(self.len() <= cap);
615
616 if self.is_heap() {
617 let old_buffer_len = self.header().buffer_len;
618 let new_buffer_len = buffer_len(cap);
619
620 let old_alloc_len = header_len() + old_buffer_len;
621 let new_alloc_len = header_len() + new_buffer_len;
622
623 let old_ptr = self.header_raw() as *mut Storage;
624 let mut v = unsafe { Vec::from_raw_parts(old_ptr, old_alloc_len, old_alloc_len) };
625 v.resize(new_alloc_len, 0);
626 v.shrink_to_fit();
627 self.data = v.as_ptr() as usize | HEAP_FLAG;
628 forget(v);
629
630 self.header_mut().buffer_len = new_buffer_len;
631 } else {
632 let old_self = replace(self, SmallBitVec::with_capacity(cap));
633 unsafe {
634 self.set_len(old_self.len());
635 for i in 0..old_self.len() {
636 self.set_unchecked(i, old_self.get_unchecked(i));
637 }
638 }
639 }
640 }
641
642 #[inline]
646 pub fn heap_ptr(&self) -> Option<*const usize> {
647 if self.is_heap() {
648 Some((self.data & !HEAP_FLAG) as *const Storage)
649 } else {
650 None
651 }
652 }
653
654 #[inline]
658 pub fn into_storage(self) -> InternalStorage {
659 if self.is_heap() {
660 let alloc_len = header_len() + self.header().buffer_len;
661 let ptr = self.header_raw() as *mut Storage;
662 let slice = unsafe { Box::from_raw(slice::from_raw_parts_mut(ptr, alloc_len)) };
663 forget(self);
664 InternalStorage::Spilled(slice)
665 } else {
666 InternalStorage::Inline(self.data)
667 }
668 }
669
670 pub unsafe fn from_storage(storage: InternalStorage) -> SmallBitVec {
703 match storage {
704 InternalStorage::Inline(data) => SmallBitVec { data },
705 InternalStorage::Spilled(vs) => {
706 let ptr = Box::into_raw(vs);
707 SmallBitVec {
708 data: (ptr as *mut usize as usize) | HEAP_FLAG,
709 }
710 }
711 }
712 }
713
714 #[inline]
716 fn is_inline(&self) -> bool {
717 self.data & HEAP_FLAG == 0
718 }
719
720 #[inline]
722 fn is_heap(&self) -> bool {
723 !self.is_inline()
724 }
725
726 #[inline]
728 fn header_raw(&self) -> *mut Header {
729 assert!(self.is_heap());
730 (self.data & !HEAP_FLAG) as *mut Header
731 }
732
733 #[inline]
734 fn header_mut(&mut self) -> &mut Header {
735 unsafe { &mut *self.header_raw() }
736 }
737
738 #[inline]
739 fn header(&self) -> &Header {
740 unsafe { &*self.header_raw() }
741 }
742
743 #[inline]
745 fn buffer_raw(&self) -> *mut [Storage] {
746 unsafe {
747 let header_ptr = self.header_raw();
748 let buffer_len = (*header_ptr).buffer_len;
749 let buffer_ptr = (header_ptr as *mut Storage)
750 .offset((size_of::<Header>() / size_of::<Storage>()) as isize);
751 slice::from_raw_parts_mut(buffer_ptr, buffer_len)
752 }
753 }
754
755 #[inline]
756 fn buffer_mut(&mut self) -> &mut [Storage] {
757 unsafe { &mut *self.buffer_raw() }
758 }
759
760 #[inline]
761 fn buffer(&self) -> &[Storage] {
762 unsafe { &*self.buffer_raw() }
763 }
764}
765
766impl fmt::Debug for SmallBitVec {
769 #[inline]
770 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
771 fmt.debug_list()
772 .entries(self.iter().map(|b| b as u8))
773 .finish()
774 }
775}
776
777impl Default for SmallBitVec {
778 #[inline]
779 fn default() -> Self {
780 Self::new()
781 }
782}
783
784impl PartialEq for SmallBitVec {
785 fn eq(&self, other: &Self) -> bool {
786 if self.is_inline() && other.is_inline() {
788 return self.data == other.data;
789 }
790
791 let len = self.len();
792 if len != other.len() {
793 return false;
794 }
795
796 if self.is_heap() && other.is_heap() {
798 let buf0 = self.buffer();
799 let buf1 = other.buffer();
800
801 let full_blocks = len / bits_per_storage();
802 let remainder = len % bits_per_storage();
803
804 if buf0[..full_blocks] != buf1[..full_blocks] {
805 return false;
806 }
807
808 if remainder != 0 {
809 let mask = (1 << remainder) - 1;
810 if buf0[full_blocks] & mask != buf1[full_blocks] & mask {
811 return false;
812 }
813 }
814 return true;
815 }
816
817 Iterator::eq(self.iter(), other.iter())
819 }
820}
821
822impl Eq for SmallBitVec {}
823
824impl Drop for SmallBitVec {
825 fn drop(&mut self) {
826 if self.is_heap() {
827 unsafe {
828 let header_ptr = self.header_raw();
829 let alloc_ptr = header_ptr as *mut Storage;
830 let alloc_len = header_len() + (*header_ptr).buffer_len;
831 Vec::from_raw_parts(alloc_ptr, alloc_len, alloc_len);
832 }
833 }
834 }
835}
836
837impl Clone for SmallBitVec {
838 fn clone(&self) -> Self {
839 if self.is_inline() {
840 return SmallBitVec { data: self.data };
841 }
842
843 let buffer_len = self.header().buffer_len;
844 let alloc_len = header_len() + buffer_len;
845 let ptr = self.header_raw() as *mut Storage;
846 let raw_allocation = unsafe { slice::from_raw_parts(ptr, alloc_len) };
847
848 let v = raw_allocation.to_vec();
849 let header_ptr = v.as_ptr() as *mut Header;
850 forget(v);
851 SmallBitVec {
852 data: (header_ptr as usize) | HEAP_FLAG,
853 }
854 }
855}
856
857impl Index<usize> for SmallBitVec {
858 type Output = bool;
859
860 #[inline(always)]
861 fn index(&self, i: usize) -> &bool {
862 assert!(i < self.len(), "index out of range");
863 if self.get(i).unwrap() {
864 &true
865 } else {
866 &false
867 }
868 }
869}
870
871#[cfg(feature = "malloc_size_of")]
872impl malloc_size_of::MallocSizeOf for SmallBitVec {
873 fn size_of(&self, ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
874 if let Some(ptr) = self.heap_ptr() {
875 unsafe { ops.malloc_size_of(ptr) }
876 } else {
877 0
878 }
879 }
880}
881
882impl hash::Hash for SmallBitVec {
883 #[inline]
884 fn hash<H: hash::Hasher>(&self, state: &mut H) {
885 let len = self.len();
886 len.hash(state);
887 if self.is_inline() {
888 (self.data & inline_ones(len)).reverse_bits().hash(state);
889 } else {
890 let full_blocks = len / bits_per_storage();
891 let remainder = len % bits_per_storage();
892 let buffer = self.buffer();
893 if full_blocks != 0 {
894 buffer[..full_blocks].hash(state);
895 }
896 if remainder != 0 {
897 let mask = (1 << remainder) - 1;
898 (buffer[full_blocks] & mask).hash(state);
899 }
900 }
901 }
902}
903
904impl Extend<bool> for SmallBitVec {
905 #[inline]
906 fn extend<I: IntoIterator<Item = bool>>(&mut self, iter: I) {
907 let iter = iter.into_iter();
908
909 let (min, _) = iter.size_hint();
910 assert!(min <= usize::max_value(), "capacity overflow");
911 self.reserve(min);
912
913 for element in iter {
914 self.push(element)
915 }
916 }
917}
918
919impl FromIterator<bool> for SmallBitVec {
920 #[inline]
921 fn from_iter<I: IntoIterator<Item = bool>>(iter: I) -> Self {
922 let mut v = SmallBitVec::new();
923 v.extend(iter);
924 v
925 }
926}
927
928impl IntoIterator for SmallBitVec {
929 type Item = bool;
930 type IntoIter = IntoIter;
931
932 #[inline]
933 fn into_iter(self) -> IntoIter {
934 IntoIter {
935 range: 0..self.len(),
936 vec: self,
937 }
938 }
939}
940
941impl<'a> IntoIterator for &'a SmallBitVec {
942 type Item = bool;
943 type IntoIter = Iter<'a>;
944
945 #[inline]
946 fn into_iter(self) -> Iter<'a> {
947 self.iter()
948 }
949}
950
951pub struct IntoIter {
957 vec: SmallBitVec,
958 range: Range<usize>,
959}
960
961impl Iterator for IntoIter {
962 type Item = bool;
963
964 #[inline]
965 fn next(&mut self) -> Option<bool> {
966 self.range
967 .next()
968 .map(|i| unsafe { self.vec.get_unchecked(i) })
969 }
970
971 #[inline]
972 fn size_hint(&self) -> (usize, Option<usize>) {
973 self.range.size_hint()
974 }
975}
976
977impl DoubleEndedIterator for IntoIter {
978 #[inline]
979 fn next_back(&mut self) -> Option<bool> {
980 self.range
981 .next_back()
982 .map(|i| unsafe { self.vec.get_unchecked(i) })
983 }
984}
985
986impl ExactSizeIterator for IntoIter {}
987
988pub struct Iter<'a> {
994 vec: &'a SmallBitVec,
995 range: Range<usize>,
996}
997
998impl<'a> Default for Iter<'a> {
999 #[inline]
1000 fn default() -> Self {
1001 const EMPTY: &'static SmallBitVec = &SmallBitVec::new();
1002 Self {
1003 vec: EMPTY,
1004 range: 0..0,
1005 }
1006 }
1007}
1008
1009impl<'a> Iterator for Iter<'a> {
1010 type Item = bool;
1011
1012 #[inline]
1013 fn next(&mut self) -> Option<bool> {
1014 self.range
1015 .next()
1016 .map(|i| unsafe { self.vec.get_unchecked(i) })
1017 }
1018
1019 #[inline]
1020 fn size_hint(&self) -> (usize, Option<usize>) {
1021 self.range.size_hint()
1022 }
1023}
1024
1025impl<'a> DoubleEndedIterator for Iter<'a> {
1026 #[inline]
1027 fn next_back(&mut self) -> Option<bool> {
1028 self.range
1029 .next_back()
1030 .map(|i| unsafe { self.vec.get_unchecked(i) })
1031 }
1032}
1033
1034impl<'a> ExactSizeIterator for Iter<'a> {}
1035
1036#[derive(Debug, Clone)]
1042pub struct VecRange<'a> {
1043 vec: &'a SmallBitVec,
1044 range: Range<usize>,
1045}
1046
1047impl<'a> VecRange<'a> {
1048 #[inline]
1049 pub fn iter(&self) -> Iter<'a> {
1050 Iter {
1051 vec: self.vec,
1052 range: self.range.clone(),
1053 }
1054 }
1055}
1056
1057impl<'a> Index<usize> for VecRange<'a> {
1058 type Output = bool;
1059
1060 #[inline]
1061 fn index(&self, i: usize) -> &bool {
1062 let vec_i = i + self.range.start;
1063 assert!(vec_i < self.range.end, "index out of range");
1064 &self.vec[vec_i]
1065 }
1066}
1067
1068#[cfg(test)]
1069mod overflow_tests {
1070 use super::*;
1071
1072 #[test]
1073 #[should_panic(expected = "capacity overflow")]
1074 fn test_poc1_from_elem_overflow() {
1075 let _v = SmallBitVec::from_elem(usize::MAX, false);
1076 }
1077
1078 #[test]
1079 #[should_panic(expected = "capacity overflow")]
1080 fn test_poc2_reserve_overflow() {
1081 let mut v = SmallBitVec::new();
1082 v.push(true);
1083 v.reserve(usize::MAX - 10);
1084 }
1085}