loro_internal/
arena.rs

1mod str_arena;
2use self::str_arena::StrArena;
3use crate::sync::{Mutex, MutexGuard};
4use crate::{
5    change::Lamport,
6    container::{
7        idx::ContainerIdx,
8        list::list_op::{InnerListOp, ListOp},
9        map::MapSet,
10        ContainerID,
11    },
12    id::Counter,
13    op::{InnerContent, ListSlice, Op, RawOp, RawOpContent, SliceRange},
14    LoroValue,
15};
16use append_only_bytes::BytesSlice;
17use loro_common::PeerID;
18use rustc_hash::FxHashMap;
19use std::fmt;
20use std::{
21    num::NonZeroU16,
22    ops::{Range, RangeBounds},
23    sync::Arc,
24};
25
26pub(crate) struct LoadAllFlag;
27type ParentResolver = dyn Fn(ContainerID) -> Option<ContainerID> + Send + Sync + 'static;
28
29#[derive(Default)]
30struct InnerSharedArena {
31    // The locks should not be exposed outside this file.
32    // It might be better to use RwLock in the future
33    container_idx_to_id: Mutex<Vec<ContainerID>>,
34    // 0 stands for unknown, 1 stands for root containers
35    depth: Mutex<Vec<Option<NonZeroU16>>>,
36    container_id_to_idx: Mutex<FxHashMap<ContainerID, ContainerIdx>>,
37    /// The parent of each container.
38    parents: Mutex<FxHashMap<ContainerIdx, Option<ContainerIdx>>>,
39    values: Mutex<Vec<LoroValue>>,
40    root_c_idx: Mutex<Vec<ContainerIdx>>,
41    str: Arc<Mutex<StrArena>>,
42    /// Optional resolver used when querying parent for a container that has not been registered yet.
43    /// If set, `get_parent` will try this resolver to lazily fetch and register the parent.
44    parent_resolver: Mutex<Option<Arc<ParentResolver>>>,
45}
46
47impl fmt::Debug for InnerSharedArena {
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49        f.debug_struct("InnerSharedArena")
50            .field("container_idx_to_id", &"<Mutex<_>>")
51            .field("depth", &"<Mutex<_>>")
52            .field("container_id_to_idx", &"<Mutex<_>>")
53            .field("parents", &"<Mutex<_>>")
54            .field("values", &"<Mutex<_>>")
55            .field("root_c_idx", &"<Mutex<_>>")
56            .field("str", &"<Arc<Mutex<_>>>")
57            .field("parent_resolver", &"<Mutex<Option<...>>>")
58            .finish()
59    }
60}
61
62/// This is shared between [OpLog] and [AppState].
63///
64#[derive(Debug, Clone)]
65pub struct SharedArena {
66    inner: Arc<InnerSharedArena>,
67}
68
69#[derive(Debug)]
70pub struct StrAllocResult {
71    /// unicode start
72    pub start: usize,
73    /// unicode end
74    pub end: usize,
75}
76
77pub(crate) struct ArenaGuards<'a> {
78    container_id_to_idx: MutexGuard<'a, FxHashMap<ContainerID, ContainerIdx>>,
79    container_idx_to_id: MutexGuard<'a, Vec<ContainerID>>,
80    depth: MutexGuard<'a, Vec<Option<NonZeroU16>>>,
81    parents: MutexGuard<'a, FxHashMap<ContainerIdx, Option<ContainerIdx>>>,
82    root_c_idx: MutexGuard<'a, Vec<ContainerIdx>>,
83    parent_resolver: MutexGuard<'a, Option<Arc<ParentResolver>>>,
84}
85
86impl ArenaGuards<'_> {
87    pub fn register_container(&mut self, id: &ContainerID) -> ContainerIdx {
88        if let Some(&idx) = self.container_id_to_idx.get(id) {
89            return idx;
90        }
91
92        let idx = self.container_idx_to_id.len();
93        self.container_idx_to_id.push(id.clone());
94        let idx = ContainerIdx::from_index_and_type(idx as u32, id.container_type());
95        self.container_id_to_idx.insert(id.clone(), idx);
96        if id.is_root() {
97            self.root_c_idx.push(idx);
98            self.parents.insert(idx, None);
99            self.depth.push(NonZeroU16::new(1));
100        } else {
101            self.depth.push(None);
102        }
103        idx
104    }
105
106    pub fn set_parent(&mut self, child: ContainerIdx, parent: Option<ContainerIdx>) {
107        self.parents.insert(child, parent);
108
109        match parent {
110            Some(p) => {
111                if let Some(d) = get_depth(
112                    p,
113                    &mut self.depth,
114                    &mut self.parents,
115                    &self.parent_resolver,
116                    &mut self.container_idx_to_id,
117                    &mut self.container_id_to_idx,
118                    &mut self.root_c_idx,
119                ) {
120                    self.depth[child.to_index() as usize] = NonZeroU16::new(d.get() + 1);
121                } else {
122                    self.depth[child.to_index() as usize] = None;
123                }
124            }
125            None => {
126                self.depth[child.to_index() as usize] = NonZeroU16::new(1);
127            }
128        }
129    }
130}
131
132impl SharedArena {
133    #[allow(clippy::new_without_default)]
134    pub fn new() -> Self {
135        Self {
136            inner: Arc::new(InnerSharedArena::default()),
137        }
138    }
139
140    pub fn fork(&self) -> Self {
141        Self {
142            inner: Arc::new(InnerSharedArena {
143                container_idx_to_id: Mutex::new(
144                    self.inner.container_idx_to_id.lock().unwrap().clone(),
145                ),
146                depth: Mutex::new(self.inner.depth.lock().unwrap().clone()),
147                container_id_to_idx: Mutex::new(
148                    self.inner.container_id_to_idx.lock().unwrap().clone(),
149                ),
150                parents: Mutex::new(self.inner.parents.lock().unwrap().clone()),
151                values: Mutex::new(self.inner.values.lock().unwrap().clone()),
152                root_c_idx: Mutex::new(self.inner.root_c_idx.lock().unwrap().clone()),
153                str: self.inner.str.clone(),
154                parent_resolver: Mutex::new(self.inner.parent_resolver.lock().unwrap().clone()),
155            }),
156        }
157    }
158
159    pub(crate) fn with_guards(&self, f: impl FnOnce(&mut ArenaGuards)) {
160        let mut guards = self.get_arena_guards();
161        f(&mut guards);
162    }
163
164    fn get_arena_guards(&self) -> ArenaGuards<'_> {
165        ArenaGuards {
166            container_id_to_idx: self.inner.container_id_to_idx.lock().unwrap(),
167            container_idx_to_id: self.inner.container_idx_to_id.lock().unwrap(),
168            depth: self.inner.depth.lock().unwrap(),
169            parents: self.inner.parents.lock().unwrap(),
170            root_c_idx: self.inner.root_c_idx.lock().unwrap(),
171            parent_resolver: self.inner.parent_resolver.lock().unwrap(),
172        }
173    }
174
175    pub fn register_container(&self, id: &ContainerID) -> ContainerIdx {
176        let mut container_id_to_idx = self.inner.container_id_to_idx.lock().unwrap();
177        if let Some(&idx) = container_id_to_idx.get(id) {
178            return idx;
179        }
180
181        let mut container_idx_to_id = self.inner.container_idx_to_id.lock().unwrap();
182        let idx = container_idx_to_id.len();
183        container_idx_to_id.push(id.clone());
184        let idx = ContainerIdx::from_index_and_type(idx as u32, id.container_type());
185        container_id_to_idx.insert(id.clone(), idx);
186        if id.is_root() {
187            self.inner.root_c_idx.lock().unwrap().push(idx);
188            self.inner.parents.lock().unwrap().insert(idx, None);
189            self.inner.depth.lock().unwrap().push(NonZeroU16::new(1));
190        } else {
191            self.inner.depth.lock().unwrap().push(None);
192        }
193        idx
194    }
195
196    pub fn get_container_id(&self, idx: ContainerIdx) -> Option<ContainerID> {
197        let lock = self.inner.container_idx_to_id.lock().unwrap();
198        lock.get(idx.to_index() as usize).cloned()
199    }
200
201    /// Fast map from `ContainerID` to `ContainerIdx` for containers already registered
202    /// in the arena.
203    ///
204    /// Important: This is not an existence check. Absence here does not imply that a
205    /// container does not exist, since registration can be lazy and containers may
206    /// be persisted only in the state KV store until first use.
207    ///
208    /// For existence-aware lookup that consults persisted state and performs lazy
209    /// registration, prefer `DocState::resolve_idx`.
210    pub fn id_to_idx(&self, id: &ContainerID) -> Option<ContainerIdx> {
211        self.inner
212            .container_id_to_idx
213            .lock()
214            .unwrap()
215            .get(id)
216            .copied()
217    }
218
219    #[inline]
220    pub fn idx_to_id(&self, id: ContainerIdx) -> Option<ContainerID> {
221        let lock = self.inner.container_idx_to_id.lock().unwrap();
222        lock.get(id.to_index() as usize).cloned()
223    }
224
225    #[inline]
226    pub fn with_idx_to_id<R>(&self, f: impl FnOnce(&Vec<ContainerID>) -> R) -> R {
227        let lock = self.inner.container_idx_to_id.lock().unwrap();
228        f(&lock)
229    }
230
231    pub fn alloc_str(&self, str: &str) -> StrAllocResult {
232        let mut text_lock = self.inner.str.lock().unwrap();
233        _alloc_str(&mut text_lock, str)
234    }
235
236    /// return slice and unicode index
237    pub fn alloc_str_with_slice(&self, str: &str) -> (BytesSlice, StrAllocResult) {
238        let mut text_lock = self.inner.str.lock().unwrap();
239        _alloc_str_with_slice(&mut text_lock, str)
240    }
241
242    /// alloc str without extra info
243    pub fn alloc_str_fast(&self, bytes: &[u8]) {
244        let mut text_lock = self.inner.str.lock().unwrap();
245        text_lock.alloc(std::str::from_utf8(bytes).unwrap());
246    }
247
248    #[inline]
249    pub fn utf16_len(&self) -> usize {
250        self.inner.str.lock().unwrap().len_utf16()
251    }
252
253    #[inline]
254    pub fn alloc_value(&self, value: LoroValue) -> usize {
255        let mut values_lock = self.inner.values.lock().unwrap();
256        _alloc_value(&mut values_lock, value)
257    }
258
259    #[inline]
260    pub fn alloc_values(&self, values: impl Iterator<Item = LoroValue>) -> std::ops::Range<usize> {
261        let mut values_lock = self.inner.values.lock().unwrap();
262        _alloc_values(&mut values_lock, values)
263    }
264
265    #[inline]
266    pub fn set_parent(&self, child: ContainerIdx, parent: Option<ContainerIdx>) {
267        let mut parents = self.inner.parents.lock().unwrap();
268        parents.insert(child, parent);
269        let mut depth = self.inner.depth.lock().unwrap();
270        let mut root_c_idx = self.inner.root_c_idx.lock().unwrap();
271
272        match parent {
273            Some(p) => {
274                // Acquire the two maps as mutable guards so we can lazily register
275                // unknown parents while computing depth.
276                let mut idx_to_id_guard = self.inner.container_idx_to_id.lock().unwrap();
277                let mut id_to_idx_guard = self.inner.container_id_to_idx.lock().unwrap();
278                let parent_resolver_guard = self.inner.parent_resolver.lock().unwrap();
279                if let Some(d) = get_depth(
280                    p,
281                    &mut depth,
282                    &mut parents,
283                    &parent_resolver_guard,
284                    &mut idx_to_id_guard,
285                    &mut id_to_idx_guard,
286                    &mut root_c_idx,
287                ) {
288                    depth[child.to_index() as usize] = NonZeroU16::new(d.get() + 1);
289                } else {
290                    depth[child.to_index() as usize] = None;
291                }
292            }
293            None => {
294                depth[child.to_index() as usize] = NonZeroU16::new(1);
295            }
296        }
297    }
298
299    pub fn log_hierarchy(&self) {
300        if cfg!(debug_assertions) {
301            for (c, p) in self.inner.parents.lock().unwrap().iter() {
302                tracing::info!(
303                    "container {:?} {:?} {:?}",
304                    c,
305                    self.get_container_id(*c),
306                    p.map(|x| self.get_container_id(x))
307                );
308            }
309        }
310    }
311
312    pub fn log_all_containers(&self) {
313        self.inner
314            .container_id_to_idx
315            .lock()
316            .unwrap()
317            .iter()
318            .for_each(|(id, idx)| {
319                tracing::info!("container {:?} {:?}", id, idx);
320            });
321        self.inner
322            .container_idx_to_id
323            .lock()
324            .unwrap()
325            .iter()
326            .enumerate()
327            .for_each(|(i, id)| {
328                tracing::info!("container {} {:?}", i, id);
329            });
330    }
331
332    pub fn get_parent(&self, child: ContainerIdx) -> Option<ContainerIdx> {
333        if self.get_container_id(child).unwrap().is_root() {
334            // TODO: PERF: we can speed this up by use a special bit in ContainerIdx to indicate
335            // whether the target is a root container
336            return None;
337        }
338
339        // Try fast path first
340        if let Some(p) = self.inner.parents.lock().unwrap().get(&child).copied() {
341            return p;
342        }
343
344        // Fallback: try to resolve parent lazily via the resolver if provided.
345        let resolver = self.inner.parent_resolver.lock().unwrap().clone();
346        if let Some(resolver) = resolver {
347            let child_id = self.get_container_id(child).unwrap();
348            if let Some(parent_id) = resolver(child_id.clone()) {
349                let parent_idx = self.register_container(&parent_id);
350                self.set_parent(child, Some(parent_idx));
351                return Some(parent_idx);
352            }
353        }
354
355        panic!("InternalError: Parent is not registered")
356    }
357
358    /// Call `f` on each ancestor of `container`, including `container` itself.
359    ///
360    /// f(ContainerIdx, is_first)
361    pub fn with_ancestors(&self, container: ContainerIdx, mut f: impl FnMut(ContainerIdx, bool)) {
362        let mut container = Some(container);
363        let mut is_first = true;
364        while let Some(c) = container {
365            f(c, is_first);
366            is_first = false;
367            container = self.get_parent(c)
368        }
369    }
370
371    #[inline]
372    pub fn slice_by_unicode(&self, range: impl RangeBounds<usize>) -> BytesSlice {
373        self.inner.str.lock().unwrap().slice_by_unicode(range)
374    }
375
376    #[inline]
377    pub fn slice_by_utf8(&self, range: impl RangeBounds<usize>) -> BytesSlice {
378        self.inner.str.lock().unwrap().slice_bytes(range)
379    }
380
381    #[inline]
382    pub fn slice_str_by_unicode_range(&self, range: Range<usize>) -> String {
383        let mut s = self.inner.str.lock().unwrap();
384        let s: &mut StrArena = &mut s;
385        let mut ans = String::with_capacity(range.len());
386        ans.push_str(s.slice_str_by_unicode(range));
387        ans
388    }
389
390    #[inline]
391    pub fn with_text_slice(&self, range: Range<usize>, mut f: impl FnMut(&str)) {
392        f(self.inner.str.lock().unwrap().slice_str_by_unicode(range))
393    }
394
395    #[inline]
396    pub fn get_value(&self, idx: usize) -> Option<LoroValue> {
397        self.inner.values.lock().unwrap().get(idx).cloned()
398    }
399
400    #[inline]
401    pub fn get_values(&self, range: Range<usize>) -> Vec<LoroValue> {
402        (self.inner.values.lock().unwrap()[range]).to_vec()
403    }
404
405    pub fn convert_single_op(
406        &self,
407        container: &ContainerID,
408        peer: PeerID,
409        counter: Counter,
410        lamport: Lamport,
411        content: RawOpContent,
412    ) -> Op {
413        let container = self.register_container(container);
414        self.inner_convert_op(content, peer, counter, lamport, container)
415    }
416
417    pub fn can_import_snapshot(&self) -> bool {
418        self.inner.str.lock().unwrap().is_empty() && self.inner.values.lock().unwrap().is_empty()
419    }
420
421    fn inner_convert_op(
422        &self,
423        content: RawOpContent<'_>,
424        _peer: PeerID,
425        counter: i32,
426        _lamport: Lamport,
427        container: ContainerIdx,
428    ) -> Op {
429        match content {
430            crate::op::RawOpContent::Map(MapSet { key, value }) => Op {
431                counter,
432                container,
433                content: crate::op::InnerContent::Map(MapSet { key, value }),
434            },
435            crate::op::RawOpContent::List(list) => match list {
436                ListOp::Insert { slice, pos } => match slice {
437                    ListSlice::RawData(values) => {
438                        let range = self.alloc_values(values.iter().cloned());
439                        Op {
440                            counter,
441                            container,
442                            content: crate::op::InnerContent::List(InnerListOp::Insert {
443                                slice: SliceRange::from(range.start as u32..range.end as u32),
444                                pos,
445                            }),
446                        }
447                    }
448                    ListSlice::RawStr { str, unicode_len } => {
449                        let (slice, info) = self.alloc_str_with_slice(&str);
450                        Op {
451                            counter,
452                            container,
453                            content: crate::op::InnerContent::List(InnerListOp::InsertText {
454                                slice,
455                                unicode_start: info.start as u32,
456                                unicode_len: unicode_len as u32,
457                                pos: pos as u32,
458                            }),
459                        }
460                    }
461                },
462                ListOp::Delete(span) => Op {
463                    counter,
464                    container,
465                    content: crate::op::InnerContent::List(InnerListOp::Delete(span)),
466                },
467                ListOp::StyleStart {
468                    start,
469                    end,
470                    info,
471                    key,
472                    value,
473                } => Op {
474                    counter,
475                    container,
476                    content: InnerContent::List(InnerListOp::StyleStart {
477                        start,
478                        end,
479                        key,
480                        info,
481                        value,
482                    }),
483                },
484                ListOp::StyleEnd => Op {
485                    counter,
486                    container,
487                    content: InnerContent::List(InnerListOp::StyleEnd),
488                },
489                ListOp::Move {
490                    from,
491                    to,
492                    elem_id: from_id,
493                } => Op {
494                    counter,
495                    container,
496                    content: InnerContent::List(InnerListOp::Move {
497                        from,
498                        to,
499                        elem_id: from_id,
500                    }),
501                },
502                ListOp::Set { elem_id, value } => Op {
503                    counter,
504                    container,
505                    content: InnerContent::List(InnerListOp::Set { elem_id, value }),
506                },
507            },
508            crate::op::RawOpContent::Tree(tree) => Op {
509                counter,
510                container,
511                content: crate::op::InnerContent::Tree(tree.clone()),
512            },
513            #[cfg(feature = "counter")]
514            crate::op::RawOpContent::Counter(c) => Op {
515                counter,
516                container,
517                content: crate::op::InnerContent::Future(crate::op::FutureInnerContent::Counter(c)),
518            },
519            crate::op::RawOpContent::Unknown { prop, value } => Op {
520                counter,
521                container,
522                content: crate::op::InnerContent::Future(crate::op::FutureInnerContent::Unknown {
523                    prop,
524                    value: Box::new(value),
525                }),
526            },
527        }
528    }
529
530    #[inline]
531    pub fn convert_raw_op(&self, op: &RawOp) -> Op {
532        self.inner_convert_op(
533            op.content.clone(),
534            op.id.peer,
535            op.id.counter,
536            op.lamport,
537            op.container,
538        )
539    }
540
541    #[inline]
542    pub fn export_containers(&self) -> Vec<ContainerID> {
543        self.inner.container_idx_to_id.lock().unwrap().clone()
544    }
545
546    pub fn export_parents(&self) -> Vec<Option<ContainerIdx>> {
547        let parents = self.inner.parents.lock().unwrap();
548        let containers = self.inner.container_idx_to_id.lock().unwrap();
549        containers
550            .iter()
551            .enumerate()
552            .map(|(x, id)| {
553                let idx = ContainerIdx::from_index_and_type(x as u32, id.container_type());
554                let parent_idx = parents.get(&idx)?;
555                *parent_idx
556            })
557            .collect()
558    }
559
560    /// Returns all the possible root containers of the docs
561    ///
562    /// We need to load all the cached kv in DocState before we can ensure all root contains are covered.
563    /// So we need the flag type here.
564    #[inline]
565    pub(crate) fn root_containers(&self, _f: LoadAllFlag) -> Vec<ContainerIdx> {
566        self.inner.root_c_idx.lock().unwrap().clone()
567    }
568
569    // TODO: this can return a u16 directly now, since the depths are always valid
570    pub(crate) fn get_depth(&self, container: ContainerIdx) -> Option<NonZeroU16> {
571        let mut depth_guard = self.inner.depth.lock().unwrap();
572        let mut parents_guard = self.inner.parents.lock().unwrap();
573        let mut root_c_idx_guard = self.inner.root_c_idx.lock().unwrap();
574        let resolver_guard = self.inner.parent_resolver.lock().unwrap();
575        let mut idx_to_id_guard = self.inner.container_idx_to_id.lock().unwrap();
576        let mut id_to_idx_guard = self.inner.container_id_to_idx.lock().unwrap();
577        get_depth(
578            container,
579            &mut depth_guard,
580            &mut parents_guard,
581            &resolver_guard,
582            &mut idx_to_id_guard,
583            &mut id_to_idx_guard,
584            &mut root_c_idx_guard,
585        )
586    }
587
588    pub(crate) fn iter_value_slice(
589        &self,
590        range: Range<usize>,
591    ) -> impl Iterator<Item = LoroValue> + '_ {
592        let values = self.inner.values.lock().unwrap();
593        range
594            .into_iter()
595            .map(move |i| values.get(i).unwrap().clone())
596    }
597
598    pub(crate) fn get_root_container_idx_by_key(
599        &self,
600        root_index: &loro_common::InternalString,
601    ) -> Option<ContainerIdx> {
602        let inner = self.inner.container_id_to_idx.lock().unwrap();
603        for t in loro_common::ContainerType::ALL_TYPES.iter() {
604            let cid = ContainerID::Root {
605                name: root_index.clone(),
606                container_type: *t,
607            };
608            if let Some(idx) = inner.get(&cid) {
609                return Some(*idx);
610            }
611        }
612        None
613    }
614
615    #[allow(unused)]
616    pub(crate) fn log_all_values(&self) {
617        let values = self.inner.values.lock().unwrap();
618        for (i, v) in values.iter().enumerate() {
619            loro_common::debug!("value {} {:?}", i, v);
620        }
621    }
622}
623
624fn _alloc_str_with_slice(
625    text_lock: &mut MutexGuard<'_, StrArena>,
626    str: &str,
627) -> (BytesSlice, StrAllocResult) {
628    let start = text_lock.len_bytes();
629    let ans = _alloc_str(text_lock, str);
630    (text_lock.slice_bytes(start..), ans)
631}
632
633fn _alloc_values(
634    values_lock: &mut MutexGuard<'_, Vec<LoroValue>>,
635    values: impl Iterator<Item = LoroValue>,
636) -> Range<usize> {
637    values_lock.reserve(values.size_hint().0);
638    let start = values_lock.len();
639    for value in values {
640        values_lock.push(value);
641    }
642
643    start..values_lock.len()
644}
645
646fn _alloc_value(values_lock: &mut MutexGuard<'_, Vec<LoroValue>>, value: LoroValue) -> usize {
647    values_lock.push(value);
648    values_lock.len() - 1
649}
650
651fn _alloc_str(text_lock: &mut MutexGuard<'_, StrArena>, str: &str) -> StrAllocResult {
652    let start = text_lock.len_unicode();
653    text_lock.alloc(str);
654    StrAllocResult {
655        start,
656        end: text_lock.len_unicode(),
657    }
658}
659
660fn _slice_str(range: Range<usize>, s: &mut StrArena) -> String {
661    let mut ans = String::with_capacity(range.len());
662    ans.push_str(s.slice_str_by_unicode(range));
663    ans
664}
665
666fn get_depth(
667    target: ContainerIdx,
668    depth: &mut Vec<Option<NonZeroU16>>,
669    parents: &mut FxHashMap<ContainerIdx, Option<ContainerIdx>>,
670    parent_resolver: &Option<Arc<ParentResolver>>,
671    idx_to_id: &mut Vec<ContainerID>,
672    id_to_idx: &mut FxHashMap<ContainerID, ContainerIdx>,
673    root_c_idx: &mut Vec<ContainerIdx>,
674) -> Option<NonZeroU16> {
675    let mut d = depth[target.to_index() as usize];
676    if d.is_some() {
677        return d;
678    }
679
680    let parent: Option<ContainerIdx> = if let Some(p) = parents.get(&target) {
681        *p
682    } else {
683        let id = idx_to_id.get(target.to_index() as usize).unwrap();
684        if id.is_root() {
685            None
686        } else if let Some(parent_resolver) = parent_resolver.as_ref() {
687            let parent_id = parent_resolver(id.clone())?;
688            let parent_is_root = parent_id.is_root();
689            let mut parent_was_new = false;
690            // If the parent is not registered yet, register it lazily instead of unwrapping.
691            let parent_idx = if let Some(idx) = id_to_idx.get(&parent_id).copied() {
692                idx
693            } else {
694                let new_index = idx_to_id.len();
695                idx_to_id.push(parent_id.clone());
696                let new_idx =
697                    ContainerIdx::from_index_and_type(new_index as u32, parent_id.container_type());
698                id_to_idx.insert(parent_id.clone(), new_idx);
699                // Keep depth vector in sync with containers list.
700                if parent_is_root {
701                    depth.push(NonZeroU16::new(1));
702                } else {
703                    depth.push(None);
704                }
705                parent_was_new = true;
706                new_idx
707            };
708
709            if parent_is_root {
710                if parent_was_new {
711                    parents.insert(parent_idx, None);
712                    root_c_idx.push(parent_idx);
713                } else {
714                    parents.entry(parent_idx).or_insert(None);
715                }
716                if depth[parent_idx.to_index() as usize].is_none() {
717                    depth[parent_idx.to_index() as usize] = NonZeroU16::new(1);
718                }
719            }
720
721            Some(parent_idx)
722        } else {
723            return None;
724        }
725    };
726
727    match parent {
728        Some(p) => {
729            d = NonZeroU16::new(
730                get_depth(
731                    p,
732                    depth,
733                    parents,
734                    parent_resolver,
735                    idx_to_id,
736                    id_to_idx,
737                    root_c_idx,
738                )?
739                .get()
740                    + 1,
741            );
742            depth[target.to_index() as usize] = d;
743        }
744        None => {
745            depth[target.to_index() as usize] = NonZeroU16::new(1);
746            d = NonZeroU16::new(1);
747        }
748    }
749
750    d
751}
752
753impl SharedArena {
754    /// Register or clear a resolver to lazily determine a container's parent when missing.
755    ///
756    /// - The resolver receives the child `ContainerIdx` and returns an optional `ContainerID` of its parent.
757    /// - If the resolver returns `Some`, `SharedArena` will register the parent in the arena and link it.
758    /// - If the resolver is `None` or returns `None`, `get_parent` will panic for non-root containers as before.
759    pub fn set_parent_resolver<F>(&self, resolver: Option<F>)
760    where
761        F: Fn(ContainerID) -> Option<ContainerID> + Send + Sync + 'static,
762    {
763        let mut slot = self.inner.parent_resolver.lock().unwrap();
764        *slot = resolver.map(|f| Arc::new(f) as Arc<ParentResolver>);
765    }
766}