1 files changed, 31 insertions, 22 deletions

diff --git a/src/document/replica.rs b/src/document/replica.rs
index 5a87016..20c02e5 100644
--- a/src/document/replica.rs
+++ b/src/document/replica.rs
@@ -10,7 +10,7 @@ use std::fmt;
 use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
-use automerge::{AutoCommit, ChangeHash};
+use super::{ChangeHash, Document, Transaction};
 use tokio::sync::watch;
 
 /// One generation of the swarm document.
@@ -95,12 +95,11 @@ impl fmt::Debug for Epoch {
 
 /// The document and its epoch, always locked and swapped together.
 ///
-/// Derefs to the document, so code that only needs the `AutoCommit` keeps working
-/// with a guard of `Replica`.
+/// Derefs to the document, so code that only needs the document keeps working with a
+/// guard of `Replica`.
 ///
-/// The heads are cached: `AutoCommit::get_heads` takes `&mut self` (it closes a pending
-/// transaction), so without the cache every keepalive and status read would need the
-/// write lock. Every path that changes the document under the write lock calls
+/// The heads are cached, sorted, so keepalives and status reads (and the stamp below)
+/// don't recompute them. Every path that changes the document under the write lock calls
 /// [`Replica::refresh_heads`] before releasing it (`AppState::mutate_doc`,
 /// `switch_epoch`, the sync apply); readers take the read lock and use
 /// [`Replica::heads`].
@@ -113,7 +112,7 @@ pub(crate) struct Replica {
     /// How many times this node changed epoch (switch, rejoin, abandoned rejoin) in this
     /// run: messages carry it, so a late one can't flip a peer's view of our epoch back.
     pub(crate) epoch_seq: u32,
-    pub(crate) doc: AutoCommit,
+    pub(crate) doc: Document,
     heads: Vec<ChangeHash>,
     stamp: watch::Sender<Stamp>,
 }
@@ -129,7 +128,7 @@ pub(crate) struct Stamp {
 }
 
 impl Replica {
-    pub(crate) fn new(epoch: Epoch, doc: AutoCommit) -> Self {
+    pub(crate) fn new(epoch: Epoch, doc: Document) -> Self {
         let (stamp, _) = watch::channel(Stamp {
             epoch,
             epoch_seq: 0,
@@ -168,30 +167,40 @@ impl Replica {
         self.stamp.subscribe()
     }
 
-    /// The writes since the last commit, as one change stamped `at`: receivers tell by
-    /// it how long the change took to reach them (`swarm::stats`).
-    pub(crate) fn commit(&mut self, at: crate::clock::UnixMillis) {
-        self.doc.commit_with(
-            automerge::transaction::CommitOptions::default().with_time(at.as_u64() as i64),
-        );
-    }
-
-    /// The writes since the last commit, discarded.
-    pub(crate) fn rollback(&mut self) {
-        self.doc.rollback();
+    /// The writes `f` makes, as one change stamped `at` if it succeeds (receivers tell
+    /// by the stamp how long the change took to reach them, `swarm::stats`), dropped if
+    /// not. Doesn't refresh the heads: the caller does, under the same lock.
+    pub(crate) fn transact<T, E>(
+        &mut self,
+        at: crate::clock::UnixMillis,
+        f: impl FnOnce(&mut Transaction<'_>) -> Result<T, E>,
+    ) -> Result<T, E> {
+        let mut transaction = self.doc.transaction();
+        let out = f(&mut transaction);
+        match &out {
+            Ok(_) => {
+                transaction.commit_with(
+                    automerge::transaction::CommitOptions::default().with_time(at.as_u64() as i64),
+                );
+            }
+            Err(_) => {
+                transaction.rollback();
+            }
+        }
+        out
     }
 }
 
 impl Deref for Replica {
-    type Target = AutoCommit;
+    type Target = Document;
 
-    fn deref(&self) -> &AutoCommit {
+    fn deref(&self) -> &Document {
         &self.doc
     }
 }
 
 impl DerefMut for Replica {
-    fn deref_mut(&mut self) -> &mut AutoCommit {
+    fn deref_mut(&mut self) -> &mut Document {
         &mut self.doc
     }
 }