typed_store/
util.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::StorageType;
5use bincode::Options;
6use serde::Serialize;
7use std::ops::{Bound, RangeBounds};
8use std::path::Path;
9
10#[inline]
11pub fn be_fix_int_ser<S>(t: &S) -> Vec<u8>
12where
13    S: ?Sized + serde::Serialize,
14{
15    bincode::DefaultOptions::new()
16        .with_big_endian()
17        .with_fixint_encoding()
18        .serialize(t)
19        .expect("failed to serialize via be_fix_int_ser method")
20}
21
22pub(crate) fn iterator_bounds<K>(
23    lower_bound: Option<K>,
24    upper_bound: Option<K>,
25) -> (Option<Vec<u8>>, Option<Vec<u8>>)
26where
27    K: Serialize,
28{
29    (
30        lower_bound.map(|b| be_fix_int_ser(&b)),
31        upper_bound.map(|b| be_fix_int_ser(&b)),
32    )
33}
34
35pub(crate) fn iterator_bounds_with_range<K>(
36    range: impl RangeBounds<K>,
37) -> (Option<Vec<u8>>, Option<Vec<u8>>)
38where
39    K: Serialize,
40{
41    let iterator_lower_bound = match range.start_bound() {
42        Bound::Included(lower_bound) => {
43            // Rocksdb lower bound is inclusive by default so nothing to do
44            Some(be_fix_int_ser(&lower_bound))
45            // readopts.set_iterate_lower_bound(key_buf);
46        }
47        Bound::Excluded(lower_bound) => {
48            let mut key_buf = be_fix_int_ser(&lower_bound);
49
50            // Since we want exclusive, we need to increment the key to exclude the previous
51            big_endian_saturating_add_one(&mut key_buf);
52            Some(key_buf)
53            // readopts.set_iterate_lower_bound(key_buf);
54        }
55        Bound::Unbounded => None,
56    };
57    let iterator_upper_bound = match range.end_bound() {
58        Bound::Included(upper_bound) => {
59            let mut key_buf = be_fix_int_ser(&upper_bound);
60
61            // If the key is already at the limit, there's nowhere else to go, so no upper bound
62            if !is_max(&key_buf) {
63                // Since we want exclusive, we need to increment the key to get the upper bound
64                big_endian_saturating_add_one(&mut key_buf);
65                // readopts.set_iterate_upper_bound(key_buf);
66            }
67            Some(key_buf)
68        }
69        Bound::Excluded(upper_bound) => {
70            // Rocksdb upper bound is inclusive by default so nothing to do
71            Some(be_fix_int_ser(&upper_bound))
72            // readopts.set_iterate_upper_bound(key_buf);
73        }
74        Bound::Unbounded => None,
75    };
76    (iterator_lower_bound, iterator_upper_bound)
77}
78
79/// Given a vec<u8>, find the value which is one more than the vector
80/// if the vector was a big endian number.
81/// If the vector is already minimum, don't change it.
82fn big_endian_saturating_add_one(v: &mut [u8]) {
83    if is_max(v) {
84        return;
85    }
86    for i in (0..v.len()).rev() {
87        if v[i] == u8::MAX {
88            v[i] = 0;
89        } else {
90            v[i] += 1;
91            break;
92        }
93    }
94}
95
96/// Check if all the bytes in the vector are 0xFF
97fn is_max(v: &[u8]) -> bool {
98    v.iter().all(|&x| x == u8::MAX)
99}
100
101pub(crate) fn ensure_database_type<P: AsRef<Path>>(
102    path: P,
103    storage_type: StorageType,
104) -> std::io::Result<()> {
105    if !path.as_ref().exists() {
106        return Ok(());
107    }
108    for entry in std::fs::read_dir(path)? {
109        let filepath = entry?.path();
110        if filepath.extension().is_some_and(|ext| ext == "sst")
111            && storage_type != StorageType::Rocks
112        {
113            panic!(
114                "DB type mismatch: expected {:?}, found RocksDB",
115                storage_type
116            );
117        }
118        if filepath
119            .file_name()
120            .is_some_and(|n| n.to_string_lossy().starts_with("wal_"))
121            && storage_type != StorageType::TideHunter
122        {
123            panic!(
124                "DB type mismatch: expected {:?}, found TideHunter",
125                storage_type
126            );
127        }
128    }
129    Ok(())
130}
131
132#[allow(clippy::assign_op_pattern, clippy::manual_div_ceil)]
133#[test]
134fn test_helpers() {
135    let v = vec![];
136    assert!(is_max(&v));
137
138    fn check_add(v: Vec<u8>) {
139        let mut v = v;
140        let num = Num32::from_big_endian(&v);
141        big_endian_saturating_add_one(&mut v);
142        assert!(num + 1 == Num32::from_big_endian(&v));
143    }
144
145    uint::construct_uint! {
146        // 32 byte number
147        struct Num32(4);
148    }
149
150    let mut v = vec![255; 32];
151    big_endian_saturating_add_one(&mut v);
152    assert!(Num32::MAX == Num32::from_big_endian(&v));
153
154    check_add(vec![1; 32]);
155    check_add(vec![6; 32]);
156    check_add(vec![254; 32]);
157
158    // TBD: More tests coming with randomized arrays
159}
typed_store/util.rs

typed_store/
util.rs
