use crate::{ctx::*, DekuError, DekuRead, DekuWrite};
use bitvec::prelude::*;
use std::collections::HashMap;
use std::hash::{BuildHasher, Hash};
#[allow(clippy::type_complexity)]
fn read_hashmap_with_predicate<
'a,
K: DekuRead<'a, Ctx> + Eq + Hash,
V: DekuRead<'a, Ctx>,
S: BuildHasher + Default,
Ctx: Copy,
Predicate: FnMut(usize, &(K, V)) -> bool,
>(
input: &'a BitSlice<u8, Msb0>,
capacity: Option<usize>,
ctx: Ctx,
mut predicate: Predicate,
) -> Result<(&'a BitSlice<u8, Msb0>, HashMap<K, V, S>), DekuError> {
let mut res = HashMap::with_capacity_and_hasher(capacity.unwrap_or(0), S::default());
let mut rest = input;
let mut found_predicate = false;
while !found_predicate {
let (new_rest, kv) = <(K, V)>::read(rest, ctx)?;
found_predicate = predicate(
unsafe { new_rest.as_bitptr().offset_from(input.as_bitptr()) } as usize,
&kv,
);
res.insert(kv.0, kv.1);
rest = new_rest;
}
Ok((rest, res))
}
impl<
'a,
K: DekuRead<'a, Ctx> + Eq + Hash,
V: DekuRead<'a, Ctx>,
S: BuildHasher + Default,
Ctx: Copy,
Predicate: FnMut(&(K, V)) -> bool,
> DekuRead<'a, (Limit<(K, V), Predicate>, Ctx)> for HashMap<K, V, S>
{
fn read(
input: &'a BitSlice<u8, Msb0>,
(limit, inner_ctx): (Limit<(K, V), Predicate>, Ctx),
) -> Result<(&'a BitSlice<u8, Msb0>, Self), DekuError>
where
Self: Sized,
{
match limit {
Limit::Count(mut count) => {
if count == 0 {
return Ok((input, HashMap::<K, V, S>::default()));
}
read_hashmap_with_predicate(input, Some(count), inner_ctx, move |_, _| {
count -= 1;
count == 0
})
}
Limit::Until(mut predicate, _) => {
read_hashmap_with_predicate(input, None, inner_ctx, move |_, kv| predicate(kv))
}
Limit::BitSize(size) => {
let bit_size = size.0;
read_hashmap_with_predicate(input, None, inner_ctx, move |read_bits, _| {
read_bits == bit_size
})
}
Limit::ByteSize(size) => {
let bit_size = size.0 * 8;
read_hashmap_with_predicate(input, None, inner_ctx, move |read_bits, _| {
read_bits == bit_size
})
}
}
}
}
impl<
'a,
K: DekuRead<'a> + Eq + Hash,
V: DekuRead<'a>,
S: BuildHasher + Default,
Predicate: FnMut(&(K, V)) -> bool,
> DekuRead<'a, Limit<(K, V), Predicate>> for HashMap<K, V, S>
{
fn read(
input: &'a BitSlice<u8, Msb0>,
limit: Limit<(K, V), Predicate>,
) -> Result<(&'a BitSlice<u8, Msb0>, Self), DekuError>
where
Self: Sized,
{
Self::read(input, (limit, ()))
}
}
impl<K: DekuWrite<Ctx>, V: DekuWrite<Ctx>, S, Ctx: Copy> DekuWrite<Ctx> for HashMap<K, V, S> {
fn write(&self, output: &mut BitVec<u8, Msb0>, inner_ctx: Ctx) -> Result<(), DekuError> {
for kv in self {
kv.write(output, inner_ctx)?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use rustc_hash::FxHashMap;
macro_rules! fxhashmap(
{ $($key:expr => $value:expr),+ } => {
{
let mut m = FxHashMap::default();
$(
m.insert($key, $value);
)+
m
}
};
);
#[rstest(input, endian, bit_size, limit, expected, expected_rest,
case::count_0([0xAA].as_ref(), Endian::Little, Some(8), 0.into(), FxHashMap::default(), bits![u8, Msb0; 1, 0, 1, 0, 1, 0, 1, 0]),
case::count_1([0x01, 0xAA, 0x02, 0xBB].as_ref(), Endian::Little, Some(8), 1.into(), fxhashmap!{0x01 => 0xAA}, bits![u8, Msb0; 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
case::count_2([0x01, 0xAA, 0x02, 0xBB, 0xBB].as_ref(), Endian::Little, Some(8), 2.into(), fxhashmap!{0x01 => 0xAA, 0x02 => 0xBB}, bits![u8, Msb0; 1, 0, 1, 1, 1, 0, 1, 1]),
case::until_null([0x01, 0xAA, 0, 0, 0xBB].as_ref(), Endian::Little, None, (|kv: &(u8, u8)| kv.0 == 0u8 && kv.1 == 0u8).into(), fxhashmap!{0x01 => 0xAA, 0 => 0}, bits![u8, Msb0; 1, 0, 1, 1, 1, 0, 1, 1]),
case::until_bits([0x01, 0xAA, 0xBB].as_ref(), Endian::Little, None, BitSize(16).into(), fxhashmap!{0x01 => 0xAA}, bits![u8, Msb0; 1, 0, 1, 1, 1, 0, 1, 1]),
case::bits_6([0b0000_0100, 0b1111_0000, 0b1000_0000].as_ref(), Endian::Little, Some(6), 2.into(), fxhashmap!{0x01 => 0x0F, 0x02 => 0}, bits![u8, Msb0;]),
#[should_panic(expected = "Parse(\"too much data: container of 8 bits cannot hold 9 bits\")")]
case::not_enough_data([].as_ref(), Endian::Little, Some(9), 1.into(), FxHashMap::default(), bits![u8, Msb0;]),
#[should_panic(expected = "Parse(\"too much data: container of 8 bits cannot hold 9 bits\")")]
case::not_enough_data([0xAA].as_ref(), Endian::Little, Some(9), 1.into(), FxHashMap::default(), bits![u8, Msb0;]),
#[should_panic(expected = "Incomplete(NeedSize { bits: 8 })")]
case::not_enough_data([0xAA].as_ref(), Endian::Little, Some(8), 2.into(), FxHashMap::default(), bits![u8, Msb0;]),
#[should_panic(expected = "Incomplete(NeedSize { bits: 8 })")]
case::not_enough_data_until([0xAA].as_ref(), Endian::Little, Some(8), (|_: &(u8, u8)| false).into(), FxHashMap::default(), bits![u8, Msb0;]),
#[should_panic(expected = "Incomplete(NeedSize { bits: 8 })")]
case::not_enough_data_bits([0xAA].as_ref(), Endian::Little, Some(8), (BitSize(16)).into(), FxHashMap::default(), bits![u8, Msb0;]),
#[should_panic(expected = "Parse(\"too much data: container of 8 bits cannot hold 9 bits\")")]
case::too_much_data([0xAA, 0xBB].as_ref(), Endian::Little, Some(9), 1.into(), FxHashMap::default(), bits![u8, Msb0;]),
)]
fn test_hashmap_read<Predicate: FnMut(&(u8, u8)) -> bool>(
input: &[u8],
endian: Endian,
bit_size: Option<usize>,
limit: Limit<(u8, u8), Predicate>,
expected: FxHashMap<u8, u8>,
expected_rest: &BitSlice<u8, Msb0>,
) {
let bit_slice = input.view_bits::<Msb0>();
let (rest, res_read) = match bit_size {
Some(bit_size) => {
FxHashMap::<u8, u8>::read(bit_slice, (limit, (endian, BitSize(bit_size)))).unwrap()
}
None => FxHashMap::<u8, u8>::read(bit_slice, (limit, (endian))).unwrap(),
};
assert_eq!(expected, res_read);
assert_eq!(expected_rest, rest);
}
#[rstest(input, endian, expected,
case::normal(fxhashmap!{0x11u8 => 0xAABBu16, 0x23u8 => 0xCCDDu16}, Endian::Little, vec![0x11, 0xBB, 0xAA, 0x23, 0xDD, 0xCC]),
)]
fn test_hashmap_write(input: FxHashMap<u8, u16>, endian: Endian, expected: Vec<u8>) {
let mut res_write = bitvec![u8, Msb0;];
input.write(&mut res_write, endian).unwrap();
assert_eq!(expected, res_write.into_vec());
}
#[rstest(input, endian, limit, expected, expected_rest, expected_write,
case::normal_le([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Little, 2.into(), fxhashmap!{0xBBAA => 0, 0xDDCC => 0}, bits![u8, Msb0;], vec![0xCC, 0xDD, 0, 0xAA, 0xBB, 0]),
case::normal_be([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Big, 2.into(), fxhashmap!{0xAABB => 0, 0xCCDD => 0}, bits![u8, Msb0;], vec![0xCC, 0xDD, 0, 0xAA, 0xBB, 0]),
case::predicate_le([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Little, (|kv: &(u16, u8)| kv.0 == 0xBBAA && kv.1 == 0).into(), fxhashmap!{0xBBAA => 0}, bits![u8, Msb0; 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], vec![0xAA, 0xBB, 0]),
case::predicate_be([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Big, (|kv: &(u16, u8)| kv.0 == 0xAABB && kv.1 == 0).into(), fxhashmap!{0xAABB => 0}, bits![u8, Msb0; 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], vec![0xAA, 0xBB, 0]),
case::bytes_le([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Little, BitSize(24).into(), fxhashmap!{0xBBAA => 0}, bits![u8, Msb0; 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], vec![0xAA, 0xBB, 0]),
case::bytes_be([0xAA, 0xBB, 0, 0xCC, 0xDD, 0].as_ref(), Endian::Big, BitSize(24).into(), fxhashmap!{0xAABB => 0}, bits![u8, Msb0; 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], vec![0xAA, 0xBB, 0]),
)]
fn test_hashmap_read_write<Predicate: FnMut(&(u16, u8)) -> bool>(
input: &[u8],
endian: Endian,
limit: Limit<(u16, u8), Predicate>,
expected: FxHashMap<u16, u8>,
expected_rest: &BitSlice<u8, Msb0>,
expected_write: Vec<u8>,
) {
let bit_slice = input.view_bits::<Msb0>();
let (rest, res_read) = FxHashMap::<u16, u8>::read(bit_slice, (limit, endian)).unwrap();
assert_eq!(expected, res_read);
assert_eq!(expected_rest, rest);
let mut res_write = bitvec![u8, Msb0;];
res_read.write(&mut res_write, endian).unwrap();
assert_eq!(expected_write, res_write.into_vec());
}
}