fix: reduce heap alloc

.gitignore

@@ -3,3 +3,4 @@
 *.ipynb
 flamegraph.svg
 target
+dhat-heap.json

Cargo.lock

@@ -551,38 +551,6 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4217ad341ebadf8d8e724e264f13e593e0648f5b3e94b3896a5df283be015ecc"
 
-[[package]]
-name = "jemalloc-ctl"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c1891c671f3db85d8ea8525dd43ab147f9977041911d24a03e5a36187a7bfde9"
-dependencies = [
- "jemalloc-sys",
- "libc",
- "paste",
-]
-
-[[package]]
-name = "jemalloc-sys"
-version = "0.5.2+5.3.0-patched"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "134163979b6eed9564c98637b710b40979939ba351f59952708234ea11b5f3f8"
-dependencies = [
- "cc",
- "fs_extra",
- "libc",
-]
-
-[[package]]
-name = "jemallocator"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "16c2514137880c52b0b4822b563fadd38257c1f380858addb74a400889696ea6"
-dependencies = [
- "jemalloc-sys",
- "libc",
-]
-
 [[package]]
 name = "js-sys"
 version = "0.3.60"
@@ -639,9 +607,6 @@ dependencies = [
  "flate2",
  "fxhash",
  "im",
- "jemalloc-ctl",
- "jemalloc-sys",
- "jemallocator",
  "js-sys",
  "num",
  "owning_ref",
@@ -659,6 +624,9 @@ dependencies = [
  "string_cache",
  "tabled",
  "thiserror",
+ "tikv-jemalloc-ctl",
+ "tikv-jemalloc-sys",
+ "tikv-jemallocator",
  "wasm-bindgen",
 ]
 
@@ -1462,6 +1430,38 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820"
 
+[[package]]
+name = "tikv-jemalloc-ctl"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e37706572f4b151dff7a0146e040804e9c26fe3a3118591112f05cf12a4216c1"
+dependencies = [
+ "libc",
+ "paste",
+ "tikv-jemalloc-sys",
+]
+
+[[package]]
+name = "tikv-jemalloc-sys"
+version = "0.5.2+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec45c14da997d0925c7835883e4d5c181f196fa142f8c19d7643d1e9af2592c3"
+dependencies = [
+ "cc",
+ "fs_extra",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20612db8a13a6c06d57ec83953694185a367e16945f66565e8028d2c0bd76979"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
 [[package]]
 name = "tinytemplate"
 version = "1.2.1"

crates/loro-core/Cargo.toml

@@ -40,15 +40,18 @@ ctor = "0.1.23"
 criterion = "0.4.0"
 serde_json = "1.0.87"
 flate2 = "1.0.24"
-jemalloc-ctl = "0.5.0"
-jemalloc-sys = { version = "0.5.0+0.5.3", features = [
+tikv-jemalloc-ctl = "0.5.0"
+tikv-jemalloc-sys = { version = "0.5.0", features = [
+  "stats",
+  "profiling",
+  "unprefixed_malloc_on_supported_platforms",
+] }
+tikv-jemallocator = { version = "0.5.0", features = [
   "stats",
   "profiling",
   "unprefixed_malloc_on_supported_platforms",
 ] }
 
-[target.'cfg(not(target_env = "msvc"))'.dev-dependencies]
-jemallocator = "0.5.0"
 
 # See https://matklad.github.io/2021/02/27/delete-cargo-integration-tests.html
 [lib]

crates/loro-core/examples/mem.rs

@@ -1,11 +1,11 @@
-#[cfg(not(target_env = "msvc"))]
-use jemallocator::Jemalloc;
+// use tikv_jemallocator::Jemalloc;
+// #[global_allocator]
+// static GLOBAL: Jemalloc = Jemalloc;
 
-#[cfg(not(target_env = "msvc"))]
 #[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+static ALLOC: dhat::Alloc = dhat::Alloc;
 
-use jemalloc_ctl::{epoch, stats};
+use tikv_jemalloc_ctl::{epoch, stats, Access, AsName};
 const RAW_DATA: &[u8; 901823] = include_bytes!("../benches/automerge-paper.json.gz");
 
 use std::{io::Read, time::Instant};
@@ -15,7 +15,7 @@ use loro_core::LoroCore;
 use serde_json::Value;
 
 pub fn main() {
-    let alloc_stats = stats::allocated::mib().unwrap();
+    // let alloc_stats = stats::allocated::mib().unwrap();
     let mut d = GzDecoder::new(&RAW_DATA[..]);
     let mut s = String::new();
     d.read_to_string(&mut s).unwrap();
@@ -24,6 +24,7 @@ pub fn main() {
     let txns = json.as_object().unwrap().get("txns");
     let e = epoch::mib().unwrap();
     let start = Instant::now();
+    let profiler = dhat::Profiler::builder().trim_backtraces(None).build();
     let mut loro = LoroCore::default();
     let mut text = loro.get_or_create_root_text("text").unwrap();
     for i in 0..1 {
@@ -42,13 +43,18 @@ pub fn main() {
                 text.delete(pos, del_here);
                 text.insert(pos, ins_content);
             }
+
+            if start.elapsed().as_secs() > 10 {
+                break;
+            }
         }
     }
     drop(json);
     drop(d);
-    e.advance().unwrap();
-    let new_new_heap = alloc_stats.read().unwrap();
+    drop(profiler);
+    // e.advance().unwrap();
+    // let new_new_heap = alloc_stats.read().unwrap();
     println!("Apply Automerge Dataset 1X");
-    println!("Mem: {} MB", new_new_heap as f64 / 1024. / 1024.);
+    // println!("Mem: {} MB", new_new_heap as f64 / 1024. / 1024.);
     println!("Used: {} ms", start.elapsed().as_millis());
 }

crates/loro-core/justfile

@@ -27,5 +27,5 @@ quick-fuzz:
 flame:
   cargo flamegraph --example test --features=fuzzing --root
 
-bench:
-  cargo bench --features fuzzing
+bench *FLAGS:
+  cargo bench --features fuzzing {{FLAGS}}

crates/loro-core/src/container/map/map_container.rs

@@ -55,7 +55,7 @@ impl MapContainer {
 
         let id = store.next_id_for(client_id);
         let counter = id.counter;
-        store.append_local_ops(vec![Op {
+        store.append_local_ops(&[Op {
             id,
             container: self_id,
             content: OpContent::Normal {

crates/loro-core/src/container/text/text_container.rs

@@ -75,7 +75,7 @@ impl TextContainer {
             self.id.clone(),
         );
         let last_id = op.id_last();
-        store.append_local_ops(vec![op]);
+        store.append_local_ops(&[op]);
         self.head = smallvec![last_id];
         self.vv.set_last(last_id);
 
@@ -99,7 +99,7 @@ impl TextContainer {
         );
 
         let last_id = op.id_last();
-        store.append_local_ops(vec![op]);
+        store.append_local_ops(&[op]);
         self.state.delete_range(Some(pos), Some(pos + len));
         self.head = smallvec![last_id];
         self.vv.set_last(last_id);

crates/loro-core/src/dag/mermaid.rs

@@ -86,7 +86,7 @@ fn break_points_to_output(input: BreakPoints) -> Output {
         })
         .collect();
     for (client_id, break_points) in breaks.iter() {
-        let mut spans = vec![];
+        let mut spans = Vec::with_capacity(break_points.len());
         for (from, to) in break_points.iter().zip(break_points.iter().skip(1)) {
             spans.push(IdSpan::new(*client_id, *from, *to));
         }

crates/loro-core/src/log_store.rs

@@ -217,7 +217,7 @@ impl LogStore {
     }
 
     /// this method would not get the container and apply op
-    pub fn append_local_ops(&mut self, ops: Vec<Op>) {
+    pub fn append_local_ops(&mut self, ops: &[Op]) {
         if ops.is_empty() {
             return;
         }

crates/rle/src/rle_tree.rs

@@ -9,6 +9,7 @@ use fxhash::FxHashMap;
 use num::FromPrimitive;
 use ouroboros::self_referencing;
 
+use smallvec::SmallVec;
 pub use tree_trait::Position;
 use tree_trait::RleTreeTrait;
 
@@ -263,7 +264,8 @@ impl<T: Rle, A: RleTreeTrait<T>> RleTree<T, A> {
         U: FnMut(&mut T),
         F: FnMut(&T, *mut LeafNode<T, A>),
     {
-        let mut updates_map: HashMap<NonNull<_>, Vec<(usize, Vec<T>)>, _> = FxHashMap::default();
+        let mut updates_map: HashMap<NonNull<_>, Vec<(usize, SmallVec<[T; 2]>)>, _> =
+            FxHashMap::default();
         for cursor in cursors {
             // SAFETY: we has the exclusive reference to the tree and the cursor is valid
             let updates = unsafe {
@@ -305,7 +307,8 @@ impl<T: Rle, A: RleTreeTrait<T>> RleTree<T, A> {
                 .push((cursor, arg));
         }
 
-        let mut updates_map: HashMap<NonNull<_>, Vec<(usize, Vec<T>)>, _> = FxHashMap::default();
+        let mut updates_map: HashMap<NonNull<_>, Vec<(usize, SmallVec<[T; 2]>)>, _> =
+            FxHashMap::default();
         for ((mut leaf, index), args) in cursor_map.iter() {
             // SAFETY: we has the exclusive reference to the tree and the cursor is valid
             let leaf = unsafe { leaf.as_mut() };
@@ -332,7 +335,7 @@ impl<T: Rle, A: RleTreeTrait<T>> RleTree<T, A> {
     // TODO: perf, use smallvec
     fn update_with_gathered_map<F, M>(
         &mut self,
-        iter: HashMap<NonNull<LeafNode<T, A>>, Vec<(usize, Vec<T>)>, M>,
+        iter: HashMap<NonNull<LeafNode<T, A>>, Vec<(usize, SmallVec<[T; 2]>)>, M>,
         notify: &mut F,
     ) where
         F: FnMut(&T, *mut LeafNode<T, A>),

crates/rle/src/rle_tree/node/internal_impl.rs

@@ -160,7 +160,7 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
             to.map_or((self.children.len(), None), |x| self._delete_end(x));
         let deleted_len = direct_delete_end as isize - direct_delete_start as isize;
         // TODO: maybe we can simplify this insertions logic
-        let mut insertions = vec![];
+        let mut insertions: SmallVec<[(usize, &mut Node<T, A>); 2]> = smallvec::smallvec![];
         {
             // handle removing at the end point
             let mut handled = false;
@@ -516,6 +516,7 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
 
 impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
     /// this can only invoke from root
+    /// TODO: need to speed this method up. maybe remove hashset here? use a miniset instead
     #[inline]
     pub(crate) fn delete<F>(&mut self, start: Option<A::Int>, end: Option<A::Int>, notify: &mut F)
     where

crates/rle/src/rle_tree/node/leaf_impl.rs

@@ -344,11 +344,11 @@ impl<'bump, T: Rle, A: RleTreeTrait<T>> LeafNode<'bump, T, A> {
         offset: usize,
         len: usize,
         update_fn: &mut U,
-    ) -> Option<Vec<T>>
+    ) -> Option<SmallVec<[T; 2]>>
     where
         U: FnMut(&mut T),
     {
-        let mut ans = vec![];
+        let mut ans = smallvec::smallvec![];
         if len == 0 {
             return None;
         }
@@ -434,7 +434,7 @@ impl<'bump, T: Rle, A: RleTreeTrait<T>> LeafNode<'bump, T, A> {
 
     pub(crate) fn apply_updates<F>(
         &mut self,
-        mut updates: Vec<(usize, Vec<T>)>,
+        mut updates: Vec<(usize, SmallVec<[T; 2]>)>,
         notify: &mut F,
     ) -> Result<(), Vec<&'bump mut Node<'bump, T, A>>>
     where

crates/rle/src/rle_vec.rs

@@ -374,6 +374,20 @@ where
     }
 }
 
+impl<A: Array> From<&[A::Item]> for RleVec<A>
+where
+    A::Item: Mergable + HasLength + Clone,
+{
+    fn from(value: &[A::Item]) -> Self {
+        let mut ans: RleVec<A> = RleVec::with_capacity(value.len());
+        for v in value.iter() {
+            ans.push(v.clone());
+        }
+        ans.vec.shrink_to_fit();
+        ans
+    }
+}
+
 impl<A: Array> From<SmallVec<A>> for RleVec<A> {
     fn from(value: SmallVec<A>) -> Self {
         RleVec { vec: value }