use either::Either;
use enum_as_inner::EnumAsInner;
use fractional_index::FractionalIndex;
use fxhash::{FxHashMap, FxHashSet};
use itertools::Itertools;
use loro_common::{
ContainerID, IdFull, IdLp, LoroError, LoroResult, LoroTreeError, LoroValue, PeerID, TreeID,
DELETED_TREE_ROOT, ID,
};
use rand::SeedableRng;
use rle::HasLength;
use serde::Serialize;
use std::collections::VecDeque;
use std::fmt::Debug;
use std::ops::{Deref, DerefMut};
use std::sync::Weak;
use super::{ApplyLocalOpReturn, ContainerState, DiffApplyContext};
use crate::configure::Configure;
use crate::container::idx::ContainerIdx;
use crate::delta::{TreeDiff, TreeDiffItem, TreeExternalDiff};
use crate::diff_calc::DiffMode;
use crate::encoding::{EncodeMode, StateSnapshotDecodeContext, StateSnapshotEncoder};
use crate::event::InternalDiff;
use crate::op::Op;
use crate::{
container::tree::tree_op::TreeOp,
delta::TreeInternalDiff,
event::{Diff, Index},
op::RawOp,
};
use crate::{DocState, LoroDocInner};
#[derive(Clone, Debug, EnumAsInner)]
pub enum TreeFractionalIndexConfigInner {
GenerateFractionalIndex {
jitter: u8,
rng: Box<rand::rngs::StdRng>,
},
MoveDisabled,
}
#[derive(Debug, Clone)]
pub struct TreeState {
idx: ContainerIdx,
trees: FxHashMap<TreeID, TreeStateNode>,
children: TreeChildrenCache,
fractional_index_config: TreeFractionalIndexConfigInner,
peer_id: PeerID,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct TreeStateNode {
pub parent: TreeParentId,
pub position: Option<FractionalIndex>,
pub last_move_op: IdFull,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct NodePosition {
pub(crate) position: FractionalIndex,
pub(crate) idlp: IdLp,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumAsInner, Serialize)]
pub enum TreeParentId {
Node(TreeID),
Root,
Deleted,
Unexist,
}
impl From<Option<TreeID>> for TreeParentId {
fn from(id: Option<TreeID>) -> Self {
match id {
Some(id) => {
if TreeID::is_deleted_root(&id) {
TreeParentId::Deleted
} else {
TreeParentId::Node(id)
}
}
None => TreeParentId::Root,
}
}
}
impl From<TreeID> for TreeParentId {
fn from(id: TreeID) -> Self {
if TreeID::is_deleted_root(&id) {
TreeParentId::Deleted
} else {
TreeParentId::Node(id)
}
}
}
impl From<&TreeID> for TreeParentId {
fn from(id: &TreeID) -> Self {
if TreeID::is_deleted_root(id) {
TreeParentId::Deleted
} else {
TreeParentId::Node(*id)
}
}
}
impl TreeParentId {
pub fn tree_id(&self) -> Option<TreeID> {
match self {
TreeParentId::Node(id) => Some(*id),
TreeParentId::Root => None,
TreeParentId::Deleted => Some(DELETED_TREE_ROOT),
TreeParentId::Unexist => unreachable!(),
}
}
}
#[derive(Debug, Clone)]
enum NodeChildren {
Vec(Vec<(NodePosition, TreeID)>),
BTree(btree::ChildTree),
}
impl Default for NodeChildren {
fn default() -> Self {
NodeChildren::Vec(vec![])
}
}
impl NodeChildren {
const MAX_SIZE_FOR_ARRAY: usize = 16;
fn get_index_by_child_id(&self, target: &TreeID) -> Option<usize> {
match self {
NodeChildren::Vec(v) => v.iter().position(|(_, id)| id == target),
NodeChildren::BTree(btree) => btree.id_to_index(target),
}
}
fn get_last_insert_index_by_position(
&self,
node_position: &NodePosition,
) -> Result<usize, usize> {
match self {
NodeChildren::Vec(v) => v.binary_search_by_key(&node_position, |x| &x.0),
NodeChildren::BTree(btree) => btree.get_index_by_node_position(node_position),
}
}
fn get_node_position_at(&self, pos: usize) -> Option<&NodePosition> {
match self {
NodeChildren::Vec(v) => v.get(pos).map(|x| &x.0),
NodeChildren::BTree(btree) => btree.get_elem_at(pos).map(|x| x.pos.as_ref()),
}
}
fn get_elem_at(&self, pos: usize) -> Option<(&NodePosition, &TreeID)> {
match self {
NodeChildren::Vec(v) => v.get(pos).map(|x| (&x.0, &x.1)),
NodeChildren::BTree(btree) => btree.get_elem_at(pos).map(|x| (x.pos.as_ref(), &x.id)),
}
}
fn generate_fi_at(&self, pos: usize, target: &TreeID) -> FractionalIndexGenResult {
let mut reset_ids = vec![];
let mut left = None;
let mut next_right = None;
{
let mut right = None;
let children_num = self.len();
if children_num == 0 {
return FractionalIndexGenResult::Ok(FractionalIndex::default());
}
if pos > 0 {
left = self.get_node_position_at(pos - 1);
}
if pos < children_num {
right = self.get_elem_at(pos);
}
let left_fi = left.map(|x| &x.position);
if let Some(left_fi) = left_fi {
if Some(left_fi) == right.map(|x| &x.0.position) {
reset_ids.push(*right.unwrap().1);
for i in (pos + 1)..children_num {
let this_position =
self.get_node_position_at(i).map(|x| &x.position).unwrap();
if this_position == left_fi {
reset_ids.push(*self.get_elem_at(i).unwrap().1);
} else {
next_right = Some(this_position.clone());
break;
}
}
}
}
if reset_ids.is_empty() {
return FractionalIndexGenResult::Ok(
FractionalIndex::new(left.map(|x| &x.position), right.map(|x| &x.0.position))
.unwrap(),
);
}
}
let positions = FractionalIndex::generate_n_evenly(
left.map(|x| &x.position),
next_right.as_ref(),
reset_ids.len() + 1,
)
.unwrap();
FractionalIndexGenResult::Rearrange(
Some(*target)
.into_iter()
.chain(reset_ids)
.zip(positions)
.collect(),
)
}
fn get_id_at(&self, pos: usize) -> Option<TreeID> {
match self {
NodeChildren::Vec(v) => v.get(pos).map(|x| x.1),
NodeChildren::BTree(btree) => btree.get_elem_at(pos).map(|x| x.id),
}
}
fn delete_child(&mut self, target: &TreeID) {
match self {
NodeChildren::Vec(v) => {
v.retain(|(_, id)| id != target);
}
NodeChildren::BTree(v) => {
v.delete_child(target);
}
}
}
fn upgrade(&mut self) {
match self {
NodeChildren::Vec(v) => {
let mut btree = btree::ChildTree::new();
for (pos, id) in v.drain(..) {
btree.insert_child(pos, id);
}
*self = NodeChildren::BTree(btree);
}
NodeChildren::BTree(_) => unreachable!(),
}
}
fn insert_child(&mut self, pos: NodePosition, id: TreeID) {
match self {
NodeChildren::Vec(v) => {
if v.len() >= Self::MAX_SIZE_FOR_ARRAY {
self.upgrade();
return self.insert_child(pos, id);
}
let r = v.binary_search_by(|(target, _)| target.cmp(&pos));
match r {
Ok(_) => unreachable!(),
Err(i) => {
v.insert(i, (pos, id));
}
}
}
NodeChildren::BTree(v) => {
v.insert_child(pos, id);
}
}
}
fn push_child_in_order(&mut self, pos: NodePosition, id: TreeID) {
match self {
NodeChildren::Vec(v) => {
if v.len() >= Self::MAX_SIZE_FOR_ARRAY {
self.upgrade();
return self.push_child_in_order(pos, id);
}
if let Some(last) = v.last() {
assert!(last.0 < pos);
}
v.push((pos, id));
}
NodeChildren::BTree(v) => {
v.push_child_in_order(pos, id);
}
}
}
fn len(&self) -> usize {
match self {
NodeChildren::Vec(v) => v.len(),
NodeChildren::BTree(v) => v.len(),
}
}
fn has_child(&self, node_position: &NodePosition) -> bool {
match self {
NodeChildren::Vec(v) => v
.binary_search_by(|(target, _)| target.cmp(node_position))
.is_ok(),
NodeChildren::BTree(v) => v.has_child(node_position),
}
}
fn iter(&self) -> impl Iterator<Item = (&NodePosition, &TreeID)> {
match self {
NodeChildren::Vec(v) => Either::Left(v.iter().map(|x| (&x.0, &x.1))),
NodeChildren::BTree(t) => Either::Right(t.iter()),
}
}
}
#[derive(Clone, Default)]
struct TreeChildrenCache(FxHashMap<TreeParentId, NodeChildren>);
mod btree {
use std::{cmp::Ordering, ops::Range, sync::Arc};
use fxhash::FxHashMap;
use generic_btree::{
rle::{CanRemove, HasLength, Mergeable, Sliceable, TryInsert},
BTree, BTreeTrait, Cursor, FindResult, LeafIndex, LengthFinder, Query, UseLengthFinder,
};
use loro_common::TreeID;
use super::NodePosition;
struct ChildTreeTrait;
#[derive(Debug, Clone)]
pub(super) struct ChildTree {
tree: BTree<ChildTreeTrait>,
id_to_leaf_index: FxHashMap<TreeID, LeafIndex>,
}
impl ChildTree {
pub(super) fn new() -> Self {
Self {
tree: BTree::new(),
id_to_leaf_index: FxHashMap::default(),
}
}
pub(super) fn insert_child(&mut self, pos: NodePosition, id: TreeID) {
let (c, _) = self.tree.insert::<KeyQuery>(
&pos,
Elem {
pos: Arc::new(pos.clone()),
id,
},
);
self.id_to_leaf_index.insert(id, c.leaf);
}
pub(super) fn push_child_in_order(&mut self, pos: NodePosition, id: TreeID) {
let c = self.tree.push(Elem {
pos: Arc::new(pos.clone()),
id,
});
self.id_to_leaf_index.insert(id, c.leaf);
}
pub(super) fn delete_child(&mut self, id: &TreeID) {
if let Some(leaf) = self.id_to_leaf_index.remove(id) {
self.tree.remove_leaf(Cursor { leaf, offset: 0 });
} else {
panic!("The id is not in the tree");
}
}
pub(super) fn has_child(&self, pos: &NodePosition) -> bool {
match self.tree.query::<KeyQuery>(pos) {
Some(r) => r.found,
None => false,
}
}
pub(super) fn iter(&self) -> impl Iterator<Item = (&NodePosition, &TreeID)> {
self.tree.iter().map(|x| (&*x.pos, &x.id))
}
pub(super) fn len(&self) -> usize {
self.tree.root_cache().len
}
pub(super) fn get_elem_at(&self, pos: usize) -> Option<&Elem> {
let result = self.tree.query::<LengthFinder>(&pos)?;
if !result.found {
return None;
}
self.tree.get_elem(result.leaf())
}
pub(super) fn id_to_index(&self, id: &TreeID) -> Option<usize> {
let leaf_index = self.id_to_leaf_index.get(id)?;
let mut ans = 0;
self.tree.visit_previous_caches(
Cursor {
leaf: *leaf_index,
offset: 0,
},
|prev| match prev {
generic_btree::PreviousCache::NodeCache(c) => {
ans += c.len;
}
generic_btree::PreviousCache::PrevSiblingElem(_) => {
ans += 1;
}
generic_btree::PreviousCache::ThisElemAndOffset { .. } => {}
},
);
Some(ans)
}
pub(super) fn get_index_by_node_position(
&self,
node_position: &NodePosition,
) -> Result<usize, usize> {
let Some(res) = self.tree.query::<KeyQuery>(node_position) else {
return Ok(0);
};
let mut ans = 0;
self.tree
.visit_previous_caches(res.cursor, |prev| match prev {
generic_btree::PreviousCache::NodeCache(c) => {
ans += c.len;
}
generic_btree::PreviousCache::PrevSiblingElem(_) => {
ans += 1;
}
generic_btree::PreviousCache::ThisElemAndOffset { elem: _, offset } => {
ans += offset;
}
});
if res.found {
Ok(ans)
} else {
Err(ans)
}
}
}
#[derive(Clone, Debug)]
pub(super) struct Elem {
pub(super) pos: Arc<NodePosition>,
pub(super) id: TreeID,
}
impl Mergeable for Elem {
fn can_merge(&self, _rhs: &Self) -> bool {
false
}
fn merge_right(&mut self, _rhs: &Self) {
unreachable!()
}
fn merge_left(&mut self, _left: &Self) {
unreachable!()
}
}
impl HasLength for Elem {
fn rle_len(&self) -> usize {
1
}
}
impl Sliceable for Elem {
fn _slice(&self, range: std::ops::Range<usize>) -> Self {
assert!(range.len() == 1);
self.clone()
}
}
impl CanRemove for Elem {
fn can_remove(&self) -> bool {
false
}
}
impl TryInsert for Elem {
fn try_insert(&mut self, _pos: usize, elem: Self) -> Result<(), Self>
where
Self: Sized,
{
Err(elem)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct Cache {
range: Option<Range<Arc<NodePosition>>>,
len: usize,
}
impl BTreeTrait for ChildTreeTrait {
type Elem = Elem;
type Cache = Cache;
type CacheDiff = ();
const USE_DIFF: bool = false;
fn calc_cache_internal(
cache: &mut Self::Cache,
caches: &[generic_btree::Child<Self>],
) -> Self::CacheDiff {
if caches.is_empty() {
*cache = Default::default();
return;
}
*cache = Cache {
range: Some(
caches[0].cache.range.as_ref().unwrap().start.clone()
..caches
.last()
.unwrap()
.cache
.range
.as_ref()
.unwrap()
.end
.clone(),
),
len: caches.iter().map(|x| x.cache.len).sum(),
};
}
fn apply_cache_diff(_cache: &mut Self::Cache, _diff: &Self::CacheDiff) {
unreachable!()
}
fn merge_cache_diff(_diff1: &mut Self::CacheDiff, _diff2: &Self::CacheDiff) {}
fn get_elem_cache(elem: &Self::Elem) -> Self::Cache {
Cache {
range: Some(elem.pos.clone()..elem.pos.clone()),
len: 1,
}
}
fn new_cache_to_diff(_cache: &Self::Cache) -> Self::CacheDiff {}
fn sub_cache(_cache_lhs: &Self::Cache, _cache_rhs: &Self::Cache) -> Self::CacheDiff {}
}
struct KeyQuery;
impl Query<ChildTreeTrait> for KeyQuery {
type QueryArg = NodePosition;
#[inline(always)]
fn init(_target: &Self::QueryArg) -> Self {
KeyQuery
}
#[inline]
fn find_node(
&mut self,
target: &Self::QueryArg,
caches: &[generic_btree::Child<ChildTreeTrait>],
) -> FindResult {
let result = caches.binary_search_by(|x| {
let range = x.cache.range.as_ref().unwrap();
if target < &range.start {
core::cmp::Ordering::Greater
} else if target > &range.end {
core::cmp::Ordering::Less
} else {
core::cmp::Ordering::Equal
}
});
match result {
Ok(i) => FindResult::new_found(i, 0),
Err(i) => FindResult::new_missing(
i.min(caches.len() - 1),
if i == caches.len() { 1 } else { 0 },
),
}
}
#[inline(always)]
fn confirm_elem(
&mut self,
q: &Self::QueryArg,
elem: &<ChildTreeTrait as BTreeTrait>::Elem,
) -> (usize, bool) {
match q.cmp(&elem.pos) {
Ordering::Less => (0, false),
Ordering::Equal => (0, true),
Ordering::Greater => (1, false),
}
}
}
impl UseLengthFinder<ChildTreeTrait> for ChildTreeTrait {
fn get_len(cache: &<ChildTreeTrait as BTreeTrait>::Cache) -> usize {
cache.len
}
}
}
impl Debug for TreeChildrenCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
writeln!(f, "TreeChildrenCache {{")?;
for (parent, children) in self.0.iter() {
writeln!(f, " {:?}:{{", parent)?;
for (position, id) in children.iter() {
writeln!(f, " {:?} -> {:?}", position, id)?;
}
writeln!(f, " }}")?;
}
writeln!(f, "}}")
}
}
impl Deref for TreeChildrenCache {
type Target = FxHashMap<TreeParentId, NodeChildren>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl DerefMut for TreeChildrenCache {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl NodePosition {
fn new(position: FractionalIndex, idlp: IdLp) -> Self {
Self { position, idlp }
}
}
pub(crate) enum FiIfNotConfigured {
UseJitterZero,
#[allow(unused)]
Throw,
#[allow(unused)]
Zero,
}
impl TreeState {
pub fn new(idx: ContainerIdx, peer_id: PeerID) -> Self {
Self {
idx,
trees: FxHashMap::default(),
children: Default::default(),
fractional_index_config: TreeFractionalIndexConfigInner::GenerateFractionalIndex {
jitter: 0,
rng: Box::new(rand::rngs::StdRng::seed_from_u64(0)),
},
peer_id,
}
}
pub fn mov(
&mut self,
target: TreeID,
parent: TreeParentId,
id: IdFull,
position: Option<FractionalIndex>,
with_check: bool,
) -> Result<(), LoroError> {
if with_check {
if let TreeParentId::Node(parent) = parent {
if !self.trees.contains_key(&parent) {
return Err(LoroTreeError::TreeNodeParentNotFound(parent).into());
}
}
if self.is_ancestor_of(&target, &parent) {
return Err(LoroTreeError::CyclicMoveError.into());
}
}
if let Some(old_parent) = self.trees.get(&target).map(|x| x.parent) {
self.delete_position(&old_parent, &target);
}
let entry = self.children.entry(parent).or_default();
let node_position = NodePosition::new(position.clone().unwrap_or_default(), id.idlp());
debug_assert!(!entry.has_child(&node_position));
entry.insert_child(node_position, target);
self.trees.insert(
target,
TreeStateNode {
parent,
position,
last_move_op: id,
},
);
Ok(())
}
fn _init_push_tree_node_in_order(
&mut self,
target: TreeID,
parent: TreeParentId,
last_move_op: IdFull,
position: Option<FractionalIndex>,
) -> Result<(), LoroError> {
debug_assert!(!self.trees.contains_key(&target));
let entry = self.children.entry(parent).or_default();
let node_position =
NodePosition::new(position.clone().unwrap_or_default(), last_move_op.idlp());
debug_assert!(!entry.has_child(&node_position));
entry.push_child_in_order(node_position, target);
self.trees.insert(
target,
TreeStateNode {
parent,
position,
last_move_op,
},
);
Ok(())
}
#[inline(never)]
fn is_ancestor_of(&self, maybe_ancestor: &TreeID, node_id: &TreeParentId) -> bool {
if !self.trees.contains_key(maybe_ancestor) {
return false;
}
if let TreeParentId::Node(id) = node_id {
if id == maybe_ancestor {
return true;
}
}
match node_id {
TreeParentId::Node(id) => {
let parent = &self.trees.get(id).unwrap().parent;
if parent == node_id {
panic!("is_ancestor_of loop")
}
self.is_ancestor_of(maybe_ancestor, parent)
}
TreeParentId::Deleted | TreeParentId::Root => false,
TreeParentId::Unexist => unreachable!(),
}
}
pub fn parent(&self, target: &TreeID) -> Option<TreeParentId> {
self.trees.get(target).map(|x| x.parent)
}
pub(crate) fn is_node_deleted(&self, target: &TreeID) -> Option<bool> {
match self.trees.get(target).map(|x| x.parent)? {
TreeParentId::Deleted => Some(true),
TreeParentId::Root => Some(false),
TreeParentId::Node(p) => self.is_node_deleted(&p),
TreeParentId::Unexist => unreachable!(),
}
}
pub(crate) fn is_node_unexist(&self, target: &TreeID) -> bool {
!self.trees.contains_key(target)
}
pub(crate) fn tree_nodes(&self) -> Vec<TreeNode> {
self.get_all_tree_nodes_under(TreeParentId::Root)
}
#[allow(unused)]
pub(crate) fn deleted_tree_nodes(&self) -> Vec<TreeNode> {
self.get_all_tree_nodes_under(TreeParentId::Deleted)
}
pub(crate) fn get_all_tree_nodes_under(&self, root: TreeParentId) -> Vec<TreeNode> {
let mut ans = vec![];
let children = self.children.get(&root);
let mut q = children
.map(|x| VecDeque::from_iter(x.iter().enumerate().zip(std::iter::repeat(root))))
.unwrap_or_default();
while let Some(((index, (position, &target)), parent)) = q.pop_front() {
ans.push(TreeNode {
id: target,
parent,
fractional_index: position.position.clone(),
index,
last_move_op: self.trees.get(&target).map(|x| x.last_move_op).unwrap(),
});
if let Some(children) = self.children.get(&TreeParentId::Node(target)) {
q.extend(
children
.iter()
.enumerate()
.map(|(index, (position, this_target))| {
((index, (position, this_target)), TreeParentId::Node(target))
}),
);
}
}
ans
}
pub(crate) fn get_all_hierarchy_nodes_under(
&self,
root: TreeParentId,
) -> Vec<TreeNodeWithChildren> {
let mut ans = vec![];
let Some(children) = self.children.get(&root) else {
return ans;
};
for (index, (position, &target)) in children.iter().enumerate() {
ans.push(TreeNodeWithChildren {
id: target,
parent: root,
fractional_index: position.position.clone(),
index,
children: self.get_all_hierarchy_nodes_under(TreeParentId::Node(target)),
})
}
ans
}
fn bfs_all_alive_nodes_for_fast_snapshot(&self) -> Vec<TreeNode> {
let mut ans = vec![];
self._bfs_all_nodes(TreeParentId::Root, &mut ans);
ans
}
fn bfs_all_deleted_nodes_for_fast_snapshot(&self) -> Vec<TreeNode> {
let mut ans = vec![];
self._bfs_all_nodes(TreeParentId::Deleted, &mut ans);
ans
}
fn _bfs_all_nodes(&self, root: TreeParentId, ans: &mut Vec<TreeNode>) {
let children = self.children.get(&root);
if let Some(children) = children {
for (index, (position, target)) in children.iter().enumerate() {
ans.push(TreeNode {
id: *target,
parent: root,
fractional_index: position.position.clone(),
index,
last_move_op: self.trees.get(target).map(|x| x.last_move_op).unwrap(),
});
}
for (_, id) in children.iter() {
self._bfs_all_nodes(TreeParentId::Node(*id), ans);
}
}
}
pub fn nodes(&self) -> Vec<TreeID> {
self.trees.keys().copied().collect::<Vec<_>>()
}
#[cfg(feature = "test_utils")]
pub fn max_counter(&self) -> i32 {
self.trees
.keys()
.filter(|&k| !self.is_node_deleted(k).unwrap())
.map(|k| k.counter)
.max()
.unwrap_or(0)
}
pub fn get_children<'a>(
&'a self,
parent: &TreeParentId,
) -> Option<impl Iterator<Item = TreeID> + 'a> {
self.children.get(parent).map(|x| x.iter().map(|x| *x.1))
}
pub fn children_num(&self, parent: &TreeParentId) -> Option<usize> {
self.children.get(parent).map(|x| x.len())
}
pub fn is_parent(&self, target: &TreeID, parent: &TreeParentId) -> bool {
self.trees.get(target).is_some_and(|x| x.parent == *parent)
}
pub(crate) fn delete_position(&mut self, parent: &TreeParentId, target: &TreeID) {
if let Some(x) = self.children.get_mut(parent) {
x.delete_child(target);
}
}
pub(crate) fn generate_position_at(
&mut self,
target: &TreeID,
parent: &TreeParentId,
index: usize,
cfg: FiIfNotConfigured,
) -> FractionalIndexGenResult {
match &mut self.fractional_index_config {
TreeFractionalIndexConfigInner::GenerateFractionalIndex { jitter, rng } => {
if *jitter == 0 {
self.children
.entry(*parent)
.or_default()
.generate_fi_at(index, target)
} else {
self.children
.entry(*parent)
.or_default()
.generate_fi_at_jitter(index, target, rng.as_mut(), *jitter)
}
}
TreeFractionalIndexConfigInner::MoveDisabled => match cfg {
FiIfNotConfigured::Throw => FractionalIndexGenResult::NotConfigured,
FiIfNotConfigured::UseJitterZero => self
.children
.entry(*parent)
.or_default()
.generate_fi_at(index, target),
FiIfNotConfigured::Zero => FractionalIndexGenResult::Ok(FractionalIndex::default()),
},
}
}
pub(crate) fn is_fractional_index_enabled(&self) -> bool {
!matches!(
self.fractional_index_config,
TreeFractionalIndexConfigInner::MoveDisabled
)
}
pub(crate) fn enable_generate_fractional_index(&mut self, jitter: u8) {
if let TreeFractionalIndexConfigInner::GenerateFractionalIndex {
jitter: old_jitter, ..
} = &mut self.fractional_index_config
{
*old_jitter = jitter;
return;
}
self.fractional_index_config = TreeFractionalIndexConfigInner::GenerateFractionalIndex {
jitter,
rng: Box::new(rand::rngs::StdRng::seed_from_u64(self.peer_id)),
};
}
pub(crate) fn disable_generate_fractional_index(&mut self) {
self.fractional_index_config = TreeFractionalIndexConfigInner::MoveDisabled;
}
pub(crate) fn get_position(&self, target: &TreeID) -> Option<FractionalIndex> {
self.trees.get(target).and_then(|x| x.position.clone())
}
pub(crate) fn get_index_by_tree_id(&self, target: &TreeID) -> Option<usize> {
let parent = self.parent(target)?;
(!parent.is_deleted())
.then(|| {
self.children
.get(&parent)
.and_then(|x| x.get_index_by_child_id(target))
})
.flatten()
}
pub(crate) fn get_index_by_position(
&self,
parent: &TreeParentId,
node_position: &NodePosition,
) -> Option<usize> {
self.children.get(parent).map(|c| {
match c.get_last_insert_index_by_position(node_position) {
Ok(i) => i,
Err(i) => i,
}
})
}
pub(crate) fn get_id_by_index(&self, parent: &TreeParentId, index: usize) -> Option<TreeID> {
(!parent.is_deleted())
.then(|| self.children.get(parent).and_then(|x| x.get_id_at(index)))
.flatten()
}
#[allow(unused)]
pub(crate) fn check_tree_integrity(&self) {
let mut parent_children_map: FxHashMap<TreeParentId, FxHashSet<TreeID>> =
FxHashMap::default();
for (id, node) in self.trees.iter() {
let parent = node.parent;
parent_children_map.entry(parent).or_default().insert(*id);
}
for (parent, children) in parent_children_map.iter() {
let cached_children = self.get_children(parent).expect("parent not found");
let cached_children = cached_children.collect::<FxHashSet<_>>();
if &cached_children != children {
panic!(
"tree integrity broken: children set of node {:?} is not consistent",
parent
);
}
}
}
pub fn is_empty(&self) -> bool {
match self.children.get(&TreeParentId::Root) {
Some(c) => c.len() == 0,
None => true,
}
}
pub fn get_last_move_id(&self, id: &TreeID) -> Option<ID> {
self.trees.get(id).map(|x| x.last_move_op.id())
}
}
pub(crate) enum FractionalIndexGenResult {
Ok(FractionalIndex),
Rearrange(Vec<(TreeID, FractionalIndex)>),
NotConfigured,
}
impl ContainerState for TreeState {
fn container_idx(&self) -> crate::container::idx::ContainerIdx {
self.idx
}
fn estimate_size(&self) -> usize {
self.trees.len() * (std::mem::size_of::<(TreeID, TreeStateNode)>())
}
fn is_state_empty(&self) -> bool {
self.nodes().is_empty()
}
fn apply_diff_and_convert(
&mut self,
diff: crate::event::InternalDiff,
ctx: DiffApplyContext,
) -> Diff {
let need_check = !matches!(ctx.mode, DiffMode::Checkout | DiffMode::Linear);
let mut ans = vec![];
if let InternalDiff::Tree(tree) = &diff {
for diff in tree.diff.iter() {
let last_move_op = diff.last_effective_move_op_id;
let target = diff.target;
match &diff.action {
TreeInternalDiff::Create { parent, position } => {
self.mov(target, *parent, last_move_op, Some(position.clone()), false)
.unwrap();
if let TreeParentId::Node(p) = parent {
if self.is_node_deleted(p).unwrap() {
continue;
}
}
let index = self.get_index_by_tree_id(&target).unwrap();
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Create {
parent: *parent,
index,
position: position.clone(),
},
});
}
TreeInternalDiff::Move { parent, position } => {
let old_parent = self.trees.get(&target).unwrap().parent;
let old_index = self.get_index_by_tree_id(&target);
let was_alive = !self.is_node_deleted(&target).unwrap();
if need_check {
if self
.mov(target, *parent, last_move_op, Some(position.clone()), true)
.is_ok()
{
if self.is_node_deleted(&target).unwrap() {
if was_alive {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Delete {
old_parent,
old_index: old_index.unwrap(),
},
});
}
} else if was_alive {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Move {
parent: *parent,
index: self.get_index_by_tree_id(&target).unwrap(),
position: position.clone(),
old_parent,
old_index: old_index.unwrap(),
},
});
} else {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Create {
parent: *parent,
index: self.get_index_by_tree_id(&target).unwrap(),
position: position.clone(),
},
});
}
}
} else {
self.mov(target, *parent, last_move_op, Some(position.clone()), false)
.unwrap();
if let TreeParentId::Node(p) = parent {
if self.is_node_deleted(p).unwrap() {
continue;
}
}
let index = self.get_index_by_tree_id(&target).unwrap();
match was_alive {
true => {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Move {
parent: *parent,
index,
position: position.clone(),
old_parent,
old_index: old_index.unwrap(),
},
});
}
false => {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Create {
parent: *parent,
index,
position: position.clone(),
},
});
}
}
};
}
TreeInternalDiff::Delete { parent, position } => {
let mut send_event = true;
if self.is_node_deleted(&target).unwrap() {
send_event = false;
}
if send_event {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Delete {
old_parent: self.trees.get(&target).unwrap().parent,
old_index: self.get_index_by_tree_id(&target).unwrap(),
},
});
}
self.mov(target, *parent, last_move_op, position.clone(), false)
.unwrap();
}
TreeInternalDiff::MoveInDelete { parent, position } => {
self.mov(target, *parent, last_move_op, position.clone(), false)
.unwrap();
}
TreeInternalDiff::UnCreate => {
if !self.is_node_deleted(&target).unwrap() {
ans.push(TreeDiffItem {
target,
action: TreeExternalDiff::Delete {
old_parent: self.trees.get(&target).unwrap().parent,
old_index: self.get_index_by_tree_id(&target).unwrap(),
},
});
}
let node = self.trees.remove(&target);
if let Some(node) = node {
if !node.parent.is_deleted() {
self.children
.get_mut(&node.parent)
.unwrap()
.delete_child(&target);
}
}
continue;
}
};
}
}
Diff::Tree(TreeDiff { diff: ans })
}
fn apply_diff(&mut self, diff: InternalDiff, ctx: DiffApplyContext) {
if let InternalDiff::Tree(tree) = &diff {
let need_check = !matches!(ctx.mode, DiffMode::Checkout | DiffMode::Linear);
for diff in tree.diff.iter() {
let last_move_op = diff.last_effective_move_op_id;
let target = diff.target;
match &diff.action {
TreeInternalDiff::Create { parent, position }
| TreeInternalDiff::Move {
parent, position, ..
} => {
if need_check {
self.mov(target, *parent, last_move_op, Some(position.clone()), true)
.unwrap_or_default();
} else {
self.mov(target, *parent, last_move_op, Some(position.clone()), false)
.unwrap();
}
}
TreeInternalDiff::Delete { parent, position } => {
self.mov(target, *parent, last_move_op, position.clone(), false)
.unwrap();
}
TreeInternalDiff::MoveInDelete { parent, position } => {
self.mov(target, *parent, last_move_op, position.clone(), false)
.unwrap();
}
TreeInternalDiff::UnCreate => {
let parent = self.trees.remove(&target);
if let Some(parent) = parent {
if !parent.parent.is_deleted() {
self.children
.get_mut(&parent.parent)
.unwrap()
.delete_child(&target);
}
}
continue;
}
};
}
}
}
fn apply_local_op(&mut self, raw_op: &RawOp, _op: &Op) -> LoroResult<ApplyLocalOpReturn> {
let mut deleted_containers = vec![];
match &raw_op.content {
crate::op::RawOpContent::Tree(tree) => match &**tree {
TreeOp::Create {
target,
parent,
position,
}
| TreeOp::Move {
target,
parent,
position,
} => {
let parent = TreeParentId::from(*parent);
self.mov(
*target,
parent,
raw_op.id_full(),
Some(position.clone()),
true,
)?;
}
TreeOp::Delete { target } => {
let parent = TreeParentId::Deleted;
deleted_containers.push(ContainerID::new_normal(
target.id(),
loro_common::ContainerType::Map,
));
self.mov(*target, parent, raw_op.id_full(), None, true)?;
}
},
_ => unreachable!(),
}
Ok(ApplyLocalOpReturn { deleted_containers })
}
fn to_diff(&mut self, _doc: &Weak<LoroDocInner>) -> Diff {
let mut diffs = vec![];
let Some(roots) = self.children.get(&TreeParentId::Root) else {
return Diff::Tree(TreeDiff { diff: vec![] });
};
let mut q = VecDeque::from_iter(roots.iter());
let mut index = 0;
let mut parent = TreeParentId::Root;
while let Some((position, node)) = q.pop_front() {
let node_parent = self.trees.get(node).unwrap().parent;
if node_parent != parent {
index = 0;
parent = node_parent;
}
let diff = TreeDiffItem {
target: *node,
action: TreeExternalDiff::Create {
parent: node_parent,
index,
position: position.position.clone(),
},
};
index += 1;
diffs.push(diff);
if let Some(children) = self.children.get(&TreeParentId::Node(*node)) {
q.extend(children.iter());
}
}
Diff::Tree(TreeDiff { diff: diffs })
}
fn get_value(&mut self) -> LoroValue {
self.get_all_hierarchy_nodes_under(TreeParentId::Root)
.into_iter()
.map(|x| x.into_value())
.collect::<Vec<_>>()
.into()
}
fn get_child_index(&self, id: &ContainerID) -> Option<Index> {
let id = id.as_normal().unwrap();
let tree_id = TreeID {
peer: *id.0,
counter: *id.1,
};
if !self.trees.contains_key(&tree_id) || self.is_node_deleted(&tree_id).unwrap() {
None
} else {
Some(Index::Node(tree_id))
}
}
fn contains_child(&self, id: &ContainerID) -> bool {
let id = id.as_normal().unwrap();
let tree_id = TreeID {
peer: *id.0,
counter: *id.1,
};
self.trees.contains_key(&tree_id) && !self.is_node_deleted(&tree_id).unwrap()
}
fn get_child_containers(&self) -> Vec<ContainerID> {
self.trees
.keys()
.map(|n| n.associated_meta_container())
.collect_vec()
}
#[doc = " Get a list of ops that can be used to restore the state to the current state"]
fn encode_snapshot(&self, mut encoder: StateSnapshotEncoder) -> Vec<u8> {
for node in self.trees.values() {
if node.last_move_op == IdFull::NONE_ID {
continue;
}
encoder.encode_op(node.last_move_op.idlp().into(), || unimplemented!());
}
Vec::new()
}
#[doc = " Restore the state to the state represented by the ops that exported by `get_snapshot_ops`"]
fn import_from_snapshot_ops(&mut self, ctx: StateSnapshotDecodeContext) -> LoroResult<()> {
assert_eq!(ctx.mode, EncodeMode::OutdatedSnapshot);
for op in ctx.ops {
assert_eq!(op.op.atom_len(), 1);
let content = op.op.content.as_tree().unwrap();
match &**content {
TreeOp::Create {
target,
parent,
position,
}
| TreeOp::Move {
target,
parent,
position,
} => {
let parent = TreeParentId::from(*parent);
self.mov(*target, parent, op.id_full(), Some(position.clone()), false)
.unwrap()
}
TreeOp::Delete { target } => {
let parent = TreeParentId::Deleted;
self.mov(*target, parent, op.id_full(), None, false)
.unwrap()
}
};
}
Ok(())
}
fn fork(&self, _config: &Configure) -> Self {
self.clone()
}
}
#[allow(clippy::ptr_arg)]
pub(crate) fn get_meta_value(nodes: &mut Vec<LoroValue>, state: &mut DocState) {
for node in nodes.iter_mut() {
let map = node.as_map_mut().unwrap().make_mut();
let meta = map.get_mut("meta").unwrap();
let id = meta.as_container().unwrap();
*meta = state.get_container_deep_value(state.arena.register_container(id));
let children = map.get_mut("children").unwrap().as_list_mut().unwrap();
get_meta_value(children.make_mut(), state);
}
}
#[derive(Debug, Clone)]
pub struct TreeNode {
pub id: TreeID,
pub parent: TreeParentId,
pub fractional_index: FractionalIndex,
pub index: usize,
pub last_move_op: IdFull,
}
#[derive(Debug, Clone)]
pub struct TreeNodeWithChildren {
pub id: TreeID,
pub parent: TreeParentId,
pub fractional_index: FractionalIndex,
pub index: usize,
pub children: Vec<TreeNodeWithChildren>,
}
impl TreeNodeWithChildren {
fn into_value(self) -> LoroValue {
let mut t = FxHashMap::default();
t.insert("id".to_string(), self.id.to_string().into());
let p = self
.parent
.tree_id()
.map(|p| p.to_string().into())
.unwrap_or(LoroValue::Null);
t.insert("parent".to_string(), p);
t.insert(
"meta".to_string(),
self.id.associated_meta_container().into(),
);
t.insert("index".to_string(), (self.index as i64).into());
t.insert(
"fractional_index".to_string(),
self.fractional_index.to_string().into(),
);
t.insert(
"children".to_string(),
self.children
.into_iter()
.map(|x| x.into_value())
.collect::<Vec<_>>()
.into(),
);
t.into()
}
}
mod jitter {
use super::{FractionalIndexGenResult, NodeChildren};
use fractional_index::FractionalIndex;
use loro_common::TreeID;
use rand::Rng;
impl NodeChildren {
pub(super) fn generate_fi_at_jitter(
&self,
pos: usize,
target: &TreeID,
rng: &mut impl Rng,
jitter: u8,
) -> FractionalIndexGenResult {
let mut reset_ids = vec![];
let mut left = None;
let mut next_right = None;
{
let mut right = None;
let children_num = self.len();
if children_num == 0 {
return FractionalIndexGenResult::Ok(FractionalIndex::jitter_default(
rng, jitter,
));
}
if pos > 0 {
left = self.get_node_position_at(pos - 1);
}
if pos < children_num {
right = self.get_elem_at(pos);
}
let left_fi = left.map(|x| &x.position);
if let Some(left_fi) = left_fi {
if Some(left_fi) == right.map(|x| &x.0.position) {
reset_ids.push(*right.unwrap().1);
for i in (pos + 1)..children_num {
let this_position = &self.get_node_position_at(i).unwrap().position;
if this_position == left_fi {
reset_ids.push(*self.get_elem_at(i).unwrap().1);
} else {
next_right = Some(this_position.clone());
break;
}
}
}
}
if reset_ids.is_empty() {
return FractionalIndexGenResult::Ok(
FractionalIndex::new_jitter(
left.map(|x| &x.position),
right.map(|x| &x.0.position),
rng,
jitter,
)
.unwrap(),
);
}
}
let positions = FractionalIndex::generate_n_evenly_jitter(
left.map(|x| &x.position),
next_right.as_ref(),
reset_ids.len() + 1,
rng,
jitter,
)
.unwrap();
FractionalIndexGenResult::Rearrange(
Some(*target)
.into_iter()
.chain(reset_ids)
.zip(positions)
.collect(),
)
}
}
}
mod snapshot {
use std::{borrow::Cow, collections::BTreeSet, io::Read};
use fractional_index::FractionalIndex;
use fxhash::FxHashMap;
use itertools::Itertools;
use loro_common::{IdFull, Lamport, PeerID, TreeID};
use serde_columnar::columnar;
use crate::{
encoding::{arena::PositionArena, value_register::ValueRegister},
state::FastStateSnapshot,
};
use super::{TreeNode, TreeParentId, TreeState};
#[columnar(vec, ser, de, iterable)]
#[derive(Debug, Clone)]
struct EncodedTreeNodeId {
#[columnar(strategy = "DeltaRle")]
peer_idx: usize,
#[columnar(strategy = "DeltaRle")]
counter: i32,
}
#[columnar(vec, ser, de, iterable)]
#[derive(Debug, Clone)]
struct EncodedTreeNode {
#[columnar(strategy = "DeltaRle")]
parent_idx_plus_two: usize,
#[columnar(strategy = "DeltaRle")]
last_set_peer_idx: usize,
#[columnar(strategy = "DeltaRle")]
last_set_counter: i32,
#[columnar(strategy = "DeltaRle")]
last_set_lamport_sub_counter: i32,
fractional_index_idx: usize,
}
#[columnar(ser, de)]
struct EncodedTree<'a> {
#[columnar(class = "vec", iter = "EncodedTreeNodeId")]
node_ids: Vec<EncodedTreeNodeId>,
#[columnar(class = "vec", iter = "EncodedTreeNode")]
nodes: Vec<EncodedTreeNode>,
#[columnar(borrow)]
fractional_indexes: Cow<'a, [u8]>,
#[columnar(borrow)]
reserved_has_effect_bool_rle: Cow<'a, [u8]>,
}
fn encode(
state: &TreeState,
input: Vec<TreeNode>,
deleted_nodes: Vec<TreeNode>,
) -> (ValueRegister<PeerID>, EncodedTree) {
let mut peers: ValueRegister<PeerID> = ValueRegister::new();
let mut position_set = BTreeSet::default();
let mut nodes = Vec::with_capacity(input.len());
let mut node_ids = Vec::with_capacity(input.len());
let mut position_register = ValueRegister::new();
let mut id_to_idx = FxHashMap::default();
for node in input.iter().chain(deleted_nodes.iter()) {
position_set.insert(node.fractional_index.clone());
let idx = node_ids.len();
node_ids.push(EncodedTreeNodeId {
peer_idx: peers.register(&node.id.peer),
counter: node.id.counter,
});
id_to_idx.insert(node.id, idx);
}
for p in position_set {
position_register.register(&p);
}
for node in input {
let n = state.trees.get(&node.id).unwrap();
let last_set_id = n.last_move_op;
nodes.push(EncodedTreeNode {
parent_idx_plus_two: match node.parent {
TreeParentId::Deleted => 1,
TreeParentId::Root => 0,
TreeParentId::Node(id) => id_to_idx.get(&id).unwrap() + 2,
TreeParentId::Unexist => unreachable!(),
},
last_set_peer_idx: peers.register(&last_set_id.peer),
last_set_counter: last_set_id.counter,
last_set_lamport_sub_counter: last_set_id.lamport as i32 - last_set_id.counter,
fractional_index_idx: position_register.register(&node.fractional_index),
})
}
for node in deleted_nodes {
let n = state.trees.get(&node.id).unwrap();
let last_set_id = n.last_move_op;
nodes.push(EncodedTreeNode {
parent_idx_plus_two: match node.parent {
TreeParentId::Deleted => 1,
TreeParentId::Root => 0,
TreeParentId::Node(id) => id_to_idx.get(&id).unwrap() + 2,
TreeParentId::Unexist => unreachable!(),
},
last_set_peer_idx: peers.register(&last_set_id.peer),
last_set_counter: last_set_id.counter,
last_set_lamport_sub_counter: last_set_id.lamport as i32 - last_set_id.counter,
fractional_index_idx: position_register.register(&node.fractional_index),
})
}
let position_vec = position_register.unwrap_vec();
let positions = PositionArena::from_positions(position_vec.iter().map(|p| p.as_bytes()));
(
peers,
EncodedTree {
node_ids,
nodes,
fractional_indexes: positions.encode().into(),
reserved_has_effect_bool_rle: vec![].into(),
},
)
}
impl FastStateSnapshot for TreeState {
fn encode_snapshot_fast<W: std::io::prelude::Write>(&mut self, mut w: W) {
let all_alive_nodes = self.bfs_all_alive_nodes_for_fast_snapshot();
let all_deleted_nodes = self.bfs_all_deleted_nodes_for_fast_snapshot();
let (peers, encoded) = encode(self, all_alive_nodes, all_deleted_nodes);
let peers = peers.unwrap_vec();
leb128::write::unsigned(&mut w, peers.len() as u64).unwrap();
for peer in peers {
w.write_all(&peer.to_le_bytes()).unwrap();
}
let vec = serde_columnar::to_vec(&encoded).unwrap();
w.write_all(&vec).unwrap();
}
fn decode_value(_: &[u8]) -> loro_common::LoroResult<(loro_common::LoroValue, &[u8])> {
unreachable!()
}
fn decode_snapshot_fast(
idx: crate::container::idx::ContainerIdx,
(_, mut bytes): (loro_common::LoroValue, &[u8]),
ctx: crate::state::ContainerCreationContext,
) -> loro_common::LoroResult<Self>
where
Self: Sized,
{
let peer_num = leb128::read::unsigned(&mut bytes).unwrap() as usize;
let mut peers = Vec::with_capacity(peer_num);
for _ in 0..peer_num {
let mut buf = [0u8; 8];
bytes.read_exact(&mut buf).unwrap();
peers.push(PeerID::from_le_bytes(buf));
}
let mut tree = TreeState::new(idx, ctx.peer);
let encoded: EncodedTree = serde_columnar::from_bytes(bytes)?;
let fractional_indexes = PositionArena::decode(&encoded.fractional_indexes).unwrap();
let fractional_indexes = fractional_indexes.parse_to_positions();
let node_ids = encoded
.node_ids
.iter()
.map(|x| TreeID::new(peers[x.peer_idx], x.counter))
.collect_vec();
for (node_id, node) in node_ids.iter().zip(encoded.nodes.into_iter()) {
tree._init_push_tree_node_in_order(
*node_id,
match node.parent_idx_plus_two {
0 => TreeParentId::Root,
1 => TreeParentId::Deleted,
n => {
let id = node_ids[n - 2];
TreeParentId::from(Some(id))
}
},
IdFull::new(
peers[node.last_set_peer_idx],
node.last_set_counter,
(node.last_set_lamport_sub_counter + node.last_set_counter) as Lamport,
),
Some(FractionalIndex::from_bytes(
fractional_indexes[node.fractional_index_idx].clone(),
)),
)
.unwrap();
}
Ok(tree)
}
}
#[cfg(test)]
mod test {
use loro_common::LoroValue;
use crate::{
container::idx::ContainerIdx,
state::{ContainerCreationContext, ContainerState},
LoroDoc,
};
use super::*;
#[test]
fn test_tree_state_snapshot() {
let doc = LoroDoc::new();
doc.set_peer_id(0).unwrap();
doc.start_auto_commit();
let tree = doc.get_tree("tree");
tree.enable_fractional_index(0);
let a = tree.create(TreeParentId::Root).unwrap();
let b = tree.create(TreeParentId::Root).unwrap();
let _c = tree.create(TreeParentId::Root).unwrap();
tree.mov(b, TreeParentId::Node(a)).unwrap();
let (bytes, value) = {
let mut doc_state = doc.app_state().try_lock().unwrap();
let tree_state = doc_state.get_tree("tree").unwrap();
let value = tree_state.get_value();
let mut bytes = Vec::new();
tree_state.encode_snapshot_fast(&mut bytes);
(bytes, value)
};
assert!(bytes.len() == 55, "{}", bytes.len());
let mut new_tree_state = TreeState::decode_snapshot_fast(
ContainerIdx::from_index_and_type(0, loro_common::ContainerType::Tree),
(LoroValue::Null, &bytes),
ContainerCreationContext {
configure: &Default::default(),
peer: 0,
},
)
.unwrap();
let mut doc_state = doc.app_state().try_lock().unwrap();
let tree_state = doc_state.get_tree("tree").unwrap();
assert_eq!(&tree_state.trees, &new_tree_state.trees);
let new_v = new_tree_state.get_value();
assert_eq!(value, new_v);
}
}
}