Cache site id -> ordinal mapping during tx

somtochiama · somtochiama · commit 7f682af595f9 · 2025-10-31T13:52:05.000+01:00
Signed-off-by: Somtochi Onyekwere &lt;somtochionyekwere@gmail.com&gt;
diff --git a/core/rs/core/src/c.rs b/core/rs/core/src/c.rs
@@ -73,6 +73,7 @@ pub struct crsql_ExtData {
     pub pSelectClockTablesStmt: *mut sqlite::stmt,
     pub mergeEqualValues: ::core::ffi::c_int,
     pub timestamp: ::core::ffi::c_ulonglong,
+    pub ordinalMap: *mut ::core::ffi::c_void,
 }
 
 #[repr(C)]
@@ -268,7 +269,7 @@ fn bindgen_test_layout_crsql_ExtData() {
     let ptr = UNINIT.as_ptr();
     assert_eq!(
         ::core::mem::size_of::<crsql_ExtData>(),
-        160usize,
+        168usize,
         concat!("Size of: ", stringify!(crsql_ExtData))
     );
     assert_eq!(
@@ -498,4 +499,14 @@ fn bindgen_test_layout_crsql_ExtData() {
             stringify!(timestamp)
         )
     );
+    assert_eq!(
+        unsafe { ::core::ptr::addr_of!((*ptr).ordinalMap) as usize - ptr as usize },
+        160usize,
+        concat!(
+            "Offset of field: ",
+            stringify!(crsql_ExtData),
+            "::",
+            stringify!(ordinalMap)
+        )
+    );
 }
diff --git a/core/rs/core/src/commit.rs b/core/rs/core/src/commit.rs
@@ -1,5 +1,7 @@
+use alloc::{boxed::Box, collections::BTreeMap, vec::Vec};
 use core::{
     ffi::{c_int, c_void},
+    mem,
     ptr::null,
 };
 
@@ -31,4 +33,9 @@ pub unsafe fn commit_or_rollback_reset(ext_data: *mut crsql_ExtData) {
     (*ext_data).seq = 0;
     (*ext_data).timestamp = 0;
     (*ext_data).updatedTableInfosThisTx = 0;
+
+    let mut ordinals: mem::ManuallyDrop<Box<BTreeMap<Vec<u8>, i64>>> = mem::ManuallyDrop::new(
+        Box::from_raw((*ext_data).ordinalMap as *mut BTreeMap<Vec<u8>, i64>),
+    );
+    ordinals.clear();
 }
diff --git a/core/rs/core/src/db_version.rs b/core/rs/core/src/db_version.rs
@@ -187,6 +187,12 @@ pub extern "C" fn crsql_init_last_db_versions_map(ext_data: *mut crsql_ExtData)
     unsafe { (*ext_data).lastDbVersions = Box::into_raw(Box::new(map)) as *mut c_void }
 }
 
+#[no_mangle]
+pub extern "C" fn crsql_init_ordinal_map(ext_data: *mut crsql_ExtData) {
+    let map: BTreeMap<Vec<u8>, i64> = BTreeMap::new();
+    unsafe { (*ext_data).ordinalMap = Box::into_raw(Box::new(map)) as *mut c_void }
+}
+
 #[no_mangle]
 pub extern "C" fn crsql_drop_last_db_versions_map(ext_data: *mut crsql_ExtData) {
     unsafe {
@@ -196,6 +202,15 @@ pub extern "C" fn crsql_drop_last_db_versions_map(ext_data: *mut crsql_ExtData)
     }
 }
 
