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
22/// Serialize `t` in big-endian fixed-int encoding directly into `buf`,
23/// returning the number of bytes written.
24#[inline]
25pub fn be_fix_int_ser_into<S>(buf: &mut Vec<u8>, t: &S) -> usize
26where
27    S: ?Sized + serde::Serialize,
28{
29    let before = buf.len();
30    bincode::DefaultOptions::new()
31        .with_big_endian()
32        .with_fixint_encoding()
33        .serialize_into(&mut *buf, t)
34        .expect("failed to serialize via be_fix_int_ser_into method");
35    buf.len() - before
36}
37
38pub(crate) fn iterator_bounds<K>(
39    lower_bound: Option<K>,
40    upper_bound: Option<K>,
41) -> (Option<Vec<u8>>, Option<Vec<u8>>)
42where
43    K: Serialize,
44{
45    (
46        lower_bound.map(|b| be_fix_int_ser(&b)),
47        upper_bound.map(|b| be_fix_int_ser(&b)),
48    )
49}
50
51pub(crate) fn iterator_bounds_with_range<K>(
52    range: impl RangeBounds<K>,
53) -> (Option<Vec<u8>>, Option<Vec<u8>>)
54where
55    K: Serialize,
56{
57    let iterator_lower_bound = match range.start_bound() {
58        Bound::Included(lower_bound) => {
59            // Rocksdb lower bound is inclusive by default so nothing to do
60            Some(be_fix_int_ser(&lower_bound))
61            // readopts.set_iterate_lower_bound(key_buf);
62        }
63        Bound::Excluded(lower_bound) => {
64            let mut key_buf = be_fix_int_ser(&lower_bound);
65
66            if is_max(&key_buf) {
67                // No representable key strictly greater than the maximum at this byte
68                // length. Append a zero byte so the lower bound is lexicographically
69                // greater than any same-length key, ensuring the iterator yields nothing
70                // -- matching the user's intent of excluding the max key.
71                key_buf.push(0);
72            } else {
73                // Since we want exclusive, we need to increment the key to exclude the previous
74                big_endian_saturating_add_one(&mut key_buf);
75            }
76            Some(key_buf)
77        }
78        Bound::Unbounded => None,
79    };
80    let iterator_upper_bound = match range.end_bound() {
81        Bound::Included(upper_bound) => {
82            let mut key_buf = be_fix_int_ser(&upper_bound);
83
84            if is_max(&key_buf) {
85                // The key is already at the limit, so the inclusive upper bound covers
86                // everything; leaving the rocksdb upper bound unset is the only way to
87                // include the max key (rocksdb upper bound is exclusive).
88                None
89            } else {
90                // Since rocksdb upper bound is exclusive, increment the key by one so the
91                // user-supplied inclusive bound is included.
92                big_endian_saturating_add_one(&mut key_buf);
93                Some(key_buf)
94            }
95        }
96        Bound::Excluded(upper_bound) => {
97            // Rocksdb upper bound is inclusive by default so nothing to do
98            Some(be_fix_int_ser(&upper_bound))
99            // readopts.set_iterate_upper_bound(key_buf);
100        }
101        Bound::Unbounded => None,
102    };
103    (iterator_lower_bound, iterator_upper_bound)
104}
105
106/// Given a vec<u8>, find the value which is one more than the vector
107/// if the vector was a big endian number.
108/// If the vector is already minimum, don't change it.
109fn big_endian_saturating_add_one(v: &mut [u8]) {
110    if is_max(v) {
111        return;
112    }
113    for i in (0..v.len()).rev() {
114        if v[i] == u8::MAX {
115            v[i] = 0;
116        } else {
117            v[i] += 1;
118            break;
119        }
120    }
121}
122
123/// Check if all the bytes in the vector are 0xFF
124fn is_max(v: &[u8]) -> bool {
125    v.iter().all(|&x| x == u8::MAX)
126}
127
128pub(crate) fn ensure_database_type<P: AsRef<Path>>(
129    path: P,
130    storage_type: StorageType,
131) -> std::io::Result<()> {
132    if !path.as_ref().exists() {
133        return Ok(());
134    }
135    for entry in std::fs::read_dir(path)? {
136        let filepath = entry?.path();
137        if filepath.extension().is_some_and(|ext| ext == "sst")
138            && storage_type != StorageType::Rocks
139        {
140            panic!(
141                "DB type mismatch: expected {:?}, found RocksDB",
142                storage_type
143            );
144        }
145        if filepath
146            .file_name()
147            .is_some_and(|n| n.to_string_lossy().starts_with("wal_"))
148            && storage_type != StorageType::TideHunter
149        {
150            panic!(
151                "DB type mismatch: expected {:?}, found TideHunter",
152                storage_type
153            );
154        }
155    }
156    Ok(())
157}
158
159#[allow(
160    clippy::assign_op_pattern,
161    clippy::manual_div_ceil,
162    clippy::disallowed_methods
163)] // Intentional zip: external construct_uint! macro uses .zip() internally
164#[test]
165fn test_helpers() {
166    let v = vec![];
167    assert!(is_max(&v));
168
169    fn check_add(v: Vec<u8>) {
170        let mut v = v;
171        let num = Num32::from_big_endian(&v);
172        big_endian_saturating_add_one(&mut v);
173        assert!(num + 1 == Num32::from_big_endian(&v));
174    }
175
176    uint::construct_uint! {
177        // 32 byte number
178        struct Num32(4);
179    }
180
181    let mut v = vec![255; 32];
182    big_endian_saturating_add_one(&mut v);
183    assert!(Num32::MAX == Num32::from_big_endian(&v));
184
185    check_add(vec![1; 32]);
186    check_add(vec![6; 32]);
187    check_add(vec![254; 32]);
188
189    // TBD: More tests coming with randomized arrays
190}
typed_store/util.rs

typed_store/
util.rs
