use crate::RevitFile;
use crate::partition_element_records::bbox_marker;
use std::collections::{BTreeMap, BTreeSet};
const MAX_ENTRIES: usize = 64;
const MAX_BLOCK_LEN: usize = 200_000;
const VALUE_BLOCK_MARK: [u8; 59] = {
let mut mark = [0xffu8; 59];
mark[56] = 0;
mark[57] = 0;
mark[58] = 0;
mark
};
const VALUE_BLOCK_MARK_2023: [u8; 31] = {
let mut mark = [0xffu8; 31];
mark[28] = 0;
mark[29] = 0;
mark[30] = 0;
mark
};
#[derive(Clone, Copy)]
struct BlockLayout {
mark: &'static [u8],
id_len: usize,
first_map_only: bool,
}
const BLOCK_LAYOUT: BlockLayout = BlockLayout {
mark: &VALUE_BLOCK_MARK,
id_len: 8,
first_map_only: false,
};
const BLOCK_LAYOUT_2023: BlockLayout = BlockLayout {
mark: &VALUE_BLOCK_MARK_2023,
id_len: 4,
first_map_only: true,
};
fn id_at(buf: &[u8], at: usize, len: usize) -> Option<u32> {
let bytes = buf.get(at..at.checked_add(len)?)?;
match len {
4 => Some(u32::from_le_bytes(bytes.try_into().ok()?)),
_ => u32::try_from(u64::from_le_bytes(bytes.try_into().ok()?)).ok(),
}
}
pub fn scan_partition(
buf: &[u8],
declared_ids: &BTreeSet<u32>,
materials: &BTreeSet<u32>,
) -> BTreeMap<u32, Vec<u32>> {
scan_partition_with(buf, declared_ids, materials, BLOCK_LAYOUT)
}
fn scan_partition_with(
buf: &[u8],
declared_ids: &BTreeSet<u32>,
materials: &BTreeSet<u32>,
layout: BlockLayout,
) -> BTreeMap<u32, Vec<u32>> {
let blocks: Vec<(usize, u32)> = memchr::memmem::find_iter(buf, layout.mark)
.filter(|&at| at >= layout.id_len && buf[at - 1] != 0xff)
.filter_map(|at| {
let id = id_at(buf, at - layout.id_len, layout.id_len)?;
declared_ids.contains(&id).then_some((at, id))
})
.collect();
let mut out: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
for (index, &(start, owner)) in blocks.iter().enumerate() {
let end = blocks
.get(index + 1)
.map_or(buf.len(), |next| next.0.saturating_sub(layout.id_len))
.min(start + MAX_BLOCK_LEN)
.min(buf.len());
let list = out.entry(owner).or_default();
let mut at = start + layout.mark.len();
while at + 4 <= end {
if let Some(found) = map_at(buf, at, end, materials, layout.id_len) {
for material in found.materials {
if !list.contains(&material) {
list.push(material);
}
}
if layout.first_map_only {
break;
}
at = found.end;
} else {
at += 1;
}
}
}
out
}
struct MaterialMap {
materials: Vec<u32>,
end: usize,
}
fn map_at(
buf: &[u8],
at: usize,
end: usize,
materials: &BTreeSet<u32>,
id_len: usize,
) -> Option<MaterialMap> {
let count = u32::from_le_bytes(buf.get(at..at + 4)?.try_into().ok()?) as usize;
if !(1..=MAX_ENTRIES).contains(&count) {
return None;
}
let entry_len = 4 + id_len;
let map_end = at + 4 + count * entry_len;
if map_end > end {
return None;
}
let mut found = Vec::with_capacity(count);
for entry in 0..count {
let value_at = at + 4 + entry * entry_len + 4;
let material = id_at(buf, value_at, id_len).filter(|m| materials.contains(m))?;
found.push(material);
}
Some(MaterialMap {
materials: found,
end: map_end,
})
}
pub fn type_material_names(rf: &mut RevitFile, revit_version: u32) -> BTreeMap<u32, Vec<String>> {
if bbox_marker(revit_version).is_none() {
return BTreeMap::new();
}
let Ok(records) = crate::elem_table::parse_records(rf) else {
return BTreeMap::new();
};
let declared = crate::elem_table::declared_ids(&records);
let Ok(names) = crate::partition_materials::scan_material_names(rf, revit_version, &declared)
else {
return BTreeMap::new();
};
let materials: BTreeSet<u32> = names.keys().copied().collect();
let mut by_owner: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
for stream in rf.partition_stream_names() {
let Ok(inflated) = rf.inflated_partition(&stream) else {
continue;
};
for (owner, found) in scan_partition(inflated.bytes(), &declared, &materials) {
let list = by_owner.entry(owner).or_default();
for material in found {
if !list.contains(&material) {
list.push(material);
}
}
}
}
by_owner
.into_iter()
.filter_map(|(owner, ids)| {
let resolved: Vec<String> =
ids.iter().filter_map(|id| names.get(id).cloned()).collect();
(ids.is_empty() || !resolved.is_empty()).then_some((owner, resolved))
})
.collect()
}
pub fn type_material_names_2023(rf: &mut RevitFile) -> BTreeMap<u32, Vec<String>> {
let Ok(records) = crate::elem_table::parse_records(rf) else {
return BTreeMap::new();
};
let declared = crate::elem_table::declared_ids(&records);
let (tagged, names) = crate::partition_materials::scan_materials_2023(rf, &declared);
let streams = rf.partition_stream_names();
let mut by_owner: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
for stream in &streams {
let Ok(inflated) = rf.inflated_partition(stream) else {
continue;
};
for (owner, found) in
scan_partition_with(inflated.bytes(), &declared, &tagged, BLOCK_LAYOUT_2023)
{
let list = by_owner.entry(owner).or_default();
for material in found {
if !list.contains(&material) {
list.push(material);
}
}
}
}
by_owner
.into_iter()
.filter(|(_, ids)| !ids.is_empty())
.filter_map(|(owner, ids)| {
let resolved: Option<Vec<String>> =
ids.iter().map(|id| names.get(id).cloned()).collect();
resolved.map(|names| (owner, names))
})
.collect()
}