typed_store/
util.rs

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