diff --git a/crates/papyrus_storage/src/db/table_types/dup_sort_tables.rs b/crates/papyrus_storage/src/db/table_types/dup_sort_tables.rs index 0b3a400dcc..6d49b3c420 100644 --- a/crates/papyrus_storage/src/db/table_types/dup_sort_tables.rs +++ b/crates/papyrus_storage/src/db/table_types/dup_sort_tables.rs @@ -346,6 +346,47 @@ impl<'env, K: KeyTrait + Debug, V: ValueSerde + Debug, T: DupSortTableType + Dup } } +// TODO(dvir): consider adding unchecked version of the append function. +#[allow(private_bounds)] +impl<'env, K: KeyTrait + Debug, V: ValueSerde + Debug, T: DupSortTableType + DupSortUtils> + TableHandle<'env, K, V, T> +{ + // Append a new value to the given key. The sub key must be bigger than the last sub key for the + // given main key, otherwise an error will be returned. + // In contrast to the append function in the Table trait, this function will return an error if + // The sub key is equal to the last sub key of the given main key. + #[allow(dead_code)] + pub(crate) fn append_greater_sub_key( + &'env self, + txn: &DbTransaction<'env, RW>, + key: &K, + value: &::Value, + ) -> DbResult<()> { + let main_key = T::get_main_key(key)?; + let sub_key_and_value = T::get_sub_key_and_value(key, value)?; + + let mut cursor = txn.txn.cursor(&self.database)?; + cursor.put(&main_key, &sub_key_and_value, WriteFlags::APPEND_DUP).map_err( + |err| match err { + libmdbx::Error::KeyMismatch => DbError::Append, + _ => err.into(), + }, + )?; + + // This checks the case where the the sub key is already the last in the sub tree; in this + // case, we revert the last put and return an error. + if let Some(prev) = cursor.prev_dup::, DbValueType<'_>>()? { + if prev.1.starts_with(&T::get_sub_key(key)?) { + cursor.next_dup::, DbValueType<'_>>()?; + cursor.del(WriteFlags::empty())?; + return Err(DbError::Append); + } + } + + Ok(()) + } +} + impl< 'txn, Mode: TransactionKind, diff --git a/crates/papyrus_storage/src/db/table_types/dup_sort_tables_test.rs b/crates/papyrus_storage/src/db/table_types/dup_sort_tables_test.rs index 255ddbbb81..c8c01a9275 100644 --- a/crates/papyrus_storage/src/db/table_types/dup_sort_tables_test.rs +++ b/crates/papyrus_storage/src/db/table_types/dup_sort_tables_test.rs @@ -1,7 +1,12 @@ +use assert_matches::assert_matches; + +use super::{DupSortTableType, DupSortUtils}; use crate::db::db_test::get_test_env; +use crate::db::serialization::NoVersionValueWrapper; use crate::db::table_types::dup_sort_tables::add_one; -use crate::db::table_types::test_utils::{random_table_test, table_test}; -use crate::db::DbWriter; +use crate::db::table_types::test_utils::{random_table_test, table_test, TableKey, TableValue}; +use crate::db::table_types::Table; +use crate::db::{DbError, DbResult, DbWriter, TableIdentifier}; #[test] fn common_prefix_table() { @@ -18,6 +23,59 @@ fn common_prefix_compare_with_simple_table_random() { random_table_test(simple_table, common_prefix_table, &reader, &mut writer); } +#[test] +fn common_prefix_append_greater_sub_key() { + append_greater_sub_key_test(DbWriter::create_common_prefix_table); +} + +#[allow(clippy::type_complexity)] +fn append_greater_sub_key_test( + create_table: fn( + &mut DbWriter, + &'static str, + ) -> DbResult>, +) where + T: DupSortTableType + DupSortUtils<(u32, u32), NoVersionValueWrapper>, +{ + let ((_reader, mut writer), _temp_dir) = get_test_env(); + let table_id = create_table(&mut writer, "table").unwrap(); + + let txn = writer.begin_rw_txn().unwrap(); + + let handle = txn.open_table(&table_id).unwrap(); + handle.append_greater_sub_key(&txn, &(2, 2), &22).unwrap(); + handle.append_greater_sub_key(&txn, &(2, 3), &23).unwrap(); + handle.append_greater_sub_key(&txn, &(1, 1), &11).unwrap(); + handle.append_greater_sub_key(&txn, &(3, 0), &30).unwrap(); + + // For DupSort tables append with key that already exists should fail. Try append with smaller + // bigger and equal values. + let result = handle.append_greater_sub_key(&txn, &(2, 2), &0); + assert_matches!(result, Err(DbError::Append)); + + let result = handle.append_greater_sub_key(&txn, &(2, 2), &22); + assert_matches!(result, Err(DbError::Append)); + + let result = handle.append_greater_sub_key(&txn, &(2, 2), &100); + assert_matches!(result, Err(DbError::Append)); + + // As before, but for the last main key. + let result = handle.append_greater_sub_key(&txn, &(3, 0), &0); + assert_matches!(result, Err(DbError::Append)); + + let result = handle.append_greater_sub_key(&txn, &(3, 0), &30); + assert_matches!(result, Err(DbError::Append)); + + let result = handle.append_greater_sub_key(&txn, &(3, 0), &100); + assert_matches!(result, Err(DbError::Append)); + + // Check the final database. + assert_eq!(handle.get(&txn, &(2, 2)).unwrap(), Some(22)); + assert_eq!(handle.get(&txn, &(2, 3)).unwrap(), Some(23)); + assert_eq!(handle.get(&txn, &(1, 1)).unwrap(), Some(11)); + assert_eq!(handle.get(&txn, &(3, 0)).unwrap(), Some(30)); +} + #[test] fn add_one_test() { let mut bytes;