use std::collections::BTreeMap;
use std::path::Path;
use crate::container::{Descriptor, ObjectSource};
use crate::error::Result;
use crate::store::{EmbeddedStore, Id, ObjectStore, externalize};
pub const KIND_OBJECT: &str = "object";
pub const KIND_CHANNEL_PAYLOAD: &str = "channel_payload";
pub const KIND_CHANNEL_HEADER: &str = "channel_header";
pub const KIND_MODEL: &str = "model";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ShareUnit {
pub kind: &'static str,
pub id: [u8; 32],
pub len: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ShareReport {
pub total_bytes: u64,
pub unique_bytes: u64,
pub unit_count: u64,
pub unique_count: u64,
pub by_kind: Vec<(String, u64, u64)>,
}
impl ShareReport {
pub fn to_json(&self) -> String {
let kinds: Vec<String> = self
.by_kind
.iter()
.map(|(k, t, u)| format!("{{\"kind\":\"{k}\",\"total\":{t},\"unique\":{u}}}"))
.collect();
format!(
concat!(
"{{",
"\"total_bytes\":{},",
"\"unique_bytes\":{},",
"\"unit_count\":{},",
"\"unique_count\":{},",
"\"by_kind\":[{}]",
"}}"
),
self.total_bytes,
self.unique_bytes,
self.unit_count,
self.unique_count,
kinds.join(",")
)
}
}
fn unit(kind: &'static str, bytes: &[u8]) -> ShareUnit {
ShareUnit {
kind,
id: *Id::of(bytes).as_bytes(),
len: bytes.len() as u64,
}
}
pub fn extract_units(descriptor: &Descriptor) -> Result<Vec<ShareUnit>> {
let mut units = Vec::new();
for model in &descriptor.models {
let bytes = model.encode()?;
units.push(unit(KIND_MODEL, &bytes));
}
for channel in &descriptor.channels {
units.push(unit(KIND_CHANNEL_HEADER, &channel.header_bytes()?));
units.push(unit(KIND_CHANNEL_PAYLOAD, &channel.payload));
}
for obj in &descriptor.objects {
match obj {
ObjectSource::Inline(bytes) => units.push(unit(KIND_OBJECT, bytes)),
ObjectSource::External { id, len } => units.push(ShareUnit {
kind: KIND_OBJECT,
id: *id.as_bytes(),
len: *len,
}),
}
}
Ok(units)
}
pub fn cohort_report(descriptors: &[Descriptor]) -> Result<ShareReport> {
let mut report = ShareReport::default();
let mut seen: BTreeMap<[u8; 32], u64> = BTreeMap::new();
let mut kind_total: BTreeMap<&'static str, u64> = BTreeMap::new();
let mut kind_unique: BTreeMap<&'static str, BTreeMap<[u8; 32], u64>> = BTreeMap::new();
for descriptor in descriptors {
for u in extract_units(descriptor)? {
report.total_bytes += u.len;
report.unit_count += 1;
*kind_total.entry(u.kind).or_default() += u.len;
kind_unique.entry(u.kind).or_default().insert(u.id, u.len);
if seen.insert(u.id, u.len).is_none() {
report.unique_bytes += u.len;
report.unique_count += 1;
}
}
}
report.by_kind = kind_total
.into_iter()
.map(|(kind, total)| {
let unique: u64 = kind_unique.get(kind).map(|m| m.values().sum()).unwrap_or(0);
(kind.to_string(), total, unique)
})
.collect();
Ok(report)
}
pub fn share_store(field_root: &Path) -> Result<EmbeddedStore> {
EmbeddedStore::open(field_root.join("share"))
}
pub fn externalize_objects(field_root: &Path, descriptor: &mut Descriptor) -> Result<u64> {
let mut store = share_store(field_root)?;
let resolver = store.clone();
externalize(descriptor, &resolver, &mut store)?;
Ok(descriptor.objects.len() as u64)
}
pub fn store_units(field_root: &Path, descriptor: &Descriptor) -> Result<u64> {
let mut store = share_store(field_root)?;
let mut offered: u64 = 0;
for model in &descriptor.models {
store.put(&model.encode()?)?;
offered += 1;
}
for channel in &descriptor.channels {
store.put(&channel.header_bytes()?)?;
store.put(&channel.payload)?;
offered += 2;
}
for obj in &descriptor.objects {
if let ObjectSource::Inline(bytes) = obj {
store.put(bytes)?;
offered += 1;
}
}
Ok(offered)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::container::UNIVERSE;
use crate::dra::Program;
use crate::entropy::codec::{
CODER_ORDER0_BYTE_RANS, CODER_VERSION_1, EntropyChannelDescriptor,
};
use crate::entropy::model::EntropyModel;
use crate::{SOURCE_FORMAT_OPAQUE, integrity};
fn channel(
payload: Vec<u8>,
initial_state: u32,
decoded_length: u64,
) -> EntropyChannelDescriptor {
EntropyChannelDescriptor {
coder: CODER_ORDER0_BYTE_RANS,
coder_version: CODER_VERSION_1,
scale_bits: 8,
lane_count: 1,
model_id: 0,
symbol_count: decoded_length,
decoded_length,
initial_state,
payload,
}
}
fn descriptor(
objects: Vec<ObjectSource>,
channels: Vec<EntropyChannelDescriptor>,
) -> Descriptor {
Descriptor {
universe: UNIVERSE.to_string(),
source_format: SOURCE_FORMAT_OPAQUE,
format_basis: "opaque:share-test".to_string(),
models: vec![EntropyModel::uniform(8).unwrap()],
channels,
objects,
program: Program::new(vec![]),
observation_index: None,
seek_directory: false,
source_sha256: [0u8; 32],
source_len: 0,
}
}
#[test]
fn repeated_objects_are_counted_twice_and_unique_once() {
let payload = b"the same object bytes".to_vec();
let distinct = b"distinct".to_vec();
let d = descriptor(
vec![
ObjectSource::Inline(payload.clone()),
ObjectSource::Inline(payload.clone()),
ObjectSource::Inline(distinct.clone()),
],
vec![],
);
let report = cohort_report(std::slice::from_ref(&d)).unwrap();
assert_eq!(report.unit_count, 4);
let model_len = EntropyModel::uniform(8).unwrap().encode().unwrap().len() as u64;
let distinct_len = distinct.len() as u64;
assert_eq!(
report.total_bytes,
model_len + 2 * payload.len() as u64 + distinct_len
);
assert_eq!(report.unique_count, 3);
assert_eq!(
report.unique_bytes,
model_len + payload.len() as u64 + distinct_len
);
assert!(report.unique_bytes <= report.total_bytes);
}
#[test]
fn identical_channel_payloads_share_the_payload_not_the_header() {
let payload = vec![7u8; 64];
let a = channel(payload.clone(), 1, 64);
let b = channel(payload, 2, 64);
let d = descriptor(vec![], vec![a, b]);
let report = cohort_report(std::slice::from_ref(&d)).unwrap();
let by = |kind: &str| {
report
.by_kind
.iter()
.find(|(k, _, _)| k == kind)
.cloned()
.unwrap()
};
let (_, payload_total, payload_unique) = by(KIND_CHANNEL_PAYLOAD);
assert_eq!(payload_total, 128);
assert_eq!(payload_unique, 64, "identical payloads share once");
let (_, header_total, header_unique) = by(KIND_CHANNEL_HEADER);
assert_eq!(header_total, 2 * 33);
assert_eq!(
header_unique,
2 * 33,
"distinct initial_state -> distinct header"
);
}
#[test]
fn identical_headers_share_when_only_payload_differs() {
let a = channel(vec![1u8; 16], 9, 16);
let b = channel(vec![2u8; 16], 9, 16);
let d = descriptor(vec![], vec![a, b]);
let report = cohort_report(std::slice::from_ref(&d)).unwrap();
let header = report
.by_kind
.iter()
.find(|(k, _, _)| k == KIND_CHANNEL_HEADER)
.unwrap();
assert_eq!(header.1, 66);
assert_eq!(header.2, 33, "equal-length payloads -> one shared header");
}
#[test]
fn report_is_deterministic_and_order_independent() {
let d1 = descriptor(
vec![ObjectSource::Inline(b"alpha".to_vec())],
vec![channel(vec![1, 2, 3], 1, 3)],
);
let d2 = descriptor(
vec![ObjectSource::Inline(b"alpha".to_vec())],
vec![channel(vec![4, 5, 6], 2, 3)],
);
let a = cohort_report(&[d1.clone(), d2.clone()]).unwrap();
let b = cohort_report(&[d2, d1]).unwrap();
assert_eq!(a, b);
assert!(a.unique_bytes <= a.total_bytes);
}
#[test]
fn external_objects_contribute_id_and_length_without_a_store() {
let inline = descriptor(vec![ObjectSource::Inline(b"0123456789".to_vec())], vec![]);
let id = Id::of(b"0123456789");
let external = descriptor(vec![ObjectSource::External { id, len: 10 }], vec![]);
assert_eq!(
extract_units(&inline).unwrap(),
extract_units(&external).unwrap()
);
}
#[test]
fn object_ids_and_source_sha_are_distinct_namespaces() {
let bytes = b"namespace check";
let u = unit(KIND_OBJECT, bytes);
assert_eq!(u.id, *Id::of(bytes).as_bytes());
assert_ne!(u.id, integrity::sha256(bytes));
}
#[test]
fn store_units_persists_every_unique_unit_exactly_once() {
let root = std::env::temp_dir().join(format!(
"vole-share-store-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
let shared_payload = vec![9u8; 32];
let a = descriptor(
vec![ObjectSource::Inline(b"shared-object".to_vec())],
vec![channel(shared_payload.clone(), 1, 32)],
);
let b = descriptor(
vec![ObjectSource::Inline(b"shared-object".to_vec())],
vec![channel(shared_payload, 1, 32)],
);
let offered = store_units(&root, &a).unwrap() + store_units(&root, &b).unwrap();
assert_eq!(
offered, 8,
"2 descriptors x (model + header + payload + object)"
);
let report = cohort_report(&[a, b]).unwrap();
let stats = share_store(&root).unwrap().stats().unwrap();
assert_eq!(stats.stored_bytes, report.unique_bytes);
assert_eq!(stats.object_count, report.unique_count);
std::fs::remove_dir_all(&root).ok();
}
#[test]
fn externalize_objects_resolves_and_stays_exact() {
let root = std::env::temp_dir().join(format!(
"vole-share-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
let source = b"externalized exact bytes";
let mut d = Descriptor {
universe: UNIVERSE.to_string(),
source_format: SOURCE_FORMAT_OPAQUE,
format_basis: "opaque:share-test".to_string(),
models: vec![],
channels: vec![],
objects: vec![
ObjectSource::Inline(source.to_vec()),
ObjectSource::Inline(source.to_vec()),
],
program: Program::new(vec![crate::dra::Op::EmitObject { object_id: 0 }]),
observation_index: None,
seek_directory: false,
source_sha256: integrity::sha256(source),
source_len: source.len() as u64,
};
let offered = store_units(&root, &d).unwrap();
assert_eq!(offered, 2);
externalize_objects(&root, &mut d).unwrap();
assert!(
d.objects
.iter()
.all(|o| matches!(o, ObjectSource::External { .. }))
);
let (bytes, _) = d.serialize().unwrap();
let parsed = Descriptor::parse(&bytes, crate::limits::Limits::DEFAULT).unwrap();
let store = share_store(&root).unwrap();
let out =
crate::materialize::materialize_with(&parsed, &store, crate::limits::Limits::DEFAULT)
.unwrap();
assert_eq!(out, source);
std::fs::remove_dir_all(&root).ok();
}
}