use std::collections::HashMap;
use crate::format::chunk_index::btree_v2::build_index;
use crate::format::fractal_heap::HeapParams;
use crate::format::fractal_heap_write::{plan_heap, HeapBlock};
use crate::format::sohm::{
encode_list, list_size, message_hash, record_size, SohmIndexHeader, SohmMasterTable,
SohmRecord, SohmRecordLocation, BT2_TYPE_SOHM_INDEX, SOHM_B2_NODE_SIZE, SOHM_HEAP_ID_LEN,
SOHM_INDEX_BTREE, SOHM_INDEX_LIST,
};
use crate::format::{FormatContext, FormatError, FormatResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SohmIndexSpec {
pub mesg_types: u16,
pub min_mesg_size: u32,
pub list_max: u16,
pub btree_min: u16,
}
pub type SharedKey = (u8, Vec<u8>);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct NestedShare {
pub heap_id_at: usize,
pub target: SharedKey,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SharedMessage {
pub msg_type: u8,
pub body: Vec<u8>,
pub nested: Vec<NestedShare>,
pub ref_count: u32,
pub ohdr_addr: Option<u64>,
}
impl SharedMessage {
pub fn kept_in_ohdr(&self) -> Option<u64> {
self.ohdr_addr.filter(|_| self.ref_count <= 1)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SohmIndexContent {
pub spec: SohmIndexSpec,
pub messages: Vec<SharedMessage>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BuiltSharedMessages {
pub table_addr: u64,
pub blocks: Vec<HeapBlock>,
pub heap_ids: HashMap<(u8, Vec<u8>), [u8; SOHM_HEAP_ID_LEN]>,
}
pub fn build_shared_messages(
indexes: &[SohmIndexContent],
ctx: &FormatContext,
alloc: &mut dyn FnMut(u64) -> u64,
) -> FormatResult<BuiltSharedMessages> {
let mut blocks = Vec::new();
let mut heap_ids: HashMap<SharedKey, [u8; SOHM_HEAP_ID_LEN]> = HashMap::new();
let mut headers = Vec::with_capacity(indexes.len());
let mut plans = Vec::with_capacity(indexes.len());
for index in indexes {
let lengths: Vec<usize> = index
.messages
.iter()
.filter(|m| m.kept_in_ohdr().is_none())
.map(|m| m.body.len())
.collect();
plans.push(plan_heap(
&HeapParams::object_header(),
ctx,
&lengths,
alloc,
)?);
}
let mut collected: HashMap<SharedKey, [u8; SOHM_HEAP_ID_LEN]> = HashMap::new();
for (index, plan) in indexes.iter().zip(&plans) {
let heaped = index.messages.iter().filter(|m| m.kept_in_ohdr().is_none());
for (message, id) in heaped.zip(plan.ids()) {
collected.insert((message.msg_type, message.body.clone()), heap_id(id)?);
}
}
for (index, plan) in indexes.iter().zip(plans) {
let num_messages = u16::try_from(index.messages.len()).map_err(|_| {
FormatError::InvalidData(format!(
"shared-message index holds {} messages, more than the count field takes",
index.messages.len()
))
})?;
let mut bodies = Vec::with_capacity(index.messages.len());
for message in &index.messages {
bodies.push(resolve_nested(message, &collected)?);
}
let heap_addr = plan.header_addr();
let ids = plan.ids().to_vec();
let heaped: Vec<Vec<u8>> = index
.messages
.iter()
.zip(&bodies)
.filter(|(m, _)| m.kept_in_ohdr().is_none())
.map(|(_, body)| body.clone())
.collect();
blocks.extend(plan.finish(&heaped)?.blocks);
let mut records = Vec::with_capacity(index.messages.len());
let mut heaped_at = 0usize;
for (message, body) in index.messages.iter().zip(bodies) {
let location = match message.kept_in_ohdr() {
Some(oh_addr) => SohmRecordLocation::InObjectHeader {
msg_type: message.msg_type,
index: 0,
oh_addr,
},
None => {
let id = heap_id(&ids[heaped_at])?;
heaped_at += 1;
heap_ids.insert((message.msg_type, body.clone()), id);
SohmRecordLocation::InHeap {
ref_count: message.ref_count,
heap_id: id,
}
}
};
records.push((
SohmRecord {
hash: message_hash(&body, message.msg_type),
location,
},
body,
));
}
let (index_type, index_addr) = if is_btree(&index.spec, num_messages) {
(
SOHM_INDEX_BTREE,
build_btree(&mut records, ctx, alloc, &mut blocks),
)
} else {
(
SOHM_INDEX_LIST,
build_list(&records, &index.spec, ctx, alloc, &mut blocks),
)
};
headers.push(SohmIndexHeader {
index_type,
mesg_types: index.spec.mesg_types,
min_mesg_size: index.spec.min_mesg_size,
list_max: index.spec.list_max,
btree_min: index.spec.btree_min,
num_messages,
index_addr,
heap_addr,
});
}
let table = SohmMasterTable { indexes: headers };
let nindexes = u8::try_from(indexes.len()).map_err(|_| {
FormatError::InvalidData(format!("{} shared-message indexes", indexes.len()))
})?;
let table_addr = alloc(SohmMasterTable::encoded_size(ctx, nindexes) as u64);
let image = table.encode(ctx);
blocks.push(HeapBlock {
addr: table_addr,
len: image.len() as u64,
image,
});
Ok(BuiltSharedMessages {
table_addr,
blocks,
heap_ids,
})
}
fn is_btree(spec: &SohmIndexSpec, num_messages: u16) -> bool {
spec.list_max == 0 || num_messages > spec.list_max
}
fn heap_id(id: &[u8]) -> FormatResult<[u8; SOHM_HEAP_ID_LEN]> {
id.try_into().map_err(|_| {
FormatError::InvalidData(format!(
"shared-message heap returned a {}-byte id, expected {SOHM_HEAP_ID_LEN}",
id.len()
))
})
}
fn resolve_nested(
message: &SharedMessage,
placed: &HashMap<SharedKey, [u8; SOHM_HEAP_ID_LEN]>,
) -> FormatResult<Vec<u8>> {
if message.nested.is_empty() {
return Ok(message.body.clone());
}
let mut body = message.body.clone();
for nested in &message.nested {
let Some(heap_id) = placed.get(&nested.target) else {
return Err(FormatError::InvalidData(
"a shared message points at a body no index holds".into(),
));
};
let at = nested.heap_id_at;
if body.len() < at + SOHM_HEAP_ID_LEN {
return Err(FormatError::InvalidData(
"a shared message's nested pointer runs past the body holding it".into(),
));
}
body[at..at + SOHM_HEAP_ID_LEN].copy_from_slice(heap_id);
}
Ok(body)
}
fn build_list(
records: &[(SohmRecord, Vec<u8>)],
spec: &SohmIndexSpec,
ctx: &FormatContext,
alloc: &mut dyn FnMut(u64) -> u64,
blocks: &mut Vec<HeapBlock>,
) -> u64 {
let len = list_size(ctx, spec.list_max) as u64;
let addr = alloc(len);
let entries: Vec<SohmRecord> = records.iter().map(|(r, _)| *r).collect();
blocks.push(HeapBlock {
addr,
len,
image: encode_list(&entries, ctx),
});
addr
}
fn build_btree(
records: &mut [(SohmRecord, Vec<u8>)],
ctx: &FormatContext,
alloc: &mut dyn FnMut(u64) -> u64,
blocks: &mut Vec<HeapBlock>,
) -> u64 {
records.sort_by(|a, b| a.0.hash.cmp(&b.0.hash).then_with(|| a.1.cmp(&b.1)));
let mut image = Vec::with_capacity(records.len() * record_size(ctx));
for (record, _) in records.iter() {
image.extend_from_slice(&record.encode(ctx));
}
let (addr, nodes) = build_index(
BT2_TYPE_SOHM_INDEX,
record_size(ctx) as u16,
SOHM_B2_NODE_SIZE,
&image,
ctx,
alloc,
);
blocks.extend(nodes.into_iter().map(|(addr, image)| HeapBlock {
addr,
len: image.len() as u64,
image,
}));
addr
}
#[cfg(test)]
mod tests {
use super::*;
use crate::format::chunk_index::btree_v2::{collect_btree_v2_records, Bt2Header};
use crate::format::fractal_heap::{
collect_managed_blocks, read_heap_object, FractalHeapHeader, HeapId,
};
use crate::format::messages::{MSG_ATTRIBUTE, MSG_DATASPACE, MSG_DATATYPE};
use crate::format::sohm::{SharedLocation, SharedMessagePointer, SOHM_IN_HEAP};
use crate::format::{BlockReader, UNDEF_ADDR};
struct MemFile {
bytes: Vec<u8>,
}
impl MemFile {
fn new() -> Self {
Self { bytes: vec![0; 16] }
}
fn alloc(&mut self, len: u64) -> u64 {
let addr = self.bytes.len() as u64;
self.bytes.resize(self.bytes.len() + len as usize, 0);
addr
}
}
impl BlockReader for MemFile {
fn read_block(&mut self, offset: u64, len: usize) -> FormatResult<Vec<u8>> {
let start = offset as usize;
if start > self.bytes.len() {
return Err(FormatError::BufferTooShort {
needed: start,
available: self.bytes.len(),
});
}
let end = (start + len).min(self.bytes.len());
Ok(self.bytes[start..end].to_vec())
}
}
fn ctx() -> FormatContext {
FormatContext::default_v3()
}
fn spec(list_max: u16) -> SohmIndexSpec {
SohmIndexSpec {
mesg_types: (1 << MSG_DATASPACE) | (1 << MSG_DATATYPE) | (1 << MSG_ATTRIBUTE),
min_mesg_size: 0,
list_max,
btree_min: 40,
}
}
fn message(msg_type: u8, seed: u8, len: usize, ref_count: u32) -> SharedMessage {
SharedMessage {
msg_type,
body: (0..len).map(|i| seed.wrapping_add(i as u8)).collect(),
nested: Vec::new(),
ref_count,
ohdr_addr: None,
}
}
fn lay_out(indexes: &[SohmIndexContent]) -> (MemFile, BuiltSharedMessages) {
let mut file = MemFile::new();
let built = build_shared_messages(indexes, &ctx(), &mut |len| file.alloc(len)).unwrap();
for block in &built.blocks {
assert!(
block.image.len() as u64 <= block.len,
"block image overruns its allocation"
);
let at = block.addr as usize;
file.bytes[at..at + block.image.len()].copy_from_slice(&block.image);
}
(file, built)
}
fn read_records(file: &mut MemFile, header: &SohmIndexHeader) -> Vec<SohmRecord> {
let size = record_size(&ctx());
let raw = if header.index_type == SOHM_INDEX_LIST {
let buf = file
.read_block(header.index_addr, 4 + size * header.num_messages as usize)
.unwrap();
assert_eq!(&buf[..4], b"SMLI");
buf[4..].to_vec()
} else {
let bt2 = Bt2Header::decode(&file.read_block(header.index_addr, 256).unwrap(), &ctx())
.unwrap();
assert_eq!(bt2.record_type, BT2_TYPE_SOHM_INDEX);
assert_eq!(bt2.node_size, SOHM_B2_NODE_SIZE);
assert_eq!(bt2.record_size as usize, size);
collect_btree_v2_records(&bt2, &ctx(), file).unwrap()
};
raw.chunks_exact(size)
.map(|r| SohmRecord {
hash: u32::from_le_bytes(r[1..5].try_into().unwrap()),
location: if r[0] == SOHM_IN_HEAP {
SohmRecordLocation::InHeap {
ref_count: u32::from_le_bytes(r[5..9].try_into().unwrap()),
heap_id: r[9..17].try_into().unwrap(),
}
} else {
SohmRecordLocation::InObjectHeader {
msg_type: r[6],
index: u16::from_le_bytes(r[7..9].try_into().unwrap()),
oh_addr: u64::from_le_bytes(r[9..17].try_into().unwrap()),
}
},
})
.collect()
}
fn in_heap(record: &SohmRecord) -> (u32, [u8; SOHM_HEAP_ID_LEN]) {
match record.location {
SohmRecordLocation::InHeap { ref_count, heap_id } => (ref_count, heap_id),
other => panic!("expected a heap record, got {other:?}"),
}
}
fn read_body(file: &mut MemFile, heap_addr: u64, record: &SohmRecord) -> Vec<u8> {
let heap =
FractalHeapHeader::decode(&file.read_block(heap_addr, 512).unwrap(), &ctx()).unwrap();
let blocks = collect_managed_blocks(&heap, &ctx(), file).unwrap();
let id = HeapId::parse(&in_heap(record).1, &heap, &ctx()).unwrap();
read_heap_object(&id, &heap, &ctx(), &blocks, file).unwrap()
}
#[test]
fn a_list_index_round_trips_every_body() {
let messages = vec![
message(MSG_DATASPACE, 1, 24, 5),
message(MSG_DATATYPE, 40, 20, 1),
message(MSG_DATASPACE, 90, 24, 1),
message(MSG_ATTRIBUTE, 7, 56, 4),
];
let (mut file, built) = lay_out(&[SohmIndexContent {
spec: spec(50),
messages: messages.clone(),
}]);
let table = SohmMasterTable::decode(
&file
.read_block(built.table_addr, SohmMasterTable::encoded_size(&ctx(), 1))
.unwrap(),
&ctx(),
1,
)
.unwrap();
let header = &table.indexes[0];
assert_eq!(header.index_type, SOHM_INDEX_LIST);
assert_eq!(header.num_messages, 4);
assert_ne!(header.index_addr, UNDEF_ADDR);
let records = read_records(&mut file, header);
assert_eq!(records.len(), 4);
for (message, record) in messages.iter().zip(&records) {
assert_eq!(record.hash, message_hash(&message.body, message.msg_type));
assert_eq!(in_heap(record).0, message.ref_count);
assert_eq!(read_body(&mut file, header.heap_addr, record), message.body);
}
for message in &messages {
let id = built.heap_ids[&(message.msg_type, message.body.clone())];
let pointer =
SharedMessagePointer::decode(&SharedMessagePointer::encode_sohm(id), &ctx())
.unwrap();
assert_eq!(pointer.location, SharedLocation::Sohm);
assert_eq!(pointer.heap_id, id);
}
}
#[test]
fn a_zero_list_max_index_is_a_btree() {
let messages = vec![
message(MSG_DATASPACE, 1, 24, 5),
message(MSG_DATATYPE, 40, 20, 1),
message(MSG_ATTRIBUTE, 7, 56, 4),
];
let (mut file, built) = lay_out(&[SohmIndexContent {
spec: SohmIndexSpec {
list_max: 0,
btree_min: 0,
..spec(0)
},
messages: messages.clone(),
}]);
let table = SohmMasterTable::decode(
&file
.read_block(built.table_addr, SohmMasterTable::encoded_size(&ctx(), 1))
.unwrap(),
&ctx(),
1,
)
.unwrap();
let header = &table.indexes[0];
assert_eq!(header.index_type, SOHM_INDEX_BTREE);
let records = read_records(&mut file, header);
assert_eq!(records.len(), 3);
assert!(
records.windows(2).all(|w| w[0].hash <= w[1].hash),
"records are not hash-ordered: {records:?}"
);
for message in &messages {
let hash = message_hash(&message.body, message.msg_type);
let record = records.iter().find(|r| r.hash == hash).unwrap();
assert_eq!(in_heap(record).0, message.ref_count);
assert_eq!(read_body(&mut file, header.heap_addr, record), message.body);
}
}
#[test]
fn an_index_past_its_list_maximum_is_a_btree() {
let messages: Vec<SharedMessage> = (0..200u32)
.map(|i| SharedMessage {
msg_type: MSG_DATASPACE,
body: i.to_le_bytes().repeat(6),
nested: Vec::new(),
ref_count: i + 1,
ohdr_addr: None,
})
.collect();
let (mut file, built) = lay_out(&[SohmIndexContent {
spec: spec(50),
messages: messages.clone(),
}]);
let table = SohmMasterTable::decode(
&file
.read_block(built.table_addr, SohmMasterTable::encoded_size(&ctx(), 1))
.unwrap(),
&ctx(),
1,
)
.unwrap();
let header = &table.indexes[0];
assert_eq!(header.index_type, SOHM_INDEX_BTREE);
assert_eq!(header.num_messages, 200);
let bt2 =
Bt2Header::decode(&file.read_block(header.index_addr, 256).unwrap(), &ctx()).unwrap();
assert!(bt2.depth > 0, "expected a multi-level index, got one leaf");
let records = read_records(&mut file, header);
assert_eq!(records.len(), 200);
for message in &messages {
let hash = message_hash(&message.body, message.msg_type);
let record = records.iter().find(|r| r.hash == hash).unwrap();
assert_eq!(read_body(&mut file, header.heap_addr, record), message.body);
}
}
#[test]
fn each_index_gets_its_own_heap() {
let (mut file, built) = lay_out(&[
SohmIndexContent {
spec: SohmIndexSpec {
mesg_types: 1 << MSG_ATTRIBUTE,
..spec(50)
},
messages: vec![message(MSG_ATTRIBUTE, 3, 40, 2)],
},
SohmIndexContent {
spec: SohmIndexSpec {
mesg_types: (1 << MSG_DATATYPE) | (1 << MSG_DATASPACE),
..spec(50)
},
messages: vec![message(MSG_DATATYPE, 9, 20, 3)],
},
]);
let table = SohmMasterTable::decode(
&file
.read_block(built.table_addr, SohmMasterTable::encoded_size(&ctx(), 2))
.unwrap(),
&ctx(),
2,
)
.unwrap();
assert_eq!(table.indexes.len(), 2);
assert_ne!(table.indexes[0].heap_addr, table.indexes[1].heap_addr);
assert_eq!(
table.heap_addr(MSG_ATTRIBUTE),
Some(table.indexes[0].heap_addr)
);
assert_eq!(
table.heap_addr(MSG_DATASPACE),
Some(table.indexes[1].heap_addr)
);
}
#[test]
fn an_empty_index_is_still_laid_out() {
let (mut file, built) = lay_out(&[SohmIndexContent {
spec: spec(50),
messages: Vec::new(),
}]);
let table = SohmMasterTable::decode(
&file
.read_block(built.table_addr, SohmMasterTable::encoded_size(&ctx(), 1))
.unwrap(),
&ctx(),
1,
)
.unwrap();
assert_eq!(table.indexes[0].num_messages, 0);
assert_ne!(table.indexes[0].heap_addr, UNDEF_ADDR);
assert!(built.heap_ids.is_empty());
}
}