catcolab_document_types/
automerge_util.rs

1//! Utilities for copying Automerge document state between heads.
2
3use automerge::transaction::Transactable;
4
5/// Overwrite the document root with the state at `target_heads`.
6///
7/// Unlike the `hydrate_to_json` + `populate_automerge_from_json` round-trip,
8/// this preserves rich-text marks and block markers on Text objects by using
9/// `marks_at` / `mark` instead of going through a plain-string intermediary.
10pub fn copy_doc_at_heads<'a>(
11    tx: &mut automerge::transaction::Transaction<'a>,
12    target_heads: &[automerge::ChangeHash],
13) -> Result<(), automerge::AutomergeError> {
14    use automerge::ReadDoc;
15
16    let current_keys: Vec<String> = tx.keys(automerge::ROOT).collect();
17    for key in &current_keys {
18        tx.delete(automerge::ROOT, key.as_str())?;
19    }
20
21    let target_entries: Vec<_> = {
22        let keys: Vec<String> = tx.keys_at(automerge::ROOT, target_heads).collect();
23        keys.into_iter()
24            .filter_map(|key| {
25                tx.get_at(automerge::ROOT, key.as_str(), target_heads)
26                    .ok()
27                    .flatten()
28                    .map(|(v, id)| (key, v.to_owned(), id))
29            })
30            .collect()
31    };
32    for (key, value, source_id) in &target_entries {
33        put_value_into_map(tx, &automerge::ROOT, key, value, source_id, target_heads)?;
34    }
35    Ok(())
36}
37
38fn collect_children_map(
39    tx: &automerge::transaction::Transaction<'_>,
40    source_id: &automerge::ObjId,
41    heads: &[automerge::ChangeHash],
42) -> Vec<(String, automerge::Value<'static>, automerge::ObjId)> {
43    use automerge::ReadDoc;
44    let keys: Vec<String> = tx.keys_at(source_id, heads).collect();
45    keys.into_iter()
46        .filter_map(|key| {
47            tx.get_at(source_id, key.as_str(), heads)
48                .ok()
49                .flatten()
50                .map(|(v, id)| (key, v.to_owned(), id))
51        })
52        .collect()
53}
54
55fn collect_children_list(
56    tx: &automerge::transaction::Transaction<'_>,
57    source_id: &automerge::ObjId,
58    heads: &[automerge::ChangeHash],
59) -> Vec<(automerge::Value<'static>, automerge::ObjId)> {
60    use automerge::ReadDoc;
61    let len = tx.length_at(source_id, heads);
62    (0..len)
63        .filter_map(|i| {
64            tx.get_at(source_id, i, heads).ok().flatten().map(|(v, id)| (v.to_owned(), id))
65        })
66        .collect()
67}
68
69fn copy_text_spans<'a>(
70    tx: &mut automerge::transaction::Transaction<'a>,
71    new_id: &automerge::ObjId,
72    source_id: &automerge::ObjId,
73    heads: &[automerge::ChangeHash],
74) -> Result<(), automerge::AutomergeError> {
75    use automerge::ReadDoc;
76    let spans: Vec<automerge::Span> = tx.spans_at(source_id, heads)?.collect();
77    tx.update_spans(new_id, automerge::marks::UpdateSpansConfig::default(), spans)
78}
79
80fn put_value_into_map<'a>(
81    tx: &mut automerge::transaction::Transaction<'a>,
82    parent: &automerge::ObjId,
83    key: &str,
84    value: &automerge::Value<'_>,
85    source_id: &automerge::ObjId,
86    heads: &[automerge::ChangeHash],
87) -> Result<(), automerge::AutomergeError> {
88    use automerge::ObjType;
89
90    match value {
91        automerge::Value::Object(ObjType::Text) => {
92            let new_id = tx.put_object(parent, key, ObjType::Text)?;
93            copy_text_spans(tx, &new_id, source_id, heads)?;
94        }
95        automerge::Value::Object(ObjType::Map) => {
96            let new_id = tx.put_object(parent, key, ObjType::Map)?;
97            let children = collect_children_map(tx, source_id, heads);
98            for (child_key, child_val, child_src) in &children {
99                put_value_into_map(tx, &new_id, child_key, child_val, child_src, heads)?;
100            }
101        }
102        automerge::Value::Object(ObjType::List) => {
103            let new_id = tx.put_object(parent, key, ObjType::List)?;
104            let children = collect_children_list(tx, source_id, heads);
105            for (i, (child_val, child_src)) in children.iter().enumerate() {
106                insert_value_into_list_from_doc(tx, &new_id, i, child_val, child_src, heads)?;
107            }
108        }
109        automerge::Value::Object(obj_type) => {
110            tx.put_object(parent, key, *obj_type)?;
111        }
112        automerge::Value::Scalar(s) => {
113            tx.put(parent, key, s.as_ref().clone())?;
114        }
115    }
116    Ok(())
117}
118
119fn insert_value_into_list_from_doc<'a>(
120    tx: &mut automerge::transaction::Transaction<'a>,
121    parent: &automerge::ObjId,
122    index: usize,
123    value: &automerge::Value<'_>,
124    source_id: &automerge::ObjId,
125    heads: &[automerge::ChangeHash],
126) -> Result<(), automerge::AutomergeError> {
127    use automerge::ObjType;
128
129    match value {
130        automerge::Value::Object(ObjType::Text) => {
131            let new_id = tx.insert_object(parent, index, ObjType::Text)?;
132            copy_text_spans(tx, &new_id, source_id, heads)?;
133        }
134        automerge::Value::Object(ObjType::Map) => {
135            let new_id = tx.insert_object(parent, index, ObjType::Map)?;
136            let children = collect_children_map(tx, source_id, heads);
137            for (child_key, child_val, child_src) in &children {
138                put_value_into_map(tx, &new_id, child_key, child_val, child_src, heads)?;
139            }
140        }
141        automerge::Value::Object(ObjType::List) => {
142            let new_id = tx.insert_object(parent, index, ObjType::List)?;
143            let children = collect_children_list(tx, source_id, heads);
144            for (i, (child_val, child_src)) in children.iter().enumerate() {
145                insert_value_into_list_from_doc(tx, &new_id, i, child_val, child_src, heads)?;
146            }
147        }
148        automerge::Value::Object(obj_type) => {
149            tx.insert_object(parent, index, *obj_type)?;
150        }
151        automerge::Value::Scalar(s) => {
152            tx.insert(parent, index, s.as_ref().clone())?;
153        }
154    }
155    Ok(())
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use crate::common_test::{doc_from_json, doc_to_json};
162    use automerge::{Automerge, ObjType, ReadDoc};
163    use serde_json::json;
164
165    #[test]
166    fn copy_restores_scalar_fields() {
167        let mut doc = doc_from_json(&json!({
168            "name": "alice",
169            "age": 30,
170            "active": true
171        }));
172        let heads_v1 = doc.get_heads();
173
174        // Mutate the doc to a different state.
175        doc.transact(|tx| {
176            let name_id = tx.put_object(automerge::ROOT, "name", ObjType::Text)?;
177            tx.splice_text(&name_id, 0, 0, "bob")?;
178            tx.put(automerge::ROOT, "age", 99_i64)?;
179            tx.put(automerge::ROOT, "active", false)?;
180            Ok::<_, automerge::AutomergeError>(())
181        })
182        .unwrap();
183
184        // Restore to v1.
185        doc.transact(|tx| {
186            copy_doc_at_heads(tx, &heads_v1)?;
187            Ok::<_, automerge::AutomergeError>(())
188        })
189        .unwrap();
190
191        let result = doc_to_json(&doc);
192        assert_eq!(result["name"], "alice");
193        assert_eq!(result["age"], 30);
194        assert_eq!(result["active"], true);
195    }
196
197    #[test]
198    fn copy_restores_nested_maps() {
199        let mut doc = doc_from_json(&json!({
200            "config": {
201                "theme": "dark",
202                "settings": {
203                    "fontSize": 14
204                }
205            }
206        }));
207        let heads_v1 = doc.get_heads();
208
209        // Overwrite with different nested structure.
210        doc.transact(|tx| {
211            tx.delete(automerge::ROOT, "config")?;
212            let config = tx.put_object(automerge::ROOT, "config", ObjType::Map)?;
213            let theme = tx.put_object(&config, "theme", ObjType::Text)?;
214            tx.splice_text(&theme, 0, 0, "light")?;
215            Ok::<_, automerge::AutomergeError>(())
216        })
217        .unwrap();
218
219        doc.transact(|tx| {
220            copy_doc_at_heads(tx, &heads_v1)?;
221            Ok::<_, automerge::AutomergeError>(())
222        })
223        .unwrap();
224
225        let result = doc_to_json(&doc);
226        assert_eq!(result["config"]["theme"], "dark");
227        assert_eq!(result["config"]["settings"]["fontSize"], 14);
228    }
229
230    #[test]
231    fn copy_restores_lists() {
232        let mut doc = doc_from_json(&json!({
233            "items": ["a", "b", "c"]
234        }));
235        let heads_v1 = doc.get_heads();
236
237        // Replace with different list.
238        doc.transact(|tx| {
239            tx.delete(automerge::ROOT, "items")?;
240            let list = tx.put_object(automerge::ROOT, "items", ObjType::List)?;
241            let x = tx.insert_object(&list, 0, ObjType::Text)?;
242            tx.splice_text(&x, 0, 0, "x")?;
243            Ok::<_, automerge::AutomergeError>(())
244        })
245        .unwrap();
246
247        doc.transact(|tx| {
248            copy_doc_at_heads(tx, &heads_v1)?;
249            Ok::<_, automerge::AutomergeError>(())
250        })
251        .unwrap();
252
253        let result = doc_to_json(&doc);
254        let items: Vec<&str> = result["items"]
255            .as_array()
256            .unwrap()
257            .iter()
258            .map(|v| v.as_str().unwrap())
259            .collect();
260        assert_eq!(items, vec!["a", "b", "c"]);
261    }
262
263    #[test]
264    fn copy_removes_keys_not_in_target() {
265        let mut doc = doc_from_json(&json!({
266            "keep": "yes"
267        }));
268        let heads_v1 = doc.get_heads();
269
270        // Add extra keys.
271        doc.transact(|tx| {
272            let extra = tx.put_object(automerge::ROOT, "extra", ObjType::Text)?;
273            tx.splice_text(&extra, 0, 0, "should be gone")?;
274            tx.put(automerge::ROOT, "another", 42_i64)?;
275            Ok::<_, automerge::AutomergeError>(())
276        })
277        .unwrap();
278
279        doc.transact(|tx| {
280            copy_doc_at_heads(tx, &heads_v1)?;
281            Ok::<_, automerge::AutomergeError>(())
282        })
283        .unwrap();
284
285        let result = doc_to_json(&doc);
286        assert_eq!(result["keep"], "yes");
287        assert!(result.get("extra").is_none());
288        assert!(result.get("another").is_none());
289    }
290
291    #[test]
292    fn copy_preserves_rich_text_marks() {
293        let mut doc = Automerge::new();
294
295        // Create text with a bold mark.
296        doc.transact(|tx| {
297            let text_id = tx.put_object(automerge::ROOT, "content", ObjType::Text)?;
298            tx.splice_text(&text_id, 0, 0, "hello world")?;
299            tx.mark(
300                &text_id,
301                automerge::marks::Mark::new("bold".into(), true, 0, 5),
302                automerge::marks::ExpandMark::After,
303            )?;
304            Ok::<_, automerge::AutomergeError>(())
305        })
306        .unwrap();
307        let heads_with_mark = doc.get_heads();
308
309        // Overwrite with plain text.
310        doc.transact(|tx| {
311            tx.delete(automerge::ROOT, "content")?;
312            let text_id = tx.put_object(automerge::ROOT, "content", ObjType::Text)?;
313            tx.splice_text(&text_id, 0, 0, "replaced")?;
314            Ok::<_, automerge::AutomergeError>(())
315        })
316        .unwrap();
317
318        // Restore.
319        doc.transact(|tx| {
320            copy_doc_at_heads(tx, &heads_with_mark)?;
321            Ok::<_, automerge::AutomergeError>(())
322        })
323        .unwrap();
324
325        // Verify text content.
326        let result = doc_to_json(&doc);
327        assert_eq!(result["content"], "hello world");
328
329        // Verify mark was preserved.
330        let (_, content_id) = doc.get(automerge::ROOT, "content").unwrap().unwrap();
331        let marks = doc.marks(&content_id).unwrap();
332        assert!(!marks.is_empty(), "bold mark should be preserved");
333        assert_eq!(marks[0].name(), "bold");
334        assert_eq!(marks[0].start, 0);
335        assert_eq!(marks[0].end, 5);
336    }
337
338    #[test]
339    fn copy_works_on_empty_doc() {
340        let mut doc = Automerge::new();
341        let heads_empty = doc.get_heads();
342
343        // Add some data.
344        doc.transact(|tx| {
345            tx.put(automerge::ROOT, "key", 1_i64)?;
346            Ok::<_, automerge::AutomergeError>(())
347        })
348        .unwrap();
349
350        // Restore to empty.
351        doc.transact(|tx| {
352            copy_doc_at_heads(tx, &heads_empty)?;
353            Ok::<_, automerge::AutomergeError>(())
354        })
355        .unwrap();
356
357        let keys: Vec<String> = doc.keys(automerge::ROOT).collect();
358        assert!(keys.is_empty());
359    }
360}
361
362#[cfg(all(test, feature = "property-tests"))]
363mod property_tests {
364    use super::*;
365    use crate::common_test::{doc_from_json, doc_to_json};
366    use crate::v1::notebook::ModelNotebook;
367    use automerge::ReadDoc;
368    use test_strategy::proptest;
369
370    /// After mutating a doc and then restoring via `copy_doc_at_heads`, the
371    /// JSON representation matches the original.
372    #[proptest(cases = 64)]
373    fn copy_doc_at_heads_restores_model_notebook(notebook: ModelNotebook) {
374        let json = serde_json::to_value(&notebook.0).expect("serialize to JSON");
375        let mut doc = doc_from_json(&json);
376        let original_heads = doc.get_heads();
377
378        // Mutate: clear all keys from root.
379        doc.transact(|tx| {
380            let keys: Vec<String> = tx.keys(automerge::ROOT).collect();
381            for key in keys {
382                tx.delete(automerge::ROOT, key.as_str())?;
383            }
384            Ok::<_, automerge::AutomergeError>(())
385        })
386        .unwrap();
387
388        // Restore to original heads.
389        doc.transact(|tx| {
390            copy_doc_at_heads(tx, &original_heads)?;
391            Ok::<_, automerge::AutomergeError>(())
392        })
393        .unwrap();
394
395        let result = doc_to_json(&doc);
396        proptest::prop_assert_eq!(json, result);
397    }
398
399    /// Restoring to empty heads after populating yields an empty document.
400    #[proptest(cases = 64)]
401    fn copy_doc_at_heads_restores_to_empty(notebook: ModelNotebook) {
402        let json = serde_json::to_value(&notebook.0).expect("serialize to JSON");
403        let mut doc = automerge::Automerge::new();
404        let empty_heads = doc.get_heads();
405
406        // Populate the doc.
407        doc.transact(|tx| {
408            crate::automerge_json::populate_automerge_from_json(tx, automerge::ROOT, &json)
409                .unwrap();
410            Ok::<_, automerge::AutomergeError>(())
411        })
412        .unwrap();
413
414        // Restore to empty.
415        doc.transact(|tx| {
416            copy_doc_at_heads(tx, &empty_heads)?;
417            Ok::<_, automerge::AutomergeError>(())
418        })
419        .unwrap();
420
421        let keys: Vec<String> = doc.keys(automerge::ROOT).collect();
422        proptest::prop_assert!(
423            keys.is_empty(),
424            "doc should be empty after restoring to empty heads"
425        );
426    }
427}