perf(storage): add append_sub_key functionallity (#1900)

* perf(storage): add append functionallity to tables

* CR fix

* perf(storage): add append_sub_key functionallity

* CR fix

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<K, V>>
+    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: &<V as ValueSerde>::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::<DbKeyType<'_>, DbValueType<'_>>()? {
+            if prev.1.starts_with(&T::get_sub_key(key)?) {
+                cursor.next_dup::<DbKeyType<'_>, DbValueType<'_>>()?;
+                cursor.del(WriteFlags::empty())?;
+                return Err(DbError::Append);
+            }
+        }
+
+        Ok(())
+    }
+}
+
 impl<
     'txn,
     Mode: TransactionKind,

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<T>(
+    create_table: fn(
+        &mut DbWriter,
+        &'static str,
+    ) -> DbResult<TableIdentifier<TableKey, TableValue, T>>,
+) where
+    T: DupSortTableType + DupSortUtils<(u32, u32), NoVersionValueWrapper<u32>>,
+{
+    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;