Skip to main content

uqa_graph/memory_store/
graphid.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! AGE graphid composition and per-graph label allocation state.
8
9use super::{BTreeMap, BTreeSet, Deserialize, GraphStoreError, GraphStoreResult, Serialize};
10
11use crate::age_names::{EDGE_DEFAULT_LABEL_NAME, VERTEX_DEFAULT_LABEL_NAME};
12
13/// Number of bits reserved for the per-label sequence inside an AGE
14/// `graphid`. The label id occupies the remaining high 16 bits.
15pub const GRAPHID_LABEL_SHIFT: u32 = 48;
16
17/// Reserved AGE label id for unlabeled vertices (`_ag_label_vertex`).
18pub const VERTEX_DEFAULT_LABEL_ID: u32 = 1;
19
20/// Reserved AGE label id for unlabeled edges (`_ag_label_edge`).
21pub const EDGE_DEFAULT_LABEL_ID: u32 = 2;
22
23/// First label id available to user labels.
24pub const FIRST_USER_LABEL_ID: u32 = 3;
25
26/// The largest label id whose AGE graphid remains representable as a signed
27/// 64-bit agtype integer.
28pub const MAX_GRAPHID_LABEL_ID: u32 = 32_767;
29
30pub(super) const MAX_GRAPHID_SEQUENCE: u64 = (1_u64 << GRAPHID_LABEL_SHIFT) - 1;
31const MAX_EXACT_F64_INTEGER: u64 = 9_007_199_254_740_992;
32
33pub(super) fn usize_to_f64_exact(value: usize, context: &str) -> GraphStoreResult<f64> {
34    if u64::try_from(value).is_ok_and(|value| value <= MAX_EXACT_F64_INTEGER) {
35        Ok(value as f64)
36    } else {
37        Err(GraphStoreError::InvalidMutation(format!(
38            "{context} {value} exceeds the exact f64 integer range"
39        )))
40    }
41}
42
43/// Compose an AGE `graphid` from a label id and per-label sequence.
44pub fn make_graphid(label_id: u32, sequence: u64) -> GraphStoreResult<u64> {
45    if label_id > MAX_GRAPHID_LABEL_ID {
46        return Err(GraphStoreError::IdExhausted(format!(
47            "label id {label_id} exceeds {MAX_GRAPHID_LABEL_ID}"
48        )));
49    }
50    if sequence == 0 || sequence > MAX_GRAPHID_SEQUENCE {
51        return Err(GraphStoreError::IdExhausted(format!(
52            "sequence {sequence} is outside 1..={MAX_GRAPHID_SEQUENCE}"
53        )));
54    }
55    Ok((u64::from(label_id) << GRAPHID_LABEL_SHIFT) | sequence)
56}
57
58/// Label id component of an AGE `graphid`.
59#[must_use]
60pub fn graphid_label_id(id: u64) -> u32 {
61    let bytes = id.to_be_bytes();
62    u32::from(u16::from_be_bytes([bytes[0], bytes[1]]))
63}
64
65/// Sequence component of an AGE `graphid`.
66#[must_use]
67pub fn graphid_sequence(id: u64) -> u64 {
68    id & ((1 << GRAPHID_LABEL_SHIFT) - 1)
69}
70
71/// AGE label kind: the `ag_label.kind` catalog value.
72#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
73pub enum LabelKind {
74    /// A vertex label (`ag_label.kind = 'v'`).
75    #[serde(rename = "v")]
76    Vertex,
77    /// An edge label (`ag_label.kind = 'e'`).
78    #[serde(rename = "e")]
79    Edge,
80}
81
82impl LabelKind {
83    /// The `ag_label.kind` character.
84    #[must_use]
85    pub fn as_char(self) -> char {
86        match self {
87            Self::Vertex => 'v',
88            Self::Edge => 'e',
89        }
90    }
91
92    /// The reserved label id used for unlabeled entities of this kind.
93    #[must_use]
94    pub fn default_label_id(self) -> u32 {
95        match self {
96            Self::Vertex => VERTEX_DEFAULT_LABEL_ID,
97            Self::Edge => EDGE_DEFAULT_LABEL_ID,
98        }
99    }
100
101    /// The reserved AGE default label name for this kind.
102    #[must_use]
103    pub fn default_label_name(self) -> &'static str {
104        match self {
105            Self::Vertex => VERTEX_DEFAULT_LABEL_NAME,
106            Self::Edge => EDGE_DEFAULT_LABEL_NAME,
107        }
108    }
109
110    fn entity_noun(self) -> &'static str {
111        match self {
112            Self::Vertex => "vertices",
113            Self::Edge => "edges",
114        }
115    }
116}
117
118/// One `ag_label` catalog entry of a graph.
119#[derive(Debug, Clone, PartialEq, Eq)]
120pub struct GraphLabelInfo {
121    /// Label name; the default labels use the reserved AGE names.
122    pub name: String,
123    /// AGE label id (the high 16 bits of every graphid under the label).
124    pub id: u32,
125    /// Vertex or edge label.
126    pub kind: LabelKind,
127    /// Last allocated per-label sequence value (0 when nothing was
128    /// allocated yet).
129    pub last_sequence: u64,
130}
131
132/// Per-graph AGE label registry: label name -> label id plus the
133/// per-label id sequences. Serializable so engines can persist it in
134/// catalog metadata and restore deterministic id allocation.
135#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
136#[serde(default)]
137pub struct GraphLabelRegistry {
138    /// Label name -> AGE label id. Vertex and edge labels share the
139    /// namespace-wide counter; the reserved names for ids 1 / 2 are
140    /// not stored here (empty labels map onto them implicitly).
141    pub labels: BTreeMap<String, u32>,
142    /// Label name -> vertex or edge kind. Registries persisted before
143    /// kinds were recorded fill this map from the stored entities.
144    pub kinds: BTreeMap<String, LabelKind>,
145    /// Label id -> last allocated per-label sequence value.
146    pub sequences: BTreeMap<u32, u64>,
147    /// AGE label ids whose relations were removed through `drop_label`.
148    /// Tombstones are persisted because edge rows can outlive the vertex
149    /// label relations that owned their endpoints. Registries written before
150    /// this field existed deserialize as an empty set.
151    pub dropped_label_ids: BTreeSet<u32>,
152    /// Next label id handed to a previously unseen label.
153    pub next_label_id: u32,
154}
155
156impl Default for GraphLabelRegistry {
157    fn default() -> Self {
158        Self {
159            labels: BTreeMap::new(),
160            kinds: BTreeMap::new(),
161            sequences: BTreeMap::new(),
162            dropped_label_ids: BTreeSet::new(),
163            next_label_id: FIRST_USER_LABEL_ID,
164        }
165    }
166}
167
168impl GraphLabelRegistry {
169    /// Resolve the label id for an entity of `kind`, allocating a new
170    /// user label on first use. Empty labels map onto the reserved
171    /// default label of the kind. Using a label registered for the
172    /// other kind fails exactly like AGE's `CREATE` transform.
173    pub(super) fn label_id(&mut self, label: &str, kind: LabelKind) -> GraphStoreResult<u32> {
174        if label.is_empty() {
175            self.require_default_label(kind)?;
176            return Ok(kind.default_label_id());
177        }
178        // The reserved AGE names always denote the default labels, so they
179        // resolve to the reserved ids instead of allocating a user label.
180        for reserved in [LabelKind::Vertex, LabelKind::Edge] {
181            if label == reserved.default_label_name() {
182                Self::require_kind(label, reserved, kind)?;
183                self.require_default_label(reserved)?;
184                return Ok(reserved.default_label_id());
185            }
186        }
187        if let Some(existing) = self.kinds.get(label).copied() {
188            Self::require_kind(label, existing, kind)?;
189        }
190        if let Some(id) = self.labels.get(label) {
191            if *id > MAX_GRAPHID_LABEL_ID {
192                return Err(GraphStoreError::IdExhausted(format!(
193                    "persisted label id {id} exceeds {MAX_GRAPHID_LABEL_ID}"
194                )));
195            }
196            self.kinds.entry(label.to_string()).or_insert(kind);
197            return Ok(*id);
198        }
199        self.require_default_label(kind)?;
200        let id = self.allocate_label_id()?;
201        self.labels.insert(label.to_string(), id);
202        self.kinds.insert(label.to_string(), kind);
203        Ok(id)
204    }
205
206    fn require_default_label(&self, kind: LabelKind) -> GraphStoreResult<()> {
207        if self.dropped_label_ids.contains(&kind.default_label_id()) {
208            return Err(GraphStoreError::InvalidMutation(format!(
209                "default label {} does not exist",
210                kind.default_label_name()
211            )));
212        }
213        Ok(())
214    }
215
216    fn require_kind(
217        label: &str,
218        existing: LabelKind,
219        requested: LabelKind,
220    ) -> GraphStoreResult<()> {
221        if existing == requested {
222            return Ok(());
223        }
224        Err(GraphStoreError::InvalidMutation(format!(
225            "label {label} is for {}, not {}",
226            existing.entity_noun(),
227            requested.entity_noun()
228        )))
229    }
230
231    fn allocate_label_id(&mut self) -> GraphStoreResult<u32> {
232        let id = self.next_label_id;
233        if id > MAX_GRAPHID_LABEL_ID {
234            return Err(GraphStoreError::IdExhausted(format!(
235                "label id {id} exceeds {MAX_GRAPHID_LABEL_ID}"
236            )));
237        }
238        self.next_label_id = id
239            .checked_add(1)
240            .ok_or_else(|| GraphStoreError::IdExhausted("label id counter overflow".to_string()))?;
241        Ok(id)
242    }
243
244    /// Whether `label` names a registered user label or a reserved
245    /// default label.
246    #[must_use]
247    pub fn contains_label(&self, label: &str) -> bool {
248        if label == VERTEX_DEFAULT_LABEL_NAME {
249            return !self
250                .dropped_label_ids
251                .contains(&LabelKind::Vertex.default_label_id());
252        }
253        if label == EDGE_DEFAULT_LABEL_NAME {
254            return !self
255                .dropped_label_ids
256                .contains(&LabelKind::Edge.default_label_id());
257        }
258        self.labels.contains_key(label)
259    }
260
261    /// The kind of a registered or default label. A user label persisted
262    /// before kinds were recorded, and whose entities are all gone, reports
263    /// as a vertex label until its next use records the kind.
264    #[must_use]
265    pub fn label_kind(&self, label: &str) -> Option<LabelKind> {
266        if label == VERTEX_DEFAULT_LABEL_NAME {
267            return self.contains_label(label).then_some(LabelKind::Vertex);
268        }
269        if label == EDGE_DEFAULT_LABEL_NAME {
270            return self.contains_label(label).then_some(LabelKind::Edge);
271        }
272        if !self.labels.contains_key(label) {
273            return None;
274        }
275        Some(self.kinds.get(label).copied().unwrap_or(LabelKind::Vertex))
276    }
277
278    /// Register an empty user label ahead of any entity, as
279    /// `create_vlabel` / `create_elabel` do. Returns the new label id;
280    /// `None` when the name is already a label of this graph.
281    pub fn register_label(
282        &mut self,
283        label: &str,
284        kind: LabelKind,
285    ) -> GraphStoreResult<Option<u32>> {
286        if self.contains_label(label) {
287            return Ok(None);
288        }
289        self.require_default_label(kind)?;
290        if label == VERTEX_DEFAULT_LABEL_NAME || label == EDGE_DEFAULT_LABEL_NAME {
291            return Err(GraphStoreError::InvalidMutation(format!(
292                "default label {label} cannot be recreated without recreating the graph"
293            )));
294        }
295        let id = self.allocate_label_id()?;
296        self.labels.insert(label.to_string(), id);
297        self.kinds.insert(label.to_string(), kind);
298        Ok(Some(id))
299    }
300
301    /// Forget a label. Default labels leave a durable tombstone so the graph
302    /// can continue to exist without their AGE relations. Returns the
303    /// released label id, or `None` when the label is not registered.
304    pub fn remove_label(&mut self, label: &str) -> Option<u32> {
305        for kind in [LabelKind::Vertex, LabelKind::Edge] {
306            if label == kind.default_label_name() {
307                if !self.dropped_label_ids.insert(kind.default_label_id()) {
308                    return None;
309                }
310                self.sequences.remove(&kind.default_label_id());
311                return Some(kind.default_label_id());
312            }
313        }
314        let id = self.labels.remove(label)?;
315        self.kinds.remove(label);
316        self.sequences.remove(&id);
317        self.dropped_label_ids.insert(id);
318        Some(id)
319    }
320
321    /// Every present label of the graph in `ag_label` order: surviving
322    /// defaults first, then user labels by ascending label id.
323    #[must_use]
324    pub fn labels(&self) -> Vec<GraphLabelInfo> {
325        let mut out = Vec::new();
326        for kind in [LabelKind::Vertex, LabelKind::Edge] {
327            if !self.dropped_label_ids.contains(&kind.default_label_id()) {
328                out.push(GraphLabelInfo {
329                    name: kind.default_label_name().to_string(),
330                    id: kind.default_label_id(),
331                    kind,
332                    last_sequence: self
333                        .sequences
334                        .get(&kind.default_label_id())
335                        .copied()
336                        .unwrap_or(0),
337                });
338            }
339        }
340        let mut user: Vec<GraphLabelInfo> = self
341            .labels
342            .iter()
343            .map(|(name, id)| GraphLabelInfo {
344                name: name.clone(),
345                id: *id,
346                kind: self.label_kind(name).unwrap_or(LabelKind::Vertex),
347                last_sequence: self.sequences.get(id).copied().unwrap_or(0),
348            })
349            .collect();
350        user.sort_by_key(|label| label.id);
351        out.extend(user);
352        out
353    }
354
355    pub(super) fn next_sequence(&mut self, label_id: u32) -> GraphStoreResult<u64> {
356        let current = self.sequences.get(&label_id).copied().unwrap_or(0);
357        let next = current.checked_add(1).ok_or_else(|| {
358            GraphStoreError::IdExhausted(format!(
359                "sequence counter overflow for label id {label_id}"
360            ))
361        })?;
362        if next > MAX_GRAPHID_SEQUENCE {
363            return Err(GraphStoreError::IdExhausted(format!(
364                "sequence {next} exceeds {MAX_GRAPHID_SEQUENCE} for label id {label_id}"
365            )));
366        }
367        self.sequences.insert(label_id, next);
368        Ok(next)
369    }
370
371    /// Fold an existing entity id back into the registry so restored
372    /// graphs never re-issue an id that is already in use.
373    pub(super) fn observe(&mut self, label: &str, id: u64, kind: LabelKind) {
374        let label_id = graphid_label_id(id);
375        if label_id == 0 {
376            // Pre-AGE id (plain counter) - nothing to learn.
377            return;
378        }
379        if !label.is_empty() && label_id >= FIRST_USER_LABEL_ID {
380            self.labels.entry(label.to_string()).or_insert(label_id);
381            self.kinds.entry(label.to_string()).or_insert(kind);
382        }
383        self.dropped_label_ids.remove(&label_id);
384        let seq = graphid_sequence(id);
385        let entry = self.sequences.entry(label_id).or_insert(0);
386        if seq > *entry {
387            *entry = seq;
388        }
389        if label_id >= self.next_label_id {
390            self.next_label_id = label_id + 1;
391        }
392    }
393
394    /// Merge another registry (e.g. persisted metadata) into this one,
395    /// keeping the larger sequence values and label id watermark.
396    pub fn merge(&mut self, other: &GraphLabelRegistry) {
397        for (label, id) in &other.labels {
398            self.labels.entry(label.clone()).or_insert(*id);
399        }
400        for (label, kind) in &other.kinds {
401            self.kinds.entry(label.clone()).or_insert(*kind);
402        }
403        for (label_id, seq) in &other.sequences {
404            let entry = self.sequences.entry(*label_id).or_insert(0);
405            if *seq > *entry {
406                *entry = *seq;
407            }
408        }
409        self.dropped_label_ids
410            .extend(other.dropped_label_ids.iter().copied());
411        if other.next_label_id > self.next_label_id {
412            self.next_label_id = other.next_label_id;
413        }
414    }
415}