Skip to main content

lora_store/memory/
snapshot.rs

1//! Snapshot payload helpers for the in-memory graph.
2//!
3//! `lora-store` no longer ships its own on-disk codec. The byte-level
4//! columnar format lives in `lora-snapshot`; this module just bridges
5//! between [`InMemoryGraph`] and the portable [`SnapshotPayload`]
6//! vocabulary.
7
8use std::collections::BTreeSet;
9use std::sync::Arc;
10
11use crate::dict::{Dicts, ValueDictionary};
12use crate::encoded::{self, StoredNode, StoredRel};
13use crate::{
14    Name, NodeId, NodeRecord, NodeRef, PropertyValue, RelRef, RelationshipId, RelationshipRecord,
15    SnapshotError, SnapshotMeta, SnapshotPayload, ValueRef,
16};
17
18use super::distinct_stats::DistinctBuilder;
19use super::graph::DeferredData;
20use super::index_catalog::StoredIndexEntity;
21use super::record_slab::BuiltLeaves;
22use super::{ConstraintDefinition, InMemoryGraph, IndexDefinition, VectorIndexSnapshot};
23use crate::scratch::ScratchVec;
24
25/// Format-version stamp surfaced through [`SnapshotMeta::format_version`]
26/// for payloads produced via the inherent helpers below. Kept stable
27/// across `lora-snapshot` codec versions because the payload shape
28/// itself has not changed; only the on-disk encoding has.
29pub(super) const PAYLOAD_FORMAT_VERSION: u32 = 1;
30
31impl InMemoryGraph {
32    /// Return the portable graph-state payload. Callers downstream of
33    /// `lora-store` (typically `lora-database`) feed this into
34    /// `lora-snapshot` for byte-level encoding.
35    pub fn snapshot_payload(&self) -> SnapshotPayload {
36        let mut vector_indexes = self
37            .vector_indexes_read(StoredIndexEntity::Node)
38            .to_snapshots(StoredIndexEntity::Node);
39        vector_indexes.extend(
40            self.vector_indexes_read(StoredIndexEntity::Relationship)
41                .to_snapshots(StoredIndexEntity::Relationship),
42        );
43        SnapshotPayload {
44            next_node_id: self.next_node_id,
45            next_rel_id: self.next_rel_id,
46            nodes: self.iter_node_refs().map(|r| r.to_record()).collect(),
47            relationships: self.iter_rel_refs().map(|r| r.to_record()).collect(),
48            indexes: self.index_catalog_read().list(),
49            constraints: self.constraint_catalog_read().list(),
50            vector_indexes,
51        }
52    }
53
54    /// Replace the graph from a portable graph-state payload, preserving the
55    /// currently installed mutation recorder across the swap.
56    pub fn load_snapshot_payload(
57        &mut self,
58        payload: SnapshotPayload,
59    ) -> Result<SnapshotMeta, SnapshotError> {
60        let meta = SnapshotMeta {
61            format_version: PAYLOAD_FORMAT_VERSION,
62            node_count: payload.nodes.len(),
63            relationship_count: payload.relationships.len(),
64            wal_lsn: None,
65        };
66
67        validate_payload_ids(&payload)?;
68
69        let mut builder = SnapshotGraphBuilder::new(payload.next_node_id, payload.next_rel_id)?;
70        for node in payload.nodes {
71            builder.push_node(node)?;
72        }
73        for rel in payload.relationships {
74            builder.push_relationship(rel)?;
75        }
76        let mut rebuilt =
77            builder.finish(payload.indexes, payload.constraints, payload.vector_indexes)?;
78
79        // Preserve the existing recorder across the swap — observers of the
80        // store's identity should not be silently detached by a restore,
81        // same policy as `clear()`.
82        rebuilt.recorder = self.recorder.take();
83        *self = rebuilt;
84
85        Ok(meta)
86    }
87
88    /// Replace the graph with `restored`, one a snapshot decoder built
89    /// (`lora_snapshot::read_graph_snapshot`), preserving the currently
90    /// installed mutation recorder as [`Self::load_snapshot_payload`] does.
91    pub fn replace_with_restored(&mut self, mut restored: InMemoryGraph) {
92        restored.recorder = self.recorder.take();
93        *self = restored;
94    }
95
96    /// The node records, in id order, as borrowed views. With
97    /// [`Self::snapshot_relationships`] and [`Self::snapshot_catalogs`] this
98    /// is what a snapshot encoder reads instead of a [`SnapshotPayload`],
99    /// which would hold a decoded copy of every record at once. Each call
100    /// is a fresh pass, so an encoder that writes column by column can
101    /// walk the graph once per column.
102    pub fn snapshot_nodes(&self) -> impl Iterator<Item = NodeRef<'_>> + '_ {
103        self.iter_node_refs()
104    }
105
106    /// The relationship records, in id order, as borrowed views.
107    pub fn snapshot_relationships(&self) -> impl Iterator<Item = RelRef<'_>> + '_ {
108        self.iter_rel_refs()
109    }
110
111    /// Everything a snapshot holds besides the records.
112    pub fn snapshot_catalogs(&self) -> SnapshotCatalogs {
113        let mut vector_indexes = self
114            .vector_indexes_read(StoredIndexEntity::Node)
115            .to_snapshots(StoredIndexEntity::Node);
116        vector_indexes.extend(
117            self.vector_indexes_read(StoredIndexEntity::Relationship)
118                .to_snapshots(StoredIndexEntity::Relationship),
119        );
120        SnapshotCatalogs {
121            next_node_id: self.next_node_id,
122            next_rel_id: self.next_rel_id,
123            indexes: self.index_catalog_read().list(),
124            constraints: self.constraint_catalog_read().list(),
125            vector_indexes,
126        }
127    }
128}
129
130/// The parts of a snapshot that are not records: the id counters and the
131/// declared catalogs. Small next to the graph.
132#[derive(Debug, Clone)]
133pub struct SnapshotCatalogs {
134    pub next_node_id: NodeId,
135    pub next_rel_id: RelationshipId,
136    pub indexes: Vec<IndexDefinition>,
137    pub constraints: Vec<ConstraintDefinition>,
138    pub vector_indexes: Vec<VectorIndexSnapshot>,
139}
140
141/// Builds an [`InMemoryGraph`] from snapshot records one at a time.
142///
143/// A decoder that streams a snapshot pushes each record as it decodes it,
144/// so the records are never all held at once beside the graph they fill.
145/// Push every node before the relationships that reference it, then call
146/// [`Self::finish`] with the catalogs, which registers the indexes and
147/// constraints over the loaded data.
148pub struct SnapshotGraphBuilder {
149    graph: InMemoryGraph,
150    node_distinct: DistinctBuilder,
151    rel_distinct: DistinctBuilder,
152}
153
154/// A node record names one label twice: no write makes one, and the label
155/// index would hold the node twice.
156fn repeated_label(id: NodeId) -> SnapshotError {
157    SnapshotError::Decode(format!("node {id} has a repeated label in snapshot"))
158}
159
160impl SnapshotGraphBuilder {
161    pub fn new(next_node_id: NodeId, next_rel_id: RelationshipId) -> Result<Self, SnapshotError> {
162        validate_next_id("node", next_node_id)?;
163        validate_next_id("relationship", next_rel_id)?;
164        let mut graph = InMemoryGraph::new();
165        graph.next_node_id = next_node_id;
166        graph.next_rel_id = next_rel_id;
167        Ok(Self {
168            graph,
169            node_distinct: DistinctBuilder::default(),
170            rel_distinct: DistinctBuilder::default(),
171        })
172    }
173
174    pub fn push_node(&mut self, node: NodeRecord) -> Result<(), SnapshotError> {
175        let id = node.id;
176        validate_entity_id("node", id, self.graph.next_node_id)?;
177        if node
178            .labels
179            .iter()
180            .enumerate()
181            .any(|(i, label)| node.labels[..i].contains(label))
182        {
183            return Err(repeated_label(id));
184        }
185        // Distinct-value sketches are not persisted: count the record as a
186        // live create would, through the bulk builder.
187        for label in &node.labels {
188            self.node_distinct.add_all(label, &node.properties);
189        }
190        if self.graph.has_node_at(id) {
191            return Err(SnapshotError::Decode(format!(
192                "duplicate node id {id} in snapshot payload"
193            )));
194        }
195        self.graph
196            .put_node_checked(id, &node)
197            .map_err(SnapshotError::Decode)?;
198        for label in &node.labels {
199            self.graph.insert_node_label_index(id, label);
200        }
201        Ok(())
202    }
203
204    pub fn push_relationship(&mut self, rel: RelationshipRecord) -> Result<(), SnapshotError> {
205        validate_entity_id("relationship", rel.id, self.graph.next_rel_id)?;
206        validate_slot_id("relationship source node", rel.src)?;
207        validate_slot_id("relationship target node", rel.dst)?;
208        if self.graph.has_rel_at(rel.id) {
209            return Err(SnapshotError::Decode(format!(
210                "duplicate relationship id {} in snapshot payload",
211                rel.id
212            )));
213        }
214        if !self.graph.has_node_at(rel.src) {
215            return Err(SnapshotError::Decode(format!(
216                "relationship {} references missing source node {}",
217                rel.id, rel.src
218            )));
219        }
220        if !self.graph.has_node_at(rel.dst) {
221            return Err(SnapshotError::Decode(format!(
222                "relationship {} references missing target node {}",
223                rel.id, rel.dst
224            )));
225        }
226        self.graph
227            .put_rel_checked(rel.id, &rel)
228            .map_err(SnapshotError::Decode)?;
229        self.graph.attach_relationship(&rel);
230        self.rel_distinct.add_all(&rel.rel_type, &rel.properties);
231        Ok(())
232    }
233
234    pub fn finish(
235        self,
236        indexes: Vec<IndexDefinition>,
237        constraints: Vec<ConstraintDefinition>,
238        vector_indexes: Vec<VectorIndexSnapshot>,
239    ) -> Result<InMemoryGraph, SnapshotError> {
240        let Self {
241            graph: mut rebuilt,
242            node_distinct,
243            rel_distinct,
244        } = self;
245        rebuilt.distinct_stats.nodes = node_distinct.finish();
246        rebuilt.distinct_stats.relationships = rel_distinct.finish();
247        // No hash property index is built here. Restart reproduces the
248        // state of a fresh process: the declared RANGE indexes and the
249        // uniqueness/key constraints below activate (and backfill) the
250        // hash indexes they need through `register_index` /
251        // `register_constraint`. Implicit, lookup-activated indexes are
252        // rebuilt lazily by the first lookup that needs them.
253
254        let constraint_owned_indexes: BTreeSet<String> = constraints
255            .iter()
256            .filter_map(|def| {
257                def.owned_index
258                    .clone()
259                    .or_else(|| def.kind.requires_backing_index().then(|| def.name.clone()))
260            })
261            .collect();
262
263        // Re-register every user-visible index in the catalog. Going through
264        // `register_index` re-populates RANGE buckets and keeps the
265        // `populate_index_data` invariant aligned with the catalog —
266        // skipping it would leave RANGE indexes registered but never populated.
267        // Constraint-owned backing indexes are restored by re-registering the
268        // owning constraint below, which keeps catalog ownership explicit.
269        //
270        // Registration only fills the catalogs, in order; the indexes'
271        // data and the constraint scans are left in `deferred` and done
272        // together at the end, on every core.
273        let mut deferred = DeferredData::default();
274        for def in indexes {
275            if constraint_owned_indexes.contains(&def.name) {
276                continue;
277            }
278            // Errors here would mean the snapshot itself is corrupt or
279            // ambiguous; map them into Decode rather than panicking.
280            rebuilt
281                .register_index_deferring(
282                    crate::memory::IndexRequest {
283                        explicit_name: Some(def.name.clone()),
284                        kind: def.kind,
285                        entity: def.entity,
286                        label: def.label.clone(),
287                        additional_labels: def.additional_labels.clone(),
288                        properties: def.properties.clone(),
289                        options: def.options.clone(),
290                    },
291                    /*if_not_exists*/ true,
292                    /*record_event*/ true,
293                    Some(&mut deferred),
294                )
295                .map_err(|e| SnapshotError::Decode(format!("index `{}`: {e}", def.name)))?;
296        }
297
298        // Overlay persisted HNSW snapshots over the freshly-registered
299        // (and freshly-backfilled) vector indexes. This is the
300        // post-step that gives Phase 5 its raison d'être: instead of
301        // paying O(n log n) to re-insert every vector through the
302        // HNSW algorithm, we install the graph topology byte-for-byte.
303        // Snapshots from versions before this trailer round-trip with
304        // `vector_indexes = []` so the fallback path (the backfill
305        // that already ran inside `register_index`) handles them
306        // correctly.
307        for snap in vector_indexes {
308            let entity = snap.entity;
309            let mut registry = rebuilt.vector_indexes_write(entity);
310            if !registry.restore_snapshot(snap) {
311                // Catalog/snapshot mismatch — registry already
312                // contains the populate-built backend, which is the
313                // safe fallback. No further action.
314            }
315        }
316
317        // Re-register constraints. Uniqueness / key constraints recreate
318        // their own backing indexes as part of registration.
319        for def in constraints {
320            rebuilt
321                .register_constraint_deferring(
322                    crate::memory::ConstraintRequest {
323                        name: def.name.clone(),
324                        kind: def.kind.clone(),
325                        entity: def.entity,
326                        label: def.label.clone(),
327                        properties: def.properties.clone(),
328                        backing: def.backing,
329                    },
330                    /*if_not_exists*/ true,
331                    Some(&mut deferred),
332                )
333                .map_err(|e| SnapshotError::Decode(format!("constraint `{}`: {e}", def.name)))?;
334        }
335
336        rebuilt
337            .run_deferred(deferred)
338            .map_err(SnapshotError::Decode)?;
339        Ok(rebuilt)
340    }
341}
342
343fn validate_payload_ids(payload: &SnapshotPayload) -> Result<(), SnapshotError> {
344    validate_next_id("node", payload.next_node_id)?;
345    validate_next_id("relationship", payload.next_rel_id)?;
346
347    for node in &payload.nodes {
348        validate_entity_id("node", node.id, payload.next_node_id)?;
349    }
350    for rel in &payload.relationships {
351        validate_entity_id("relationship", rel.id, payload.next_rel_id)?;
352        validate_slot_id("relationship source node", rel.src)?;
353        validate_slot_id("relationship target node", rel.dst)?;
354    }
355
356    Ok(())
357}
358
359fn validate_next_id(kind: &str, next_id: u64) -> Result<(), SnapshotError> {
360    validate_slot_id(&format!("next {kind} id"), next_id)?;
361    if next_id == u64::MAX {
362        return Err(SnapshotError::Decode(format!(
363            "next {kind} id {next_id} leaves no allocatable id"
364        )));
365    }
366    Ok(())
367}
368
369fn validate_entity_id(kind: &str, id: u64, next_id: u64) -> Result<(), SnapshotError> {
370    validate_slot_id(kind, id)?;
371    if id >= next_id {
372        return Err(SnapshotError::Decode(format!(
373            "{kind} id {id} is not below next {kind} id {next_id}"
374        )));
375    }
376    Ok(())
377}
378
379fn validate_slot_id(label: &str, id: u64) -> Result<(), SnapshotError> {
380    let idx = usize::try_from(id).map_err(|_| {
381        SnapshotError::Decode(format!(
382            "{label} {id} does not fit in usize on this platform"
383        ))
384    })?;
385    idx.checked_add(1)
386        .ok_or_else(|| SnapshotError::Decode(format!("{label} {id} leaves no valid slab slot")))?;
387    Ok(())
388}
389
390// ---------------------------------------------------------------------------
391// Encoding on several threads
392// ---------------------------------------------------------------------------
393
394/// The dictionaries a graph numbers its records with, in number order.
395/// A snapshot writer names records by these numbers, and a
396/// [`SnapshotGraphBuilder`] primed with them ([`SnapshotGraphBuilder::prime`])
397/// encodes records on any thread.
398#[derive(Debug, Clone, Default, PartialEq, Eq)]
399pub struct SnapshotNames {
400    pub labels: Vec<Name>,
401    pub types: Vec<Name>,
402    pub keys: Vec<Name>,
403    /// Which short string values the graph stores by number.
404    pub values: ValueDictionary,
405}
406
407/// A stored node as a snapshot writer reads it: names by number in
408/// [`SnapshotNames`], values borrowed.
409#[derive(Clone, Copy)]
410pub struct SnapshotNode<'a> {
411    stored: StoredNode<'a>,
412}
413
414impl<'a> SnapshotNode<'a> {
415    pub fn id(self) -> NodeId {
416        self.stored.id
417    }
418
419    /// Numbers in [`SnapshotNames::labels`], in the node's order.
420    pub fn label_ids(self) -> impl ExactSizeIterator<Item = u32> + 'a {
421        let labels = self.stored.labels();
422        let count = labels.len();
423        ExactLen {
424            inner: labels.ids(),
425            remaining: count,
426        }
427    }
428
429    /// `(number in` [`SnapshotNames::keys`]`, value)` in key-name order.
430    pub fn properties(self) -> impl Iterator<Item = (u32, ValueRef<'a>)> + 'a {
431        self.stored.properties().iter_ids()
432    }
433
434    pub fn property_count(self) -> usize {
435        self.stored.properties().len()
436    }
437}
438
439/// A stored relationship as a snapshot writer reads it.
440#[derive(Clone, Copy)]
441pub struct SnapshotRelationship<'a> {
442    stored: StoredRel<'a>,
443}
444
445impl<'a> SnapshotRelationship<'a> {
446    pub fn id(self) -> RelationshipId {
447        self.stored.id
448    }
449
450    pub fn src(self) -> NodeId {
451        self.stored.src
452    }
453
454    pub fn dst(self) -> NodeId {
455        self.stored.dst
456    }
457
458    /// Number in [`SnapshotNames::types`].
459    pub fn type_id(self) -> u32 {
460        self.stored.type_id()
461    }
462
463    /// `(number in` [`SnapshotNames::keys`]`, value)` in key-name order.
464    pub fn properties(self) -> impl Iterator<Item = (u32, ValueRef<'a>)> + 'a {
465        self.stored.properties().iter_ids()
466    }
467
468    pub fn property_count(self) -> usize {
469        self.stored.properties().len()
470    }
471}
472
473/// An iterator of known length.
474struct ExactLen<I> {
475    inner: I,
476    remaining: usize,
477}
478
479impl<I: Iterator> Iterator for ExactLen<I> {
480    type Item = I::Item;
481
482    fn next(&mut self) -> Option<I::Item> {
483        let item = self.inner.next()?;
484        self.remaining -= 1;
485        Some(item)
486    }
487
488    fn size_hint(&self) -> (usize, Option<usize>) {
489        (self.remaining, Some(self.remaining))
490    }
491}
492
493impl<I: Iterator> ExactSizeIterator for ExactLen<I> {}
494
495impl InMemoryGraph {
496    /// The graph's dictionaries. Every name a record holds is in them,
497    /// and so may names no record holds any more.
498    pub fn snapshot_names(&self) -> SnapshotNames {
499        let names = |dict: &crate::dict::NameDict| {
500            (0..dict.len() as u32)
501                .map(|id| dict.name(id).clone())
502                .collect()
503        };
504        SnapshotNames {
505            labels: names(&self.dicts.labels),
506            types: names(&self.dicts.types),
507            keys: names(&self.dicts.keys),
508            values: self
509                .dicts
510                .values
511                .export(|key| self.dicts.keys.name(key).clone()),
512        }
513    }
514
515    /// The nodes with an id in `ids`, in id order, numbered as
516    /// [`Self::snapshot_names`] numbers them. Writers on several threads
517    /// each take a range.
518    pub fn snapshot_nodes_in(
519        &self,
520        ids: std::ops::Range<NodeId>,
521    ) -> impl Iterator<Item = SnapshotNode<'_>> + '_ {
522        let dicts = &self.dicts;
523        let start = usize::try_from(ids.start).unwrap_or(usize::MAX);
524        let end = ids.end;
525        self.nodes
526            .iter_from(start)
527            .map(|(i, bytes)| (i as NodeId, bytes))
528            .take_while(move |(id, _)| *id < end)
529            .map(move |(id, bytes)| SnapshotNode {
530                stored: StoredNode::new(id, bytes, dicts),
531            })
532    }
533
534    /// The relationships with an id in `ids`, in id order.
535    pub fn snapshot_relationships_in(
536        &self,
537        ids: std::ops::Range<RelationshipId>,
538    ) -> impl Iterator<Item = SnapshotRelationship<'_>> + '_ {
539        let dicts = &self.dicts;
540        let start = usize::try_from(ids.start).unwrap_or(usize::MAX);
541        let end = ids.end;
542        self.relationships
543            .iter_from(start)
544            .map(|(i, bytes)| (i as RelationshipId, bytes))
545            .take_while(move |(id, _)| *id < end)
546            .map(move |(id, bytes)| SnapshotRelationship {
547                stored: StoredRel::new(id, bytes, dicts),
548            })
549    }
550
551    /// One past the highest node slot (an upper bound on node ids).
552    pub fn snapshot_node_slots(&self) -> u64 {
553        self.nodes.len() as u64
554    }
555
556    /// One past the highest relationship slot.
557    pub fn snapshot_relationship_slots(&self) -> u64 {
558        self.relationships.len() as u64
559    }
560}
561
562/// Encodes snapshot records into a graph's stored form on any thread,
563/// for the [`SnapshotGraphBuilder`] that [primed](SnapshotGraphBuilder::prime)
564/// it. Cheap to clone.
565#[derive(Clone)]
566pub struct SnapshotEncoder {
567    inner: Arc<EncoderTables>,
568}
569
570struct EncoderTables {
571    /// The builder's dictionaries when it was primed. Encoding reads them
572    /// only; a value they do not number yet is stored in full.
573    dicts: Dicts,
574    /// Snapshot label / type / key number -> graph number.
575    labels: Vec<u32>,
576    types: Vec<u32>,
577    keys: Vec<u32>,
578    /// The snapshot's key names, and each one's place in name order:
579    /// records keep their properties in that order.
580    key_names: Vec<Arc<str>>,
581    key_ranks: Vec<u32>,
582    label_names: Vec<Name>,
583    type_names: Vec<Name>,
584    next_node_id: NodeId,
585    next_rel_id: RelationshipId,
586}
587
588/// Nodes encoded off the builder's thread, for
589/// [`SnapshotGraphBuilder::push_nodes`]. Fill with [`SnapshotEncoder::push_node`].
590pub struct NodeBatch {
591    /// The encoded records, in whole leaves of the node slab.
592    records: BuiltLeaves,
593    by_label: Vec<(u32, Vec<NodeId>)>,
594    distinct: DistinctBuilder,
595    /// Reused per record.
596    scratch: Vec<u8>,
597    labels: Vec<u32>,
598}
599
600/// Relationships encoded off the builder's thread, for
601/// [`SnapshotGraphBuilder::push_relationships`].
602pub struct RelationshipBatch {
603    /// The encoded records, in whole leaves of the relationship slab.
604    records: BuiltLeaves,
605    /// `(id, source, target, graph type number)` per relationship.
606    links: ScratchVec<(RelationshipId, NodeId, NodeId, u32)>,
607    by_type: Vec<(u32, Vec<RelationshipId>)>,
608    distinct: DistinctBuilder,
609    scratch: Vec<u8>,
610}
611
612impl NodeBatch {
613    pub fn len(&self) -> usize {
614        self.records.len()
615    }
616
617    pub fn is_empty(&self) -> bool {
618        self.records.len() == 0
619    }
620}
621
622impl RelationshipBatch {
623    pub fn len(&self) -> usize {
624        self.records.len()
625    }
626
627    pub fn is_empty(&self) -> bool {
628        self.records.len() == 0
629    }
630}
631
632/// Add record `id` to a batch's leaves. Ids ascend within a batch, so
633/// a repeated or falling one is refused here.
634fn push_record(records: &mut BuiltLeaves, id: u64, bytes: &[u8]) -> Result<(), SnapshotError> {
635    let slot = usize::try_from(id)
636        .map_err(|_| SnapshotError::Decode(format!("id {id} does not fit in usize")))?;
637    if records.last_slot().is_some_and(|last| last >= slot) {
638        return Err(SnapshotError::Decode(format!(
639            "record ids in a snapshot batch must ascend (id {id})"
640        )));
641    }
642    records.push(slot, bytes);
643    Ok(())
644}
645
646/// `ids` for `number` in a batch's by-name lists, which hold few names.
647fn ids_for(lists: &mut Vec<(u32, Vec<u64>)>, number: u32) -> &mut Vec<u64> {
648    let at = match lists.iter().rposition(|(n, _)| *n == number) {
649        Some(at) => at,
650        None => {
651            lists.push((number, Vec::new()));
652            lists.len() - 1
653        }
654    };
655    &mut lists[at].1
656}
657
658impl SnapshotEncoder {
659    /// A batch for about `capacity` records.
660    pub fn node_batch(&self, capacity: usize) -> NodeBatch {
661        let _ = capacity;
662        NodeBatch {
663            records: BuiltLeaves::default(),
664            by_label: Vec::new(),
665            distinct: DistinctBuilder::default(),
666            scratch: Vec::new(),
667            labels: Vec::new(),
668        }
669    }
670
671    /// A batch for about `capacity` records. `capacity` comes from the
672    /// snapshot being read, so at most a frame's worth is reserved up front.
673    pub fn relationship_batch(&self, capacity: usize) -> RelationshipBatch {
674        RelationshipBatch {
675            records: BuiltLeaves::default(),
676            links: ScratchVec::with_capacity(capacity.min(1 << 16)),
677            by_type: Vec::new(),
678            distinct: DistinctBuilder::default(),
679            scratch: Vec::new(),
680        }
681    }
682
683    /// Put `properties` (snapshot key numbers) in the order a record keeps
684    /// them, the last of a repeated key winning as in [`Properties`], and
685    /// renumber the keys for the graph.
686    fn order_properties(
687        &self,
688        properties: &mut Vec<(u32, PropertyValue)>,
689    ) -> Result<(), SnapshotError> {
690        let tables = &*self.inner;
691        let mut sorted = true;
692        let mut last = None;
693        for (key, _) in properties.iter() {
694            let rank = *tables
695                .key_ranks
696                .get(*key as usize)
697                .ok_or_else(|| SnapshotError::Decode("invalid property key number".into()))?;
698            if last.is_some_and(|last| last >= rank) {
699                sorted = false;
700            }
701            last = Some(rank);
702        }
703        if !sorted {
704            // Stable, so of two equal keys the later stays later; keep it.
705            properties.sort_by_key(|(key, _)| tables.key_ranks[*key as usize]);
706            let mut i = properties.len();
707            while i > 1 {
708                i -= 1;
709                let same = tables.key_ranks[properties[i].0 as usize]
710                    == tables.key_ranks[properties[i - 1].0 as usize];
711                if same {
712                    properties.remove(i - 1);
713                }
714            }
715        }
716        Ok(())
717    }
718
719    /// Encode a node: `labels` and the keys of `properties` are numbers in
720    /// the snapshot's [`SnapshotNames`]. Takes the values out of
721    /// `properties`, leaving it empty for the next record.
722    pub fn push_node(
723        &self,
724        batch: &mut NodeBatch,
725        id: NodeId,
726        labels: &[u32],
727        properties: &mut Vec<(u32, PropertyValue)>,
728    ) -> Result<(), SnapshotError> {
729        let tables = &*self.inner;
730        validate_entity_id("node", id, tables.next_node_id)?;
731        self.order_properties(properties)?;
732        batch.labels.clear();
733        for &label in labels {
734            let number = *tables
735                .labels
736                .get(label as usize)
737                .ok_or_else(|| SnapshotError::Decode("invalid label dictionary id".into()))?;
738            if batch.labels.contains(&number) {
739                return Err(repeated_label(id));
740            }
741            batch.labels.push(number);
742            ids_for(&mut batch.by_label, number).push(id);
743        }
744        for &label in labels {
745            let name = &tables.label_names[label as usize];
746            batch.distinct.add_all(
747                name,
748                properties
749                    .iter()
750                    .map(|(key, value)| (&tables.key_names[*key as usize], value)),
751            );
752        }
753        for (key, _) in properties.iter_mut() {
754            *key = tables.keys[*key as usize];
755        }
756        encoded::encode_node_ids(&batch.labels, properties, &tables.dicts, &mut batch.scratch);
757        properties.clear();
758        push_record(&mut batch.records, id, &batch.scratch)
759    }
760
761    /// Encode a relationship: `rel_type` and the keys of `properties` are
762    /// numbers in the snapshot's [`SnapshotNames`].
763    pub fn push_relationship(
764        &self,
765        batch: &mut RelationshipBatch,
766        id: RelationshipId,
767        src: NodeId,
768        dst: NodeId,
769        rel_type: u32,
770        properties: &mut Vec<(u32, PropertyValue)>,
771    ) -> Result<(), SnapshotError> {
772        let tables = &*self.inner;
773        validate_entity_id("relationship", id, tables.next_rel_id)?;
774        validate_slot_id("relationship source node", src)?;
775        validate_slot_id("relationship target node", dst)?;
776        let type_id = *tables
777            .types
778            .get(rel_type as usize)
779            .ok_or_else(|| SnapshotError::Decode("invalid relationship type id".into()))?;
780        self.order_properties(properties)?;
781        batch.distinct.add_all(
782            &tables.type_names[rel_type as usize],
783            properties
784                .iter()
785                .map(|(key, value)| (&tables.key_names[*key as usize], value)),
786        );
787        for (key, _) in properties.iter_mut() {
788            *key = tables.keys[*key as usize];
789        }
790        encoded::encode_rel_ids(
791            src,
792            dst,
793            type_id,
794            properties,
795            &tables.dicts,
796            &mut batch.scratch,
797        );
798        properties.clear();
799        push_record(&mut batch.records, id, &batch.scratch)?;
800        batch.links.push((id, src, dst, type_id));
801        ids_for(&mut batch.by_type, type_id).push(id);
802        Ok(())
803    }
804}
805
806impl SnapshotGraphBuilder {
807    /// Number the snapshot's names in the graph being built, and return
808    /// an encoder that turns records named by those numbers into the
809    /// graph's stored form on any thread. Call before pushing records.
810    pub fn prime(&mut self, names: &SnapshotNames) -> Result<SnapshotEncoder, SnapshotError> {
811        let dicts = &mut self.graph.dicts;
812        let labels = names
813            .labels
814            .iter()
815            .map(|name| dicts.labels.id_or_insert(name))
816            .collect();
817        let types = names
818            .types
819            .iter()
820            .map(|name| dicts.types.id_or_insert(name))
821            .collect();
822        let keys: Vec<u32> = names
823            .keys
824            .iter()
825            .map(|name| dicts.keys.id_or_insert(name))
826            .collect();
827        let mut by_name: Vec<u32> = (0..names.keys.len() as u32).collect();
828        by_name.sort_by(|a, b| {
829            names.keys[*a as usize]
830                .as_str()
831                .cmp(names.keys[*b as usize].as_str())
832        });
833        let mut key_ranks = vec![0u32; names.keys.len()];
834        let mut rank = 0u32;
835        for (i, &key) in by_name.iter().enumerate() {
836            if i > 0
837                && names.keys[by_name[i - 1] as usize].as_str() != names.keys[key as usize].as_str()
838            {
839                rank += 1;
840            }
841            key_ranks[key as usize] = rank;
842        }
843        let keys_dict = &mut dicts.keys;
844        dicts
845            .values
846            .import(&names.values, |name| keys_dict.id_or_insert(name))
847            .map_err(SnapshotError::Decode)?;
848        Ok(SnapshotEncoder {
849            inner: Arc::new(EncoderTables {
850                dicts: dicts.clone(),
851                labels,
852                types,
853                keys,
854                key_names: names
855                    .keys
856                    .iter()
857                    .map(|name| name.as_arc().clone())
858                    .collect(),
859                key_ranks,
860                label_names: names.labels.clone(),
861                type_names: names.types.clone(),
862                next_node_id: self.graph.next_node_id,
863                next_rel_id: self.graph.next_rel_id,
864            }),
865        })
866    }
867
868    /// Store nodes an encoder from [`Self::prime`] encoded.
869    pub fn push_nodes(&mut self, batch: NodeBatch) -> Result<(), SnapshotError> {
870        self.graph
871            .put_encoded_node_leaves(batch.records)
872            .map_err(SnapshotError::Decode)?;
873        for (label, ids) in &batch.by_label {
874            self.graph.extend_node_label_index(*label, ids);
875        }
876        self.node_distinct.merge(batch.distinct);
877        Ok(())
878    }
879
880    /// Store relationships an encoder from [`Self::prime`] encoded. Their
881    /// endpoints must already be pushed.
882    pub fn push_relationships(&mut self, batch: RelationshipBatch) -> Result<(), SnapshotError> {
883        self.push_relationship_batches(vec![batch])
884    }
885
886    /// [`Self::push_relationships`] for several batches, in order: their
887    /// adjacency is linked in one pass, on every core when the build
888    /// allows.
889    pub fn push_relationship_batches(
890        &mut self,
891        batches: Vec<RelationshipBatch>,
892    ) -> Result<(), SnapshotError> {
893        let mut links = ScratchVec::with_capacity(batches.iter().map(|b| b.links.len()).sum());
894        for batch in batches {
895            self.graph
896                .put_encoded_relationship_leaves(batch.records, batch.links.as_slice())
897                .map_err(SnapshotError::Decode)?;
898            for (rel_type, ids) in &batch.by_type {
899                self.graph.extend_relationship_type_index(*rel_type, ids);
900            }
901            self.rel_distinct.merge(batch.distinct);
902            links.extend_from(&batch.links);
903        }
904        self.graph.link_relationships(links.as_slice());
905        Ok(())
906    }
907}