Use stable memory

paulyoung · paulyoung · commit 61c0b172ccca · 2022-09-12T16:44:40.000-07:00
diff --git a/crates/ic-sqlite/src/lib.rs b/crates/ic-sqlite/src/lib.rs
@@ -1,27 +1,26 @@
-use std::cell::RefCell;
-use std::ffi::CString;
-use std::os::raw::c_char;
-use std::ptr::NonNull;
-
+use crate::vfs::PagesVfs;
+use icfs::StableMemory;
 use rusqlite::{params_from_iter, OpenFlags, Row, Rows};
 use serde::ser::Serializer;
 use serde::Serialize;
 use serde_json::Value as JsonValue;
 use sqlite_vfs::{register, RegisterError};
-
-pub use crate::vfs::PagesVfs;
+use std::cell::RefCell;
+use std::ffi::CString;
+use std::os::raw::c_char;
+use std::ptr::NonNull;
 
 mod vfs;
 
 // TODO: icfs::stable_memory::WASM_PAGE_SIZE_IN_BYTES;
 
-extern "C" {
-    pub fn page_count() -> u32;
-    pub fn get_page(ix: u32, ptr: *mut u8);
-    pub fn put_page(ix: u32, ptr: *const u8);
-    pub fn del_page(ix: u32);
-    pub fn conn_sleep(ms: u32);
-}
+// extern "C" {
+//     pub fn page_count() -> u32;
+//     pub fn get_page(ix: u32, ptr: *mut u8);
+//     pub fn put_page(ix: u32, ptr: *const u8);
+//     pub fn del_page(ix: u32);
+//     pub fn conn_sleep(ms: u32);
+// }
 
 // TODO: is there any way to provide this method for SQLite, but not export it as part of the WASM
 // module?
@@ -49,7 +48,7 @@ pub struct Connection {
 
 #[no_mangle]
 pub unsafe extern "C" fn conn_new() -> *mut Connection {
-    let is_new = page_count() == 0;
+    let is_new = StableMemory::capacity() == 0;
 
     let conn = rusqlite::Connection::open_with_flags_and_vfs(
         &vfs::DB_NAME,
diff --git a/crates/ic-sqlite/src/vfs.rs b/crates/ic-sqlite/src/vfs.rs
@@ -1,7 +1,7 @@
 // use ic_cdk::api::stable;
-// use icfs::StableMemory;
+use icfs::StableMemory;
 use std::convert::TryInto;
-use std::io::{self, ErrorKind};
+use std::io::{self, ErrorKind, Read, Seek, SeekFrom, Write};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
@@ -10,6 +10,11 @@ use sqlite_vfs::{LockKind, OpenKind, OpenOptions, Vfs};
 pub const DB_NAME: &str = "main.db";
 pub const VFS_NAME: &str = "ic-sqlite";
 
+thread_local! {
+    static STABLE_MEMORY: std::cell::RefCell<StableMemory>
+        = std::cell::RefCell::new(StableMemory::default());
+}
+
 #[derive(Default)]
 pub struct PagesVfs<const PAGE_SIZE: usize> {
     lock_state: Arc<Mutex<LockState>>,
@@ -69,6 +74,10 @@ impl<const PAGE_SIZE: usize> Vfs for PagesVfs<PAGE_SIZE> {
     fn random(&self, buffer: &mut [i8]) {
         // Calling `raw_rand` would be preferable, but this method can't be async.
         // let raw_rand: Vec<u8> = call(Principal::management_canister(), "raw_rand", ()).await;
+        //
+        // We can't write and register a custom getrandom function either because that can't be async:
+        // * https://github.com/rust-random/rand/blob/b73640705d6714509f8ceccc49e8df996fa19f51/README.md#wasm-support
+        // * https://docs.rs/getrandom/0.2.7/getrandom/macro.register_custom_getrandom.html
         todo!();
         // rand::Rng::fill(&mut rand::thread_rng(), buffer);
     }
@@ -187,27 +196,39 @@ impl<const PAGE_SIZE: usize> sqlite_vfs::DatabaseHandle for Connection<PAGE_SIZE
     }
 }
 
+fn put_page<const PAGE_SIZE: usize>(ix: u32, data: &[u8; PAGE_SIZE]) {
+    STABLE_MEMORY.with(|stable_memory| {
+        let mut stable_memory = *stable_memory.borrow();
+        let page_size: u64 = PAGE_SIZE.try_into().unwrap();
+        let offset: u64 = ix as u64 * page_size;
+        stable_memory.seek(SeekFrom::Start(offset)).unwrap();
+        stable_memory.write(data).unwrap();
+    })
+}
+
 impl<const PAGE_SIZE: usize> Connection<PAGE_SIZE> {
     fn get_page(ix: u32) -> [u8; PAGE_SIZE] {
-        let mut data = [0u8; PAGE_SIZE];
-        unsafe { crate::get_page(ix, data.as_mut_ptr()) };
-        data
+        STABLE_MEMORY.with(|stable_memory| {
+            let mut stable_memory = *stable_memory.borrow();
+            let mut data = [0u8; PAGE_SIZE];
+            let page_size: u64 = PAGE_SIZE.try_into().unwrap();
+            let offset: u64 = ix as u64 * page_size;
+            stable_memory.seek(SeekFrom::Start(offset)).unwrap();
+            stable_memory.read(&mut data).unwrap();
+            data
+        })
     }
 
     fn put_page(ix: u32, data: &[u8; PAGE_SIZE]) {
-        unsafe {
-            crate::put_page(ix, data.as_ptr());
-        }
+        put_page(ix, data)
     }
 
     fn del_page(ix: u32) {
-        unsafe {
-            crate::del_page(ix);
-        }
+        put_page(ix, &[0; PAGE_SIZE])
     }
 
     fn page_count() -> usize {
-        unsafe { crate::page_count() as usize }
+        StableMemory::size().try_into().unwrap()
     }
 
     fn lock(&mut self, to: LockKind) -> bool {
