Skip to main content

rudb_encoding/
bitpack.rs

1//! Bit packing in the FastLanes unified transposed layout.
2//!
3//! Packing N values of a fixed bit width into a dense buffer is the bottom of every integer
4//! encoding in `spec/06-compression.md` section 6.2. FOR subtracts a base and packs. DELTA
5//! differences and packs. DICT produces codes and packs them. So this is the one kernel that runs
6//! over more bytes than anything else in the system, and the layout it uses decides whether the
7//! decoder can be data parallel or has to walk a dependency chain.
8//!
9//! The obvious layout writes value 0 in the low bits of the first word, value 1 above it, and so
10//! on. Unpacking that requires knowing where the previous value ended, which is a sequential
11//! dependency, and a SIMD implementation has to fight it with shuffles that differ per width and
12//! per instruction set. FastLanes takes the other road. The 1024 values of a vector are seen as a
13//! matrix of `T` rows by `1024 / T` lanes, where `T` is the bit width of the type, and the packing
14//! runs down the rows of every lane at once. Every lane has the same bit schedule, so unpacking
15//! lane 0 and lane 31 is the same instruction sequence with no cross lane data movement at all.
16//! That is what makes one scalar reference implementation and one AVX-512 implementation and one
17//! NEON implementation agree bit for bit, and it is why the vector size is 1024 rather than a
18//! rounder number.
19//!
20//! The price is that the values come out permuted within the vector. Row `r` lane `l` is not the
21//! `r * lanes + l`th value of the input. Section 6.2 says why that is acceptable: an operator
22//! working inside one vector does not care what order the rows are in, so the permutation only
23//! has to be undone when a vector is materialized in row order. [`transpose`] and [`untranspose`]
24//! are that step, and they are deliberately separate from [`pack_transposed`] and
25//! [`unpack_transposed`] so that the engine can keep data permuted through a whole pipeline and
26//! pay for the reordering once at the end rather than twice per operator.
27//!
28//! There is a second layout in here, in [`pack_tail`], and it is the sequential one this module
29//! opens by arguing against. The transposed layout is all or nothing: a value lives at a row and a
30//! lane, the lanes are interleaved through the whole buffer, and no prefix of a packed unit holds a
31//! prefix of the values. So a unit holding 3 values costs exactly what a unit holding 1024 costs,
32//! and a cascade is full of short arrays. A five entry dictionary, a run length array, an exception
33//! list. Storing three numbers in 5 KB is not a compressed format. The tail packer handles anything
34//! shorter than a unit, it has the dependency chain the transposed layout exists to avoid, and that
35//! is affordable there and nowhere else, because a tail is at most 1023 values and is decoded once
36//! while a full unit is on the hot path of every scan in the system.
37//!
38//! The permutation itself is a fixed shuffle of the eight bit groups of a row index, in the order
39//! 0, 4, 2, 6, 1, 5, 3, 7. That order is not arbitrary. It is the one that makes an eight way
40//! interleave of the rows land back in sequence under the pairwise unpacking pattern the paper
41//! uses, and the important property for us is only that it is a bijection that both directions
42//! agree on.
43
44use rudb_common::{Error, Result};
45
46/// How many values a packed unit holds. One vector, per `spec/06-compression.md` section 6.2.
47pub const VALUES: usize = 1024;
48
49/// The interleaving order of the eight row groups. See the module documentation.
50const 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
60/// An unsigned integer type that can be bit packed.
61///
62/// Sealed, because the layout constants are only correct for the four widths that divide 1024 into
63/// a whole number of lanes, and because every kernel here does its arithmetic in `u64` and relies
64/// on every implementor fitting in one.
65pub trait Packable: sealed::Sealed + Copy + Ord + std::fmt::Debug {
66    /// Width of the type in bits. `T` in the module documentation.
67    const WIDTH: usize;
68    /// How many of these fit in the 1024 bit virtual register, which is how many lanes there are.
69    const LANES: usize = VALUES / Self::WIDTH;
70
71    /// Widens to the type the packing arithmetic is done in.
72    fn to_u64(self) -> u64;
73    /// Narrows back. The high bits are already known to be zero.
74    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/// A mask of the low `bits` bits, correct at 0 and at 64 where the shift would overflow.
98#[inline]
99const fn low_mask(bits: usize) -> u64 {
100    if bits >= 64 { u64::MAX } else { (1u64 << bits) - 1 }
101}
102
103/// A right shift that saturates to zero at 64 rather than overflowing.
104#[inline]
105const fn shift_right(value: u64, bits: usize) -> u64 {
106    if bits >= 64 { 0 } else { value >> bits }
107}
108
109/// Where the value at row `row` lane `lane` of the transposed matrix came from in the input.
110///
111/// The row index is split into a group and an offset within the group, the group is permuted by
112/// the fixed order in the module documentation, and the two are recombined with the offset as the
113/// high part. The lane index is untouched, which is the property that makes the layout lane
114/// parallel.
115///
116/// # Panics
117///
118/// If `row` is not below `T::WIDTH` or `lane` is not below `T::LANES`.
119#[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
130/// Rewrites 1024 values from row order into the transposed layout.
131///
132/// # Errors
133///
134/// If either slice is not exactly [`VALUES`] long.
135pub 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
146/// Rewrites 1024 values from the transposed layout back into row order.
147///
148/// # Errors
149///
150/// If either slice is not exactly [`VALUES`] long.
151pub 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/// How many words of `T` a packed vector of the given width occupies.
163///
164/// Every lane contributes `width` words, which is the same `width * 1024` bits the naive layout
165/// would use. The layout costs nothing in space.
166#[must_use]
167pub fn packed_len<T: Packable>(width: usize) -> usize {
168    width * T::LANES
169}
170
171/// The smallest bit width that can hold every value in the slice. Zero for an empty slice or a
172/// slice of zeros, which [`pack_transposed`] handles as the degenerate case that stores nothing.
173#[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
179/// Packs a transposed vector at a fixed bit width.
180///
181/// The input is 1024 values already in the layout [`transpose`] produces, and the output is
182/// [`packed_len`] words. Every lane is packed independently and the loop over lanes is the one a
183/// SIMD implementation replaces with a single register.
184///
185/// # Errors
186///
187/// If the input is not [`VALUES`] long, if the output is not [`packed_len`] long, if `width`
188/// exceeds the width of the type, or if a value does not fit in `width` bits.
189pub 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        // Nothing is stored. The caller has already established that every value is zero, either
201        // by asking for `required_width` or by being the CONSTANT encoding, and the check below
202        // enforces it rather than trusting it.
203        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        // Bits already sitting in `accumulator`, always below `T::WIDTH` between iterations.
210        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                // The only value that can straddle the word boundary is the one just written, so
224                // the carry is a shift of it rather than anything kept from earlier rows.
225                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
235/// Unpacks into the transposed layout. The inverse of [`pack_transposed`].
236///
237/// # Errors
238///
239/// If the input is not [`packed_len`] long, if the output is not [`VALUES`] long, or if `width`
240/// exceeds the width of the type.
241pub 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        // Bits of the current word not yet handed out, right aligned in `buffer`.
260        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/// The buffer [`pack_with`] transposes through, kept so it can be reused.
285///
286/// Going between row order and the transposed layout needs somewhere to put the other order, and
287/// that somewhere is [`VALUES`] values, which is 8 KB for a `u64`. Allocating it per call is not the
288/// expensive part. Zeroing it is, because the allocator hands back a page it has to clear and the
289/// transpose then writes every element of it anyway. On a scan of a packed integer column that is
290/// once per 1024 rows, and it showed up as the largest single item in a ClickBench profile, larger
291/// than the unpacking it was making room for.
292///
293/// So a caller that packs more than one unit should make one of these and pass it in. The unpacking
294/// side does not need one at all any more: see [`unpack`].
295///
296/// It starts empty and grows on the first unit that needs it, because a caller holds one for a whole
297/// decode and most chunks are not bit packed at all. Making the buffer in the constructor was tried
298/// and was worse than what it replaced, by more than the zeroing it saved.
299#[derive(Debug)]
300pub struct Scratch<T: Packable> {
301    transposed: Vec<T>,
302}
303
304impl<T: Packable> Scratch<T> {
305    /// A scratch buffer that has not made room for anything yet.
306    #[must_use]
307    pub const fn new() -> Self {
308        Self { transposed: Vec::new() }
309    }
310
311    /// Makes room for one unit. A no op every time after the first.
312    fn ready(&mut self) {
313        if self.transposed.len() != VALUES {
314            self.transposed.resize(VALUES, T::from_u64(0));
315        }
316    }
317}
318
319impl<T: Packable> Default for Scratch<T> {
320    fn default() -> Self {
321        Self::new()
322    }
323}
324
325/// Packs a vector given in row order, transposing it first.
326///
327/// The engine does not use this. Data written by the storage layer is transposed once on the way
328/// in and stays that way, per the module documentation. This exists for tests, for the format lab,
329/// and for the one place that has to hand back a vector in the order the user gave it.
330///
331/// # Errors
332///
333/// As [`pack_transposed`].
334pub fn pack<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
335    pack_with(input, width, output, &mut Scratch::new())
336}
337
338/// As [`pack`], through a buffer the caller keeps rather than one allocated per call.
339///
340/// # Errors
341///
342/// As [`pack_transposed`].
343pub fn pack_with<T: Packable>(
344    input: &[T],
345    width: usize,
346    output: &mut [T],
347    scratch: &mut Scratch<T>,
348) -> Result<()> {
349    check_vector_len(input.len(), "input")?;
350    scratch.ready();
351    transpose(input, &mut scratch.transposed)?;
352    pack_transposed(&scratch.transposed, width, output)
353}
354
355/// Unpacks into row order. The inverse of [`pack`].
356///
357/// This is what every scan of a packed integer column goes through, so it is written as one pass
358/// rather than as [`unpack_transposed`] followed by [`untranspose`]. Those two are still here and
359/// still the definition of the layout, and the test below checks this agrees with them at every
360/// width, but running them in sequence costs three things this does not. A 1024 value buffer to
361/// hold the middle, a second read of all of it, and a scatter: `untranspose` walks its input in
362/// order and writes all over its output, which is a store that misses and a loop no compiler will
363/// turn into wider instructions.
364///
365/// The fused form works because a row has the same bit schedule in every lane. That is the whole
366/// point of the layout. Row `r` of every lane takes bits `r * width` to `(r + 1) * width` of that
367/// lane's stream, so which word to read and how far to shift it are decided once for the row, and
368/// what is left for the lanes is a load, a shift, an or, a mask and a store with no branch and no
369/// carry from the lane before. The lanes of a row are next to each other in both the packed words
370/// and the output, so that inner loop reads and writes straight lines. Where a row lands in the
371/// output is the permutation `untranspose` was applying, and since the lane index is the low part
372/// of it, it comes out as a base address for the row and costs nothing.
373///
374/// # Errors
375///
376/// As [`unpack_transposed`].
377pub fn unpack<T: Packable>(input: &[T], width: usize, output: &mut [T]) -> Result<()> {
378    check_width::<T>(width)?;
379    check_vector_len(output.len(), "output")?;
380    if input.len() != packed_len::<T>(width) {
381        return Err(Error::internal(format!(
382            "a {width} bit packed vector is {} words, not {}",
383            packed_len::<T>(width),
384            input.len()
385        )));
386    }
387    if width == 0 {
388        output.fill(T::from_u64(0));
389        return Ok(());
390    }
391
392    let mask = low_mask(width);
393    let lanes = T::LANES;
394    let group_size = T::WIDTH / 8;
395    for row in 0..T::WIDTH {
396        let bit = row * width;
397        let word = bit / T::WIDTH;
398        let shift = bit % T::WIDTH;
399        // The same arithmetic as `source_index` with the lane left off, because the lane is the low
400        // part of it and the lanes of a row are consecutive from here.
401        let base = ((row % group_size) * 8 + ORDER[row / group_size]) * lanes;
402        let low = &input[word * lanes..(word + 1) * lanes];
403        let into = &mut output[base..base + lanes];
404        if shift + width <= T::WIDTH {
405            for lane in 0..lanes {
406                into[lane] = T::from_u64((low[lane].to_u64() >> shift) & mask);
407            }
408        } else {
409            // The value straddles two words, so `shift` is above zero, the carry in from the word
410            // above is a left shift by less than the word width, and neither shift can overflow.
411            // There is a word above to read: a value that straddles into word `word + 1` is one the
412            // packer wrote there, and it wrote `width` words a lane.
413            let carried = T::WIDTH - shift;
414            let high = &input[(word + 1) * lanes..(word + 2) * lanes];
415            for lane in 0..lanes {
416                let value = (low[lane].to_u64() >> shift) | (high[lane].to_u64() << carried);
417                into[lane] = T::from_u64(value & mask);
418            }
419        }
420    }
421    Ok(())
422}
423
424/// How many bytes [`pack_tail`] writes for `count` values at `width` bits.
425#[must_use]
426pub fn tail_len(count: usize, width: usize) -> usize {
427    (count * width).div_ceil(8)
428}
429
430/// Packs fewer than [`VALUES`] values, sequentially and to a byte boundary.
431///
432/// The transposed layout is all or nothing. A value lives at a row and a lane, the lanes are
433/// interleaved through the whole buffer, and there is no prefix of a packed unit that holds a
434/// prefix of the values. So a unit holding 3 values costs the same as a unit holding 1024, which is
435/// 5 KB to store three numbers, and every nested array in a cascade is short: a dictionary of five
436/// entries, a run length array, an exception list.
437///
438/// This is the other layout for exactly those. It is the obvious sequential one, value 0 in the low
439/// bits, and it has the dependency chain the transposed layout was chosen to avoid. That is
440/// affordable here and only here: a tail is at most 1023 values and is decoded once, so the chain
441/// is bounded by a number that does not grow with the data, while a full unit is on the hot path of
442/// every scan in the system.
443///
444/// # Errors
445///
446/// If `count` is not below [`VALUES`], if `width` exceeds 64, or if a value does not fit.
447pub fn pack_tail(values: &[u64], width: usize, output: &mut Vec<u8>) -> Result<()> {
448    check_tail(values.len(), width)?;
449    if width == 0 {
450        return check_all_zero(values);
451    }
452    let mask = low_mask(width);
453    // 128 bits, because the accumulator holds up to 7 bits left over from the previous value plus a
454    // whole 64 bit one.
455    let mut accumulator: u128 = 0;
456    let mut filled = 0usize;
457    for value in values {
458        if value & !mask != 0 {
459            return Err(Error::internal(format!("value {value} does not fit in {width} bits")));
460        }
461        accumulator |= u128::from(*value) << filled;
462        filled += width;
463        while filled >= 8 {
464            output.push((accumulator & 0xff) as u8);
465            accumulator >>= 8;
466            filled -= 8;
467        }
468    }
469    if filled > 0 {
470        output.push((accumulator & 0xff) as u8);
471    }
472    Ok(())
473}
474
475/// Unpacks what [`pack_tail`] wrote.
476///
477/// The writer has a dependency chain because it has to know how many bits are left over from the
478/// value before, but the reader does not, and this does not carry one. Value `index` occupies the
479/// `width` bits starting at bit `index * width`, so its position is arithmetic rather than history,
480/// and since it begins at most seven bits into a byte and runs at most sixty four, it always lies
481/// inside sixteen bytes read from that byte. One unaligned load, one shift and one mask.
482///
483/// That matters more than the module documentation lets on. The argument there is that a tail is at
484/// most 1023 values and so is bounded by a number that does not grow with the data, which is true
485/// per call and misleading in aggregate, because a cascade puts a short array in every chunk and a
486/// scan reads every chunk. ClickBench 9 is where it showed. UserID is nearly unique, so its
487/// dictionary holds about a thousand sixty four bit values per part and lands one value short of a
488/// full unit, which sends the whole column down this path: nine hundred and seventy four parts,
489/// about a million values, and the byte at a time version fed eight bytes through a `u128` for each
490/// one. That was fifty five percent of the instructions of a scan of that column on its own.
491///
492/// # Errors
493///
494/// If `count` is not below [`VALUES`], if `width` exceeds 64, or if the input is shorter than
495/// [`tail_len`].
496pub fn unpack_tail(input: &[u8], width: usize, count: usize) -> Result<Vec<u64>> {
497    check_tail(count, width)?;
498    if width == 0 {
499        return Ok(vec![0; count]);
500    }
501    if input.len() < tail_len(count, width) {
502        return Err(Error::internal(format!(
503            "{count} values at {width} bits need {} bytes and there are {}",
504            tail_len(count, width),
505            input.len()
506        )));
507    }
508    let mask = u128::from(low_mask(width));
509    let mut values = Vec::with_capacity(count);
510    let read = |window: u128, bit: usize| ((window >> bit) & mask) as u64;
511    // A buffer shorter than a window is one load for the whole call, because everything it holds is
512    // inside it. Short arrays are most of what a cascade stores, so this is the common case by
513    // count of calls even though it is the rare one by count of values.
514    if input.len() < WINDOW {
515        let mut window = [0u8; WINDOW];
516        window[..input.len()].copy_from_slice(input);
517        let word = u128::from_le_bytes(window);
518        for index in 0..count {
519            values.push(read(word, index * width));
520        }
521        return Ok(values);
522    }
523    // Otherwise a value is read where it lies, until the window would run off the end.
524    let whole = (((input.len() - WINDOW) * 8) / width + 1).min(count);
525    if width <= NARROW {
526        // Half the window, because a value this wide that starts at most seven bits into a byte
527        // ends inside the eight bytes from that byte. The shift and the mask are then one
528        // instruction each where a 128 bit shift is three, and every real width is down here: the
529        // offsets a text block carries are seventeen bits and a dictionary code is fewer.
530        let mask = low_mask(width);
531        for index in 0..whole {
532            let bit = index * width;
533            let word = word_at(input, bit / 8);
534            values.push((word >> (bit % 8)) & mask);
535        }
536    } else {
537        for index in 0..whole {
538            let bit = index * width;
539            let mut window = [0u8; WINDOW];
540            window.copy_from_slice(&input[bit / 8..bit / 8 + WINDOW]);
541            values.push(read(u128::from_le_bytes(window), bit % 8));
542        }
543    }
544    if whole < count {
545        // Every value left over begins past the sixteenth byte from the end, by the definition of
546        // `whole` just above, and the buffer stops on the byte holding the top bits of the last
547        // one. So all of them lie inside the final window and one load serves the lot.
548        let base = input.len() - WINDOW;
549        let mut window = [0u8; WINDOW];
550        window.copy_from_slice(&input[base..]);
551        let word = u128::from_le_bytes(window);
552        for index in whole..count {
553            values.push(read(word, index * width - base * 8));
554        }
555    }
556    Ok(values)
557}
558
559/// One value of a run written by [`pack_tail`], read where it lies.
560///
561/// [`unpack_tail`] decodes the whole run, which is what a scan wants and what nearly every caller
562/// here is. A binary search is the other kind of caller: it wants one value out of the middle of a
563/// block, it makes about as many probes as the block has bits, and decoding the block to answer one
564/// of them would cost more than reading the value it was avoiding.
565///
566/// # Errors
567///
568/// If `width` exceeds 64, or if the value would run past the end of `input`.
569#[inline]
570pub fn tail_at(input: &[u8], width: usize, index: usize) -> Result<u64> {
571    if width > 64 {
572        return Err(Error::internal(format!("a width of {width} is past what a u64 holds")));
573    }
574    if width == 0 {
575        return Ok(0);
576    }
577    let start = index * width;
578    let end = start + width;
579    if end.div_ceil(8) > input.len() {
580        return Err(Error::internal(format!(
581            "value {index} at {width} bits ends past the {} bytes there are",
582            input.len()
583        )));
584    }
585    let first = start / 8;
586    let last = (end - 1) / 8;
587    // A value that ends inside the eight bytes it starts in is one load, one shift and one mask.
588    // The window below copies a length the compiler does not know, which is a call to `memcpy`
589    // rather than a load, and this reads one value at a time for every string a text column hands
590    // out. It was fifteen percent of ClickBench 27.
591    if first + 8 <= input.len() && last - first < 8 {
592        return Ok((word_at(input, first) >> (start % 8)) & low_mask(width));
593    }
594    let mut window = [0u8; WINDOW];
595    window[..=last - first].copy_from_slice(&input[first..=last]);
596    let word = u128::from_le_bytes(window);
597    Ok(((word >> (start % 8)) & u128::from(low_mask(width))) as u64)
598}
599
600/// Two neighbouring values of a run, read from one load where the pair fits inside it.
601///
602/// `index` is the later of the two and the answer is the pair at `index - 1` and `index`. A text
603/// column asks for exactly this once per string it hands out, because a value starts where the one
604/// before it ended. Two calls to [`tail_at`] read the same eight bytes twice and do the bounds
605/// arithmetic twice, where a pair of seventeen bit offsets, which is what a block of text carries,
606/// both lie inside one load.
607///
608/// # Errors
609///
610/// If `index` is zero, if `width` exceeds 64, or if the pair would run past the end of `input`.
611#[inline]
612pub fn tail_pair(input: &[u8], width: usize, index: usize) -> Result<(u64, u64)> {
613    let Some(before) = index.checked_sub(1) else {
614        return Err(Error::internal("a tail pair has nothing before its first value"));
615    };
616    if width == 0 {
617        return Ok((0, 0));
618    }
619    let start = before * width;
620    let shift = start % 8;
621    let first = start / 8;
622    if shift + 2 * width <= u64::BITS as usize && first + 8 <= input.len() {
623        let word = word_at(input, first) >> shift;
624        let mask = low_mask(width);
625        return Ok((word & mask, (word >> width) & mask));
626    }
627    Ok((tail_at(input, width, before)?, tail_at(input, width, index)?))
628}
629
630/// The bytes a single tail value can span, which is a shift of at most seven plus a width of at
631/// most sixty four, so seventy one bits and therefore nine bytes, rounded up to the load that
632/// covers it.
633const WINDOW: usize = 16;
634
635/// The widest value that always ends inside the eight bytes it starts in, which is sixty four bits
636/// less the seven a value can begin into its first byte.
637const NARROW: usize = 57;
638
639/// Eight bytes read where they lie, as one load.
640///
641/// The length is a constant the compiler can see, which is what makes it a load. The caller is
642/// responsible for `at + 8` being inside `input`, and the index below says so where it is not.
643#[inline]
644fn word_at(input: &[u8], at: usize) -> u64 {
645    let run: [u8; 8] = input[at..at + 8].try_into().expect("eight bytes");
646    u64::from_le_bytes(run)
647}
648
649fn check_tail(count: usize, width: usize) -> Result<()> {
650    if count >= VALUES {
651        return Err(Error::internal(format!(
652            "{count} values is a whole unit and belongs in the transposed layout"
653        )));
654    }
655    if width > 64 {
656        return Err(Error::internal(format!("{width} bits does not fit in 64")));
657    }
658    Ok(())
659}
660
661fn check_vector_len(len: usize, what: &str) -> Result<()> {
662    if len == VALUES {
663        Ok(())
664    } else {
665        Err(Error::internal(format!("{what} is {len} values, and a packed unit is {VALUES}")))
666    }
667}
668
669fn check_width<T: Packable>(width: usize) -> Result<()> {
670    if width <= T::WIDTH {
671        Ok(())
672    } else {
673        Err(Error::internal(format!("{width} bits does not fit in a {} bit type", T::WIDTH)))
674    }
675}
676
677fn check_all_zero<T: Packable>(input: &[T]) -> Result<()> {
678    match input.iter().position(|value| value.to_u64() != 0) {
679        None => Ok(()),
680        Some(index) => Err(Error::internal(format!(
681            "a zero bit vector cannot hold {:?} at {index}",
682            input[index]
683        ))),
684    }
685}
686
687#[cfg(test)]
688mod tests {
689    use super::*;
690
691    /// A xorshift, so that the test data is the same on every host and in every run without the
692    /// workspace growing a dependency for it.
693    struct Random(u64);
694
695    impl Random {
696        fn new() -> Self {
697            Self(0x2545_f491_4f6c_dd1d)
698        }
699
700        fn next(&mut self) -> u64 {
701            self.0 ^= self.0 << 13;
702            self.0 ^= self.0 >> 7;
703            self.0 ^= self.0 << 17;
704            self.0
705        }
706    }
707
708    fn sample<T: Packable>(width: usize) -> Vec<T> {
709        let mut random = Random::new();
710        (0..VALUES).map(|_| T::from_u64(random.next() & low_mask(width))).collect()
711    }
712
713    fn round_trip<T: Packable>(width: usize) {
714        let values = sample::<T>(width);
715        let mut packed = vec![T::from_u64(0); packed_len::<T>(width)];
716        pack(&values, width, &mut packed).unwrap();
717        let mut back = vec![T::from_u64(0); VALUES];
718        unpack(&packed, width, &mut back).unwrap();
719        assert_eq!(back, values, "{width} bits of a {} bit type", T::WIDTH);
720    }
721
722    #[test]
723    fn every_width_of_every_type_round_trips() {
724        for width in 0..=8 {
725            round_trip::<u8>(width);
726        }
727        for width in 0..=16 {
728            round_trip::<u16>(width);
729        }
730        for width in 0..=32 {
731            round_trip::<u32>(width);
732        }
733        for width in 0..=64 {
734            round_trip::<u64>(width);
735        }
736    }
737
738    #[test]
739    fn a_reused_scratch_gives_what_a_fresh_one_gives() {
740        // The buffer a unit transposes through is handed in so it is not zeroed per call, which is
741        // only sound if every element of it is written every time. If some were not, a narrow unit
742        // following a wide one would read whatever the wide one left behind, so the widths here go
743        // up and down rather than in order and each answer is checked against the same unit packed
744        // through a buffer nothing has touched.
745        let mut scratch = Scratch::<u64>::new();
746        for width in [64, 1, 33, 7, 64, 0, 17, 60, 3] {
747            let values = sample::<u64>(width);
748            let mut reused = vec![0u64; packed_len::<u64>(width)];
749            pack_with(&values, width, &mut reused, &mut scratch).unwrap();
750            let mut fresh = vec![0u64; packed_len::<u64>(width)];
751            pack(&values, width, &mut fresh).unwrap();
752            assert_eq!(reused, fresh, "at {width} bits after a wider unit");
753            let mut back = vec![0u64; VALUES];
754            unpack(&reused, width, &mut back).unwrap();
755            assert_eq!(back, values, "at {width} bits");
756        }
757    }
758
759    #[test]
760    fn the_one_pass_unpack_gives_what_the_two_passes_give() {
761        // `unpack` is the fused form of `unpack_transposed` followed by `untranspose`, and those two
762        // are the definition of the layout. So this checks the fast one against the slow one at
763        // every width of every type rather than against a remembered answer, which is the check that
764        // would catch the fused one getting a shift or a row base wrong at one width out of sixty
765        // five.
766        fn agree<T: Packable>() {
767            for width in 0..=T::WIDTH {
768                let values = sample::<T>(width);
769                let mut transposed = vec![T::from_u64(0); VALUES];
770                transpose(&values, &mut transposed).unwrap();
771                let mut packed = vec![T::from_u64(0); packed_len::<T>(width)];
772                pack_transposed(&transposed, width, &mut packed).unwrap();
773
774                let mut middle = vec![T::from_u64(0); VALUES];
775                unpack_transposed(&packed, width, &mut middle).unwrap();
776                let mut slow = vec![T::from_u64(0); VALUES];
777                untranspose(&middle, &mut slow).unwrap();
778
779                let mut fast = vec![T::from_u64(0); VALUES];
780                unpack(&packed, width, &mut fast).unwrap();
781
782                assert_eq!(fast, slow, "{} bit type at {width} bits", T::WIDTH);
783                assert_eq!(fast, values, "{} bit type at {width} bits round trip", T::WIDTH);
784            }
785        }
786        agree::<u8>();
787        agree::<u16>();
788        agree::<u32>();
789        agree::<u64>();
790    }
791
792    #[test]
793    fn the_transposed_form_also_round_trips_without_being_reordered() {
794        // What the engine actually does: transpose once, then pack and unpack any number of times
795        // without ever going back to row order.
796        let values = sample::<u32>(19);
797        let mut transposed = vec![0u32; VALUES];
798        transpose(&values, &mut transposed).unwrap();
799        let mut packed = vec![0u32; packed_len::<u32>(19)];
800        pack_transposed(&transposed, 19, &mut packed).unwrap();
801        let mut back = vec![0u32; VALUES];
802        unpack_transposed(&packed, 19, &mut back).unwrap();
803        assert_eq!(back, transposed);
804    }
805
806    #[test]
807    fn the_permutation_is_a_bijection() {
808        // Every value has to land somewhere and no two may land in the same place, or a round trip
809        // would silently drop rows. Checked for all four widths because the group size changes.
810        fn check<T: Packable>() {
811            let mut seen = vec![false; VALUES];
812            for row in 0..T::WIDTH {
813                for lane in 0..T::LANES {
814                    let index = source_index::<T>(row, lane);
815                    assert!(!seen[index], "{index} is written twice for {} bits", T::WIDTH);
816                    seen[index] = true;
817                }
818            }
819            assert!(seen.into_iter().all(|hit| hit));
820        }
821        check::<u8>();
822        check::<u16>();
823        check::<u32>();
824        check::<u64>();
825    }
826
827    #[test]
828    fn transposing_is_not_the_identity() {
829        // If it were, the test above would be passing on a layout that is not the FastLanes one.
830        let values: Vec<u32> = (0..VALUES).map(|index| index as u32).collect();
831        let mut transposed = vec![0u32; VALUES];
832        transpose(&values, &mut transposed).unwrap();
833        assert_ne!(transposed, values);
834        let mut back = vec![0u32; VALUES];
835        untranspose(&transposed, &mut back).unwrap();
836        assert_eq!(back, values);
837    }
838
839    #[test]
840    fn a_full_width_pack_is_the_data_itself() {
841        // 64 bits of a 64 bit type has no packing to do, and the loop that handles the general case
842        // has to get the degenerate one right rather than shifting by 64 and wrapping.
843        let values = sample::<u64>(64);
844        let mut transposed = vec![0u64; VALUES];
845        transpose(&values, &mut transposed).unwrap();
846        let mut packed = vec![0u64; packed_len::<u64>(64)];
847        pack_transposed(&transposed, 64, &mut packed).unwrap();
848        assert_eq!(packed, transposed);
849    }
850
851    #[test]
852    fn a_zero_width_vector_stores_nothing_and_reads_back_as_zeros() {
853        let values = vec![0u32; VALUES];
854        assert_eq!(required_width(&values), 0);
855        let mut packed = Vec::new();
856        pack(&values, 0, &mut packed).unwrap();
857        let mut back = vec![7u32; VALUES];
858        unpack(&packed, 0, &mut back).unwrap();
859        assert_eq!(back, values);
860    }
861
862    #[test]
863    fn required_width_is_the_bits_of_the_largest_value() {
864        assert_eq!(required_width::<u32>(&[]), 0);
865        assert_eq!(required_width::<u32>(&[0, 0]), 0);
866        assert_eq!(required_width::<u32>(&[1]), 1);
867        assert_eq!(required_width::<u32>(&[255, 3]), 8);
868        assert_eq!(required_width::<u32>(&[256]), 9);
869        assert_eq!(required_width::<u64>(&[u64::MAX]), 64);
870    }
871
872    #[test]
873    fn a_value_too_wide_for_the_width_is_an_error_rather_than_silent_truncation() {
874        let mut values = vec![0u32; VALUES];
875        values[500] = 8;
876        let mut transposed = vec![0u32; VALUES];
877        transpose(&values, &mut transposed).unwrap();
878        let mut packed = vec![0u32; packed_len::<u32>(3)];
879        let error = pack_transposed(&transposed, 3, &mut packed).unwrap_err();
880        assert!(error.message().contains("does not fit in 3 bits"), "{error}");
881    }
882
883    #[test]
884    fn a_wrong_sized_buffer_is_an_error() {
885        let values = vec![0u32; VALUES];
886        let mut packed = vec![0u32; 3];
887        let error = pack(&values, 5, &mut packed).unwrap_err();
888        assert!(error.message().contains("words"), "{error}");
889
890        let short = vec![0u32; 7];
891        let mut output = vec![0u32; VALUES];
892        let error = unpack(&short, 5, &mut output).unwrap_err();
893        assert!(error.message().contains("words"), "{error}");
894    }
895
896    #[test]
897    fn a_nonzero_value_at_zero_width_is_an_error() {
898        let mut values = vec![0u32; VALUES];
899        values[9] = 1;
900        let mut packed = Vec::new();
901        let error = pack(&values, 0, &mut packed).unwrap_err();
902        assert!(error.message().contains("zero bit vector"), "{error}");
903    }
904
905    #[test]
906    fn packing_at_a_width_the_type_cannot_hold_is_an_error() {
907        let values = vec![0u16; VALUES];
908        let mut packed = vec![0u16; 17 * 64];
909        let error = pack(&values, 17, &mut packed).unwrap_err();
910        assert!(error.message().contains("16 bit type"), "{error}");
911    }
912
913    #[test]
914    fn a_tail_round_trips_at_every_width_and_every_length() {
915        let mut random = Random::new();
916        for width in 0..=64usize {
917            for count in [0usize, 1, 2, 7, 8, 9, 100, 1023] {
918                let values: Vec<u64> =
919                    (0..count).map(|_| random.next() & low_mask(width)).collect();
920                let mut bytes = Vec::new();
921                pack_tail(&values, width, &mut bytes).unwrap();
922                assert_eq!(bytes.len(), tail_len(count, width), "{count} at {width}");
923                assert_eq!(
924                    unpack_tail(&bytes, width, count).unwrap(),
925                    values,
926                    "{count} at {width}"
927                );
928            }
929        }
930    }
931
932    /// Reading one value where it lies agrees with decoding the whole run.
933    ///
934    /// Every width and every position, since the point of it is the arithmetic that finds the bytes
935    /// a value straddles, and that is what is off by one.
936    ///
937    /// A value that runs off the buffer is an error. The buffer stops on a byte boundary and a value
938    /// does not, so an index a little past the count can still lie inside the padding of the last
939    /// byte and that reads rather than complains. It is the caller that knows how many values it
940    /// wrote, the same way it does for `unpack_tail`.
941    #[test]
942    fn one_value_of_a_tail_reads_the_same_as_the_whole_of_it() {
943        let mut random = Random::new();
944        for width in 0..=64usize {
945            let count = 37;
946            let values: Vec<u64> = (0..count).map(|_| random.next() & low_mask(width)).collect();
947            let mut bytes = Vec::new();
948            pack_tail(&values, width, &mut bytes).unwrap();
949            for (index, value) in values.iter().enumerate() {
950                assert_eq!(tail_at(&bytes, width, index).unwrap(), *value, "{index} at {width}");
951            }
952            let Some(fits) = (bytes.len() * 8).checked_div(width) else { continue };
953            assert!(tail_at(&bytes, width, fits + 1).is_err(), "past the end at {width}");
954        }
955    }
956
957    /// The two halves of the reader agree with each other.
958    ///
959    /// A value is read with one sixteen byte load, which the values near the end of the buffer
960    /// cannot have because the buffer stops on the byte holding the top bits of the last one. Those
961    /// go through a zero padded copy instead, and the split between the two is arithmetic on
962    /// lengths, which is the kind of thing that is off by one. Handing the same bytes to the reader
963    /// twice, once exactly sized so the last values take the padded path and once with slack on the
964    /// end so every value takes the fast one, makes the two paths check each other at every width.
965    #[test]
966    fn the_padded_end_of_a_tail_reads_the_same_as_the_windowed_start() {
967        let mut random = Random::new();
968        for width in 1..=64usize {
969            for count in [1usize, 2, 3, 17, 129, 1023] {
970                let values: Vec<u64> =
971                    (0..count).map(|_| random.next() & low_mask(width)).collect();
972                let mut exact = Vec::new();
973                pack_tail(&values, width, &mut exact).unwrap();
974                let mut slack = exact.clone();
975                slack.extend_from_slice(&[0u8; WINDOW]);
976                assert_eq!(
977                    unpack_tail(&exact, width, count).unwrap(),
978                    values,
979                    "{count} at {width}"
980                );
981                assert_eq!(
982                    unpack_tail(&slack, width, count).unwrap(),
983                    values,
984                    "{count} at {width}"
985                );
986            }
987        }
988    }
989
990    /// The pair read agrees with two single reads, at every width and every position.
991    ///
992    /// The pair has its own arithmetic for the case where both values fit one load, so the thing to
993    /// check is that it falls back to the same answer everywhere that does not hold, which is every
994    /// width past thirty two and every value near the end of the buffer.
995    #[test]
996    fn a_pair_of_tail_values_reads_the_same_as_the_two_of_them_apart() {
997        let mut random = Random::new();
998        for width in 0..=64usize {
999            let count = 37;
1000            let values: Vec<u64> = (0..count).map(|_| random.next() & low_mask(width)).collect();
1001            let mut bytes = Vec::new();
1002            pack_tail(&values, width, &mut bytes).unwrap();
1003            for index in 1..count {
1004                assert_eq!(
1005                    tail_pair(&bytes, width, index).unwrap(),
1006                    (values[index - 1], values[index]),
1007                    "{index} at {width}"
1008                );
1009            }
1010            assert!(tail_pair(&bytes, width, 0).is_err(), "nothing before the first at {width}");
1011        }
1012    }
1013
1014    #[test]
1015    fn a_tail_costs_its_own_values_and_not_a_whole_unit() {
1016        // The reason it exists. Three 40 bit values in the transposed layout is a 5 KB buffer.
1017        let values = vec![(1u64 << 39) + 1; 3];
1018        let mut bytes = Vec::new();
1019        pack_tail(&values, 40, &mut bytes).unwrap();
1020        assert_eq!(bytes.len(), 15);
1021        assert_eq!(packed_len::<u64>(40) * 8, 5120);
1022    }
1023
1024    #[test]
1025    fn a_whole_unit_is_refused_by_the_tail_packer() {
1026        let values = vec![0u64; VALUES];
1027        let error = pack_tail(&values, 4, &mut Vec::new()).unwrap_err();
1028        assert!(error.message().contains("whole unit"), "{error}");
1029    }
1030
1031    #[test]
1032    fn a_short_tail_buffer_is_an_error() {
1033        let error = unpack_tail(&[0, 0], 8, 5).unwrap_err();
1034        assert!(error.message().contains("need 5 bytes"), "{error}");
1035    }
1036
1037    #[test]
1038    fn the_packed_size_is_the_same_as_the_naive_layout() {
1039        for width in 0..=32 {
1040            assert_eq!(packed_len::<u32>(width) * 32, width * VALUES);
1041        }
1042    }
1043}