use proptest::prelude::*;
use rstest::rstest;
use crate::MltError;
use crate::decoder::stream::header01;
use crate::decoder::{
DictionaryType, IntEncoding, LengthType, LogicalEncoding, LogicalValue, Morton, OffsetType,
PhysicalEncoding, RawStream, RleMeta, StreamMeta, StreamType,
};
use crate::encoder::model::StreamCtx;
use crate::encoder::{
Codecs, EncodedStream, Encoder, EncoderConfig, ExplicitEncoder, IntEncoder, PhysicalEncoder,
};
use crate::test_helpers::{assert_empty, dec, parser};
use crate::utils::BinarySerializer as _;
fn roundtrip_stream<'a>(buffer: &'a mut Vec<u8>, stream: &EncodedStream) -> RawStream<'a> {
buffer.clear();
buffer.write_stream(stream).unwrap();
assert_empty(header01::parse_stream(buffer, &mut parser()))
}
fn roundtrip_stream_u32s(wire: &[u8]) -> Vec<u32> {
let parsed_stream = assert_empty(header01::parse_stream(wire, &mut parser()));
let mut decoder = dec();
let values = parsed_stream.decode_ints::<u32>(&mut decoder).unwrap();
if !values.is_empty() {
assert!(
decoder.consumed() > 0,
"decoder should consume bytes after decode"
);
}
values
}
fn make_logical_val(logical_encoding: LogicalEncoding, num_values: usize) -> LogicalValue {
LogicalValue::new(
StreamMeta::new2(
StreamType::Data(DictionaryType::None),
logical_encoding,
PhysicalEncoding::VarInt,
num_values,
)
.unwrap(),
)
}
#[derive(Debug)]
struct StreamTestCase {
meta: StreamMeta,
data: &'static [u8],
expected_u32_logical_value: Option<Vec<u32>>,
}
fn generate_stream_test_cases() -> Vec<StreamTestCase> {
vec![
StreamTestCase {
meta: StreamMeta::new(
StreamType::Data(DictionaryType::None),
IntEncoding::new(LogicalEncoding::None, PhysicalEncoding::VarInt),
4,
),
data: &[0x04, 0x03, 0x02, 0x01],
expected_u32_logical_value: Some(vec![4, 3, 2, 1]),
},
StreamTestCase {
meta: StreamMeta::new(
StreamType::Data(DictionaryType::None),
IntEncoding::none(),
1,
),
data: &[0x04, 0x03, 0x02, 0x01],
expected_u32_logical_value: Some(vec![0x0102_0304]),
},
]
}
fn create_stream_from_test_case(test_case: &StreamTestCase) -> RawStream<'_> {
RawStream::new(test_case.meta, test_case.data)
}
#[test]
fn test_decode_bits_u32() {
let test_cases = generate_stream_test_cases();
for test_case in test_cases {
if let Some(expected_buf) = &test_case.expected_u32_logical_value {
let stream = create_stream_from_test_case(&test_case);
let mut buf = Vec::new();
stream
.decode_bits::<u32>(&mut buf, &mut dec())
.expect("Should successfully decode u32 values");
assert_eq!(
&buf, expected_buf,
"Should produce decoded u32 values correctly"
);
}
}
}
#[rstest]
#[case::componentwise_delta(LogicalEncoding::ComponentwiseDelta, vec![0u32, 0, 2, 4, 2, 4], vec![0i32, 0, 1, 2, 2, 4]
)]
#[case::delta(LogicalEncoding::Delta, vec![0u32, 1, 2, 1, 2], vec![0i32, -1, 0, -1, 0])]
#[case::delta_rle(LogicalEncoding::DeltaRle(RleMeta::Split { runs: 2, num_rle_values: 5 }), vec![3u32, 2, 0, 2], vec![0i32, 0, 0, 1, 2]
)]
#[case::delta_empty(LogicalEncoding::Delta, vec![], vec![])]
fn test_decode_i32(
#[case] logical_encoding: LogicalEncoding,
#[case] input_data: Vec<u32>,
#[case] expected: Vec<i32>,
) {
let result =
make_logical_val(logical_encoding, input_data.len()).decode_i32(&input_data, &mut dec());
assert!(result.is_ok(), "should decode successfully");
assert_eq!(result.unwrap(), expected, "should match expected output");
}
#[rstest]
#[case::empty(LogicalEncoding::None, vec![], vec![])]
#[case::new_encoded(LogicalEncoding::None, vec![10u32, 20, 30, 40], vec![10u32, 20, 30, 40])]
#[case::rle(LogicalEncoding::Rle(RleMeta::Split { runs: 3, num_rle_values: 6 }), vec![3u32, 2, 1, 10, 20, 30], vec![10u32, 10, 10, 20, 20, 30]
)]
#[case::delta(LogicalEncoding::Delta, vec![0u32, 2, 2, 2, 2], vec![0u32, 1, 2, 3, 4])]
fn test_decode_u32(
#[case] logical_encoding: LogicalEncoding,
#[case] input_data: Vec<u32>,
#[case] expected: Vec<u32>,
) {
let result =
make_logical_val(logical_encoding, input_data.len()).decode_u32(&input_data, &mut dec());
assert!(result.is_ok(), "should decode successfully");
assert_eq!(result.unwrap(), expected, "should match expected output");
}
#[rstest]
#[case::basic(vec![1, 2, 3, 4, 5, 100, 1000])]
#[case::large(vec![1_000_000; 256])]
#[case::edge_values(vec![0, 1, 2, 4, 8, 16, 1024, 65535, 1_000_000_000, u32::MAX])]
#[case::empty(vec![])]
fn test_fastpfor_roundtrip(#[case] values: Vec<u32>) {
let mut enc = Encoder::with_explicit(
EncoderConfig::default(),
ExplicitEncoder::all(IntEncoder::fastpfor()),
);
let codecs = &mut Codecs::default();
let ctx = StreamCtx::prop_data("test");
codecs.write_int_stream(&values, &ctx, &mut enc).unwrap();
let decoded_values = roundtrip_stream_u32s(enc.data());
assert_eq!(decoded_values, values);
}
fn auto_physical(values: &[u32], cfg: EncoderConfig) -> PhysicalEncoding {
let mut enc = Encoder::new(cfg);
let codecs = &mut Codecs::default();
let ctx = StreamCtx::prop_data("test");
codecs.write_int_stream(values, &ctx, &mut enc).unwrap();
let parsed = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
parsed.meta.encoding.physical
}
#[rstest]
#[case::u32_zero(0u32, PhysicalEncoding::VarInt)]
#[case::u32_small(5u32, PhysicalEncoding::VarInt)]
#[case::u32_two_bytes(1000u32, PhysicalEncoding::VarInt)]
#[case::u32_boundary_lo((1u32 << 28) - 1, PhysicalEncoding::VarInt)]
#[case::u32_boundary_hi(1u32 << 28, PhysicalEncoding::None)]
#[case::u32_max(u32::MAX, PhysicalEncoding::None)]
fn single_value_u32_picks_smaller_physical(
#[case] v: u32,
#[case] expected_physical: PhysicalEncoding,
) {
let mut enc = Encoder::new(EncoderConfig::default());
let codecs = &mut Codecs::default();
let ctx = StreamCtx::prop_data("test");
codecs.write_int_stream(&[v], &ctx, &mut enc).unwrap();
let parsed = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
assert_eq!(parsed.meta.encoding.logical, LogicalEncoding::None);
assert_eq!(parsed.meta.encoding.physical, expected_physical);
assert_eq!(parsed.meta.num_values, 1);
assert_eq!(roundtrip_stream_u32s(enc.data()), vec![v]);
}
#[rstest]
#[case::i32_zero(0i32, PhysicalEncoding::VarInt)]
#[case::i32_neg_one(-1i32, PhysicalEncoding::VarInt)]
#[case::i32_small_neg(-1000i32, PhysicalEncoding::VarInt)]
#[case::i32_large_pos(i32::MAX, PhysicalEncoding::None)]
#[case::i32_large_neg(i32::MIN, PhysicalEncoding::None)]
fn single_value_i32_picks_smaller_physical(
#[case] v: i32,
#[case] expected_physical: PhysicalEncoding,
) {
let mut enc = Encoder::new(EncoderConfig::default());
let codecs = &mut Codecs::default();
let ctx = StreamCtx::prop_data("test");
codecs.write_int_stream(&[v], &ctx, &mut enc).unwrap();
let parsed = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
assert_eq!(parsed.meta.encoding.logical, LogicalEncoding::None);
assert_eq!(parsed.meta.encoding.physical, expected_physical);
}
#[test]
fn allow_fastpfor_gates_fastpfor_selection() {
let values: Vec<u32> = (0..2000u32)
.map(|i| i.wrapping_mul(2_654_435_761) % 4096)
.collect();
let on = EncoderConfig::default().with_fastpfor(true);
let off = EncoderConfig::default().with_fastpfor(false);
assert_eq!(
auto_physical(&values, on),
PhysicalEncoding::FastPFor256,
"FastPFOR should win for this data when allow_fastpfor = true"
);
assert_ne!(
auto_physical(&values, off),
PhysicalEncoding::FastPFor256,
"allow_fastpfor = false must prevent FastPFOR from being selected"
);
}
#[rstest]
#[case::new_encoded(StreamType::Data(DictionaryType::None), 2, LogicalEncoding::None, PhysicalEncoding::None, vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], false
)]
#[case::new_encoded(StreamType::Data(DictionaryType::None), 2, LogicalEncoding::ComponentwiseDelta, PhysicalEncoding::None, vec![0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00], false
)]
#[case::new_encoded(StreamType::Offset(OffsetType::Vertex), 3, LogicalEncoding::None, PhysicalEncoding::None, vec![0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00], false
)]
#[case::varint(StreamType::Data(DictionaryType::None), 4, LogicalEncoding::None, PhysicalEncoding::VarInt, vec![0x0A, 0x14, 0x1E, 0x28], false
)]
#[case::varint(StreamType::Data(DictionaryType::None), 5, LogicalEncoding::Delta, PhysicalEncoding::VarInt, vec![0x00, 0x02, 0x02, 0x02, 0x02], false
)]
#[case::varint(StreamType::Data(DictionaryType::None), 3, LogicalEncoding::PseudoDecimal, PhysicalEncoding::VarInt, vec![0x01, 0x02, 0x03], false
)]
#[case::varint(StreamType::Length(LengthType::VarBinary), 3, LogicalEncoding::Delta, PhysicalEncoding::VarInt, vec![0x00, 0x02, 0x02], false
)]
#[case::rle(StreamType::Data(DictionaryType::None), 6, LogicalEncoding::Rle(RleMeta::Split { runs: 3, num_rle_values: 6 }), PhysicalEncoding::VarInt, vec![0x03, 0x02, 0x01, 0x0A, 0x14, 0x1E], false
)]
#[case::rle(StreamType::Data(DictionaryType::None), 5, LogicalEncoding::DeltaRle(RleMeta::Split { runs: 2, num_rle_values: 5 }), PhysicalEncoding::VarInt, vec![0x03, 0x02, 0x00, 0x02], false
)]
#[case::morton(StreamType::Data(DictionaryType::Morton), 4, LogicalEncoding::Morton(Morton { bits: 16, shift: 0 }), PhysicalEncoding::VarInt, vec![0x01, 0x02, 0x03, 0x04], false
)]
#[case::boolean(StreamType::Present, 16, LogicalEncoding::Rle(RleMeta::Split { runs: 2, num_rle_values: 2 }), PhysicalEncoding::VarInt, vec![0xFF, 0x00], true
)]
fn test_stream_roundtrip(
#[case] stream_type: StreamType,
#[case] num_values: u32,
#[case] logical_encoding: LogicalEncoding,
#[case] physical_encoding: PhysicalEncoding,
#[case] data_bytes: Vec<u8>,
#[case] is_bool: bool,
) {
let stream = EncodedStream {
meta: StreamMeta::new(
stream_type,
IntEncoding::new(logical_encoding, physical_encoding),
num_values,
),
data: data_bytes,
};
let mut buffer = Vec::new();
if is_bool {
buffer.write_boolean_stream(&stream).unwrap();
} else {
buffer.write_stream(&stream).unwrap();
}
let parsed = assert_empty(if is_bool {
header01::parse_bool_stream(&buffer, &mut parser())
} else {
header01::parse_stream(&buffer, &mut parser())
});
assert_eq!(parsed.meta, stream.meta, "metadata mismatch");
assert_eq!(stream.data.as_slice(), parsed.data, "data mismatch");
}
#[test]
fn test_morton_parse_rejects_too_many_bits() {
let stream = EncodedStream {
meta: StreamMeta::new(
StreamType::Data(DictionaryType::Morton),
IntEncoding::new(
LogicalEncoding::Morton(Morton { bits: 17, shift: 0 }),
PhysicalEncoding::VarInt,
),
1,
),
data: vec![0],
};
let mut buffer = Vec::new();
buffer.write_stream(&stream).unwrap();
let err = header01::parse_stream(&buffer, &mut parser()).unwrap_err();
assert!(matches!(err, MltError::InvalidMortonBits(17)));
}
#[test]
fn test_varint_stream_huge_num_values_empty_data() {
let wire: &[u8] = &[0x00, 0x02, 0xd5, 0xff, 0xd5, 0xff, 0x03, 0x00];
let result = header01::parse_stream(wire, &mut parser());
assert!(
result.is_err(),
"parse must fail when num_values * 8 exceeds the memory budget"
);
}
#[test]
fn test_rle_num_rle_values_mismatch() {
let rle = RleMeta::Split {
runs: 1,
num_rle_values: u32::MAX,
};
let data = [1u32, 42u32];
let result = rle.decode::<u32>(&data, &mut dec());
assert!(
result.is_err(),
"must reject mismatched num_rle_values before allocating"
);
}
fn encoding_no_fastpfor() -> impl Strategy<Value = IntEncoder> {
any::<IntEncoder>().prop_filter("not fastpfor", |v| v.physical != PhysicalEncoder::FastPFOR)
}
fn dedup_and_get_parts(values: &[&str]) -> (Vec<u32>, Vec<u32>) {
use crate::encoder::stream::dedup_strings;
use crate::utils::strings_to_lengths;
let (unique, offset_indices) = dedup_strings(values).unwrap();
let lengths = strings_to_lengths(&unique).unwrap();
(offset_indices, lengths)
}
#[rstest]
#[case::with_duplicates(&["apple", "banana", "apple", "cherry", "banana", "apple"], &[0, 1, 0, 2, 1, 0], &[5, 6, 6]
)]
#[case::all_unique(&["a", "b", "c", "d"], &[0, 1, 2, 3], &[1, 1, 1, 1])]
#[case::all_same(&["same", "same", "same", "same"], &[0, 0, 0, 0], &[4])]
fn test_encode_strings_dict(
#[case] values: &[&str],
#[case] expected_offsets: &[u32],
#[case] expected_lengths: &[u32],
) {
let (offsets, lengths) = dedup_and_get_parts(values);
assert_eq!(offsets, expected_offsets);
assert_eq!(lengths, expected_lengths);
}
proptest! {
#[test]
fn test_i8_roundtrip(
values in prop::collection::vec(any::<i8>(), 0..100),
encoding in any::<IntEncoder>(),
) {
let widened: Vec<i32> = values.iter().map(|&v| i32::from(v)).collect();
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&widened, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let parsed_stream = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
let decoded_values = parsed_stream.decode_narrow::<i8, i32>(&mut dec()).unwrap();
assert_eq!(decoded_values, values);
}
#[test]
fn test_u8_roundtrip(
values in prop::collection::vec(any::<u8>(), 0..100),
encoding in any::<IntEncoder>()
) {
let widened: Vec<u32> = values.iter().map(|&v| u32::from(v)).collect();
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&widened, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let parsed_stream = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
let decoded_values = parsed_stream.decode_narrow::<u8, u32>(&mut dec()).unwrap();
assert_eq!(decoded_values, values);
}
#[test]
fn test_u32_roundtrip(
values in prop::collection::vec(any::<u32>(), 0..100),
encoding in any::<IntEncoder>()
) {
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&values, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let decoded_values = roundtrip_stream_u32s(enc.data());
assert_eq!(decoded_values, values);
}
#[test]
fn test_i32_roundtrip(
values in prop::collection::vec(any::<i32>(), 0..100),
encoding in any::<IntEncoder>(),
) {
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&values, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let parsed_stream = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
let decoded_values = parsed_stream.decode_ints::<i32>(&mut dec()).unwrap();
assert_eq!(decoded_values, values);
}
#[test]
fn test_u64_roundtrip(
values in prop::collection::vec(any::<u64>(), 0..100),
encoding in encoding_no_fastpfor()
) {
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&values, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let parsed_stream = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
let decoded_values = parsed_stream.decode_ints::<u64>(&mut dec()).unwrap();
assert_eq!(decoded_values, values);
}
#[test]
fn test_i64_roundtrip(
values in prop::collection::vec(any::<i64>(), 0..100),
encoding in encoding_no_fastpfor()
) {
let mut enc = Encoder::with_explicit(EncoderConfig::default(), ExplicitEncoder::all(encoding));
let mut codecs = Codecs::default();
codecs.write_int_stream(&values, &StreamCtx::prop_data("test"), &mut enc).unwrap();
let parsed_stream = assert_empty(header01::parse_stream(enc.data(), &mut parser()));
let decoded_values = parsed_stream.decode_ints::<i64>(&mut dec()).unwrap();
assert_eq!(decoded_values, values);
}
#[test]
fn test_f32_roundtrip(values in prop::collection::vec(any::<f32>(), 0..100)) {
let owned_stream = EncodedStream::encode_floats(&values).unwrap();
let mut buf = Vec::new();
let parsed_stream = roundtrip_stream(&mut buf, &owned_stream);
let decoded_values = parsed_stream.decode_floats::<f32>(&mut dec()).unwrap();
assert_eq!(decoded_values.len(), values.len());
for (v1, v2) in decoded_values.iter().zip(values.iter()) {
assert_eq!(
v1.to_bits(),
v2.to_bits(),
"despite being semantically equal, the values are not actually equal"
);
}
}
#[test]
fn test_f64_roundtrip(values in prop::collection::vec(any::<f64>(), 0..100)) {
let owned_stream = EncodedStream::encode_floats(&values).unwrap();
let mut buf = Vec::new();
let parsed_stream = roundtrip_stream(&mut buf, &owned_stream);
let decoded_values = parsed_stream.decode_floats::<f64>(&mut dec()).unwrap();
assert_eq!(decoded_values.len(), values.len());
for (v1, v2) in decoded_values.iter().zip(values.iter()) {
assert_eq!(
v1.to_bits(),
v2.to_bits(),
"despite being semantically equal, the values are not actually equal"
);
}
}
}