use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fmt;
use crate::text_crdt::{TextCrdt, TextOp};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeOpId {
pub counter: u64,
pub peer: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeNodeId(pub TreeOpId);
impl TreeNodeId {
pub const ROOT: TreeNodeId = TreeNodeId(TreeOpId {
counter: 0,
peer: 0,
});
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum LeafKind {
Token,
Trivia,
Raw,
Error,
}
#[derive(Debug, Clone)]
enum NodeBody {
Element { kind: String },
Leaf { kind: LeafKind, text: TextCrdt },
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct SortKey {
frac: Vec<u8>,
peer: u64,
}
#[derive(Debug, Clone)]
struct NodeRecord {
parent: Option<TreeNodeId>,
sort: SortKey,
sort_stamp: TreeOpId,
body: NodeBody,
tomb: Option<TreeOpId>,
text_head: TreeOpId,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum NodeSeed {
Element {
kind: String,
},
Leaf {
kind: LeafKind,
text: String,
},
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
enum TreeOpKind {
CreateNode {
id: TreeNodeId,
parent: TreeNodeId,
sort: SortKey,
seed: NodeSeed,
},
Tombstone {
node: TreeNodeId,
},
Reorder {
node: TreeNodeId,
sort: SortKey,
},
LeafEdit {
node: TreeNodeId,
prev: TreeOpId,
ops: Vec<TextOp>,
},
SplitLeaf {
node: TreeNodeId,
new: TreeNodeId,
sort: SortKey,
at_char: usize,
prev: TreeOpId,
},
MergeLeaves {
left: TreeNodeId,
right: TreeNodeId,
prev_left: TreeOpId,
prev_right: TreeOpId,
},
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeOp {
pub id: TreeOpId,
kind: TreeOpKind,
}
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeUpdate {
pub ops: Vec<TreeOp>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct DotRange {
contiguous: u64,
sparse: BTreeSet<u64>,
}
impl DotRange {
fn contains(&self, counter: u64) -> bool {
counter <= self.contiguous || self.sparse.contains(&counter)
}
fn observe(&mut self, counter: u64) {
if counter <= self.contiguous {
return;
}
self.sparse.insert(counter);
while self.sparse.remove(&(self.contiguous + 1)) {
self.contiguous += 1;
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TreeVersionFrontier {
dots: BTreeMap<u64, DotRange>,
}
impl TreeVersionFrontier {
pub fn contains(&self, id: TreeOpId) -> bool {
self.dots
.get(&id.peer)
.is_some_and(|r| r.contains(id.counter))
}
fn observe(&mut self, id: TreeOpId) {
self.dots.entry(id.peer).or_default().observe(id.counter);
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TreeError {
NotFound,
NotLeaf,
NonCharBoundary,
NotAdjacent,
}
impl fmt::Display for TreeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let msg = match self {
TreeError::NotFound => "node not found",
TreeError::NotLeaf => "node is not a leaf",
TreeError::NonCharBoundary => "offset out of range or not on a char boundary",
TreeError::NotAdjacent => "leaves are not adjacent live siblings",
};
f.write_str(msg)
}
}
impl std::error::Error for TreeError {}
pub type Result<T> = std::result::Result<T, TreeError>;
#[derive(Debug, Clone)]
pub struct LosslessTreeCrdt {
peer: u64,
counter: u64,
nodes: HashMap<TreeNodeId, NodeRecord>,
frontier: TreeVersionFrontier,
log: Vec<TreeOp>,
buffered: Vec<TreeOp>,
}
impl LosslessTreeCrdt {
pub fn new(peer: u64) -> Self {
let mut nodes = HashMap::new();
nodes.insert(
TreeNodeId::ROOT,
NodeRecord {
parent: None,
sort: SortKey {
frac: Vec::new(),
peer: 0,
},
sort_stamp: TreeOpId {
counter: 0,
peer: 0,
},
body: NodeBody::Element {
kind: "root".to_string(),
},
tomb: None,
text_head: TreeOpId {
counter: 0,
peer: 0,
},
},
);
Self {
peer,
counter: 0,
nodes,
frontier: TreeVersionFrontier::default(),
log: Vec::new(),
buffered: Vec::new(),
}
}
pub fn fork(&self, peer: u64) -> Self {
Self {
peer,
counter: self.counter,
nodes: self.nodes.clone(),
frontier: self.frontier.clone(),
log: self.log.clone(),
buffered: self.buffered.clone(),
}
}
fn next_op_id(&mut self) -> TreeOpId {
self.counter += 1;
TreeOpId {
counter: self.counter,
peer: self.peer,
}
}
fn live_children(&self, parent: TreeNodeId) -> Vec<TreeNodeId> {
let mut kids: Vec<(&TreeNodeId, &SortKey)> = self
.nodes
.iter()
.filter(|(_, r)| r.parent == Some(parent) && r.tomb.is_none())
.map(|(id, r)| (id, &r.sort))
.collect();
kids.sort_by(|a, b| a.1.cmp(b.1));
kids.into_iter().map(|(id, _)| *id).collect()
}
pub fn render(&self) -> String {
let mut out = String::new();
self.render_into(TreeNodeId::ROOT, &mut out);
out
}
fn render_into(&self, id: TreeNodeId, out: &mut String) {
let Some(rec) = self.nodes.get(&id) else {
return;
};
match &rec.body {
NodeBody::Leaf { text, .. } => out.push_str(&text.text()),
NodeBody::Element { .. } => {
for child in self.live_children(id) {
self.render_into(child, out);
}
}
}
}
pub fn live_node_count(&self) -> usize {
self.nodes
.iter()
.filter(|(id, r)| **id != TreeNodeId::ROOT && r.tomb.is_none())
.count()
}
pub fn frontier(&self) -> TreeVersionFrontier {
self.frontier.clone()
}
pub fn element_kind(&self, node: TreeNodeId) -> Option<&str> {
match self.nodes.get(&node).map(|r| &r.body) {
Some(NodeBody::Element { kind }) => Some(kind),
_ => None,
}
}
pub fn leaf_kind(&self, node: TreeNodeId) -> Option<LeafKind> {
match self.nodes.get(&node).map(|r| &r.body) {
Some(NodeBody::Leaf { kind, .. }) => Some(*kind),
_ => None,
}
}
pub fn children(&self, parent: TreeNodeId) -> Vec<TreeNodeId> {
self.live_children(parent)
}
pub fn leaf_text(&self, node: TreeNodeId) -> Result<String> {
match self.nodes.get(&node).map(|r| &r.body) {
Some(NodeBody::Leaf { text, .. }) => Ok(text.text()),
Some(NodeBody::Element { .. }) => Err(TreeError::NotLeaf),
None => Err(TreeError::NotFound),
}
}
fn key_after(&self, parent: TreeNodeId, after: Option<TreeNodeId>) -> SortKey {
let order = self.live_children(parent);
let (lo, hi) = match after {
None => (None, order.first().copied()),
Some(a) => {
let idx = order.iter().position(|x| *x == a);
match idx {
Some(i) => (Some(a), order.get(i + 1).copied()),
None => (order.last().copied(), None),
}
}
};
let lo_frac = lo.map(|id| self.nodes[&id].sort.frac.clone());
let hi_frac = hi.map(|id| self.nodes[&id].sort.frac.clone());
SortKey {
frac: key_between(lo_frac.as_deref(), hi_frac.as_deref()),
peer: self.peer,
}
}
pub fn create_node(
&mut self,
parent: TreeNodeId,
after: Option<TreeNodeId>,
seed: NodeSeed,
) -> Result<TreeNodeId> {
if !self.nodes.contains_key(&parent) {
return Err(TreeError::NotFound);
}
let sort = self.key_after(parent, after);
let op_id = self.next_op_id();
let node = TreeNodeId(op_id);
let op = TreeOp {
id: op_id,
kind: TreeOpKind::CreateNode {
id: node,
parent,
sort,
seed,
},
};
self.commit_local(op);
Ok(node)
}
pub fn tombstone_node(&mut self, node: TreeNodeId) -> Result<()> {
if !self.nodes.contains_key(&node) || node == TreeNodeId::ROOT {
return Err(TreeError::NotFound);
}
let op_id = self.next_op_id();
self.commit_local(TreeOp {
id: op_id,
kind: TreeOpKind::Tombstone { node },
});
Ok(())
}
pub fn reorder_child(&mut self, node: TreeNodeId, after: Option<TreeNodeId>) -> Result<()> {
let parent = self
.nodes
.get(&node)
.and_then(|r| r.parent)
.ok_or(TreeError::NotFound)?;
let sort = self.key_after(parent, after);
let op_id = self.next_op_id();
self.commit_local(TreeOp {
id: op_id,
kind: TreeOpKind::Reorder { node, sort },
});
Ok(())
}
pub fn edit_leaf(
&mut self,
node: TreeNodeId,
at_byte: usize,
delete_bytes: usize,
insert: &str,
) -> Result<()> {
let s = self.leaf_text(node)?;
let start = byte_to_char(&s, at_byte).ok_or(TreeError::NonCharBoundary)?;
let end = byte_to_char(&s, at_byte + delete_bytes).ok_or(TreeError::NonCharBoundary)?;
let delete_chars = end - start;
let editor = self.peer;
let ops = {
let text = self.leaf_text_mut(node)?;
*text = text.fork(editor);
let vv = text.version_vector();
for _ in 0..delete_chars {
text.delete(start);
}
text.insert_str(start, insert);
text.delta_since(&vv)
};
let prev = self.nodes[&node].text_head;
let op_id = self.next_op_id();
self.commit_local(TreeOp {
id: op_id,
kind: TreeOpKind::LeafEdit { node, prev, ops },
});
Ok(())
}
pub fn split_leaf(&mut self, node: TreeNodeId, at_byte: usize) -> Result<TreeNodeId> {
let s = self.leaf_text(node)?;
let at_char = byte_to_char(&s, at_byte).ok_or(TreeError::NonCharBoundary)?;
let sort = {
let parent = self.nodes[&node].parent.ok_or(TreeError::NotFound)?;
self.key_after(parent, Some(node))
};
let prev = self.nodes[&node].text_head;
let op_id = self.next_op_id();
let new = TreeNodeId(op_id);
self.commit_local(TreeOp {
id: op_id,
kind: TreeOpKind::SplitLeaf {
node,
new,
sort,
at_char,
prev,
},
});
Ok(new)
}
pub fn merge_adjacent_leaves(&mut self, left: TreeNodeId, right: TreeNodeId) -> Result<()> {
self.leaf_text(left)?;
self.leaf_text(right)?;
let parent = self.nodes[&left].parent.ok_or(TreeError::NotFound)?;
let order = self.live_children(parent);
let adjacent = order
.iter()
.position(|x| *x == left)
.and_then(|i| order.get(i + 1))
.is_some_and(|nxt| *nxt == right);
if !adjacent {
return Err(TreeError::NotAdjacent);
}
let prev_left = self.nodes[&left].text_head;
let prev_right = self.nodes[&right].text_head;
let op_id = self.next_op_id();
self.commit_local(TreeOp {
id: op_id,
kind: TreeOpKind::MergeLeaves {
left,
right,
prev_left,
prev_right,
},
});
Ok(())
}
pub fn diff(&self, their: &TreeVersionFrontier) -> TreeUpdate {
let mut ops: Vec<TreeOp> = self
.log
.iter()
.filter(|op| !their.contains(op.id))
.cloned()
.collect();
ops.sort_by_key(|op| (op.id.counter, op.id.peer));
TreeUpdate { ops }
}
pub fn apply_update(&mut self, update: &TreeUpdate) {
for op in &update.ops {
self.counter = self.counter.max(op.id.counter);
if self.frontier.contains(op.id) {
continue;
}
self.buffered.push(op.clone());
}
self.drain_buffered();
}
fn drain_buffered(&mut self) {
loop {
let mut progressed = false;
let pending = std::mem::take(&mut self.buffered);
for op in pending {
if self.frontier.contains(op.id) {
continue;
}
if self.dependencies_ready(&op) {
self.apply_op(&op);
self.record(op);
progressed = true;
} else {
self.buffered.push(op);
}
}
if !progressed {
break;
}
}
}
fn dependencies_ready(&self, op: &TreeOp) -> bool {
match &op.kind {
TreeOpKind::CreateNode { parent, .. } => self.nodes.contains_key(parent),
TreeOpKind::Tombstone { node } | TreeOpKind::Reorder { node, .. } => {
self.nodes.contains_key(node)
}
TreeOpKind::LeafEdit { node, prev, .. } | TreeOpKind::SplitLeaf { node, prev, .. } => {
self.nodes.contains_key(node) && self.frontier.contains(*prev)
}
TreeOpKind::MergeLeaves {
left,
right,
prev_left,
prev_right,
} => {
self.nodes.contains_key(left)
&& self.nodes.contains_key(right)
&& self.frontier.contains(*prev_left)
&& self.frontier.contains(*prev_right)
}
}
}
fn commit_local(&mut self, op: TreeOp) {
self.apply_op(&op);
self.record(op);
}
fn record(&mut self, op: TreeOp) {
self.frontier.observe(op.id);
self.log.push(op);
}
fn apply_op(&mut self, op: &TreeOp) {
match &op.kind {
TreeOpKind::CreateNode {
id,
parent,
sort,
seed,
} => {
if self.nodes.contains_key(id) {
return;
}
let body = match seed {
NodeSeed::Element { kind } => NodeBody::Element { kind: kind.clone() },
NodeSeed::Leaf { kind, text } => NodeBody::Leaf {
kind: *kind,
text: TextCrdt::from_str(id.0.peer, text),
},
};
self.nodes.insert(
*id,
NodeRecord {
parent: Some(*parent),
sort: sort.clone(),
sort_stamp: op.id,
body,
tomb: None,
text_head: op.id,
},
);
}
TreeOpKind::Tombstone { node } => {
if let Some(rec) = self.nodes.get_mut(node) {
rec.tomb = Some(match rec.tomb {
Some(existing) => existing.min(op.id),
None => op.id,
});
}
}
TreeOpKind::Reorder { node, sort } => {
if let Some(rec) = self.nodes.get_mut(node)
&& op.id > rec.sort_stamp
{
rec.sort = sort.clone();
rec.sort_stamp = op.id;
}
}
TreeOpKind::LeafEdit { node, ops, .. } => {
if let Some(rec) = self.nodes.get_mut(node)
&& let NodeBody::Leaf { text, .. } = &mut rec.body
{
text.apply_delta(ops);
rec.text_head = op.id;
}
}
TreeOpKind::SplitLeaf {
node,
new,
sort,
at_char,
..
} => self.apply_split(*node, *new, sort.clone(), *at_char, op.id),
TreeOpKind::MergeLeaves { left, right, .. } => self.apply_merge(*left, *right, op.id),
}
}
fn apply_split(
&mut self,
node: TreeNodeId,
new: TreeNodeId,
sort: SortKey,
at_char: usize,
op_id: TreeOpId,
) {
let Some(rec) = self.nodes.get(&node) else {
return;
};
let NodeBody::Leaf { kind, text } = &rec.body else {
return;
};
let kind = *kind;
let parent = rec.parent;
let chars: Vec<char> = text.text().chars().collect();
let clamp = at_char.min(chars.len());
let head: String = chars[..clamp].iter().collect();
let tail: String = chars[clamp..].iter().collect();
if let Some(rec) = self.nodes.get_mut(&node) {
rec.body = NodeBody::Leaf {
kind,
text: TextCrdt::from_str(node.0.peer, &head),
};
rec.text_head = op_id;
}
self.nodes.entry(new).or_insert(NodeRecord {
parent,
sort,
sort_stamp: op_id,
body: NodeBody::Leaf {
kind,
text: TextCrdt::from_str(new.0.peer, &tail),
},
tomb: None,
text_head: op_id,
});
}
fn apply_merge(&mut self, left: TreeNodeId, right: TreeNodeId, op_id: TreeOpId) {
let (Some(l), Some(r)) = (self.nodes.get(&left), self.nodes.get(&right)) else {
return;
};
let (NodeBody::Leaf { kind, text: lt }, NodeBody::Leaf { text: rt, .. }) =
(&l.body, &r.body)
else {
return;
};
let kind = *kind;
let combined = format!("{}{}", lt.text(), rt.text());
if let Some(rec) = self.nodes.get_mut(&left) {
rec.body = NodeBody::Leaf {
kind,
text: TextCrdt::from_str(left.0.peer, &combined),
};
rec.text_head = op_id;
}
if let Some(rec) = self.nodes.get_mut(&right) {
rec.tomb = Some(match rec.tomb {
Some(existing) => existing.min(op_id),
None => op_id,
});
}
}
fn leaf_text_mut(&mut self, node: TreeNodeId) -> Result<&mut TextCrdt> {
match self.nodes.get_mut(&node).map(|r| &mut r.body) {
Some(NodeBody::Leaf { text, .. }) => Ok(text),
Some(NodeBody::Element { .. }) => Err(TreeError::NotLeaf),
None => Err(TreeError::NotFound),
}
}
}
fn byte_to_char(s: &str, byte: usize) -> Option<usize> {
if byte > s.len() || !s.is_char_boundary(byte) {
return None;
}
Some(s[..byte].chars().count())
}
fn key_between(lo: Option<&[u8]>, hi: Option<&[u8]>) -> Vec<u8> {
let mut result = Vec::new();
let mut i = 0usize;
let cap = lo.map_or(0, |l| l.len()) + hi.map_or(0, |h| h.len()) + 2;
while i <= cap {
let a: u16 = lo.and_then(|l| l.get(i)).map_or(0, |&d| d as u16);
let b: u16 = match hi {
Some(h) => h.get(i).map_or(0, |&d| d as u16),
None => 256,
};
if a + 1 < b {
result.push(((a + b) / 2) as u8);
return result;
}
result.push(a as u8);
i += 1;
if a < b {
let lo_tail: Vec<u8> = lo
.map(|l| l.get(i..).unwrap_or(&[]).to_vec())
.unwrap_or_default();
result.extend(key_between(Some(&lo_tail), None));
return result;
}
}
result.push(128);
result
}
#[cfg(test)]
mod tests {
use super::*;
fn elem(kind: &str) -> NodeSeed {
NodeSeed::Element {
kind: kind.to_string(),
}
}
fn leaf(kind: LeafKind, text: &str) -> NodeSeed {
NodeSeed::Leaf {
kind,
text: text.to_string(),
}
}
fn sample() -> (LosslessTreeCrdt, TreeNodeId, [TreeNodeId; 3]) {
let mut t = LosslessTreeCrdt::new(1);
let para = t.create_node(TreeNodeId::ROOT, None, elem("para")).unwrap();
let a = t
.create_node(para, None, leaf(LeafKind::Token, "# "))
.unwrap();
let b = t
.create_node(para, Some(a), leaf(LeafKind::Raw, "héllo"))
.unwrap();
let c = t
.create_node(para, Some(b), leaf(LeafKind::Trivia, "\n"))
.unwrap();
(t, para, [a, b, c])
}
#[test]
fn render_is_exact_concatenation_including_multibyte() {
let (t, _para, _) = sample();
assert_eq!(t.render(), "# héllo\n");
assert_eq!(t.live_node_count(), 4); }
#[test]
fn edit_leaf_at_byte_offset_into_multibyte_text() {
let (mut t, _para, [_, b, _]) = sample();
t.edit_leaf(b, 3, 0, "X").unwrap();
assert_eq!(t.render(), "# héXllo\n");
}
#[test]
fn edit_leaf_rejects_non_char_boundary() {
let (mut t, _para, [_, b, _]) = sample();
assert_eq!(t.edit_leaf(b, 2, 0, "X"), Err(TreeError::NonCharBoundary));
}
#[test]
fn split_then_merge_preserves_render() {
let (mut t, _para, [_, b, _]) = sample();
let before = t.render();
let n = t.live_node_count();
let new = t.split_leaf(b, 3).unwrap(); assert_eq!(t.render(), before, "split preserves render");
assert_eq!(t.live_node_count(), n + 1);
t.merge_adjacent_leaves(b, new).unwrap();
assert_eq!(t.render(), before, "merge restores render");
assert_eq!(t.live_node_count(), n);
}
#[test]
fn reorder_child_changes_order_only() {
let (mut t, para, [a, b, c]) = sample();
t.reorder_child(c, None).unwrap(); assert_eq!(t.render(), "\n# héllo");
assert_eq!(t.live_children(para), vec![c, a, b]);
}
#[test]
fn diff_apply_converges_two_replicas() {
let (mut a, para, [_, _, _]) = sample();
let mut b = a.fork(2);
a.edit_leaf(a.live_children(para)[1], 0, 0, "!").unwrap();
b.create_node(para, None, leaf(LeafKind::Trivia, ">> "))
.unwrap();
let a_to_b = a.diff(&b.frontier());
let b_to_a = b.diff(&a.frontier());
a.apply_update(&b_to_a);
b.apply_update(&a_to_b);
assert_eq!(a.render(), b.render(), "converged");
}
#[test]
fn non_contiguous_delivery_leaves_a_recoverable_hole() {
let (mut a, para, _) = sample();
let mut b = a.fork(2);
a.create_node(para, None, leaf(LeafKind::Trivia, "1"))
.unwrap();
a.create_node(para, None, leaf(LeafKind::Trivia, "2"))
.unwrap();
a.create_node(para, None, leaf(LeafKind::Trivia, "3"))
.unwrap();
let update = a.diff(&b.frontier());
assert_eq!(update.ops.len(), 3);
let holed = TreeUpdate {
ops: vec![update.ops[0].clone(), update.ops[2].clone()],
};
b.apply_update(&holed);
assert_ne!(a.render(), b.render(), "b is missing the held-back op");
let repair = a.diff(&b.frontier());
assert_eq!(repair.ops.len(), 1, "only the hole is resent");
b.apply_update(&repair);
assert_eq!(a.render(), b.render(), "converged after repair");
}
}