Skip to main content

a3s_code_core/state_graph/
graph.rs

1use super::event::MIN_GRAPH_EVENT_SCHEMA_VERSION;
2use super::{GraphEvent, GraphEventRecord, GraphObject, GraphRelation, GRAPH_EVENT_SCHEMA_VERSION};
3use serde::{Deserialize, Serialize};
4use std::collections::{BTreeMap, BTreeSet};
5use thiserror::Error;
6
7#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
8pub struct StateGraph {
9    version: u64,
10    objects: BTreeMap<String, GraphObject>,
11    relations: BTreeMap<String, GraphRelation>,
12}
13
14impl StateGraph {
15    pub fn version(&self) -> u64 {
16        self.version
17    }
18
19    pub fn objects(&self) -> impl Iterator<Item = &GraphObject> {
20        self.objects.values()
21    }
22
23    pub fn relations(&self) -> impl Iterator<Item = &GraphRelation> {
24        self.relations.values()
25    }
26
27    pub fn object(&self, id: &str) -> Option<&GraphObject> {
28        self.objects.get(id)
29    }
30
31    pub fn relation(&self, id: &str) -> Option<&GraphRelation> {
32        self.relations.get(id)
33    }
34
35    pub fn relations_from<'a>(
36        &'a self,
37        source: &'a str,
38    ) -> impl Iterator<Item = &'a GraphRelation> + 'a {
39        self.relations.values().filter(move |r| r.source == source)
40    }
41
42    pub fn relations_to<'a>(
43        &'a self,
44        target: &'a str,
45    ) -> impl Iterator<Item = &'a GraphRelation> + 'a {
46        self.relations.values().filter(move |r| r.target == target)
47    }
48
49    pub(crate) fn apply(&mut self, event: &GraphEvent) -> Result<bool, ReplayError> {
50        let mutated = match event {
51            GraphEvent::ObjectCreated {
52                id,
53                object_type,
54                data,
55            } => {
56                if self.objects.contains_key(id) {
57                    return Err(ReplayError::InvalidMutation(format!(
58                        "object `{id}` already exists"
59                    )));
60                }
61                self.objects.insert(
62                    id.clone(),
63                    GraphObject {
64                        id: id.clone(),
65                        object_type: object_type.clone(),
66                        data: data.clone(),
67                        version: 1,
68                    },
69                );
70                true
71            }
72            GraphEvent::ObjectUpdated { id, version, data } => {
73                let object = self.objects.get_mut(id).ok_or_else(|| {
74                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
75                })?;
76                if *version != object.version + 1 {
77                    return Err(ReplayError::InvalidMutation(format!(
78                        "object `{id}` version jumped from {} to {version}",
79                        object.version
80                    )));
81                }
82                object.version = *version;
83                object.data = data.clone();
84                true
85            }
86            GraphEvent::ObjectRemoved { id, version } => {
87                let object = self.objects.get(id).ok_or_else(|| {
88                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
89                })?;
90                if *version != object.version + 1 {
91                    return Err(ReplayError::InvalidMutation(format!(
92                        "object `{id}` removal version is stale"
93                    )));
94                }
95                if self
96                    .relations
97                    .values()
98                    .any(|relation| relation.source == *id || relation.target == *id)
99                {
100                    return Err(ReplayError::InvalidMutation(format!(
101                        "object `{id}` still has relations"
102                    )));
103                }
104                self.objects.remove(id);
105                true
106            }
107            GraphEvent::RelationCreated {
108                id,
109                relation_type,
110                source,
111                target,
112                data,
113            } => {
114                if self.relations.contains_key(id) {
115                    return Err(ReplayError::InvalidMutation(format!(
116                        "relation `{id}` already exists"
117                    )));
118                }
119                if !self.objects.contains_key(source) || !self.objects.contains_key(target) {
120                    return Err(ReplayError::InvalidMutation(format!(
121                        "relation `{id}` has a missing endpoint"
122                    )));
123                }
124                self.relations.insert(
125                    id.clone(),
126                    GraphRelation {
127                        id: id.clone(),
128                        relation_type: relation_type.clone(),
129                        source: source.clone(),
130                        target: target.clone(),
131                        data: data.clone(),
132                        version: 1,
133                    },
134                );
135                true
136            }
137            GraphEvent::RelationUpdated { id, version, data } => {
138                let relation = self.relations.get_mut(id).ok_or_else(|| {
139                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
140                })?;
141                if *version != relation.version + 1 {
142                    return Err(ReplayError::InvalidMutation(format!(
143                        "relation `{id}` version is stale"
144                    )));
145                }
146                relation.version = *version;
147                relation.data = data.clone();
148                true
149            }
150            GraphEvent::RelationRemoved { id, version } => {
151                let relation = self.relations.get(id).ok_or_else(|| {
152                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
153                })?;
154                if *version != relation.version + 1 {
155                    return Err(ReplayError::InvalidMutation(format!(
156                        "relation `{id}` removal version is stale"
157                    )));
158                }
159                self.relations.remove(id);
160                true
161            }
162            _ => false,
163        };
164        if mutated {
165            self.version += 1;
166        }
167        Ok(mutated)
168    }
169
170    pub fn diff(&self, other: &Self) -> GraphDiff {
171        GraphDiff::between(self, other)
172    }
173
174    pub fn state_hash(&self) -> Result<String, ReplayError> {
175        StructuralHash::from_graph(self)?.digest(self.version)
176    }
177
178    fn legacy_state_hash(&self) -> Result<String, ReplayError> {
179        let bytes = serde_json::to_vec(self)
180            .map_err(|error| ReplayError::Serialization(error.to_string()))?;
181        Ok(sha256::digest(bytes))
182    }
183}
184
185#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
186pub(crate) struct StructuralHash([u8; 32]);
187
188impl StructuralHash {
189    pub(crate) fn from_graph(graph: &StateGraph) -> Result<Self, ReplayError> {
190        let mut hash = Self::default();
191        for object in graph.objects.values() {
192            hash.add(contribution(b"object", object)?);
193        }
194        for relation in graph.relations.values() {
195            hash.add(contribution(b"relation", relation)?);
196        }
197        Ok(hash)
198    }
199
200    pub(crate) fn apply(
201        &mut self,
202        event: &GraphEvent,
203        graph: &StateGraph,
204    ) -> Result<(), ReplayError> {
205        match event {
206            GraphEvent::ObjectCreated {
207                id,
208                object_type,
209                data,
210            } => self.add(contribution(
211                b"object",
212                &GraphObject {
213                    id: id.clone(),
214                    object_type: object_type.clone(),
215                    data: data.clone(),
216                    version: 1,
217                },
218            )?),
219            GraphEvent::ObjectUpdated { id, version, data } => {
220                let current = graph.object(id).ok_or_else(|| {
221                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
222                })?;
223                self.subtract(contribution(b"object", current)?);
224                self.add(contribution(
225                    b"object",
226                    &GraphObject {
227                        id: id.clone(),
228                        object_type: current.object_type.clone(),
229                        data: data.clone(),
230                        version: *version,
231                    },
232                )?);
233            }
234            GraphEvent::ObjectRemoved { id, .. } => {
235                let current = graph.object(id).ok_or_else(|| {
236                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
237                })?;
238                self.subtract(contribution(b"object", current)?);
239            }
240            GraphEvent::RelationCreated {
241                id,
242                relation_type,
243                source,
244                target,
245                data,
246            } => self.add(contribution(
247                b"relation",
248                &GraphRelation {
249                    id: id.clone(),
250                    relation_type: relation_type.clone(),
251                    source: source.clone(),
252                    target: target.clone(),
253                    data: data.clone(),
254                    version: 1,
255                },
256            )?),
257            GraphEvent::RelationUpdated { id, version, data } => {
258                let current = graph.relation(id).ok_or_else(|| {
259                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
260                })?;
261                self.subtract(contribution(b"relation", current)?);
262                self.add(contribution(
263                    b"relation",
264                    &GraphRelation {
265                        id: id.clone(),
266                        relation_type: current.relation_type.clone(),
267                        source: current.source.clone(),
268                        target: current.target.clone(),
269                        data: data.clone(),
270                        version: *version,
271                    },
272                )?);
273            }
274            GraphEvent::RelationRemoved { id, .. } => {
275                let current = graph.relation(id).ok_or_else(|| {
276                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
277                })?;
278                self.subtract(contribution(b"relation", current)?);
279            }
280            _ => {}
281        }
282        Ok(())
283    }
284
285    pub(crate) fn digest(&self, version: u64) -> Result<String, ReplayError> {
286        let mut bytes = b"a3s-state-graph-v2\0".to_vec();
287        bytes.extend_from_slice(&version.to_be_bytes());
288        bytes.extend_from_slice(&self.0);
289        Ok(sha256::digest(bytes))
290    }
291
292    fn add(&mut self, value: [u8; 32]) {
293        let mut carry = 0u16;
294        for index in (0..32).rev() {
295            let sum = u16::from(self.0[index]) + u16::from(value[index]) + carry;
296            self.0[index] = sum as u8;
297            carry = sum >> 8;
298        }
299    }
300
301    fn subtract(&mut self, value: [u8; 32]) {
302        let mut borrow = 0i16;
303        for index in (0..32).rev() {
304            let difference = i16::from(self.0[index]) - i16::from(value[index]) - borrow;
305            if difference < 0 {
306                self.0[index] = (difference + 256) as u8;
307                borrow = 1;
308            } else {
309                self.0[index] = difference as u8;
310                borrow = 0;
311            }
312        }
313    }
314}
315
316fn contribution<T: Serialize>(domain: &[u8], value: &T) -> Result<[u8; 32], ReplayError> {
317    let mut bytes = domain.to_vec();
318    bytes.push(0);
319    bytes.extend(
320        serde_json::to_vec(value).map_err(|error| ReplayError::Serialization(error.to_string()))?,
321    );
322    let encoded = sha256::digest(bytes);
323    let mut output = [0u8; 32];
324    for (index, byte) in output.iter_mut().enumerate() {
325        *byte = u8::from_str_radix(&encoded[index * 2..index * 2 + 2], 16)
326            .map_err(|error| ReplayError::Serialization(error.to_string()))?;
327    }
328    Ok(output)
329}
330
331#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
332pub struct GraphDiff {
333    pub objects_added: Vec<GraphObject>,
334    pub objects_removed: Vec<GraphObject>,
335    pub objects_changed: Vec<(GraphObject, GraphObject)>,
336    pub relations_added: Vec<GraphRelation>,
337    pub relations_removed: Vec<GraphRelation>,
338    pub relations_changed: Vec<(GraphRelation, GraphRelation)>,
339}
340
341impl GraphDiff {
342    fn between(left: &StateGraph, right: &StateGraph) -> Self {
343        let object_ids = left
344            .objects
345            .keys()
346            .chain(right.objects.keys())
347            .cloned()
348            .collect::<BTreeSet<_>>();
349        let relation_ids = left
350            .relations
351            .keys()
352            .chain(right.relations.keys())
353            .cloned()
354            .collect::<BTreeSet<_>>();
355        let mut diff = Self::default();
356        for id in object_ids {
357            match (left.objects.get(&id), right.objects.get(&id)) {
358                (None, Some(value)) => diff.objects_added.push(value.clone()),
359                (Some(value), None) => diff.objects_removed.push(value.clone()),
360                (Some(a), Some(b)) if a != b => diff.objects_changed.push((a.clone(), b.clone())),
361                _ => {}
362            }
363        }
364        for id in relation_ids {
365            match (left.relations.get(&id), right.relations.get(&id)) {
366                (None, Some(value)) => diff.relations_added.push(value.clone()),
367                (Some(value), None) => diff.relations_removed.push(value.clone()),
368                (Some(a), Some(b)) if a != b => diff.relations_changed.push((a.clone(), b.clone())),
369                _ => {}
370            }
371        }
372        diff
373    }
374
375    pub fn is_empty(&self) -> bool {
376        self.objects_added.is_empty()
377            && self.objects_removed.is_empty()
378            && self.objects_changed.is_empty()
379            && self.relations_added.is_empty()
380            && self.relations_removed.is_empty()
381            && self.relations_changed.is_empty()
382    }
383}
384
385#[derive(Debug, Clone, Error, PartialEq, Eq)]
386pub enum ReplayError {
387    #[error("unsupported graph event schema version {actual}; expected {expected}")]
388    UnsupportedSchema { expected: u32, actual: u32 },
389    #[error(
390        "graph event schema downgraded at sequence {sequence}: previous {previous}, got {actual}"
391    )]
392    SchemaDowngrade {
393        sequence: u64,
394        previous: u32,
395        actual: u32,
396    },
397    #[error("event sequence diverged: expected {expected}, got {actual}")]
398    SequenceDiverged { expected: u64, actual: u64 },
399    #[error("state version diverged at sequence {sequence}: expected {expected}, got {actual}")]
400    StateVersionDiverged {
401        sequence: u64,
402        expected: u64,
403        actual: u64,
404    },
405    #[error("event `{event_id}` references unknown or future cause `{causation_id}`")]
406    InvalidCausation {
407        event_id: String,
408        causation_id: String,
409    },
410    #[error("duplicate event id `{0}`")]
411    DuplicateEventId(String),
412    #[error("state hash diverged at sequence {sequence}")]
413    StateHashDiverged { sequence: u64 },
414    #[error("record hash chain diverged at sequence {sequence}")]
415    RecordHashDiverged { sequence: u64 },
416    #[error("invalid graph mutation: {0}")]
417    InvalidMutation(String),
418    #[error("failed to serialize graph replay state: {0}")]
419    Serialization(String),
420}
421
422pub(crate) fn replay_strict(records: &[GraphEventRecord]) -> Result<StateGraph, ReplayError> {
423    let mut graph = StateGraph::default();
424    let mut seen_ids = BTreeSet::new();
425    let mut structural_hash = StructuralHash::default();
426    let mut legacy_state_hash = graph.legacy_state_hash()?;
427    let mut previous_schema = MIN_GRAPH_EVENT_SCHEMA_VERSION;
428    for (expected_sequence, record) in records.iter().enumerate() {
429        if !(MIN_GRAPH_EVENT_SCHEMA_VERSION..=GRAPH_EVENT_SCHEMA_VERSION)
430            .contains(&record.schema_version)
431        {
432            return Err(ReplayError::UnsupportedSchema {
433                expected: GRAPH_EVENT_SCHEMA_VERSION,
434                actual: record.schema_version,
435            });
436        }
437        let expected_sequence = expected_sequence as u64;
438        if record.schema_version < previous_schema {
439            return Err(ReplayError::SchemaDowngrade {
440                sequence: expected_sequence,
441                previous: previous_schema,
442                actual: record.schema_version,
443            });
444        }
445        if record.sequence != expected_sequence {
446            return Err(ReplayError::SequenceDiverged {
447                expected: expected_sequence,
448                actual: record.sequence,
449            });
450        }
451        if let Some(cause) = &record.causation_id {
452            if !seen_ids.contains(cause) {
453                return Err(ReplayError::InvalidCausation {
454                    event_id: record.id.clone(),
455                    causation_id: cause.clone(),
456                });
457            }
458        }
459        if seen_ids.contains(&record.id) {
460            return Err(ReplayError::DuplicateEventId(record.id.clone()));
461        }
462        let expected_previous = expected_sequence
463            .checked_sub(1)
464            .map(|index| records[index as usize].record_hash.clone());
465        if record.previous_record_hash != expected_previous
466            || record.record_hash != record_hash(record)?
467        {
468            return Err(ReplayError::RecordHashDiverged {
469                sequence: record.sequence,
470            });
471        }
472        if record.state_version_before != graph.version {
473            return Err(ReplayError::StateVersionDiverged {
474                sequence: record.sequence,
475                expected: graph.version,
476                actual: record.state_version_before,
477            });
478        }
479        let mut next_structural_hash = structural_hash;
480        next_structural_hash.apply(&record.event, &graph)?;
481        let mutated = graph.apply(&record.event)?;
482        structural_hash = next_structural_hash;
483        if record.state_version_after != graph.version {
484            return Err(ReplayError::StateVersionDiverged {
485                sequence: record.sequence,
486                expected: graph.version,
487                actual: record.state_version_after,
488            });
489        }
490        if mutated && record.schema_version == 1 {
491            legacy_state_hash = graph.legacy_state_hash()?;
492        }
493        let expected_state_hash = if record.schema_version == 1 {
494            legacy_state_hash.clone()
495        } else {
496            legacy_state_hash = String::new();
497            structural_hash.digest(graph.version)?
498        };
499        if record.state_hash_after != expected_state_hash {
500            return Err(ReplayError::StateHashDiverged {
501                sequence: record.sequence,
502            });
503        }
504        seen_ids.insert(record.id.clone());
505        previous_schema = record.schema_version;
506    }
507    Ok(graph)
508}
509
510pub(crate) fn record_hash(record: &GraphEventRecord) -> Result<String, ReplayError> {
511    let value = serde_json::json!({
512        "schema_version": record.schema_version,
513        "id": record.id,
514        "sequence": record.sequence,
515        "timestamp_ms": record.timestamp_ms,
516        "branch_id": record.branch_id,
517        "causation_id": record.causation_id,
518        "correlation_id": record.correlation_id,
519        "state_version_before": record.state_version_before,
520        "state_version_after": record.state_version_after,
521        "state_hash_after": record.state_hash_after,
522        "previous_record_hash": record.previous_record_hash,
523        "event": record.event,
524    });
525    let bytes = serde_json::to_vec(&value)
526        .map_err(|error| ReplayError::Serialization(error.to_string()))?;
527    Ok(sha256::digest(bytes))
528}