Skip to content

Commit

Permalink
Implement RedbValue for Vec<T>
Browse files Browse the repository at this point in the history
  • Loading branch information
cberner committed Oct 13, 2023
1 parent 9089a20 commit 79aa893
Show file tree
Hide file tree
Showing 3 changed files with 157 additions and 0 deletions.
96 changes: 96 additions & 0 deletions src/complex_types.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
use crate::types::{RedbValue, TypeName};

// Encode len as a varint and store it at the end of output
fn encode_varint_len(len: usize, output: &mut Vec<u8>) {
if len < 254 {
output.push(len.try_into().unwrap());
} else if len <= u16::MAX.into() {
let u16_len: u16 = len.try_into().unwrap();
output.push(254);
output.extend_from_slice(&u16_len.to_be_bytes())
} else {
assert!(len <= u32::MAX as usize);
let u32_len: u32 = len.try_into().unwrap();
output.push(255);
output.extend_from_slice(&u32_len.to_be_bytes())
}
}

// Decode a variable length int starting at the beginning of data
// Returns (decoded length, length consumed of `data`)
#[allow(clippy::comparison_chain)]
fn decode_varint_len(data: &[u8]) -> (usize, usize) {
if data[0] < 254 {
(data[0] as usize, 1)
} else if data[0] == 254 {
(
u16::from_be_bytes(data[1..3].try_into().unwrap()) as usize,
3,
)
} else {
(
u32::from_be_bytes(data[1..5].try_into().unwrap()) as usize,
5,
)
}
}

impl<T: RedbValue> RedbValue for Vec<T> {
type SelfType<'a> = Vec<T::SelfType<'a>>
where
Self: 'a;
type AsBytes<'a> = Vec<u8>
where
Self: 'a;

fn fixed_width() -> Option<usize> {
None
}

fn from_bytes<'a>(data: &'a [u8]) -> Vec<T::SelfType<'a>>
where
Self: 'a,
{
let (elements, mut offset) = decode_varint_len(data);
let mut result = Vec::with_capacity(elements);
for _ in 0..elements {
let element_len = if let Some(len) = T::fixed_width() {
len
} else {
let (len, consumed) = decode_varint_len(&data[offset..]);
offset += consumed;
len
};
result.push(T::from_bytes(&data[offset..(offset + element_len)]));
offset += element_len;
}
assert_eq!(offset, data.len());
result
}

fn as_bytes<'a, 'b: 'a>(value: &'a Vec<T::SelfType<'b>>) -> Vec<u8>
where
Self: 'a,
Self: 'b,
{
let mut result = if let Some(width) = T::fixed_width() {
Vec::with_capacity(value.len() * width + 5)
} else {
Vec::with_capacity(value.len() * 2 + 5)
};
encode_varint_len(value.len(), &mut result);

for element in value.iter() {
let serialized = T::as_bytes(element);
if T::fixed_width().is_none() {
encode_varint_len(serialized.as_ref().len(), &mut result);
}
result.extend_from_slice(serialized.as_ref());
}
result
}

fn type_name() -> TypeName {
TypeName::internal(&format!("Vec<{}>", T::type_name().name()))
}
}
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ type Result<T = (), E = StorageError> = std::result::Result<T, E>;
#[cfg(feature = "python")]
pub use crate::python::redb;

mod complex_types;
mod db;
mod error;
mod multimap_table;
Expand Down
60 changes: 60 additions & 0 deletions tests/basic_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1043,6 +1043,66 @@ fn array_type() {
assert!(iter.next().is_none());
}

#[test]
fn vec_fixed_width_value_type() {
let tmpfile = create_tempfile();
let db = Database::create(tmpfile.path()).unwrap();

let definition: TableDefinition<u8, Vec<u64>> = TableDefinition::new("x");

let value = vec![0, 1, 2, 3];
let write_txn = db.begin_write().unwrap();
{
let mut table = write_txn.open_table(definition).unwrap();
table.insert(0, &value).unwrap();
}
write_txn.commit().unwrap();

let read_txn = db.begin_read().unwrap();
let table = read_txn.open_table(definition).unwrap();
assert_eq!(value, table.get(0).unwrap().unwrap().value());
}

#[test]
fn vec_var_width_value_type() {
let tmpfile = create_tempfile();
let db = Database::create(tmpfile.path()).unwrap();

let definition: TableDefinition<u8, Vec<&str>> = TableDefinition::new("x");

let value = vec!["hello", "world"];
let write_txn = db.begin_write().unwrap();
{
let mut table = write_txn.open_table(definition).unwrap();
table.insert(0, &value).unwrap();
}
write_txn.commit().unwrap();

let read_txn = db.begin_read().unwrap();
let table = read_txn.open_table(definition).unwrap();
assert_eq!(value, table.get(0).unwrap().unwrap().value());
}

#[test]
fn vec_vec_type() {
let tmpfile = create_tempfile();
let db = Database::create(tmpfile.path()).unwrap();

let definition: TableDefinition<u8, Vec<Vec<&str>>> = TableDefinition::new("x");

let value = vec![vec!["hello", "world"], vec!["this", "is", "a", "test"]];
let write_txn = db.begin_write().unwrap();
{
let mut table = write_txn.open_table(definition).unwrap();
table.insert(0, &value).unwrap();
}
write_txn.commit().unwrap();

let read_txn = db.begin_read().unwrap();
let table = read_txn.open_table(definition).unwrap();
assert_eq!(value, table.get(0).unwrap().unwrap().value());
}

#[test]
fn range_lifetime() {
let tmpfile = create_tempfile();
Expand Down

0 comments on commit 79aa893

Please sign in to comment.