Skip to main content

nodedb_codec/fastlanes/
block.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Block-level encode and decode for the FastLanes FOR + bit-packing codec.
4
5use super::bits::{low_mask_u8, pack_bits, unpack_bits};
6use super::header::BLOCK_SIZE;
7use crate::bounds::{checked_add, checked_mul, checked_range, encode_u32_len};
8use crate::error::CodecError;
9
10/// Per-block header: 2 bytes count + 1 byte bit_width + 8 bytes min_value.
11pub(super) const BLOCK_HEADER_SIZE: usize = 11;
12
13/// Encode a single nonempty block (up to 1024 values).
14pub(super) fn encode_block(values: &[i64], out: &mut Vec<u8>) -> Result<(), CodecError> {
15    if values.is_empty() || values.len() > BLOCK_SIZE {
16        return Err(CodecError::Corrupt {
17            detail: "FastLanes encoder received an invalid block size".into(),
18        });
19    }
20    let count_u16 = u16::try_from(values.len()).map_err(|_| CodecError::Corrupt {
21        detail: "FastLanes block count does not fit u16".into(),
22    })?;
23
24    let mut min_val = values[0];
25    let mut max_val = values[0];
26    for &value in &values[1..] {
27        min_val = min_val.min(value);
28        max_val = max_val.max(value);
29    }
30
31    let range = (max_val as u128).wrapping_sub(min_val as u128) as u64;
32    let bit_width = if range == 0 {
33        0
34    } else {
35        64 - range.leading_zeros() as u8
36    };
37    let count = usize::from(count_u16);
38    let packed_bits = checked_mul(count, usize::from(bit_width), "FastLanes packed bits")?;
39    let packed_bytes = packed_bits.div_ceil(8);
40    let header_end = checked_add(out.len(), BLOCK_HEADER_SIZE, "FastLanes block header")?;
41    let output_end = checked_add(header_end, packed_bytes, "FastLanes block output")?;
42    encode_u32_len(output_end, "FastLanes encoded output")?;
43
44    out.extend_from_slice(&count_u16.to_le_bytes());
45    out.push(bit_width);
46    out.extend_from_slice(&min_val.to_le_bytes());
47    if bit_width == 0 {
48        return Ok(());
49    }
50
51    out.resize(output_end, 0);
52    let packed = &mut out[header_end..output_end];
53    let mask = if bit_width == 64 {
54        u64::MAX
55    } else {
56        (1u64 << bit_width) - 1
57    };
58    let mut bit_offset = 0usize;
59    for &value in values {
60        let residual = (value.wrapping_sub(min_val) as u64) & mask;
61        pack_bits(packed, bit_offset, residual, bit_width);
62        bit_offset = checked_add(bit_offset, usize::from(bit_width), "FastLanes bit offset")?;
63    }
64    Ok(())
65}
66
67/// Decode a single block from the byte stream and return the next offset.
68pub(super) fn decode_block(
69    data: &[u8],
70    offset: usize,
71    values: &mut Vec<i64>,
72    block_idx: usize,
73    expected_count: usize,
74) -> Result<usize, CodecError> {
75    let header = checked_range(data, offset, BLOCK_HEADER_SIZE, "FastLanes block header")?;
76    let count = usize::from(u16::from_le_bytes([header[0], header[1]]));
77    let bit_width = header[2];
78    let min_val = i64::from_le_bytes([
79        header[3], header[4], header[5], header[6], header[7], header[8], header[9], header[10],
80    ]);
81    validate_block_shape(count, bit_width, block_idx, expected_count)?;
82    let packed_bytes = packed_byte_len(count, bit_width)?;
83    let packed_start = checked_add(offset, BLOCK_HEADER_SIZE, "FastLanes packed start")?;
84    let packed = checked_range(data, packed_start, packed_bytes, "FastLanes packed data")?;
85    validate_padding(packed, count, bit_width, block_idx)?;
86    let next_offset = checked_add(packed_start, packed_bytes, "FastLanes block end")?;
87
88    if bit_width == 0 {
89        values.extend(std::iter::repeat_n(min_val, count));
90        return Ok(next_offset);
91    }
92
93    let mask = if bit_width == 64 {
94        u64::MAX
95    } else {
96        (1u64 << bit_width) - 1
97    };
98    let mut bit_offset = 0usize;
99    for _ in 0..count {
100        let residual = unpack_bits(packed, bit_offset, bit_width) & mask;
101        values.push(min_val.wrapping_add(residual as i64));
102        bit_offset = checked_add(bit_offset, usize::from(bit_width), "FastLanes bit offset")?;
103    }
104    Ok(next_offset)
105}
106
107/// Skip a block without decoding, returning the next byte offset.
108pub(super) fn skip_block(
109    data: &[u8],
110    offset: usize,
111    block_idx: usize,
112    expected_count: usize,
113) -> Result<usize, CodecError> {
114    let header = checked_range(data, offset, BLOCK_HEADER_SIZE, "FastLanes block header")?;
115    let count = usize::from(u16::from_le_bytes([header[0], header[1]]));
116    let bit_width = header[2];
117    validate_block_shape(count, bit_width, block_idx, expected_count)?;
118    let packed_bytes = packed_byte_len(count, bit_width)?;
119    let packed_start = checked_add(offset, BLOCK_HEADER_SIZE, "FastLanes packed start")?;
120    let packed = checked_range(data, packed_start, packed_bytes, "FastLanes packed data")?;
121    validate_padding(packed, count, bit_width, block_idx)?;
122    checked_add(packed_start, packed_bytes, "FastLanes block end")
123}
124
125fn validate_block_shape(
126    count: usize,
127    bit_width: u8,
128    block_idx: usize,
129    expected_count: usize,
130) -> Result<(), CodecError> {
131    if count != expected_count || count == 0 || count > BLOCK_SIZE {
132        return Err(CodecError::Corrupt {
133            detail: format!(
134                "block {block_idx}: count {count} does not match expected {expected_count}"
135            ),
136        });
137    }
138    if bit_width > 64 {
139        return Err(CodecError::Corrupt {
140            detail: format!("block {block_idx}: invalid bit_width {bit_width}"),
141        });
142    }
143    Ok(())
144}
145
146fn packed_byte_len(count: usize, bit_width: u8) -> Result<usize, CodecError> {
147    checked_mul(count, usize::from(bit_width), "FastLanes packed bits").map(|bits| bits.div_ceil(8))
148}
149
150fn validate_padding(
151    packed: &[u8],
152    count: usize,
153    bit_width: u8,
154    block_idx: usize,
155) -> Result<(), CodecError> {
156    let meaningful_bits = checked_mul(count, usize::from(bit_width), "FastLanes packed bits")?;
157    let final_bits = meaningful_bits % 8;
158    if final_bits != 0
159        && packed
160            .last()
161            .is_some_and(|byte| byte & !low_mask_u8(final_bits) != 0)
162    {
163        return Err(CodecError::Corrupt {
164            detail: format!("block {block_idx}: non-zero FastLanes padding bits"),
165        });
166    }
167    Ok(())
168}
169
170/// Compute the minimum number of bits needed to represent the range of values.
171pub fn bit_width_for_range(min: i64, max: i64) -> u8 {
172    let range = (max as u128).wrapping_sub(min as u128) as u64;
173    if range == 0 {
174        0
175    } else {
176        64 - range.leading_zeros() as u8
177    }
178}