1use std::collections::BTreeMap;
14use std::io::{self, Read, Write};
15
16use roaring::RoaringBitmap;
17
18use crate::ordinal::{Ordinal, to_usize};
19
20const MAGIC: &[u8; 8] = b"FTATTR\0\0";
21const VERSION: u32 = 1;
22const TAG_CONCEPT: u8 = 0;
23const TAG_NUMBER: u8 = 1;
24const TAG_STRING: u8 = 2;
25
26#[derive(Debug, thiserror::Error)]
28pub enum AttributesError {
29 #[error("attribute row ({from}, {kind}) is out of range for {nodes} nodes and {types} types")]
31 OutOfRange {
32 from: u32,
34 kind: u32,
36 nodes: u32,
38 types: u32,
40 },
41 #[error("too many attribute rows")]
43 TooMany(#[source] std::num::TryFromIntError),
44 #[error("attributes I/O failed")]
46 Io(#[from] io::Error),
47 #[error("not an attributes artifact")]
49 Magic,
50 #[error("attributes layout version {found}, expected {expected}")]
52 Version {
53 found: u32,
55 expected: u32,
57 },
58 #[error("the attribute arrays are inconsistent")]
60 Inconsistent,
61 #[error("an attribute value is not UTF-8")]
63 Text(#[from] std::string::FromUtf8Error),
64}
65
66#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
68pub enum Value {
69 Concept(Ordinal),
71 Number(String),
73 String(String),
75}
76
77#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
79pub struct Edge {
80 pub source: Ordinal,
82 pub group: u32,
84 pub kind: u32,
86 pub value: Value,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum ValueRef<'a> {
93 Concept(Ordinal),
95 Number(&'a str),
97 String(&'a str),
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
103pub struct Row<'a> {
104 pub group: u32,
106 pub kind: u32,
108 pub value: ValueRef<'a>,
110}
111
112#[derive(Debug, Clone, PartialEq, Eq, Default)]
114struct Inverted {
115 targets: Vec<u32>,
117 offsets: Vec<u32>,
119 sources: Vec<u32>,
121 all_sources: RoaringBitmap,
123}
124
125impl Inverted {
126 fn build(mut pairs: Vec<(u32, u32)>, all_sources: RoaringBitmap) -> Self {
127 pairs.sort_unstable();
128 pairs.dedup();
129 let mut targets = Vec::new();
130 let mut offsets = Vec::new();
131 let mut sources = Vec::with_capacity(pairs.len());
132 for (target, source) in pairs {
133 if targets.last() != Some(&target) {
134 targets.push(target);
135 offsets.push(u32::try_from(sources.len()).unwrap_or(u32::MAX));
136 }
137 sources.push(source);
138 }
139 offsets.push(u32::try_from(sources.len()).unwrap_or(u32::MAX));
140 Self {
141 targets,
142 offsets,
143 sources,
144 all_sources,
145 }
146 }
147
148 fn sources(&self, target: u32) -> &[u32] {
149 let Ok(index) = self.targets.binary_search(&target) else {
150 return &[];
151 };
152 match (
153 self.offsets.get(index),
154 self.offsets.get(index.saturating_add(1)),
155 ) {
156 (Some(&start), Some(&end)) => self
157 .sources
158 .get(to_usize(start)..to_usize(end))
159 .unwrap_or_default(),
160 _ => &[],
161 }
162 }
163}
164
165#[derive(Debug, Clone, PartialEq, Eq, Default)]
167pub struct Attributes {
168 types: Vec<u64>,
170 offsets: Vec<u32>,
172 groups: Vec<u32>,
173 kinds: Vec<u32>,
174 tags: Vec<u8>,
175 payloads: Vec<u32>,
176 strings: Vec<String>,
177 inverted: Vec<Inverted>,
179}
180
181impl Attributes {
182 pub fn build(
190 nodes: u32,
191 types: Vec<u64>,
192 mut edges: Vec<Edge>,
193 ) -> Result<Self, AttributesError> {
194 let type_count = u32::try_from(types.len()).map_err(AttributesError::TooMany)?;
195 Self::check_range(nodes, type_count, &edges)?;
196 edges.sort_unstable();
197 edges.dedup();
198 u32::try_from(edges.len()).map_err(AttributesError::TooMany)?;
199 let mut attributes = Self::pack(nodes, types, &edges)?;
200 attributes.derive();
201 Ok(attributes)
202 }
203
204 fn check_range(nodes: u32, type_count: u32, edges: &[Edge]) -> Result<(), AttributesError> {
207 for edge in edges {
208 let target_ok = match edge.value {
209 Value::Concept(target) => target.index() < nodes,
210 Value::Number(_) | Value::String(_) => true,
211 };
212 if edge.source.index() >= nodes || edge.kind >= type_count || !target_ok {
213 return Err(AttributesError::OutOfRange {
214 from: edge.source.index(),
215 kind: edge.kind,
216 nodes,
217 types: type_count,
218 });
219 }
220 }
221 Ok(())
222 }
223
224 fn pack(nodes: u32, types: Vec<u64>, edges: &[Edge]) -> Result<Self, AttributesError> {
227 let mut interned: BTreeMap<String, u32> = BTreeMap::new();
228 let mut strings = Vec::new();
229 let mut intern = |text: &str| -> Result<u32, AttributesError> {
230 if let Some(&index) = interned.get(text) {
231 return Ok(index);
232 }
233 let index = u32::try_from(strings.len()).map_err(AttributesError::TooMany)?;
234 strings.push(text.to_owned());
235 interned.insert(text.to_owned(), index);
236 Ok(index)
237 };
238 let mut offsets = Vec::with_capacity(to_usize(nodes).saturating_add(1));
239 let (mut groups, mut kinds, mut tags, mut payloads) =
240 (Vec::new(), Vec::new(), Vec::new(), Vec::new());
241 let mut cursor = 0_usize;
242 for node in 0..nodes {
243 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
244 while let Some(edge) = edges.get(cursor) {
245 if edge.source.index() != node {
246 break;
247 }
248 groups.push(edge.group);
249 kinds.push(edge.kind);
250 let (tag, payload) = match &edge.value {
251 Value::Concept(target) => (TAG_CONCEPT, target.index()),
252 Value::Number(text) => (TAG_NUMBER, intern(text)?),
253 Value::String(text) => (TAG_STRING, intern(text)?),
254 };
255 tags.push(tag);
256 payloads.push(payload);
257 cursor = cursor.saturating_add(1);
258 }
259 }
260 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
261 Ok(Self {
262 types,
263 offsets,
264 groups,
265 kinds,
266 tags,
267 payloads,
268 strings,
269 inverted: Vec::new(),
270 })
271 }
272
273 fn derive(&mut self) {
275 let mut pairs: Vec<Vec<(u32, u32)>> = vec![Vec::new(); self.types.len()];
276 let mut all: Vec<RoaringBitmap> = vec![RoaringBitmap::new(); self.types.len()];
277 for node in 0..self.nodes() {
278 for row in self.rows(Ordinal::new(node)) {
279 let kind = to_usize(row.kind);
280 if let Some(set) = all.get_mut(kind) {
281 set.insert(node);
282 }
283 if let (ValueRef::Concept(target), Some(list)) = (row.value, pairs.get_mut(kind)) {
284 list.push((target.index(), node));
285 }
286 }
287 }
288 self.inverted = pairs
289 .into_iter()
290 .zip(all)
291 .map(|(pairs, all_sources)| Inverted::build(pairs, all_sources))
292 .collect();
293 }
294
295 #[must_use]
297 pub fn types(&self) -> &[u64] {
298 &self.types
299 }
300
301 #[must_use]
303 pub fn kind(&self, sctid: u64) -> Option<u32> {
304 self.types
305 .iter()
306 .position(|&t| t == sctid)
307 .and_then(|i| u32::try_from(i).ok())
308 }
309
310 #[must_use]
312 pub fn nodes(&self) -> u32 {
313 u32::try_from(self.offsets.len().saturating_sub(1)).unwrap_or(u32::MAX)
314 }
315
316 #[must_use]
318 pub fn edges(&self) -> usize {
319 self.groups.len()
320 }
321
322 pub fn rows(&self, source: Ordinal) -> impl Iterator<Item = Row<'_>> + '_ {
324 let index = to_usize(source.index());
325 let (start, end) = match (
326 self.offsets.get(index),
327 self.offsets.get(index.saturating_add(1)),
328 ) {
329 (Some(&s), Some(&e)) => (to_usize(s), to_usize(e)),
330 _ => (0, 0),
331 };
332 (start..end).filter_map(move |i| {
333 Some(Row {
334 group: *self.groups.get(i)?,
335 kind: *self.kinds.get(i)?,
336 value: match (*self.tags.get(i)?, *self.payloads.get(i)?) {
337 (TAG_CONCEPT, target) => ValueRef::Concept(Ordinal::new(target)),
338 (TAG_NUMBER, text) => ValueRef::Number(self.strings.get(to_usize(text))?),
339 (_, text) => ValueRef::String(self.strings.get(to_usize(text))?),
340 },
341 })
342 })
343 }
344
345 #[must_use]
347 pub fn sources(&self, kind: u32, target: Ordinal) -> &[u32] {
348 self.inverted
349 .get(to_usize(kind))
350 .map_or(&[], |inverted| inverted.sources(target.index()))
351 }
352
353 #[must_use]
355 pub fn sources_of_kind(&self, kind: u32) -> Option<&RoaringBitmap> {
356 self.inverted.get(to_usize(kind)).map(|i| &i.all_sources)
357 }
358
359 #[must_use]
361 pub fn targets_of_kind(&self, kind: u32) -> &[u32] {
362 self.inverted
363 .get(to_usize(kind))
364 .map_or(&[], |i| i.targets.as_slice())
365 }
366
367 pub fn write_to(&self, out: &mut impl Write) -> Result<(), AttributesError> {
373 let count = |len: usize| u32::try_from(len).map_err(AttributesError::TooMany);
374 out.write_all(MAGIC)?;
375 out.write_all(&VERSION.to_le_bytes())?;
376 out.write_all(&count(self.types.len())?.to_le_bytes())?;
377 for sctid in &self.types {
378 out.write_all(&sctid.to_le_bytes())?;
379 }
380 out.write_all(&self.nodes().to_le_bytes())?;
381 write_u32s(out, &self.offsets)?;
382 out.write_all(&count(self.groups.len())?.to_le_bytes())?;
383 write_u32s(out, &self.groups)?;
384 write_u32s(out, &self.kinds)?;
385 out.write_all(&self.tags)?;
386 write_u32s(out, &self.payloads)?;
387 out.write_all(&count(self.strings.len())?.to_le_bytes())?;
388 for text in &self.strings {
389 out.write_all(&count(text.len())?.to_le_bytes())?;
390 out.write_all(text.as_bytes())?;
391 }
392 Ok(())
393 }
394
395 pub fn read_from(input: &mut impl Read) -> Result<Self, AttributesError> {
402 let mut magic = [0_u8; 8];
403 input.read_exact(&mut magic)?;
404 if &magic != MAGIC {
405 return Err(AttributesError::Magic);
406 }
407 let version = read_u32(input)?;
408 if version != VERSION {
409 return Err(AttributesError::Version {
410 found: version,
411 expected: VERSION,
412 });
413 }
414 let type_count = read_u32(input)?;
415 let mut types = Vec::with_capacity(to_usize(type_count));
416 for _ in 0..type_count {
417 let mut long = [0_u8; 8];
418 input.read_exact(&mut long)?;
419 types.push(u64::from_le_bytes(long));
420 }
421 let nodes = read_u32(input)?;
422 let offsets = read_u32s(input, to_usize(nodes).saturating_add(1))?;
423 let rows = to_usize(read_u32(input)?);
424 let groups = read_u32s(input, rows)?;
425 let kinds = read_u32s(input, rows)?;
426 let mut tags = vec![0_u8; rows];
427 input.read_exact(&mut tags)?;
428 let payloads = read_u32s(input, rows)?;
429 let string_count = read_u32(input)?;
430 let mut strings = Vec::with_capacity(to_usize(string_count));
431 for _ in 0..string_count {
432 let len = to_usize(read_u32(input)?);
433 let mut bytes = vec![0_u8; len];
434 input.read_exact(&mut bytes)?;
435 strings.push(String::from_utf8(bytes)?);
436 }
437 let consistent = offsets.last().is_some_and(|&l| to_usize(l) == rows)
438 && offsets.windows(2).all(|w| w.first() <= w.get(1))
439 && kinds.iter().all(|&k| to_usize(k) < types.len())
440 && tags
441 .iter()
442 .zip(&payloads)
443 .all(|(&tag, &payload)| match tag {
444 TAG_CONCEPT => payload < nodes,
445 TAG_NUMBER | TAG_STRING => to_usize(payload) < strings.len(),
446 _ => false,
447 });
448 if !consistent {
449 return Err(AttributesError::Inconsistent);
450 }
451 let mut attributes = Self {
452 types,
453 offsets,
454 groups,
455 kinds,
456 tags,
457 payloads,
458 strings,
459 inverted: Vec::new(),
460 };
461 attributes.derive();
462 Ok(attributes)
463 }
464}
465
466fn write_u32s(out: &mut impl Write, values: &[u32]) -> Result<(), AttributesError> {
467 for value in values {
468 out.write_all(&value.to_le_bytes())?;
469 }
470 Ok(())
471}
472
473fn read_u32(input: &mut impl Read) -> Result<u32, AttributesError> {
474 let mut bytes = [0_u8; 4];
475 input.read_exact(&mut bytes)?;
476 Ok(u32::from_le_bytes(bytes))
477}
478
479fn read_u32s(input: &mut impl Read, count: usize) -> Result<Vec<u32>, AttributesError> {
480 let mut bytes = vec![0_u8; count.saturating_mul(4)];
481 input.read_exact(&mut bytes)?;
482 Ok(bytes
483 .as_chunks::<4>()
484 .0
485 .iter()
486 .map(|chunk| u32::from_le_bytes(*chunk))
487 .collect())
488}
489
490#[cfg(test)]
491mod tests {
492 use super::{Attributes, AttributesError, Edge, Value, ValueRef};
493 use crate::ordinal::Ordinal;
494
495 fn sample() -> Attributes {
496 let edge = |source: u32, group: u32, kind: u32, value: Value| Edge {
497 source: Ordinal::new(source),
498 group,
499 kind,
500 value,
501 };
502 Attributes::build(
503 5,
504 vec![100, 200],
505 vec![
506 edge(1, 1, 0, Value::Concept(Ordinal::new(3))),
507 edge(1, 1, 1, Value::Number(String::from("4"))),
508 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
509 edge(2, 2, 1, Value::Number(String::from("4"))),
510 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
511 edge(4, 0, 1, Value::String(String::from("blue"))),
512 ],
513 )
514 .expect("builds")
515 }
516
517 #[test]
518 fn rows_are_grouped_and_the_inverted_index_answers_by_type_and_target() {
519 let attributes = sample();
520 assert_eq!(attributes.edges(), 5, "the duplicate row is dropped");
521 let cat: Vec<_> = attributes.rows(Ordinal::new(1)).collect();
522 assert_eq!(cat.len(), 2);
523 assert_eq!(cat[0].group, 1);
524 assert_eq!(cat[0].value, ValueRef::Concept(Ordinal::new(3)));
525 assert_eq!(cat[1].value, ValueRef::Number("4"));
526 assert_eq!(attributes.sources(0, Ordinal::new(3)), [1, 2]);
527 assert!(attributes.sources(1, Ordinal::new(3)).is_empty());
528 assert_eq!(attributes.targets_of_kind(0), [3]);
529 assert_eq!(
530 attributes
531 .sources_of_kind(1)
532 .expect("kind")
533 .iter()
534 .collect::<Vec<_>>(),
535 [1, 2, 4]
536 );
537 assert_eq!(attributes.kind(200), Some(1));
538 assert_eq!(attributes.kind(300), None);
539 assert_eq!(
540 attributes.rows(Ordinal::new(4)).next().expect("row").value,
541 ValueRef::String("blue")
542 );
543 assert!(attributes.rows(Ordinal::new(0)).next().is_none());
544 }
545
546 #[test]
547 fn the_layout_round_trips_and_refuses_bad_input() {
548 let attributes = sample();
549 let mut bytes = Vec::new();
550 attributes.write_to(&mut bytes).expect("writes");
551 let again = Attributes::read_from(&mut bytes.as_slice()).expect("reads");
552 assert_eq!(again, attributes);
553 assert!(matches!(
554 Attributes::read_from(&mut b"nope".as_slice()),
555 Err(AttributesError::Io(_))
556 ));
557 assert!(matches!(
558 Attributes::build(
559 2,
560 vec![1],
561 vec![Edge {
562 source: Ordinal::new(1),
563 group: 0,
564 kind: 1,
565 value: Value::Concept(Ordinal::new(0)),
566 }]
567 ),
568 Err(AttributesError::OutOfRange { kind: 1, .. })
569 ));
570 let mut wrong = bytes.clone();
571 wrong[0] = b'X';
572 assert!(matches!(
573 Attributes::read_from(&mut wrong.as_slice()),
574 Err(AttributesError::Magic)
575 ));
576 }
577}