+#[no_mangle]
+pub extern "C" fn crsql_drop_ordinal_map(ext_data: *mut crsql_ExtData) {
+    unsafe {
+        drop(Box::from_raw(
+            (*ext_data).ordinalMap as *mut BTreeMap<Vec<u8>, i64>,
+        ));
+    }
+}
+
 pub fn insert_db_version(
     ext_data: *mut crsql_ExtData,
     insert_site_id: &[u8],
@@ -258,6 +273,15 @@ pub unsafe fn get_or_set_site_ordinal(
     ext_data: *mut crsql_ExtData,
     site_id: &[u8],
 ) -> Result<i64, ResultCode> {
+    // check the cache first
+    let mut ordinals: mem::ManuallyDrop<Box<BTreeMap<Vec<u8>, i64>>> = mem::ManuallyDrop::new(
+        Box::from_raw((*ext_data).ordinalMap as *mut BTreeMap<Vec<u8>, i64>),
+    );
+
+    if let Some(ordinal) = ordinals.get(site_id) {
+        return Ok(*ordinal);
+    }
+
     let bind_result =
         (*ext_data)
             .pSelectSiteIdOrdinalStmt
@@ -268,11 +292,11 @@ pub unsafe fn get_or_set_site_ordinal(
         return Err(rc);
     }
 
-    match (*ext_data).pSelectSiteIdOrdinalStmt.step() {
+    let ordinal = match (*ext_data).pSelectSiteIdOrdinalStmt.step() {
         Ok(ResultCode::ROW) => {
             let ordinal = (*ext_data).pSelectSiteIdOrdinalStmt.column_int64(0);
             reset_cached_stmt((*ext_data).pSelectSiteIdOrdinalStmt)?;
-            Ok(ordinal)
+            ordinal
         }
         Ok(_) => {
             reset_cached_stmt((*ext_data).pSelectSiteIdOrdinalStmt)?;
@@ -296,7 +320,7 @@ pub unsafe fn get_or_set_site_ordinal(
                 Ok(_) => {
                     let ordinal = (*ext_data).pSetSiteIdOrdinalStmt.column_int64(0);
                     reset_cached_stmt((*ext_data).pSetSiteIdOrdinalStmt)?;
-                    Ok(ordinal)
+                    ordinal
                 }
                 Err(rc) => {
                     reset_cached_stmt((*ext_data).pSetSiteIdOrdinalStmt)?;
@@ -308,5 +332,7 @@ pub unsafe fn get_or_set_site_ordinal(
             reset_cached_stmt((*ext_data).pSetSiteIdOrdinalStmt)?;
             return Err(rc);
         }
-    }
+    };
+    ordinals.insert(site_id.to_vec(), ordinal);
+    Ok(ordinal)
 }
diff --git a/core/rs/integration_check/src/t/test_db_version.rs b/core/rs/integration_check/src/t/test_db_version.rs
@@ -1,6 +1,6 @@
 extern crate alloc;
-use alloc::{ffi::CString, string::String};
-use core::ffi::c_char;
+use alloc::{boxed::Box, collections::BTreeMap, ffi::CString, format, string::String, vec::Vec};
+use core::{ffi::c_char, mem};
 use crsql_bundle::test_exports;
 use sqlite::{Connection, ResultCode};
 use sqlite_nostd as sqlite;
@@ -10,6 +10,11 @@ fn make_site() -> *mut c_char {
     inner_ptr
 }
 
+fn convert_to_bytes(id: *mut c_char) -> Vec<u8> {
+    let site_id_string = unsafe { CString::from_raw(id) };
+    site_id_string.into_bytes()
+}
+
 fn get_site_id(db: *mut sqlite::sqlite3) -> *mut c_char {
     let stmt = db
         .prepare_v2("SELECT crsql_site_id();")
@@ -112,8 +117,49 @@ fn test_next_db_version() -> Result<(), String> {
     Ok(())
 }
 
+fn test_get_or_set_site_ordinal() -> Result<(), ResultCode> {
+    let c = crate::opendb().expect("db opened");
+    let db = &c.db;
+    let raw_db = db.db;
+    let ext_data = unsafe { test_exports::c::crsql_newExtData(raw_db, get_site_id(raw_db)) };
+
+    let site_id = convert_to_bytes(get_site_id(raw_db));
+    let ordinal = unsafe { test_exports::db_version::get_or_set_site_ordinal(ext_data, &site_id)? };
+    assert_eq!(0, ordinal);
+
+    let mut ordinals = unsafe {
+        mem::ManuallyDrop::new(Box::from_raw(
+            (*ext_data).ordinalMap as *mut BTreeMap<Vec<u8>, i64>,
+        ))
+    };
+
+    assert_eq!(0, *(ordinals.get(&site_id).unwrap()));
+
+    // update ordinal in db but it should remain the same on the map
+    let update_stmt =
+        raw_db.prepare_v2("UPDATE crsql_site_id SET ordinal = 8 WHERE site_id = ?")?;
+    update_stmt.bind_blob(1, &site_id, sqlite::Destructor::STATIC)?;
+    update_stmt.step()?;
+
+    let ordinal = unsafe { test_exports::db_version::get_or_set_site_ordinal(ext_data, &site_id)? };
+    assert_eq!(0, ordinal);
+
+    // clear ordinals and call func again, we should get update
+    ordinals.clear();
+    let ordinal = unsafe { test_exports::db_version::get_or_set_site_ordinal(ext_data, &site_id)? };
+    assert_eq!(8, ordinal);
+
+    unsafe {
+        test_exports::c::crsql_freeExtData(ext_data);
+    };
+    Ok(())
+}
+
 pub fn run_suite() -> Result<(), String> {
     test_fetch_db_version_from_storage()?;
     test_next_db_version()?;
+    if let Err(rc) = test_get_or_set_site_ordinal() {
+        return Err(format!("test_get_or_set_site_ordinal failed: {:?}", rc));
+    }
     Ok(())
 }
diff --git a/core/src/ext-data.c b/core/src/ext-data.c
@@ -7,6 +7,8 @@
 
 void crsql_clear_stmt_cache(crsql_ExtData *pExtData);
 void crsql_init_table_info_vec(crsql_ExtData *pExtData);
+void crsql_init_ordinal_map(crsql_ExtData *pExtData);
+void crsql_drop_ordinal_map(crsql_ExtData *pExtData);
 void crsql_drop_table_info_vec(crsql_ExtData *pExtData);
 void crsql_init_last_db_versions_map(crsql_ExtData *pExtData);
 void crsql_drop_last_db_versions_map(crsql_ExtData *pExtData);
@@ -74,10 +76,12 @@ crsql_ExtData *crsql_newExtData(sqlite3 *db, unsigned char *siteIdBuffer) {
       -1, SQLITE_PREPARE_PERSISTENT, &(pExtData->pDbVersionStmt), 0);
   pExtData->tableInfos = 0;
   pExtData->lastDbVersions = 0;
+  pExtData->ordinalMap = 0;
   pExtData->rowsImpacted = 0;
   pExtData->updatedTableInfosThisTx = 0;
   crsql_init_table_info_vec(pExtData);
   crsql_init_last_db_versions_map(pExtData);
+  crsql_init_ordinal_map(pExtData);
 
   sqlite3_stmt *pStmt;
 
@@ -147,6 +151,7 @@ void crsql_freeExtData(crsql_ExtData *pExtData) {
   crsql_clear_stmt_cache(pExtData);
   crsql_drop_table_info_vec(pExtData);
   crsql_drop_last_db_versions_map(pExtData);
+  crsql_drop_ordinal_map(pExtData);
   sqlite3_free(pExtData);
 }
 
diff --git a/core/src/ext-data.h b/core/src/ext-data.h
@@ -48,6 +48,7 @@ struct crsql_ExtData {
 
   int mergeEqualValues;
   unsigned long long timestamp;
+  void *ordinalMap;
 };
 
 crsql_ExtData *crsql_newExtData(sqlite3 *db, unsigned char *siteIdBuffer);
@@ -57,5 +58,7 @@ int crsql_fetchPragmaSchemaVersion(sqlite3 *db, crsql_ExtData *pExtData,
 int crsql_fetchPragmaDataVersion(sqlite3 *db, crsql_ExtData *pExtData);
 int crsql_recreate_db_version_stmt(sqlite3 *db, crsql_ExtData *pExtData);
 void crsql_finalize(crsql_ExtData *pExtData);
+void crsql_init_ordinal_map(crsql_ExtData *pExtData);
+void crsql_drop_ordinal_map(crsql_ExtData *pExtData);
 
 #endif
