1use rudb_common::{Error, Result};
45
46pub const VALUES: usize = 1024;
48
49const ORDER: [usize; 8] = [0, 4, 2, 6, 1, 5, 3, 7];
51
52mod sealed {
53 pub trait Sealed {}
54 impl Sealed for u8 {}
55 impl Sealed for u16 {}
56 impl Sealed for u32 {}
57 impl Sealed for u64 {}
58}
59
60pub trait Packable: sealed::Sealed + Copy + Ord + std::fmt::Debug {
66 const WIDTH: usize;
68 const LANES: usize = VALUES / Self::WIDTH;
70
71 fn to_u64(self) -> u64;
73 fn from_u64(value: u64) -> Self;
75}
76
77macro_rules! impl_packable {
78 ($($ty:ty),*) => {$(
79 impl Packable for $ty {
80 const WIDTH: usize = <$ty>::BITS as usize;
81
82 #[inline]
83 fn to_u64(self) -> u64 {
84 u64::from(self)
85 }
86
87 #[inline]
88 fn from_u64(value: u64) -> Self {
89 value as $ty
90 }
91 }
92 )*};
93}
94
95impl_packable!(u8, u16, u32, u64);
96
97#[inline]
99const fn low_mask(bits: usize) -> u64 {
100 if bits >= 64 { u64::MAX } else { (1u64 << bits) - 1 }
101}
102
103#[inline]
105const fn shift_right(value: u64, bits: usize) -> u64 {
106 if bits >= 64 { 0 } else { value >> bits }
107}
108
109#[inline]
120#[must_use]
121pub fn source_index<T: Packable>(row: usize, lane: usize) -> usize {
122 assert!(row < T::WIDTH, "row {row} is outside a {} bit type", T::WIDTH);
123 assert!(lane < T::LANES, "lane {lane} is outside {} lanes", T::LANES);
124 let group_size = T::WIDTH / 8;
125 let group = row / group_size;
126 let offset = row % group_size;
127 ((offset * 8) + ORDER[group]) * T::LANES + lane
128}
129
130pub fn transpose<T: Packable>(input: &[T], output: &mut [T]) -> Result<()> {
136 check_vector_len(input.len(), "input")?;
137 check_vector_len(output.len(), "output")?;
138 for row in 0..T::WIDTH {
139 for lane in 0..T::LANES {
140 output[row * T::LANES + lane] = input[source_index::<T>(row, lane)];
141 }
142 }
143 Ok(())
144}
145
146pub fn untranspose<T: Packable>(input: &[T], output: &mut [T]) -> Result<()> {
152 check_vector_len(input.len(), "input")?;
153 check_vector_len(output.len(), "output")?;
154 for row in 0..T::WIDTH {
155 for lane in 0..T::LANES {
156 output[source_index::<T>(row, lane)] = input[row * T::LANES + lane];
157 }
158 }
159 Ok(())
160}
161
162#[must_use]
167pub fn packed_len<T: Packable>(width: usize) -> usize {
168 width * T::LANES
169}
170
171#[must_use]
174pub fn required_width<T: Packable>(values: &[T]) -> usize {
175 let max = values.iter().copied().max().map_or(0, T::to_u64);
176 (64 - max.leading_zeros()) as usize
177}
178
179pub fn pack_transposed<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
190 check_vector_len(input.len(), "input")?;
191 check_width::<T>(width)?;
192 if output.len() != packed_len::<T>(width) {
193 return Err(Error::internal(format!(
194 "a {width} bit packed vector is {} words, not {}",
195 packed_len::<T>(width),
196 output.len()
197 )));
198 }
199 if width == 0 {
200 return check_all_zero(input);
204 }
205
206 let mask = low_mask(width);
207 let lanes = T::LANES;
208 for lane in 0..lanes {
209 let mut filled = 0usize;
211 let mut accumulator = 0u64;
212 let mut word = 0usize;
213 for row in 0..T::WIDTH {
214 let value = input[row * lanes + lane].to_u64();
215 if value & !mask != 0 {
216 return Err(Error::internal(format!("value {value} does not fit in {width} bits")));
217 }
218 accumulator |= value << filled;
219 filled += width;
220 if filled >= T::WIDTH {
221 output[word * lanes + lane] = T::from_u64(accumulator & low_mask(T::WIDTH));
222 word += 1;
223 let consumed = width - (filled - T::WIDTH);
226 filled -= T::WIDTH;
227 accumulator = shift_right(value, consumed);
228 }
229 }
230 debug_assert_eq!(filled, 0, "a packed lane always ends on a word boundary");
231 }
232 Ok(())
233}
234
235pub fn unpack_transposed<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
242 check_width::<T>(width)?;
243 check_vector_len(output.len(), "output")?;
244 if input.len() != packed_len::<T>(width) {
245 return Err(Error::internal(format!(
246 "a {width} bit packed vector is {} words, not {}",
247 packed_len::<T>(width),
248 input.len()
249 )));
250 }
251 if width == 0 {
252 output.fill(T::from_u64(0));
253 return Ok(());
254 }
255
256 let mask = low_mask(width);
257 let lanes = T::LANES;
258 for lane in 0..lanes {
259 let mut available = 0usize;
261 let mut buffer = 0u64;
262 let mut word = 0usize;
263 for row in 0..T::WIDTH {
264 let value = if available >= width {
265 let value = buffer & mask;
266 buffer = shift_right(buffer, width);
267 available -= width;
268 value
269 } else {
270 let next = input[word * lanes + lane].to_u64();
271 word += 1;
272 let taken = width - available;
273 let value = buffer | ((next & low_mask(taken)) << available);
274 buffer = shift_right(next, taken);
275 available = T::WIDTH - taken;
276 value
277 };
278 output[row * lanes + lane] = T::from_u64(value);
279 }
280 }
281 Ok(())
282}
283
284#[derive(Debug)]
299pub struct Scratch<T: Packable> {
300 transposed: Vec<T>,
301}
302
303impl<T: Packable> Scratch<T> {
304 #[must_use]
306 pub const fn new() -> Self {
307 Self { transposed: Vec::new() }
308 }
309
310 fn ready(&mut self) {
312 if self.transposed.len() != VALUES {
313 self.transposed.resize(VALUES, T::from_u64(0));
314 }
315 }
316}
317
318impl<T: Packable> Default for Scratch<T> {
319 fn default() -> Self {
320 Self::new()
321 }
322}
323
324pub fn pack<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
334 pack_with(input, width, output, &mut Scratch::new())
335}
336
337pub fn pack_with<T: Packable>(
343 input: &[T],
344 width: usize,
345 output: &mut [T],
346 scratch: &mut Scratch<T>,
347) -> Result<()> {
348 check_vector_len(input.len(), "input")?;
349 scratch.ready();
350 transpose(input, &mut scratch.transposed)?;
351 pack_transposed(&scratch.transposed, width, output)
352}
353
354pub fn unpack<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
360 unpack_with(input, width, output, &mut Scratch::new())
361}
362
363pub fn unpack_with<T: Packable>(
369 input: &[T],
370 width: usize,
371 output: &mut [T],
372 scratch: &mut Scratch<T>,
373) -> Result<()> {
374 check_vector_len(output.len(), "output")?;
375 scratch.ready();
376 unpack_transposed(input, width, &mut scratch.transposed)?;
377 untranspose(&scratch.transposed, output)
378}
379
380#[must_use]
382pub fn tail_len(count: usize, width: usize) -> usize {
383 (count * width).div_ceil(8)
384}
385
386pub fn pack_tail(values: &[u64], width: usize, output: &mut Vec<u8>) -> Result<()> {
404 check_tail(values.len(), width)?;
405 if width == 0 {
406 return check_all_zero(values);
407 }
408 let mask = low_mask(width);
409 let mut accumulator: u128 = 0;
412 let mut filled = 0usize;
413 for value in values {
414 if value & !mask != 0 {
415 return Err(Error::internal(format!("value {value} does not fit in {width} bits")));
416 }
417 accumulator |= u128::from(*value) << filled;
418 filled += width;
419 while filled >= 8 {
420 output.push((accumulator & 0xff) as u8);
421 accumulator >>= 8;
422 filled -= 8;
423 }
424 }
425 if filled > 0 {
426 output.push((accumulator & 0xff) as u8);
427 }
428 Ok(())
429}
430
431pub fn unpack_tail(input: &[u8], width: usize, count: usize) -> Result<Vec<u64>> {
438 check_tail(count, width)?;
439 if width == 0 {
440 return Ok(vec![0; count]);
441 }
442 if input.len() < tail_len(count, width) {
443 return Err(Error::internal(format!(
444 "{count} values at {width} bits need {} bytes and there are {}",
445 tail_len(count, width),
446 input.len()
447 )));
448 }
449 let mask = u128::from(low_mask(width));
450 let mut values = Vec::with_capacity(count);
451 let mut accumulator: u128 = 0;
452 let mut available = 0usize;
453 let mut at = 0usize;
454 for _ in 0..count {
455 while available < width {
456 accumulator |= u128::from(input[at]) << available;
457 at += 1;
458 available += 8;
459 }
460 values.push((accumulator & mask) as u64);
461 accumulator >>= width;
462 available -= width;
463 }
464 Ok(values)
465}
466
467fn check_tail(count: usize, width: usize) -> Result<()> {
468 if count >= VALUES {
469 return Err(Error::internal(format!(
470 "{count} values is a whole unit and belongs in the transposed layout"
471 )));
472 }
473 if width > 64 {
474 return Err(Error::internal(format!("{width} bits does not fit in 64")));
475 }
476 Ok(())
477}
478
479fn check_vector_len(len: usize, what: &str) -> Result<()> {
480 if len == VALUES {
481 Ok(())
482 } else {
483 Err(Error::internal(format!("{what} is {len} values, and a packed unit is {VALUES}")))
484 }
485}
486
487fn check_width<T: Packable>(width: usize) -> Result<()> {
488 if width <= T::WIDTH {
489 Ok(())
490 } else {
491 Err(Error::internal(format!("{width} bits does not fit in a {} bit type", T::WIDTH)))
492 }
493}
494
495fn check_all_zero<T: Packable>(input: &[T]) -> Result<()> {
496 match input.iter().position(|value| value.to_u64() != 0) {
497 None => Ok(()),
498 Some(index) => Err(Error::internal(format!(
499 "a zero bit vector cannot hold {:?} at {index}",
500 input[index]
501 ))),
502 }
503}
504
505#[cfg(test)]
506mod tests {
507 use super::*;
508
509 struct Random(u64);
512
513 impl Random {
514 fn new() -> Self {
515 Self(0x2545_f491_4f6c_dd1d)
516 }
517
518 fn next(&mut self) -> u64 {
519 self.0 ^= self.0 << 13;
520 self.0 ^= self.0 >> 7;
521 self.0 ^= self.0 << 17;
522 self.0
523 }
524 }
525
526 fn sample<T: Packable>(width: usize) -> Vec<T> {
527 let mut random = Random::new();
528 (0..VALUES).map(|_| T::from_u64(random.next() & low_mask(width))).collect()
529 }
530
531 fn round_trip<T: Packable>(width: usize) {
532 let values = sample::<T>(width);
533 let mut packed = vec![T::from_u64(0); packed_len::<T>(width)];
534 pack(&values, width, &mut packed).unwrap();
535 let mut back = vec![T::from_u64(0); VALUES];
536 unpack(&packed, width, &mut back).unwrap();
537 assert_eq!(back, values, "{width} bits of a {} bit type", T::WIDTH);
538 }
539
540 #[test]
541 fn every_width_of_every_type_round_trips() {
542 for width in 0..=8 {
543 round_trip::<u8>(width);
544 }
545 for width in 0..=16 {
546 round_trip::<u16>(width);
547 }
548 for width in 0..=32 {
549 round_trip::<u32>(width);
550 }
551 for width in 0..=64 {
552 round_trip::<u64>(width);
553 }
554 }
555
556 #[test]
557 fn a_reused_scratch_gives_what_a_fresh_one_gives() {
558 let mut scratch = Scratch::<u64>::new();
564 for width in [64, 1, 33, 7, 64, 0, 17, 60, 3] {
565 let values = sample::<u64>(width);
566 let mut packed = vec![0u64; packed_len::<u64>(width)];
567 pack_with(&values, width, &mut packed, &mut scratch).unwrap();
568 let mut reused = vec![0u64; VALUES];
569 unpack_with(&packed, width, &mut reused, &mut scratch).unwrap();
570 let mut fresh = vec![0u64; VALUES];
571 unpack(&packed, width, &mut fresh).unwrap();
572 assert_eq!(reused, fresh, "at {width} bits after a wider unit");
573 assert_eq!(reused, values, "at {width} bits");
574 }
575 }
576
577 #[test]
578 fn the_transposed_form_also_round_trips_without_being_reordered() {
579 let values = sample::<u32>(19);
582 let mut transposed = vec![0u32; VALUES];
583 transpose(&values, &mut transposed).unwrap();
584 let mut packed = vec![0u32; packed_len::<u32>(19)];
585 pack_transposed(&transposed, 19, &mut packed).unwrap();
586 let mut back = vec![0u32; VALUES];
587 unpack_transposed(&packed, 19, &mut back).unwrap();
588 assert_eq!(back, transposed);
589 }
590
591 #[test]
592 fn the_permutation_is_a_bijection() {
593 fn check<T: Packable>() {
596 let mut seen = vec![false; VALUES];
597 for row in 0..T::WIDTH {
598 for lane in 0..T::LANES {
599 let index = source_index::<T>(row, lane);
600 assert!(!seen[index], "{index} is written twice for {} bits", T::WIDTH);
601 seen[index] = true;
602 }
603 }
604 assert!(seen.into_iter().all(|hit| hit));
605 }
606 check::<u8>();
607 check::<u16>();
608 check::<u32>();
609 check::<u64>();
610 }
611
612 #[test]
613 fn transposing_is_not_the_identity() {
614 let values: Vec<u32> = (0..VALUES).map(|index| index as u32).collect();
616 let mut transposed = vec![0u32; VALUES];
617 transpose(&values, &mut transposed).unwrap();
618 assert_ne!(transposed, values);
619 let mut back = vec![0u32; VALUES];
620 untranspose(&transposed, &mut back).unwrap();
621 assert_eq!(back, values);
622 }
623
624 #[test]
625 fn a_full_width_pack_is_the_data_itself() {
626 let values = sample::<u64>(64);
629 let mut transposed = vec![0u64; VALUES];
630 transpose(&values, &mut transposed).unwrap();
631 let mut packed = vec![0u64; packed_len::<u64>(64)];
632 pack_transposed(&transposed, 64, &mut packed).unwrap();
633 assert_eq!(packed, transposed);
634 }
635
636 #[test]
637 fn a_zero_width_vector_stores_nothing_and_reads_back_as_zeros() {
638 let values = vec![0u32; VALUES];
639 assert_eq!(required_width(&values), 0);
640 let mut packed = Vec::new();
641 pack(&values, 0, &mut packed).unwrap();
642 let mut back = vec![7u32; VALUES];
643 unpack(&packed, 0, &mut back).unwrap();
644 assert_eq!(back, values);
645 }
646
647 #[test]
648 fn required_width_is_the_bits_of_the_largest_value() {
649 assert_eq!(required_width::<u32>(&[]), 0);
650 assert_eq!(required_width::<u32>(&[0, 0]), 0);
651 assert_eq!(required_width::<u32>(&[1]), 1);
652 assert_eq!(required_width::<u32>(&[255, 3]), 8);
653 assert_eq!(required_width::<u32>(&[256]), 9);
654 assert_eq!(required_width::<u64>(&[u64::MAX]), 64);
655 }
656
657 #[test]
658 fn a_value_too_wide_for_the_width_is_an_error_rather_than_silent_truncation() {
659 let mut values = vec![0u32; VALUES];
660 values[500] = 8;
661 let mut transposed = vec![0u32; VALUES];
662 transpose(&values, &mut transposed).unwrap();
663 let mut packed = vec![0u32; packed_len::<u32>(3)];
664 let error = pack_transposed(&transposed, 3, &mut packed).unwrap_err();
665 assert!(error.message().contains("does not fit in 3 bits"), "{error}");
666 }
667
668 #[test]
669 fn a_wrong_sized_buffer_is_an_error() {
670 let values = vec![0u32; VALUES];
671 let mut packed = vec![0u32; 3];
672 let error = pack(&values, 5, &mut packed).unwrap_err();
673 assert!(error.message().contains("words"), "{error}");
674
675 let short = vec![0u32; 7];
676 let mut output = vec![0u32; VALUES];
677 let error = unpack(&short, 5, &mut output).unwrap_err();
678 assert!(error.message().contains("words"), "{error}");
679 }
680
681 #[test]
682 fn a_nonzero_value_at_zero_width_is_an_error() {
683 let mut values = vec![0u32; VALUES];
684 values[9] = 1;
685 let mut packed = Vec::new();
686 let error = pack(&values, 0, &mut packed).unwrap_err();
687 assert!(error.message().contains("zero bit vector"), "{error}");
688 }
689
690 #[test]
691 fn packing_at_a_width_the_type_cannot_hold_is_an_error() {
692 let values = vec![0u16; VALUES];
693 let mut packed = vec![0u16; 17 * 64];
694 let error = pack(&values, 17, &mut packed).unwrap_err();
695 assert!(error.message().contains("16 bit type"), "{error}");
696 }
697
698 #[test]
699 fn a_tail_round_trips_at_every_width_and_every_length() {
700 let mut random = Random::new();
701 for width in [0usize, 1, 3, 7, 8, 13, 31, 32, 33, 63, 64] {
702 for count in [0usize, 1, 2, 7, 8, 9, 100, 1023] {
703 let values: Vec<u64> =
704 (0..count).map(|_| random.next() & low_mask(width)).collect();
705 let mut bytes = Vec::new();
706 pack_tail(&values, width, &mut bytes).unwrap();
707 assert_eq!(bytes.len(), tail_len(count, width), "{count} at {width}");
708 assert_eq!(unpack_tail(&bytes, width, count).unwrap(), values);
709 }
710 }
711 }
712
713 #[test]
714 fn a_tail_costs_its_own_values_and_not_a_whole_unit() {
715 let values = vec![(1u64 << 39) + 1; 3];
717 let mut bytes = Vec::new();
718 pack_tail(&values, 40, &mut bytes).unwrap();
719 assert_eq!(bytes.len(), 15);
720 assert_eq!(packed_len::<u64>(40) * 8, 5120);
721 }
722
723 #[test]
724 fn a_whole_unit_is_refused_by_the_tail_packer() {
725 let values = vec![0u64; VALUES];
726 let error = pack_tail(&values, 4, &mut Vec::new()).unwrap_err();
727 assert!(error.message().contains("whole unit"), "{error}");
728 }
729
730 #[test]
731 fn a_short_tail_buffer_is_an_error() {
732 let error = unpack_tail(&[0, 0], 8, 5).unwrap_err();
733 assert!(error.message().contains("need 5 bytes"), "{error}");
734 }
735
736 #[test]
737 fn the_packed_size_is_the_same_as_the_naive_layout() {
738 for width in 0..=32 {
739 assert_eq!(packed_len::<u32>(width) * 32, width * VALUES);
740 }
741 }
742}