fix: register container id err when using shallow snapshot #822 (#823)

* fix: #822 should register container id correctly

* chore: changeset

Author: zxch3n

.changeset/curvy-maps-sleep.md

@@ -0,0 +1,5 @@
+---
+"loro-crdt": patch
+---
+
+fix: register container id err when using shallow snapshot #823

crates/loro-internal/src/arena.rs

@@ -14,8 +14,8 @@ use crate::{
     LoroValue,
 };
 use append_only_bytes::BytesSlice;
-use rustc_hash::FxHashMap;
 use loro_common::PeerID;
+use rustc_hash::FxHashMap;
 use std::fmt;
 use std::{
     num::NonZeroU16,
@@ -111,10 +111,11 @@ impl ArenaGuards<'_> {
                 if let Some(d) = get_depth(
                     p,
                     &mut self.depth,
-                    &self.parents,
+                    &mut self.parents,
                     &self.parent_resolver,
                     &mut self.container_idx_to_id,
                     &mut self.container_id_to_idx,
+                    &mut self.root_c_idx,
                 ) {
                     self.depth[child.to_index() as usize] = NonZeroU16::new(d.get() + 1);
                 } else {
@@ -263,23 +264,26 @@ impl SharedArena {
 
     #[inline]
     pub fn set_parent(&self, child: ContainerIdx, parent: Option<ContainerIdx>) {
-        let parents = &mut self.inner.parents.lock().unwrap();
+        let mut parents = self.inner.parents.lock().unwrap();
         parents.insert(child, parent);
         let mut depth = self.inner.depth.lock().unwrap();
+        let mut root_c_idx = self.inner.root_c_idx.lock().unwrap();
 
         match parent {
             Some(p) => {
                 // Acquire the two maps as mutable guards so we can lazily register
                 // unknown parents while computing depth.
                 let mut idx_to_id_guard = self.inner.container_idx_to_id.lock().unwrap();
                 let mut id_to_idx_guard = self.inner.container_id_to_idx.lock().unwrap();
+                let parent_resolver_guard = self.inner.parent_resolver.lock().unwrap();
                 if let Some(d) = get_depth(
                     p,
                     &mut depth,
-                    parents,
-                    &self.inner.parent_resolver.lock().unwrap(),
+                    &mut parents,
+                    &parent_resolver_guard,
                     &mut idx_to_id_guard,
                     &mut id_to_idx_guard,
+                    &mut root_c_idx,
                 ) {
                     depth[child.to_index() as usize] = NonZeroU16::new(d.get() + 1);
                 } else {
@@ -564,21 +568,21 @@ impl SharedArena {
 
     // TODO: this can return a u16 directly now, since the depths are always valid
     pub(crate) fn get_depth(&self, container: ContainerIdx) -> Option<NonZeroU16> {
-        {
-            let mut depth_guard = self.inner.depth.lock().unwrap();
-            let parents_guard = self.inner.parents.lock().unwrap();
-            let resolver_guard = self.inner.parent_resolver.lock().unwrap();
-            let mut idx_to_id_guard = self.inner.container_idx_to_id.lock().unwrap();
-            let mut id_to_idx_guard = self.inner.container_id_to_idx.lock().unwrap();
-            get_depth(
-                container,
-                &mut depth_guard,
-                &parents_guard,
-                &resolver_guard,
-                &mut idx_to_id_guard,
-                &mut id_to_idx_guard,
-            )
-        }
+        let mut depth_guard = self.inner.depth.lock().unwrap();
+        let mut parents_guard = self.inner.parents.lock().unwrap();
+        let mut root_c_idx_guard = self.inner.root_c_idx.lock().unwrap();
+        let resolver_guard = self.inner.parent_resolver.lock().unwrap();
+        let mut idx_to_id_guard = self.inner.container_idx_to_id.lock().unwrap();
+        let mut id_to_idx_guard = self.inner.container_id_to_idx.lock().unwrap();
+        get_depth(
+            container,
+            &mut depth_guard,
+            &mut parents_guard,
+            &resolver_guard,
+            &mut idx_to_id_guard,
+            &mut id_to_idx_guard,
+            &mut root_c_idx_guard,
+        )
     }
 
     pub(crate) fn iter_value_slice(
@@ -662,10 +666,11 @@ fn _slice_str(range: Range<usize>, s: &mut StrArena) -> String {
 fn get_depth(
     target: ContainerIdx,
     depth: &mut Vec<Option<NonZeroU16>>,
-    parents: &FxHashMap<ContainerIdx, Option<ContainerIdx>>,
+    parents: &mut FxHashMap<ContainerIdx, Option<ContainerIdx>>,
     parent_resolver: &Option<Arc<ParentResolver>>,
     idx_to_id: &mut Vec<ContainerID>,
     id_to_idx: &mut FxHashMap<ContainerID, ContainerIdx>,
+    root_c_idx: &mut Vec<ContainerIdx>,
 ) -> Option<NonZeroU16> {
     let mut d = depth[target.to_index() as usize];
     if d.is_some() {
@@ -680,6 +685,8 @@ fn get_depth(
             None
         } else if let Some(parent_resolver) = parent_resolver.as_ref() {
             let parent_id = parent_resolver(id.clone())?;
+            let parent_is_root = parent_id.is_root();
+            let mut parent_was_new = false;
             // If the parent is not registered yet, register it lazily instead of unwrapping.
             let parent_idx = if let Some(idx) = id_to_idx.get(&parent_id).copied() {
                 idx
@@ -690,13 +697,27 @@ fn get_depth(
                     ContainerIdx::from_index_and_type(new_index as u32, parent_id.container_type());
                 id_to_idx.insert(parent_id.clone(), new_idx);
                 // Keep depth vector in sync with containers list.
-                if parent_id.is_root() {
+                if parent_is_root {
                     depth.push(NonZeroU16::new(1));
                 } else {
                     depth.push(None);
                 }
+                parent_was_new = true;
                 new_idx
             };
+
+            if parent_is_root {
+                if parent_was_new {
+                    parents.insert(parent_idx, None);
+                    root_c_idx.push(parent_idx);
+                } else {
+                    parents.entry(parent_idx).or_insert(None);
+                }
+                if depth[parent_idx.to_index() as usize].is_none() {
+                    depth[parent_idx.to_index() as usize] = NonZeroU16::new(1);
+                }
+            }
+
             Some(parent_idx)
         } else {
             return None;
@@ -706,7 +727,17 @@ fn get_depth(
     match parent {
         Some(p) => {
             d = NonZeroU16::new(
-                get_depth(p, depth, parents, parent_resolver, idx_to_id, id_to_idx)?.get() + 1,
+                get_depth(
+                    p,
+                    depth,
+                    parents,
+                    parent_resolver,
+                    idx_to_id,
+                    id_to_idx,
+                    root_c_idx,
+                )?
+                .get()
+                    + 1,
             );
             depth[target.to_index() as usize] = d;
         }

crates/loro-wasm/VERSION

@@ -1 +1 @@
-1.7.1
+1.8.1

crates/loro/tests/issue.rs

@@ -1,6 +1,7 @@
 #![allow(unexpected_cfgs)]
-use loro::{LoroDoc, UndoManager};
+use loro::{EncodedBlobMode, ExportMode, LoroDoc, UndoManager};
 use std::sync::{Arc, Mutex};
+use tracing::{trace, trace_span};
 
 #[ctor::ctor]
 fn init() {
@@ -251,3 +252,55 @@ fn undo_tree_mov_between_children() {
     let doc_value_1 = doc.get_deep_value();
     assert_eq!(doc_value_0, doc_value_1);
 }
+
+#[test]
+fn issue_822_tree_shallow_snapshot_roundtrip() {
+    let snapshot_bytes = include_bytes!("./issue_822.bin");
+    let doc = LoroDoc::new();
+    doc.import(snapshot_bytes).expect("import snapshot blob");
+
+    let tree = doc.get_tree("nodes");
+    let tree_before = tree.get_value();
+    let doc_before = doc.get_value();
+
+    let snapshot_meta =
+        LoroDoc::decode_import_blob_meta(snapshot_bytes, false).expect("decode snapshot meta");
+    assert!(snapshot_meta.mode.is_snapshot());
+    let imported_is_shallow = snapshot_meta.mode == EncodedBlobMode::ShallowSnapshot;
+
+    let frontiers = doc.state_frontiers();
+    let shallow_bytes = trace_span!("EXPORT").in_scope(|| {
+        doc.export(ExportMode::shallow_snapshot(&frontiers))
+            .expect("export shallow snapshot")
+    });
+
+    let snapshot_meta_1 = LoroDoc::decode_import_blob_meta(&shallow_bytes, false).unwrap();
+    assert!(matches!(
+        snapshot_meta_1.mode,
+        EncodedBlobMode::ShallowSnapshot
+    ));
+
+    let shallow_meta =
+        LoroDoc::decode_import_blob_meta(&shallow_bytes, false).expect("decode shallow meta");
+    assert_eq!(shallow_meta.mode, EncodedBlobMode::ShallowSnapshot);
+
+    let shallow_doc = LoroDoc::new();
+    trace_span!("FINAL_IMPORT").in_scope(|| {
+        trace!("bytes.len: {}", shallow_bytes.len());
+        shallow_doc
+            .import(&shallow_bytes)
+            .expect("import shallow snapshot");
+    });
+
+    assert!(shallow_doc.is_shallow());
+    assert_eq!(doc.is_shallow(), imported_is_shallow);
+
+    let tree_after = shallow_doc.get_tree("nodes").get_value();
+    let doc_after = shallow_doc.get_value();
+
+    assert_eq!(
+        tree_before, tree_after,
+        "tree shallow value should roundtrip"
+    );
+    assert_eq!(doc_before, doc_after, "doc shallow value should roundtrip");
+}