use rmp_serde::{decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice};
use serde::{Deserialize, Serialize};
use crate::message_pack_format::envelope;
use crate::sketch_framework::eh_sketch_list::wire::{SketchState, rebuild_sketch, sketch_state};
use super::{EHBucket, ExponentialHistogram, compute_l2_mass, infer_merge_norm};
const EH_KIND: &[u8] = &[0x13, 0x00];
#[derive(Debug, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct EhMetadata {
pub(crate) metadata_version: u8,
pub(crate) window: u64,
pub(crate) k: u32,
}
pub(crate) fn eh_metadata(window: u64, k: u32) -> EhMetadata {
EhMetadata {
metadata_version: 1,
window,
k,
}
}
#[derive(Debug, Serialize, Deserialize)]
struct EhPayload {
buckets: Vec<SketchState>,
sizes: Vec<u64>,
min_times: Vec<u64>,
max_times: Vec<u64>,
prototype: SketchState,
}
fn check_bucket(index: usize, bucket: &EHBucket) -> Result<(), String> {
if bucket.size == 0 {
return Err(format!("bucket {index} has size 0"));
}
if bucket.min_time > bucket.max_time {
return Err(format!(
"bucket {index} spans [{}, {}]",
bucket.min_time, bucket.max_time
));
}
let mass = compute_l2_mass(&bucket.bucket);
if bucket.l2_mass != mass {
return Err(format!(
"bucket {index} caches l2_mass {} against its sketch's {mass}",
bucket.l2_mass
));
}
Ok(())
}
fn check_variant(index: usize, bucket_kind: &[u8], prototype_kind: &[u8]) -> Result<(), String> {
if bucket_kind != prototype_kind {
return Err(format!(
"bucket {index} carries kind_id {bucket_kind:02x?}, the prototype's is {prototype_kind:02x?}"
));
}
Ok(())
}
fn check_order(index: usize, previous_max_time: u64, min_time: u64) -> Result<(), String> {
if min_time < previous_max_time {
return Err(format!(
"buckets {} and {index} are out of order: [_, {previous_max_time}] precedes [{min_time}, _]",
index - 1
));
}
Ok(())
}
impl ExponentialHistogram {
pub fn serialize_to_bytes(&self) -> Result<Vec<u8>, RmpEncodeError> {
let fail = |problem: String| {
RmpEncodeError::Syntax(format!("ASAPv1 ExponentialHistogram envelope: {problem}"))
};
if self.k == 0 {
return Err(fail("k must be at least 1".to_string()));
}
let k = u32::try_from(self.k)
.map_err(|_| fail(format!("k {} exceeds the u32 metadata field", self.k)))?;
if self.merge_norm != infer_merge_norm(&self.type_to_clone) {
return Err(fail(format!(
"merge_norm {:?} disagrees with the prototype's",
self.merge_norm
)));
}
let prototype = sketch_state(&self.type_to_clone)?;
let mut buckets = Vec::with_capacity(self.payload.len());
let mut sizes = Vec::with_capacity(self.payload.len());
for (index, bucket) in self.payload.iter().enumerate() {
check_bucket(index, bucket).map_err(fail)?;
if index > 0 {
check_order(index, self.payload[index - 1].max_time, bucket.min_time)
.map_err(fail)?;
}
sizes.push(u64::try_from(bucket.size).map_err(|_| {
fail(format!(
"bucket {index} size {} exceeds the u64 payload field",
bucket.size
))
})?);
let state = sketch_state(&bucket.bucket)?;
check_variant(index, &state.kind_id, &prototype.kind_id).map_err(fail)?;
buckets.push(state);
}
let metadata = rmp_serde::to_vec_named(&eh_metadata(self.window, k))?;
let payload = rmp_serde::to_vec(&EhPayload {
buckets,
sizes,
min_times: self.payload.iter().map(|b| b.min_time).collect(),
max_times: self.payload.iter().map(|b| b.max_time).collect(),
prototype,
})?;
Ok(envelope::encode(EH_KIND, &metadata, &payload))
}
pub fn deserialize_from_bytes(bytes: &[u8]) -> Result<Self, RmpDecodeError> {
let (kind_id, metadata, payload) =
envelope::split(bytes).map_err(RmpDecodeError::Uncategorized)?;
if kind_id != EH_KIND {
return Err(RmpDecodeError::Uncategorized(format!(
"ExponentialHistogram kind_id mismatch: stored {kind_id:?}, expected {EH_KIND:?}"
)));
}
let meta: EhMetadata = from_slice(metadata)?;
if meta != eh_metadata(meta.window, meta.k) {
return Err(RmpDecodeError::Uncategorized(
"ASAPv1 ExponentialHistogram envelope: metadata mismatch".to_string(),
));
}
if meta.k == 0 {
return Err(RmpDecodeError::Uncategorized(
"ExponentialHistogram k must be at least 1".to_string(),
));
}
let p: EhPayload = from_slice(payload)?;
let count = p.buckets.len();
if p.sizes.len() != count || p.min_times.len() != count || p.max_times.len() != count {
return Err(RmpDecodeError::Uncategorized(format!(
"ExponentialHistogram parallel lengths (buckets {count}, sizes {}, min_times {}, max_times {}) disagree",
p.sizes.len(),
p.min_times.len(),
p.max_times.len()
)));
}
let mut decoded = Vec::with_capacity(count);
for (index, triple) in p.buckets.iter().enumerate() {
if index > 0 {
check_order(index, p.max_times[index - 1], p.min_times[index])
.map_err(RmpDecodeError::Uncategorized)?;
}
let sketch = rebuild_sketch(triple)?;
check_variant(index, &triple.kind_id, &p.prototype.kind_id).map_err(|problem| {
RmpDecodeError::Uncategorized(format!("ExponentialHistogram {problem}"))
})?;
let bucket = EHBucket {
l2_mass: compute_l2_mass(&sketch),
bucket: sketch,
size: usize::try_from(p.sizes[index]).map_err(|_| {
RmpDecodeError::Uncategorized(format!(
"ExponentialHistogram bucket {index} size {} exceeds this target's usize",
p.sizes[index]
))
})?,
min_time: p.min_times[index],
max_time: p.max_times[index],
};
check_bucket(index, &bucket).map_err(RmpDecodeError::Uncategorized)?;
decoded.push(bucket);
}
let type_to_clone = rebuild_sketch(&p.prototype)?;
Ok(ExponentialHistogram {
payload: decoded,
window: meta.window,
k: meta.k as usize,
merge_norm: infer_merge_norm(&type_to_clone),
type_to_clone,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sketch_framework::eh_sketch_list::wire::tests::{
alt_profile_triple, populated_variants, relabelled, sample_input,
};
use crate::sketch_framework::eh_sketch_list::wire::{CM_KIND, UNIFORM_KIND};
use crate::sketch_framework::eh_sketch_list::{EHSketchList, SketchNorm};
use crate::{Count, CountMin, DataInput, FastPath, Vector2D};
fn populated_eh() -> ExponentialHistogram {
let mut eh = ExponentialHistogram::new(
2,
1000,
EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8)),
);
for i in 0..6u64 {
eh.update(i * 10, &DataInput::U64(i % 3));
}
eh
}
fn ranges(eh: &ExponentialHistogram) -> Vec<(usize, u64, u64, u64)> {
eh.payload
.iter()
.map(|b| (b.size, b.min_time, b.max_time, b.l2_mass.to_bits()))
.collect()
}
fn envelope_for(payload: &EhPayload) -> Vec<u8> {
let metadata = rmp_serde::to_vec_named(&eh_metadata(1000, 2)).unwrap();
envelope::encode(EH_KIND, &metadata, &rmp_serde::to_vec(payload).unwrap())
}
#[test]
fn eh_round_trip_serialization() {
let eh = populated_eh();
let encoded = eh.serialize_to_bytes().expect("serialize EH");
assert!(encoded.starts_with(b"ASAPv1"));
assert_eq!(&encoded[7..10], &[2u8, 0x13, 0x00]);
let decoded = ExponentialHistogram::deserialize_from_bytes(&encoded).expect("deserialize");
assert_eq!(decoded.window, eh.window);
assert_eq!(decoded.k, eh.k);
assert_eq!(decoded.merge_norm, eh.merge_norm);
assert_eq!(ranges(&decoded), ranges(&eh));
assert_eq!(decoded.bucket_count(), eh.bucket_count());
}
#[test]
fn eh_every_variant_round_trips_as_a_bucket() {
for prototype in populated_variants() {
let name = prototype.sketch_type();
let key = sample_input(name);
let mut eh = ExponentialHistogram::new(3, 1000, prototype);
for i in 0..4u64 {
eh.update(i * 5, &key);
}
let encoded = eh
.serialize_to_bytes()
.unwrap_or_else(|e| panic!("serialize EH<{name}>: {e}"));
let decoded = ExponentialHistogram::deserialize_from_bytes(&encoded)
.unwrap_or_else(|e| panic!("deserialize EH<{name}>: {e}"));
assert_eq!(decoded.type_to_clone.sketch_type(), name);
assert_eq!(ranges(&decoded), ranges(&eh));
let again = decoded.serialize_to_bytes().expect("re-serialize");
assert_eq!(encoded, again, "EH<{name}> is not byte-stable");
}
}
#[test]
fn eh_empty_has_one_encoding_and_round_trips() {
let prototype =
EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let a = ExponentialHistogram::new(2, 1000, prototype.clone());
let b = ExponentialHistogram::new(2, 1000, prototype);
let bytes = a.serialize_to_bytes().expect("serialize");
assert_eq!(bytes, b.serialize_to_bytes().expect("serialize"));
let decoded = ExponentialHistogram::deserialize_from_bytes(&bytes).expect("deserialize");
assert_eq!(decoded.bucket_count(), 0);
assert_eq!(bytes, decoded.serialize_to_bytes().expect("re-serialize"));
}
#[test]
fn eh_carries_a_non_empty_prototype() {
let mut prototype =
EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
for _ in 0..7 {
prototype.insert(&DataInput::U64(9));
}
let eh = ExponentialHistogram::new(2, 1000, prototype);
let bytes = eh.serialize_to_bytes().expect("serialize");
let decoded = ExponentialHistogram::deserialize_from_bytes(&bytes).expect("deserialize");
assert_eq!(
decoded.type_to_clone.query(&DataInput::U64(9)),
eh.type_to_clone.query(&DataInput::U64(9))
);
assert!(decoded.type_to_clone.query(&DataInput::U64(9)).unwrap() >= 7.0);
}
#[test]
fn eh_decoded_re_serializes_byte_identically_and_queries_agree() {
let eh = populated_eh();
let bytes = eh.serialize_to_bytes().expect("serialize");
let decoded = ExponentialHistogram::deserialize_from_bytes(&bytes).expect("deserialize");
assert_eq!(bytes, decoded.serialize_to_bytes().expect("re-serialize"));
let key = DataInput::U64(1);
let original = eh.query_interval_merge(0, 50).expect("query");
let round_tripped = decoded.query_interval_merge(0, 50).expect("query");
assert_eq!(original.query(&key).ok(), round_tripped.query(&key).ok());
}
#[test]
fn eh_rejects_foreign_kind_ids() {
let list = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let list_bytes = list.serialize_to_bytes().expect("serialize EHSketchList");
assert!(ExponentialHistogram::deserialize_from_bytes(&list_bytes).is_err());
let cms = CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8);
let cms_bytes = cms.serialize_to_bytes().expect("serialize CMS");
assert!(ExponentialHistogram::deserialize_from_bytes(&cms_bytes).is_err());
}
#[test]
fn eh_metadata_rejects_unknown_keys() {
#[derive(Serialize)]
struct WithExtra {
metadata_version: u8,
window: u64,
k: u32,
bogus_field: u8,
}
let bytes = rmp_serde::to_vec_named(&WithExtra {
metadata_version: 1,
window: 1000,
k: 2,
bogus_field: 7,
})
.unwrap();
assert!(rmp_serde::from_slice::<EhMetadata>(&bytes).is_err());
}
#[test]
fn eh_metadata_rejects_a_missing_key() {
#[derive(Serialize)]
struct WithoutK {
metadata_version: u8,
window: u64,
}
let bytes = rmp_serde::to_vec_named(&WithoutK {
metadata_version: 1,
window: 1000,
})
.unwrap();
assert!(rmp_serde::from_slice::<EhMetadata>(&bytes).is_err());
}
#[test]
fn eh_rejects_a_zero_k() {
let mut eh = populated_eh();
eh.k = 0;
assert!(eh.serialize_to_bytes().is_err());
let good = populated_eh().serialize_to_bytes().expect("serialize");
let (_, _, payload) = envelope::split(&good).expect("split");
let metadata = rmp_serde::to_vec_named(&eh_metadata(1000, 0)).unwrap();
let bytes = envelope::encode(EH_KIND, &metadata, payload);
assert!(ExponentialHistogram::deserialize_from_bytes(&bytes).is_err());
}
#[test]
fn eh_rejects_parallel_arrays_of_unequal_length() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let triple = sketch_state(&sketch).expect("state");
let payload = EhPayload {
buckets: vec![sketch_state(&sketch).expect("state")],
sizes: vec![1; 1_000_000],
min_times: vec![0],
max_times: vec![0],
prototype: triple,
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
}
#[test]
fn eh_rejects_impossible_bucket_state() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let payload = EhPayload {
buckets: vec![sketch_state(&sketch).expect("state")],
sizes: vec![0],
min_times: vec![0],
max_times: vec![0],
prototype: sketch_state(&sketch).expect("state"),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
let payload = EhPayload {
buckets: vec![sketch_state(&sketch).expect("state")],
sizes: vec![1],
min_times: vec![9],
max_times: vec![4],
prototype: sketch_state(&sketch).expect("state"),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
let mut eh = populated_eh();
eh.payload[0].size = 0;
assert!(eh.serialize_to_bytes().is_err());
}
#[test]
fn eh_rejects_derived_fields_that_disagree() {
let mut eh = populated_eh();
eh.payload[0].l2_mass = 42.0;
assert!(eh.serialize_to_bytes().is_err());
let mut eh = populated_eh();
eh.merge_norm = SketchNorm::L2;
assert!(eh.serialize_to_bytes().is_err());
}
#[test]
fn eh_rejects_an_experimental_kind_id_in_a_bucket() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let payload = EhPayload {
buckets: vec![relabelled(&sketch, UNIFORM_KIND)],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: relabelled(&sketch, CM_KIND),
};
let message = ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload))
.expect_err("a relabelled bucket must not decode")
.to_string();
assert!(!message.is_empty());
#[cfg(not(feature = "experimental"))]
{
assert!(message.contains("UniformSampling"), "{message}");
assert!(message.contains("experimental"), "{message}");
}
}
#[test]
fn eh_rejects_a_custom_hash_profile_bucket() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let payload = EhPayload {
buckets: vec![alt_profile_triple()],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: sketch_state(&sketch).expect("state"),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
}
#[test]
fn eh_rejects_a_bucket_that_disagrees_with_the_prototype() {
let cm = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let cs = EHSketchList::CS(Count::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let payload = EhPayload {
buckets: vec![
sketch_state(&cm).expect("state"),
sketch_state(&cs).expect("state"),
],
sizes: vec![1, 1],
min_times: vec![0, 1],
max_times: vec![0, 1],
prototype: sketch_state(&cm).expect("state"),
};
let message = ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload))
.expect_err("a heterogeneous bucket list must not decode")
.to_string();
assert!(message.contains("the prototype's is"), "{message}");
let payload = EhPayload {
buckets: vec![sketch_state(&cs).expect("state")],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: sketch_state(&cm).expect("state"),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
}
#[test]
fn eh_refuses_to_serialize_a_bucket_that_disagrees_with_the_prototype() {
let mut mixed = populated_eh();
let foreign = EHSketchList::CS(Count::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let last = mixed.payload.len() - 1;
mixed.payload[last].l2_mass = compute_l2_mass(&foreign);
mixed.payload[last].bucket = foreign;
let message = mixed
.serialize_to_bytes()
.expect_err("a bucket of another variant must not serialize")
.to_string();
assert!(message.contains("the prototype's is"), "{message}");
}
#[test]
fn eh_rejects_buckets_out_of_order() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let triple = || sketch_state(&sketch).expect("state");
let payload = EhPayload {
buckets: vec![triple(), triple()],
sizes: vec![1, 1],
min_times: vec![50, 10],
max_times: vec![59, 19],
prototype: triple(),
};
let message = ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload))
.expect_err("shuffled buckets must not decode")
.to_string();
assert!(message.contains("out of order"), "{message}");
let payload = EhPayload {
buckets: vec![triple(), triple()],
sizes: vec![1, 1],
min_times: vec![10, 15],
max_times: vec![19, 25],
prototype: triple(),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
let payload = EhPayload {
buckets: vec![triple(), triple()],
sizes: vec![1, 1],
min_times: vec![10, 19],
max_times: vec![19, 25],
prototype: triple(),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_ok());
}
#[test]
fn eh_refuses_to_serialize_buckets_out_of_order() {
let mut eh = ExponentialHistogram::new(
8,
1000,
EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8)),
);
eh.update(50, &DataInput::U64(1));
eh.update(10, &DataInput::U64(2));
assert_eq!(
(eh.payload[0].min_time, eh.payload[1].min_time),
(50, 10),
"a backwards stream must be what builds the out-of-order payload"
);
let message = eh
.serialize_to_bytes()
.expect_err("an out-of-order payload must not serialize")
.to_string();
assert!(message.contains("out of order"), "{message}");
}
#[test]
fn eh_rejects_an_over_long_kind_id() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let long = || SketchState {
kind_id: vec![0x13; 256],
descriptor: Vec::new(),
state: Vec::new(),
};
let payload = EhPayload {
buckets: vec![long()],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: long(),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
let payload = EhPayload {
buckets: Vec::new(),
sizes: Vec::new(),
min_times: Vec::new(),
max_times: Vec::new(),
prototype: long(),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
let payload = EhPayload {
buckets: vec![sketch_state(&sketch).expect("state")],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: long(),
};
assert!(ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload)).is_err());
}
#[test]
fn eh_rejects_an_unknown_kind_id_in_a_bucket() {
let sketch = EHSketchList::CM(CountMin::<Vector2D<i32>, FastPath>::with_dimensions(3, 8));
let payload = EhPayload {
buckets: vec![relabelled(&sketch, &[0xff, 0xff])],
sizes: vec![1],
min_times: vec![0],
max_times: vec![0],
prototype: sketch_state(&sketch).expect("state"),
};
let message = ExponentialHistogram::deserialize_from_bytes(&envelope_for(&payload))
.expect_err("an unknown kind_id must not decode")
.to_string();
assert!(message.contains("not a wire variant"), "{message}");
}
}