From ec7774dc71617cd05f1a38c5d671f3b45a469e85 Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Fri, 13 Oct 2023 11:37:15 -0700
Subject: [PATCH] Implement RedbValue for Vec<T>

---
 src/complex_types.rs | 93 ++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs           |  1 +
 tests/basic_tests.rs | 60 ++++++++++++++++++++++++++++
 3 files changed, 154 insertions(+)
 create mode 100644 src/complex_types.rs

diff --git a/src/complex_types.rs b/src/complex_types.rs
new file mode 100644
index 00000000..39ed736a
--- /dev/null
+++ b/src/complex_types.rs
@@ -0,0 +1,93 @@
+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_le_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_le_bytes())
+    }
+}
+
+// Decode a variable length int starting at the beginning of data
+// Returns (decoded length, length consumed of `data`)
+fn decode_varint_len(data: &[u8]) -> (usize, usize) {
+    match data[0] {
+        0..=253 => (data[0] as usize, 1),
+        254 => (
+            u16::from_le_bytes(data[1..3].try_into().unwrap()) as usize,
+            3,
+        ),
+        255 => (
+            u32::from_le_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()))
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 24d5cb4e..1abf7c16 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -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;
diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
index 8d73c5a4..1f92bece 100644
--- a/tests/basic_tests.rs
+++ b/tests/basic_tests.rs
@@ -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();