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,
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 for edge in &edges {
196 let target_ok = match edge.value {
197 Value::Concept(target) => target.index() < nodes,
198 Value::Number(_) | Value::String(_) => true,
199 };
200 if edge.source.index() >= nodes || edge.kind >= type_count || !target_ok {
201 return Err(AttributesError::OutOfRange {
202 from: edge.source.index(),
203 kind: edge.kind,
204 nodes,
205 types: type_count,
206 });
207 }
208 }
209 edges.sort_unstable();
210 edges.dedup();
211 u32::try_from(edges.len()).map_err(|_| AttributesError::TooMany)?;
212 let mut interned: BTreeMap<String, u32> = BTreeMap::new();
213 let mut strings = Vec::new();
214 let mut intern = |text: &str| -> Result<u32, AttributesError> {
215 if let Some(&index) = interned.get(text) {
216 return Ok(index);
217 }
218 let index = u32::try_from(strings.len()).map_err(|_| AttributesError::TooMany)?;
219 strings.push(text.to_owned());
220 interned.insert(text.to_owned(), index);
221 Ok(index)
222 };
223 let mut offsets = Vec::with_capacity(to_usize(nodes).saturating_add(1));
224 let (mut groups, mut kinds, mut tags, mut payloads) =
225 (Vec::new(), Vec::new(), Vec::new(), Vec::new());
226 let mut cursor = 0_usize;
227 for node in 0..nodes {
228 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
229 while let Some(edge) = edges.get(cursor) {
230 if edge.source.index() != node {
231 break;
232 }
233 groups.push(edge.group);
234 kinds.push(edge.kind);
235 let (tag, payload) = match &edge.value {
236 Value::Concept(target) => (TAG_CONCEPT, target.index()),
237 Value::Number(text) => (TAG_NUMBER, intern(text)?),
238 Value::String(text) => (TAG_STRING, intern(text)?),
239 };
240 tags.push(tag);
241 payloads.push(payload);
242 cursor = cursor.saturating_add(1);
243 }
244 }
245 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
246 let mut attributes = Self {
247 types,
248 offsets,
249 groups,
250 kinds,
251 tags,
252 payloads,
253 strings,
254 inverted: Vec::new(),
255 };
256 attributes.derive();
257 Ok(attributes)
258 }
259
260 fn derive(&mut self) {
262 let mut pairs: Vec<Vec<(u32, u32)>> = vec![Vec::new(); self.types.len()];
263 let mut all: Vec<RoaringBitmap> = vec![RoaringBitmap::new(); self.types.len()];
264 for node in 0..self.nodes() {
265 for row in self.rows(Ordinal::new(node)) {
266 let kind = to_usize(row.kind);
267 if let Some(set) = all.get_mut(kind) {
268 set.insert(node);
269 }
270 if let (ValueRef::Concept(target), Some(list)) = (row.value, pairs.get_mut(kind)) {
271 list.push((target.index(), node));
272 }
273 }
274 }
275 self.inverted = pairs
276 .into_iter()
277 .zip(all)
278 .map(|(pairs, all_sources)| Inverted::build(pairs, all_sources))
279 .collect();
280 }
281
282 #[must_use]
284 pub fn types(&self) -> &[u64] {
285 &self.types
286 }
287
288 #[must_use]
290 pub fn kind(&self, sctid: u64) -> Option<u32> {
291 self.types
292 .iter()
293 .position(|&t| t == sctid)
294 .and_then(|i| u32::try_from(i).ok())
295 }
296
297 #[must_use]
299 pub fn nodes(&self) -> u32 {
300 u32::try_from(self.offsets.len().saturating_sub(1)).unwrap_or(u32::MAX)
301 }
302
303 #[must_use]
305 pub fn edges(&self) -> usize {
306 self.groups.len()
307 }
308
309 pub fn rows(&self, source: Ordinal) -> impl Iterator<Item = Row<'_>> + '_ {
311 let index = to_usize(source.index());
312 let (start, end) = match (
313 self.offsets.get(index),
314 self.offsets.get(index.saturating_add(1)),
315 ) {
316 (Some(&s), Some(&e)) => (to_usize(s), to_usize(e)),
317 _ => (0, 0),
318 };
319 (start..end).filter_map(move |i| {
320 Some(Row {
321 group: *self.groups.get(i)?,
322 kind: *self.kinds.get(i)?,
323 value: match (*self.tags.get(i)?, *self.payloads.get(i)?) {
324 (TAG_CONCEPT, target) => ValueRef::Concept(Ordinal::new(target)),
325 (TAG_NUMBER, text) => ValueRef::Number(self.strings.get(to_usize(text))?),
326 (_, text) => ValueRef::String(self.strings.get(to_usize(text))?),
327 },
328 })
329 })
330 }
331
332 #[must_use]
334 pub fn sources(&self, kind: u32, target: Ordinal) -> &[u32] {
335 self.inverted
336 .get(to_usize(kind))
337 .map_or(&[], |inverted| inverted.sources(target.index()))
338 }
339
340 #[must_use]
342 pub fn sources_of_kind(&self, kind: u32) -> Option<&RoaringBitmap> {
343 self.inverted.get(to_usize(kind)).map(|i| &i.all_sources)
344 }
345
346 #[must_use]
348 pub fn targets_of_kind(&self, kind: u32) -> &[u32] {
349 self.inverted
350 .get(to_usize(kind))
351 .map_or(&[], |i| i.targets.as_slice())
352 }
353
354 pub fn write_to(&self, out: &mut impl Write) -> Result<(), AttributesError> {
360 let count = |len: usize| u32::try_from(len).map_err(|_| AttributesError::TooMany);
361 out.write_all(MAGIC)?;
362 out.write_all(&VERSION.to_le_bytes())?;
363 out.write_all(&count(self.types.len())?.to_le_bytes())?;
364 for sctid in &self.types {
365 out.write_all(&sctid.to_le_bytes())?;
366 }
367 out.write_all(&self.nodes().to_le_bytes())?;
368 write_u32s(out, &self.offsets)?;
369 out.write_all(&count(self.groups.len())?.to_le_bytes())?;
370 write_u32s(out, &self.groups)?;
371 write_u32s(out, &self.kinds)?;
372 out.write_all(&self.tags)?;
373 write_u32s(out, &self.payloads)?;
374 out.write_all(&count(self.strings.len())?.to_le_bytes())?;
375 for text in &self.strings {
376 out.write_all(&count(text.len())?.to_le_bytes())?;
377 out.write_all(text.as_bytes())?;
378 }
379 Ok(())
380 }
381
382 pub fn read_from(input: &mut impl Read) -> Result<Self, AttributesError> {
389 let mut magic = [0_u8; 8];
390 input.read_exact(&mut magic)?;
391 if &magic != MAGIC {
392 return Err(AttributesError::Magic);
393 }
394 let version = read_u32(input)?;
395 if version != VERSION {
396 return Err(AttributesError::Version {
397 found: version,
398 expected: VERSION,
399 });
400 }
401 let type_count = read_u32(input)?;
402 let mut types = Vec::with_capacity(to_usize(type_count));
403 for _ in 0..type_count {
404 let mut long = [0_u8; 8];
405 input.read_exact(&mut long)?;
406 types.push(u64::from_le_bytes(long));
407 }
408 let nodes = read_u32(input)?;
409 let offsets = read_u32s(input, to_usize(nodes).saturating_add(1))?;
410 let rows = to_usize(read_u32(input)?);
411 let groups = read_u32s(input, rows)?;
412 let kinds = read_u32s(input, rows)?;
413 let mut tags = vec![0_u8; rows];
414 input.read_exact(&mut tags)?;
415 let payloads = read_u32s(input, rows)?;
416 let string_count = read_u32(input)?;
417 let mut strings = Vec::with_capacity(to_usize(string_count));
418 for _ in 0..string_count {
419 let len = to_usize(read_u32(input)?);
420 let mut bytes = vec![0_u8; len];
421 input.read_exact(&mut bytes)?;
422 strings.push(String::from_utf8(bytes)?);
423 }
424 let consistent = offsets.last().is_some_and(|&l| to_usize(l) == rows)
425 && offsets.windows(2).all(|w| w.first() <= w.get(1))
426 && kinds.iter().all(|&k| to_usize(k) < types.len())
427 && tags
428 .iter()
429 .zip(&payloads)
430 .all(|(&tag, &payload)| match tag {
431 TAG_CONCEPT => payload < nodes,
432 TAG_NUMBER | TAG_STRING => to_usize(payload) < strings.len(),
433 _ => false,
434 });
435 if !consistent {
436 return Err(AttributesError::Inconsistent);
437 }
438 let mut attributes = Self {
439 types,
440 offsets,
441 groups,
442 kinds,
443 tags,
444 payloads,
445 strings,
446 inverted: Vec::new(),
447 };
448 attributes.derive();
449 Ok(attributes)
450 }
451}
452
453fn write_u32s(out: &mut impl Write, values: &[u32]) -> Result<(), AttributesError> {
454 for value in values {
455 out.write_all(&value.to_le_bytes())?;
456 }
457 Ok(())
458}
459
460fn read_u32(input: &mut impl Read) -> Result<u32, AttributesError> {
461 let mut bytes = [0_u8; 4];
462 input.read_exact(&mut bytes)?;
463 Ok(u32::from_le_bytes(bytes))
464}
465
466fn read_u32s(input: &mut impl Read, count: usize) -> Result<Vec<u32>, AttributesError> {
467 let mut bytes = vec![0_u8; count.saturating_mul(4)];
468 input.read_exact(&mut bytes)?;
469 Ok(bytes
470 .as_chunks::<4>()
471 .0
472 .iter()
473 .map(|chunk| u32::from_le_bytes(*chunk))
474 .collect())
475}
476
477#[cfg(test)]
478mod tests {
479 use super::{Attributes, AttributesError, Edge, Value, ValueRef};
480 use crate::ordinal::Ordinal;
481
482 fn sample() -> Attributes {
483 let edge = |source: u32, group: u32, kind: u32, value: Value| Edge {
484 source: Ordinal::new(source),
485 group,
486 kind,
487 value,
488 };
489 Attributes::build(
490 5,
491 vec![100, 200],
492 vec![
493 edge(1, 1, 0, Value::Concept(Ordinal::new(3))),
494 edge(1, 1, 1, Value::Number(String::from("4"))),
495 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
496 edge(2, 2, 1, Value::Number(String::from("4"))),
497 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
498 edge(4, 0, 1, Value::String(String::from("blue"))),
499 ],
500 )
501 .expect("builds")
502 }
503
504 #[test]
505 fn rows_are_grouped_and_the_inverted_index_answers_by_type_and_target() {
506 let attributes = sample();
507 assert_eq!(attributes.edges(), 5, "the duplicate row is dropped");
508 let cat: Vec<_> = attributes.rows(Ordinal::new(1)).collect();
509 assert_eq!(cat.len(), 2);
510 assert_eq!(cat[0].group, 1);
511 assert_eq!(cat[0].value, ValueRef::Concept(Ordinal::new(3)));
512 assert_eq!(cat[1].value, ValueRef::Number("4"));
513 assert_eq!(attributes.sources(0, Ordinal::new(3)), [1, 2]);
514 assert!(attributes.sources(1, Ordinal::new(3)).is_empty());
515 assert_eq!(attributes.targets_of_kind(0), [3]);
516 assert_eq!(
517 attributes
518 .sources_of_kind(1)
519 .expect("kind")
520 .iter()
521 .collect::<Vec<_>>(),
522 [1, 2, 4]
523 );
524 assert_eq!(attributes.kind(200), Some(1));
525 assert_eq!(attributes.kind(300), None);
526 assert_eq!(
527 attributes.rows(Ordinal::new(4)).next().expect("row").value,
528 ValueRef::String("blue")
529 );
530 assert!(attributes.rows(Ordinal::new(0)).next().is_none());
531 }
532
533 #[test]
534 fn the_layout_round_trips_and_refuses_bad_input() {
535 let attributes = sample();
536 let mut bytes = Vec::new();
537 attributes.write_to(&mut bytes).expect("writes");
538 let again = Attributes::read_from(&mut bytes.as_slice()).expect("reads");
539 assert_eq!(again, attributes);
540 assert!(matches!(
541 Attributes::read_from(&mut b"nope".as_slice()),
542 Err(AttributesError::Io(_))
543 ));
544 assert!(matches!(
545 Attributes::build(
546 2,
547 vec![1],
548 vec![Edge {
549 source: Ordinal::new(1),
550 group: 0,
551 kind: 1,
552 value: Value::Concept(Ordinal::new(0)),
553 }]
554 ),
555 Err(AttributesError::OutOfRange { kind: 1, .. })
556 ));
557 let mut wrong = bytes.clone();
558 wrong[0] = b'X';
559 assert!(matches!(
560 Attributes::read_from(&mut wrong.as_slice()),
561 Err(AttributesError::Magic)
562 ));
563 }
564}