From b22bd98f6b55e2e2661e708d623f9458515c7919 Mon Sep 17 00:00:00 2001
From: Zixuan Chen <remch183@outlook.com>
Date: Sat, 5 Aug 2023 10:23:14 +0800
Subject: [PATCH] feat: make getting child container handler simple (#104)

---
 crates/loro-internal/src/handler.rs | 54 +++++++++++++++++++++++++++++
 crates/loro-internal/src/loro.rs    |  4 +--
 crates/loro-internal/src/state.rs   |  2 +-
 crates/loro-internal/tests/test.rs  |  6 ++++
 4 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/crates/loro-internal/src/handler.rs b/crates/loro-internal/src/handler.rs
index adef7ab5..95a6a423 100644
--- a/crates/loro-internal/src/handler.rs
+++ b/crates/loro-internal/src/handler.rs
@@ -8,6 +8,7 @@ use crate::{
     delta::MapValue,
     txn::EventHint,
 };
+use enum_as_inner::EnumAsInner;
 use loro_common::{ContainerID, ContainerType, LoroResult, LoroValue};
 use std::{
     borrow::Cow,
@@ -32,6 +33,23 @@ pub struct ListHandler {
     state: Weak<Mutex<DocState>>,
 }
 
+#[derive(Clone, EnumAsInner)]
+pub enum Handler {
+    Text(TextHandler),
+    Map(MapHandler),
+    List(ListHandler),
+}
+
+impl Handler {
+    fn new(value: ContainerIdx, state: Weak<Mutex<DocState>>) -> Self {
+        match value.get_type() {
+            ContainerType::Text => Self::Text(TextHandler::new(value, state)),
+            ContainerType::Map => Self::Map(MapHandler::new(value, state)),
+            ContainerType::List => Self::List(ListHandler::new(value, state)),
+        }
+    }
+}
+
 impl TextHandler {
     pub fn new(idx: ContainerIdx, state: Weak<Mutex<DocState>>) -> Self {
         assert_eq!(idx.get_type(), ContainerType::Text);
@@ -343,6 +361,24 @@ impl ListHandler {
         )
     }
 
+    pub fn get_child_handler(&self, index: usize) -> Handler {
+        let mutex = &self.state.upgrade().unwrap();
+        let state = mutex.lock().unwrap();
+        let container_id = state.with_state(self.container_idx, |state| {
+            state
+                .as_list_state()
+                .as_ref()
+                .unwrap()
+                .get(index)
+                .unwrap()
+                .as_container()
+                .unwrap()
+                .clone()
+        });
+        let idx = state.arena.register_container(&container_id);
+        Handler::new(idx, self.state.clone())
+    }
+
     pub fn len(&self) -> usize {
         self.state
             .upgrade()
@@ -504,6 +540,24 @@ impl MapHandler {
             .get_value_by_idx(self.container_idx)
     }
 
+    pub fn get_child_handler(&self, key: &str) -> Handler {
+        let mutex = &self.state.upgrade().unwrap();
+        let state = mutex.lock().unwrap();
+        let container_id = state.with_state(self.container_idx, |state| {
+            state
+                .as_map_state()
+                .as_ref()
+                .unwrap()
+                .get(key)
+                .unwrap()
+                .as_container()
+                .unwrap()
+                .clone()
+        });
+        let idx = state.arena.register_container(&container_id);
+        Handler::new(idx, self.state.clone())
+    }
+
     pub fn get_deep_value(&self) -> LoroValue {
         self.state
             .upgrade()
diff --git a/crates/loro-internal/src/loro.rs b/crates/loro-internal/src/loro.rs
index bf25277d..2a4542c5 100644
--- a/crates/loro-internal/src/loro.rs
+++ b/crates/loro-internal/src/loro.rs
@@ -118,10 +118,10 @@ impl LoroDoc {
     #[inline(always)]
     pub fn with_txn<F>(&self, f: F) -> LoroResult<()>
     where
-        F: Fn(&mut Transaction),
+        F: Fn(&mut Transaction) -> LoroResult<()>,
     {
         let mut txn = self.txn().unwrap();
-        f(&mut txn);
+        f(&mut txn)?;
         txn.commit()
     }
 
diff --git a/crates/loro-internal/src/state.rs b/crates/loro-internal/src/state.rs
index d616d9c5..3d7a8527 100644
--- a/crates/loro-internal/src/state.rs
+++ b/crates/loro-internal/src/state.rs
@@ -1,6 +1,5 @@
 use std::{borrow::Cow, mem::take, sync::Arc};
 
-use debug_log::debug_dbg;
 use enum_as_inner::EnumAsInner;
 use enum_dispatch::enum_dispatch;
 use fxhash::{FxHashMap, FxHashSet};
@@ -374,6 +373,7 @@ impl DocState {
             .as_text_state()
     }
 
+    #[inline(always)]
     pub(crate) fn with_state<F, R>(&self, idx: ContainerIdx, f: F) -> R
     where
         F: FnOnce(&State) -> R,
diff --git a/crates/loro-internal/tests/test.rs b/crates/loro-internal/tests/test.rs
index 7c063b9e..18cee753 100644
--- a/crates/loro-internal/tests/test.rs
+++ b/crates/loro-internal/tests/test.rs
@@ -53,11 +53,13 @@ fn map_checkout() {
     let v_empty = doc.oplog_frontiers();
     doc.with_txn(|txn| {
         meta.insert(txn, "key", 0.into()).unwrap();
+        Ok(())
     })
     .unwrap();
     let v0 = doc.oplog_frontiers();
     doc.with_txn(|txn| {
         meta.insert(txn, "key", 1.into()).unwrap();
+        Ok(())
     })
     .unwrap();
     let v1 = doc.oplog_frontiers();
@@ -80,18 +82,21 @@ fn map_concurrent_checkout() {
     doc_a
         .with_txn(|txn| {
             meta_a.insert(txn, "key", 0.into()).unwrap();
+            Ok(())
         })
         .unwrap();
     let va = doc_a.oplog_frontiers();
     doc_b
         .with_txn(|txn| {
             meta_b.insert(txn, "s", 1.into()).unwrap();
+            Ok(())
         })
         .unwrap();
     let vb_0 = doc_b.oplog_frontiers();
     doc_b
         .with_txn(|txn| {
             meta_b.insert(txn, "key", 1.into()).unwrap();
+            Ok(())
         })
         .unwrap();
     let vb_1 = doc_b.oplog_frontiers();
@@ -99,6 +104,7 @@ fn map_concurrent_checkout() {
     doc_a
         .with_txn(|txn| {
             meta_a.insert(txn, "key", 2.into()).unwrap();
+            Ok(())
         })
         .unwrap();