a3s_code_core/state_graph/
graph.rs

1use super::{GraphEvent, GraphEventRecord, GraphObject, GraphRelation, GRAPH_EVENT_SCHEMA_VERSION};
2use serde::{Deserialize, Serialize};
3use std::collections::{BTreeMap, BTreeSet};
4use thiserror::Error;
5
6#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
7pub struct StateGraph {
8    version: u64,
9    objects: BTreeMap<String, GraphObject>,
10    relations: BTreeMap<String, GraphRelation>,
11}
12
13impl StateGraph {
14    pub fn version(&self) -> u64 {
15        self.version
16    }
17
18    pub fn objects(&self) -> impl Iterator<Item = &GraphObject> {
19        self.objects.values()
20    }
21
22    pub fn relations(&self) -> impl Iterator<Item = &GraphRelation> {
23        self.relations.values()
24    }
25
26    pub fn object(&self, id: &str) -> Option<&GraphObject> {
27        self.objects.get(id)
28    }
29
30    pub fn relation(&self, id: &str) -> Option<&GraphRelation> {
31        self.relations.get(id)
32    }
33
34    pub fn relations_from<'a>(
35        &'a self,
36        source: &'a str,
37    ) -> impl Iterator<Item = &'a GraphRelation> + 'a {
38        self.relations.values().filter(move |r| r.source == source)
39    }
40
41    pub fn relations_to<'a>(
42        &'a self,
43        target: &'a str,
44    ) -> impl Iterator<Item = &'a GraphRelation> + 'a {
45        self.relations.values().filter(move |r| r.target == target)
46    }
47
48    pub(crate) fn apply(&mut self, event: &GraphEvent) -> Result<bool, ReplayError> {
49        let mutated = match event {
50            GraphEvent::ObjectCreated {
51                id,
52                object_type,
53                data,
54            } => {
55                if self.objects.contains_key(id) {
56                    return Err(ReplayError::InvalidMutation(format!(
57                        "object `{id}` already exists"
58                    )));
59                }
60                self.objects.insert(
61                    id.clone(),
62                    GraphObject {
63                        id: id.clone(),
64                        object_type: object_type.clone(),
65                        data: data.clone(),
66                        version: 1,
67                    },
68                );
69                true
70            }
71            GraphEvent::ObjectUpdated { id, version, data } => {
72                let object = self.objects.get_mut(id).ok_or_else(|| {
73                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
74                })?;
75                if *version != object.version + 1 {
76                    return Err(ReplayError::InvalidMutation(format!(
77                        "object `{id}` version jumped from {} to {version}",
78                        object.version
79                    )));
80                }
81                object.version = *version;
82                object.data = data.clone();
83                true
84            }
85            GraphEvent::ObjectRemoved { id, version } => {
86                let object = self.objects.get(id).ok_or_else(|| {
87                    ReplayError::InvalidMutation(format!("object `{id}` does not exist"))
88                })?;
89                if *version != object.version + 1 {
90                    return Err(ReplayError::InvalidMutation(format!(
91                        "object `{id}` removal version is stale"
92                    )));
93                }
94                if self
95                    .relations
96                    .values()
97                    .any(|relation| relation.source == *id || relation.target == *id)
98                {
99                    return Err(ReplayError::InvalidMutation(format!(
100                        "object `{id}` still has relations"
101                    )));
102                }
103                self.objects.remove(id);
104                true
105            }
106            GraphEvent::RelationCreated {
107                id,
108                relation_type,
109                source,
110                target,
111                data,
112            } => {
113                if self.relations.contains_key(id) {
114                    return Err(ReplayError::InvalidMutation(format!(
115                        "relation `{id}` already exists"
116                    )));
117                }
118                if !self.objects.contains_key(source) || !self.objects.contains_key(target) {
119                    return Err(ReplayError::InvalidMutation(format!(
120                        "relation `{id}` has a missing endpoint"
121                    )));
122                }
123                self.relations.insert(
124                    id.clone(),
125                    GraphRelation {
126                        id: id.clone(),
127                        relation_type: relation_type.clone(),
128                        source: source.clone(),
129                        target: target.clone(),
130                        data: data.clone(),
131                        version: 1,
132                    },
133                );
134                true
135            }
136            GraphEvent::RelationUpdated { id, version, data } => {
137                let relation = self.relations.get_mut(id).ok_or_else(|| {
138                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
139                })?;
140                if *version != relation.version + 1 {
141                    return Err(ReplayError::InvalidMutation(format!(
142                        "relation `{id}` version is stale"
143                    )));
144                }
145                relation.version = *version;
146                relation.data = data.clone();
147                true
148            }
149            GraphEvent::RelationRemoved { id, version } => {
150                let relation = self.relations.get(id).ok_or_else(|| {
151                    ReplayError::InvalidMutation(format!("relation `{id}` does not exist"))
152                })?;
153                if *version != relation.version + 1 {
154                    return Err(ReplayError::InvalidMutation(format!(
155                        "relation `{id}` removal version is stale"
156                    )));
157                }
158                self.relations.remove(id);
159                true
160            }
161            _ => false,
162        };
163        if mutated {
164            self.version += 1;
165        }
166        Ok(mutated)
167    }
168
169    pub fn diff(&self, other: &Self) -> GraphDiff {
170        GraphDiff::between(self, other)
171    }
172
173    pub fn state_hash(&self) -> Result<String, ReplayError> {
174        let bytes = serde_json::to_vec(self)
175            .map_err(|error| ReplayError::Serialization(error.to_string()))?;
176        Ok(sha256::digest(bytes))
177    }
178}
179
180#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
181pub struct GraphDiff {
182    pub objects_added: Vec<GraphObject>,
183    pub objects_removed: Vec<GraphObject>,
184    pub objects_changed: Vec<(GraphObject, GraphObject)>,
185    pub relations_added: Vec<GraphRelation>,
186    pub relations_removed: Vec<GraphRelation>,
187    pub relations_changed: Vec<(GraphRelation, GraphRelation)>,
188}
189
190impl GraphDiff {
191    fn between(left: &StateGraph, right: &StateGraph) -> Self {
192        let object_ids = left
193            .objects
194            .keys()
195            .chain(right.objects.keys())
196            .cloned()
197            .collect::<BTreeSet<_>>();
198        let relation_ids = left
199            .relations
200            .keys()
201            .chain(right.relations.keys())
202            .cloned()
203            .collect::<BTreeSet<_>>();
204        let mut diff = Self::default();
205        for id in object_ids {
206            match (left.objects.get(&id), right.objects.get(&id)) {
207                (None, Some(value)) => diff.objects_added.push(value.clone()),
208                (Some(value), None) => diff.objects_removed.push(value.clone()),
209                (Some(a), Some(b)) if a != b => diff.objects_changed.push((a.clone(), b.clone())),
210                _ => {}
211            }
212        }
213        for id in relation_ids {
214            match (left.relations.get(&id), right.relations.get(&id)) {
215                (None, Some(value)) => diff.relations_added.push(value.clone()),
216                (Some(value), None) => diff.relations_removed.push(value.clone()),
217                (Some(a), Some(b)) if a != b => diff.relations_changed.push((a.clone(), b.clone())),
218                _ => {}
219            }
220        }
221        diff
222    }
223
224    pub fn is_empty(&self) -> bool {
225        self.objects_added.is_empty()
226            && self.objects_removed.is_empty()
227            && self.objects_changed.is_empty()
228            && self.relations_added.is_empty()
229            && self.relations_removed.is_empty()
230            && self.relations_changed.is_empty()
231    }
232}
233
234#[derive(Debug, Clone, Error, PartialEq, Eq)]
235pub enum ReplayError {
236    #[error("unsupported graph event schema version {actual}; expected {expected}")]
237    UnsupportedSchema { expected: u32, actual: u32 },
238    #[error("event sequence diverged: expected {expected}, got {actual}")]
239    SequenceDiverged { expected: u64, actual: u64 },
240    #[error("state version diverged at sequence {sequence}: expected {expected}, got {actual}")]
241    StateVersionDiverged {
242        sequence: u64,
243        expected: u64,
244        actual: u64,
245    },
246    #[error("event `{event_id}` references unknown or future cause `{causation_id}`")]
247    InvalidCausation {
248        event_id: String,
249        causation_id: String,
250    },
251    #[error("duplicate event id `{0}`")]
252    DuplicateEventId(String),
253    #[error("state hash diverged at sequence {sequence}")]
254    StateHashDiverged { sequence: u64 },
255    #[error("record hash chain diverged at sequence {sequence}")]
256    RecordHashDiverged { sequence: u64 },
257    #[error("invalid graph mutation: {0}")]
258    InvalidMutation(String),
259    #[error("failed to serialize graph replay state: {0}")]
260    Serialization(String),
261}
262
263pub(crate) fn replay_strict(records: &[GraphEventRecord]) -> Result<StateGraph, ReplayError> {
264    let mut graph = StateGraph::default();
265    let mut seen_ids = BTreeSet::new();
266    for (expected_sequence, record) in records.iter().enumerate() {
267        if record.schema_version != GRAPH_EVENT_SCHEMA_VERSION {
268            return Err(ReplayError::UnsupportedSchema {
269                expected: GRAPH_EVENT_SCHEMA_VERSION,
270                actual: record.schema_version,
271            });
272        }
273        let expected_sequence = expected_sequence as u64;
274        if record.sequence != expected_sequence {
275            return Err(ReplayError::SequenceDiverged {
276                expected: expected_sequence,
277                actual: record.sequence,
278            });
279        }
280        if let Some(cause) = &record.causation_id {
281            if !seen_ids.contains(cause) {
282                return Err(ReplayError::InvalidCausation {
283                    event_id: record.id.clone(),
284                    causation_id: cause.clone(),
285                });
286            }
287        }
288        if seen_ids.contains(&record.id) {
289            return Err(ReplayError::DuplicateEventId(record.id.clone()));
290        }
291        let expected_previous = expected_sequence
292            .checked_sub(1)
293            .map(|index| records[index as usize].record_hash.clone());
294        if record.previous_record_hash != expected_previous
295            || record.record_hash != record_hash(record)?
296        {
297            return Err(ReplayError::RecordHashDiverged {
298                sequence: record.sequence,
299            });
300        }
301        if record.state_version_before != graph.version {
302            return Err(ReplayError::StateVersionDiverged {
303                sequence: record.sequence,
304                expected: graph.version,
305                actual: record.state_version_before,
306            });
307        }
308        graph.apply(&record.event)?;
309        if record.state_version_after != graph.version {
310            return Err(ReplayError::StateVersionDiverged {
311                sequence: record.sequence,
312                expected: graph.version,
313                actual: record.state_version_after,
314            });
315        }
316        if record.state_hash_after != graph.state_hash()? {
317            return Err(ReplayError::StateHashDiverged {
318                sequence: record.sequence,
319            });
320        }
321        seen_ids.insert(record.id.clone());
322    }
323    Ok(graph)
324}
325
326pub(crate) fn record_hash(record: &GraphEventRecord) -> Result<String, ReplayError> {
327    let value = serde_json::json!({
328        "schema_version": record.schema_version,
329        "id": record.id,
330        "sequence": record.sequence,
331        "timestamp_ms": record.timestamp_ms,
332        "branch_id": record.branch_id,
333        "causation_id": record.causation_id,
334        "correlation_id": record.correlation_id,
335        "state_version_before": record.state_version_before,
336        "state_version_after": record.state_version_after,
337        "state_hash_after": record.state_hash_after,
338        "previous_record_hash": record.previous_record_hash,
339        "event": record.event,
340    });
341    let bytes = serde_json::to_vec(&value)
342        .map_err(|error| ReplayError::Serialization(error.to_string()))?;
343    Ok(sha256::digest(bytes))
344}
a3s_code_core/state_graph/graph.rs

a3s_code_core/state_graph/
graph.rs
