benchmark sorting

experimental/storage/layered-map/Cargo.toml

@@ -21,4 +21,11 @@ bitvec = "1.0.1"
 once_cell = { workspace = true }
 
 [dev-dependencies]
+criterion = { workspace = true }
+itertools = { workspace = true }
 proptest = { workspace = true }
+jemallocator = { workspace = true }
+
+[[bench]]
+name = "sorting"
+harness = false

experimental/storage/layered-map/benches/sorting.rs

@@ -0,0 +1,138 @@
+// Copyright (c) Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use aptos_crypto::HashValue;
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use itertools::Itertools;
+
+#[cfg(unix)]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+fn recursive_bin_search(sorted_data: &[HashValue], depth: usize) {
+    if sorted_data.len() <= 1 {
+        return;
+    }
+
+    let pivot = sorted_data.partition_point(|key| !key.bit(depth));
+    recursive_bin_search(&sorted_data[..pivot], depth + 1);
+    recursive_bin_search(&sorted_data[pivot..], depth + 1);
+}
+
+fn partition(data: &mut [HashValue], depth: usize) -> usize {
+    if data.is_empty() {
+        return 0;
+    }
+
+    let mut zero_cur = 0;
+    let mut one_cur = data.len() - 1;
+
+    while zero_cur < one_cur {
+        while zero_cur < one_cur && !data[zero_cur].bit(depth) {
+            zero_cur += 1;
+        }
+        while one_cur > zero_cur && data[one_cur].bit(depth) {
+            one_cur -= 1;
+        }
+        if zero_cur < one_cur {
+            data.swap(zero_cur, one_cur);
+            zero_cur += 1;
+            one_cur -= 1;
+        }
+    }
+
+    if data[zero_cur].bit(depth) {
+        zero_cur
+    } else {
+        zero_cur + 1
+    }
+}
+
+fn recursive_partition(data: &mut [HashValue], depth: usize) {
+    if data.len() <= 1 {
+        return;
+    }
+
+    let pivot = partition(data, depth);
+    recursive_partition(&mut data[..pivot], depth + 1);
+    recursive_partition(&mut data[pivot..], depth + 1);
+}
+
+fn compare_sorting(c: &mut Criterion) {
+    let mut group = c.benchmark_group("sorting");
+
+    const SET_SIZE: usize = 100000;
+
+    let data = std::iter::repeat_with(HashValue::random)
+        .take(SET_SIZE)
+        .collect_vec();
+    group.throughput(criterion::Throughput::Elements(SET_SIZE as u64));
+
+    group.bench_function("sort_then_bin_search", |b| {
+        b.iter_batched(
+            || data.clone(),
+            |mut data| {
+                data.sort();
+                recursive_bin_search(&data, 0);
+                data
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("recursive_partition", |b| {
+        b.iter_batched(
+            || data.clone(),
+            |mut data| {
+                recursive_partition(&mut data, 0);
+                data
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    let mut data = data.clone();
+    data.sort();
+
+    group.bench_function("sort_then_bin_search_pre_sorted", |b| {
+        b.iter_batched(
+            || data.clone(),
+            |mut data| {
+                data.sort();
+                recursive_bin_search(&data, 0);
+                data
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("bin_search_pre_sorted", |b| {
+        b.iter_batched(
+            || data.clone(),
+            |data| {
+                recursive_bin_search(&data, 0);
+                data
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("recursive_partition_pre_sorted", |b| {
+        b.iter_batched(
+            || data.clone(),
+            |mut data| {
+                recursive_partition(&mut data, 0);
+                data
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+criterion_group!(
+    name = sorting;
+    config = Criterion::default();
+    targets = compare_sorting
+);
+
+criterion_main!(sorting);

experimental/storage/layered-map/src/key.rs

@@ -0,0 +1,24 @@
+// Copyright (c) Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use aptos_crypto::HashValue;
+use bitvec::prelude::*;
+
+/// When recursively creating a new `MapLayer` (a crit bit tree overlay), passing down `Vec<(K, Option<V>)>`
+/// That's why we require `Key: Clone` and clone the key and value only when the leaf node is
+/// created.
+pub trait Key: Clone + Eq {
+    fn iter_bits(&self) -> impl Iterator<Item = bool>;
+
+    fn bit(&self, depth: usize) -> bool;
+}
+
+impl Key for HashValue {
+    fn iter_bits(&self) -> impl Iterator<Item = bool> {
+        self.iter_bits()
+    }
+
+    fn bit(&self, depth: usize) -> bool {
+        *self.as_slice().view_bits::<Msb0>().get(depth).unwrap()
+    }
+}

experimental/storage/layered-map/src/lib.rs

@@ -10,25 +10,18 @@ use aptos_crypto::HashValue;
 use aptos_drop_helper::ArcAsyncDrop;
 use aptos_infallible::Mutex;
 use aptos_metrics_core::{IntGaugeHelper, TimerHelper};
+pub use key::Key;
 use std::sync::Arc;
 
 mod dropper;
 mod metrics;
 mod node;
 pub(crate) mod r#ref;
 
+mod key;
 #[cfg(test)]
 mod tests;
 
-/// When recursively creating a new `MapLayer` (a crit bit tree overlay), passing down `Vec<(K, Option<V>)>`
-/// That's why we require `Key: Clone` and clone the key and value only when the leaf node is
-/// created.
-pub trait Key: Clone + Eq {
-    fn iter_bits(&self) -> impl Iterator<Item = bool>;
-
-    fn bit(&self, depth: usize) -> bool;
-}
-
 /// Similar to `Key`, we require `Value: Clone`, another reason being it's tricky to figure out the
 /// lifetime if `get()` returns a reference to the value -- we simply clone the value.
 pub trait Value: Clone {}