use std::collections::BTreeMap;
use std::io::{self, Read, Write};
use roaring::RoaringBitmap;
use crate::ordinal::{Ordinal, to_usize};
const MAGIC: &[u8; 8] = b"FTMBRS\0\0";
const VERSION: u32 = 1;
const TAG_CONCEPT: u8 = 0;
const TAG_COMPONENT: u8 = 1;
const TAG_INTEGER: u8 = 2;
const TAG_STRING: u8 = 3;
#[derive(Debug, thiserror::Error)]
pub enum RefsetsError {
#[error("reference set {refset} row has {values} values for {fields} fields")]
Arity {
refset: u64,
values: usize,
fields: usize,
},
#[error("too many reference set rows")]
TooMany,
#[error("reference set members I/O failed")]
Io(#[from] io::Error),
#[error("not a reference set members artifact")]
Magic,
#[error("reference set members layout version {found}, expected {expected}")]
Version {
found: u32,
expected: u32,
},
#[error("the reference set member arrays are inconsistent")]
Inconsistent,
#[error("a reference set field is not UTF-8")]
Text(#[from] std::string::FromUtf8Error),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FieldKind {
Component,
Integer,
String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FieldValue {
Concept(Ordinal),
Component(u64),
Integer(i64),
String(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ValueRef<'a> {
Concept(Ordinal),
Component(u64),
Integer(i64),
String(&'a str),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemberRow {
pub concept: Ordinal,
pub effective_time: u32,
pub module: u64,
pub values: Vec<FieldValue>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct Table {
fields: Vec<String>,
kinds: Vec<FieldKind>,
concepts: Vec<u32>,
times: Vec<u32>,
modules: Vec<u64>,
tags: Vec<u8>,
payloads: Vec<u32>,
longs: Vec<u64>,
strings: Vec<String>,
members: RoaringBitmap,
}
impl Table {
#[must_use]
pub fn fields(&self) -> &[String] {
&self.fields
}
#[must_use]
pub fn kinds(&self) -> &[FieldKind] {
&self.kinds
}
#[must_use]
pub fn field(&self, name: &str) -> Option<usize> {
self.fields
.iter()
.position(|f| f.eq_ignore_ascii_case(name))
}
#[must_use]
pub fn len(&self) -> usize {
self.concepts.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.concepts.is_empty()
}
#[must_use]
pub fn members(&self) -> &RoaringBitmap {
&self.members
}
#[must_use]
pub fn concept(&self, row: usize) -> Option<Ordinal> {
self.concepts.get(row).map(|&c| Ordinal::new(c))
}
#[must_use]
pub fn effective_time(&self, row: usize) -> Option<u32> {
self.times.get(row).copied()
}
#[must_use]
pub fn module(&self, row: usize) -> Option<u64> {
self.modules.get(row).copied()
}
#[must_use]
pub fn value(&self, row: usize, field: usize) -> Option<ValueRef<'_>> {
if field >= self.fields.len() {
return None;
}
let index = row.checked_mul(self.fields.len())?.checked_add(field)?;
let payload = to_usize(*self.payloads.get(index)?);
Some(match *self.tags.get(index)? {
TAG_CONCEPT => ValueRef::Concept(Ordinal::new(*self.payloads.get(index)?)),
TAG_COMPONENT => ValueRef::Component(*self.longs.get(payload)?),
TAG_INTEGER => {
ValueRef::Integer(i64::from_le_bytes(self.longs.get(payload)?.to_le_bytes()))
}
_ => ValueRef::String(self.strings.get(payload)?),
})
}
pub fn rows_with(&self, field: usize, target: Ordinal) -> impl Iterator<Item = usize> + '_ {
(0..self.len())
.filter(move |&row| self.value(row, field) == Some(ValueRef::Concept(target)))
}
fn check(&self) -> Result<(), RefsetsError> {
let rows = self.concepts.len();
let cells = rows.saturating_mul(self.fields.len());
let consistent = self.kinds.len() == self.fields.len()
&& self.times.len() == rows
&& self.modules.len() == rows
&& self.tags.len() == cells
&& self.payloads.len() == cells
&& self
.tags
.iter()
.zip(&self.payloads)
.all(|(&tag, &payload)| match tag {
TAG_CONCEPT => true,
TAG_COMPONENT | TAG_INTEGER => to_usize(payload) < self.longs.len(),
TAG_STRING => to_usize(payload) < self.strings.len(),
_ => false,
});
consistent.then_some(()).ok_or(RefsetsError::Inconsistent)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct RefsetMembers {
tables: BTreeMap<u64, Table>,
}
impl RefsetMembers {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn insert(
&mut self,
refset: u64,
fields: &[(String, FieldKind)],
rows: Vec<MemberRow>,
) -> Result<(), RefsetsError> {
let mut table = Table {
fields: fields.iter().map(|(name, _)| name.clone()).collect(),
kinds: fields.iter().map(|(_, kind)| *kind).collect(),
..Table::default()
};
let mut interned: BTreeMap<String, u32> = BTreeMap::new();
let mut longs_seen: BTreeMap<u64, u32> = BTreeMap::new();
for row in rows {
if row.values.len() != table.fields.len() {
return Err(RefsetsError::Arity {
refset,
values: row.values.len(),
fields: table.fields.len(),
});
}
table.concepts.push(row.concept.index());
table.members.insert(row.concept.index());
table.times.push(row.effective_time);
table.modules.push(row.module);
for value in row.values {
let (tag, payload) = match value {
FieldValue::Concept(concept) => (TAG_CONCEPT, concept.index()),
FieldValue::Component(id) => (
TAG_COMPONENT,
intern_long(&mut table.longs, &mut longs_seen, id)?,
),
FieldValue::Integer(value) => (
TAG_INTEGER,
intern_long(
&mut table.longs,
&mut longs_seen,
u64::from_le_bytes(value.to_le_bytes()),
)?,
),
FieldValue::String(text) => {
let index = if let Some(&index) = interned.get(&text) {
index
} else {
let index = u32::try_from(table.strings.len())
.map_err(|_| RefsetsError::TooMany)?;
table.strings.push(text.clone());
interned.insert(text, index);
index
};
(TAG_STRING, index)
}
};
table.tags.push(tag);
table.payloads.push(payload);
}
}
u32::try_from(table.concepts.len()).map_err(|_| RefsetsError::TooMany)?;
self.tables.insert(refset, table);
Ok(())
}
#[must_use]
pub fn table(&self, refset: u64) -> Option<&Table> {
self.tables.get(&refset)
}
pub fn refsets(&self) -> impl Iterator<Item = u64> + '_ {
self.tables.keys().copied()
}
#[must_use]
pub fn len(&self) -> usize {
self.tables.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.tables.is_empty()
}
#[must_use]
pub fn total(&self) -> u64 {
self.tables
.values()
.map(|t| u64::try_from(t.len()).unwrap_or(u64::MAX))
.sum()
}
pub fn write_to(&self, out: &mut impl Write) -> Result<(), RefsetsError> {
let count = |len: usize| u32::try_from(len).map_err(|_| RefsetsError::TooMany);
out.write_all(MAGIC)?;
out.write_all(&VERSION.to_le_bytes())?;
out.write_all(&count(self.tables.len())?.to_le_bytes())?;
for (refset, table) in &self.tables {
out.write_all(&refset.to_le_bytes())?;
out.write_all(&count(table.fields.len())?.to_le_bytes())?;
for (name, kind) in table.fields.iter().zip(&table.kinds) {
write_text(out, name)?;
out.write_all(&[match kind {
FieldKind::Component => 0,
FieldKind::Integer => 1,
FieldKind::String => 2,
}])?;
}
out.write_all(&count(table.concepts.len())?.to_le_bytes())?;
write_u32s(out, &table.concepts)?;
write_u32s(out, &table.times)?;
for module in &table.modules {
out.write_all(&module.to_le_bytes())?;
}
out.write_all(&table.tags)?;
write_u32s(out, &table.payloads)?;
out.write_all(&count(table.longs.len())?.to_le_bytes())?;
for long in &table.longs {
out.write_all(&long.to_le_bytes())?;
}
out.write_all(&count(table.strings.len())?.to_le_bytes())?;
for text in &table.strings {
write_text(out, text)?;
}
}
Ok(())
}
pub fn read_from(input: &mut impl Read) -> Result<Self, RefsetsError> {
let mut magic = [0_u8; 8];
input.read_exact(&mut magic)?;
if &magic != MAGIC {
return Err(RefsetsError::Magic);
}
let version = read_u32(input)?;
if version != VERSION {
return Err(RefsetsError::Version {
found: version,
expected: VERSION,
});
}
let mut tables = BTreeMap::new();
for _ in 0..read_u32(input)? {
let refset = read_u64(input)?;
let field_count = to_usize(read_u32(input)?);
let mut fields = Vec::with_capacity(field_count);
let mut kinds = Vec::with_capacity(field_count);
for _ in 0..field_count {
fields.push(read_text(input)?);
let mut kind = [0_u8; 1];
input.read_exact(&mut kind)?;
kinds.push(match kind[0] {
0 => FieldKind::Component,
1 => FieldKind::Integer,
2 => FieldKind::String,
_ => return Err(RefsetsError::Inconsistent),
});
}
let rows = to_usize(read_u32(input)?);
let concepts = read_u32s(input, rows)?;
let times = read_u32s(input, rows)?;
let mut modules = Vec::with_capacity(rows);
for _ in 0..rows {
modules.push(read_u64(input)?);
}
let cells = rows.saturating_mul(field_count);
let mut tags = vec![0_u8; cells];
input.read_exact(&mut tags)?;
let payloads = read_u32s(input, cells)?;
let mut longs = Vec::new();
for _ in 0..read_u32(input)? {
longs.push(read_u64(input)?);
}
let mut strings = Vec::new();
for _ in 0..read_u32(input)? {
strings.push(read_text(input)?);
}
let members = concepts.iter().copied().collect();
let table = Table {
fields,
kinds,
concepts,
times,
modules,
tags,
payloads,
longs,
strings,
members,
};
table.check()?;
tables.insert(refset, table);
}
Ok(Self { tables })
}
}
fn intern_long(
longs: &mut Vec<u64>,
seen: &mut BTreeMap<u64, u32>,
value: u64,
) -> Result<u32, RefsetsError> {
if let Some(&index) = seen.get(&value) {
return Ok(index);
}
let index = u32::try_from(longs.len()).map_err(|_| RefsetsError::TooMany)?;
longs.push(value);
seen.insert(value, index);
Ok(index)
}
fn write_text(out: &mut impl Write, text: &str) -> Result<(), RefsetsError> {
let len = u32::try_from(text.len()).map_err(|_| RefsetsError::TooMany)?;
out.write_all(&len.to_le_bytes())?;
out.write_all(text.as_bytes())?;
Ok(())
}
fn read_text(input: &mut impl Read) -> Result<String, RefsetsError> {
let len = to_usize(read_u32(input)?);
let mut bytes = vec![0_u8; len];
input.read_exact(&mut bytes)?;
Ok(String::from_utf8(bytes)?)
}
fn write_u32s(out: &mut impl Write, values: &[u32]) -> Result<(), RefsetsError> {
for value in values {
out.write_all(&value.to_le_bytes())?;
}
Ok(())
}
fn read_u32(input: &mut impl Read) -> Result<u32, RefsetsError> {
let mut bytes = [0_u8; 4];
input.read_exact(&mut bytes)?;
Ok(u32::from_le_bytes(bytes))
}
fn read_u64(input: &mut impl Read) -> Result<u64, RefsetsError> {
let mut bytes = [0_u8; 8];
input.read_exact(&mut bytes)?;
Ok(u64::from_le_bytes(bytes))
}
fn read_u32s(input: &mut impl Read, count: usize) -> Result<Vec<u32>, RefsetsError> {
let mut bytes = vec![0_u8; count.saturating_mul(4)];
input.read_exact(&mut bytes)?;
Ok(bytes
.as_chunks::<4>()
.0
.iter()
.map(|chunk| u32::from_le_bytes(*chunk))
.collect())
}
#[cfg(test)]
mod tests {
use super::{FieldKind, FieldValue, MemberRow, RefsetMembers, RefsetsError, ValueRef};
use crate::ordinal::Ordinal;
fn sample() -> RefsetMembers {
let mut members = RefsetMembers::new();
members
.insert(
42,
&[
(String::from("mapGroup"), FieldKind::Integer),
(String::from("mapTarget"), FieldKind::String),
(String::from("correlationId"), FieldKind::Component),
],
vec![
MemberRow {
concept: Ordinal::new(2),
effective_time: 20_240_101,
module: 99,
values: vec![
FieldValue::Integer(1),
FieldValue::String(String::from("J45.9")),
FieldValue::Concept(Ordinal::new(7)),
],
},
MemberRow {
concept: Ordinal::new(3),
effective_time: 20_230_731,
module: 99,
values: vec![
FieldValue::Integer(-2),
FieldValue::String(String::from("J45.9")),
FieldValue::Component(123_456_789_012),
],
},
],
)
.expect("inserts");
members
.insert(7, &[], vec![])
.expect("an empty simple reference set");
members
}
#[test]
fn rows_answer_by_field_and_the_layout_round_trips() {
let members = sample();
let table = members.table(42).expect("table");
assert_eq!(table.len(), 2);
assert_eq!(table.field("MAPTARGET"), Some(1));
assert_eq!(table.value(0, 0), Some(ValueRef::Integer(1)));
assert_eq!(table.value(1, 0), Some(ValueRef::Integer(-2)));
assert_eq!(table.value(1, 1), Some(ValueRef::String("J45.9")));
assert_eq!(table.value(0, 2), Some(ValueRef::Concept(Ordinal::new(7))));
assert_eq!(
table.value(1, 2),
Some(ValueRef::Component(123_456_789_012))
);
assert_eq!(table.value(1, 3), None);
assert_eq!(table.effective_time(1), Some(20_230_731));
assert_eq!(table.members().iter().collect::<Vec<_>>(), [2, 3]);
assert_eq!(table.rows_with(2, Ordinal::new(7)).collect::<Vec<_>>(), [0]);
assert_eq!(members.total(), 2);
let mut bytes = Vec::new();
members.write_to(&mut bytes).expect("writes");
assert_eq!(
RefsetMembers::read_from(&mut bytes.as_slice()).expect("reads"),
members
);
assert!(matches!(
RefsetMembers::read_from(&mut b"XXXXXXXX\0\0\0\0".as_slice()),
Err(RefsetsError::Magic)
));
let mut bad = RefsetMembers::new();
assert!(matches!(
bad.insert(
1,
&[(String::from("f"), FieldKind::String)],
vec![MemberRow {
concept: Ordinal::new(0),
effective_time: 0,
module: 0,
values: Vec::new(),
}]
),
Err(RefsetsError::Arity { .. })
));
}
}