From 68eca4feed06be901d4e8b9489baf8d871f1dcab Mon Sep 17 00:00:00 2001
From: PThorpe92 <preston@unlockedlabs.org>
Date: Thu, 6 Mar 2025 15:37:02 -0500
Subject: [PATCH] Add demo vfs module to vtab kvstore

---
 extensions/core/src/lib.rs         |  4 +-
 extensions/core/src/vfs_modules.rs |  4 +-
 extensions/kvstore/src/lib.rs      | 63 +++++++++++++++++++++++++++++-
 3 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/extensions/core/src/lib.rs b/extensions/core/src/lib.rs
index 417391f84..6cdaa8a05 100644
--- a/extensions/core/src/lib.rs
+++ b/extensions/core/src/lib.rs
@@ -1,13 +1,13 @@
 mod types;
 mod vfs_modules;
-pub use limbo_macros::{register_extension, scalar, AggregateDerive, VTabModuleDerive};
+pub use limbo_macros::{register_extension, scalar, AggregateDerive, VTabModuleDerive, VfsDerive};
 use std::{
     fmt::Display,
     os::raw::{c_char, c_void},
 };
 pub use types::{ResultCode, Value, ValueType};
 use vfs_modules::RegisterVfsFn;
-pub use vfs_modules::{VfsFileImpl, VfsImpl};
+pub use vfs_modules::{VfsExtension, VfsFile, VfsFileImpl, VfsImpl};
 
 pub type ExtResult<T> = std::result::Result<T, ResultCode>;
 
diff --git a/extensions/core/src/vfs_modules.rs b/extensions/core/src/vfs_modules.rs
index 6421596d4..3b2ab3a28 100644
--- a/extensions/core/src/vfs_modules.rs
+++ b/extensions/core/src/vfs_modules.rs
@@ -2,7 +2,7 @@ use crate::{ExtResult, ResultCode};
 use std::ffi::{c_char, c_void};
 
 #[cfg(not(target_family = "wasm"))]
-pub trait VfsExtension: Default {
+pub trait VfsExtension: Default + Send + Sync {
     const NAME: &'static str;
     type File: VfsFile;
     fn open_file(&self, path: &str, flags: i32, direct: bool) -> ExtResult<Self::File>;
@@ -23,7 +23,7 @@ pub trait VfsExtension: Default {
 }
 
 #[cfg(not(target_family = "wasm"))]
-pub trait VfsFile: Sized {
+pub trait VfsFile: Send + Sync {
     fn lock(&mut self, _exclusive: bool) -> ExtResult<()> {
         Ok(())
     }
diff --git a/extensions/kvstore/src/lib.rs b/extensions/kvstore/src/lib.rs
index a9de7c71d..dbd17d5f1 100644
--- a/extensions/kvstore/src/lib.rs
+++ b/extensions/kvstore/src/lib.rs
@@ -1,8 +1,11 @@
 use lazy_static::lazy_static;
 use limbo_ext::{
-    register_extension, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, Value,
+    register_extension, scalar, ExtResult, ResultCode, VTabCursor, VTabKind, VTabModule,
+    VTabModuleDerive, Value, VfsDerive, VfsExtension, VfsFile,
 };
 use std::collections::BTreeMap;
+use std::fs::{File, OpenOptions};
+use std::io::{Read, Seek, SeekFrom, Write};
 use std::sync::Mutex;
 
 lazy_static! {
@@ -145,3 +148,61 @@ impl VTabCursor for KVStoreCursor {
         <KVStoreVTab as VTabModule>::next(self)
     }
 }
+
+struct TestFile {
+    file: File,
+}
+
+#[derive(VfsDerive, Default)]
+struct TestFS;
+
+// Test that we can have additional extension types in the same file
+// and still register the vfs at comptime if linking staticly
+#[scalar(name = "test_scalar")]
+fn test_scalar(_args: limbo_ext::Value) -> limbo_ext::Value {
+    limbo_ext::Value::from_integer(42)
+}
+
+impl VfsExtension for TestFS {
+    const NAME: &'static str = "testvfs";
+    type File = TestFile;
+    fn open_file(&self, path: &str, flags: i32, _direct: bool) -> ExtResult<Self::File> {
+        let file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(flags & 1 != 0)
+            .open(path)
+            .map_err(|_| ResultCode::Error)?;
+        Ok(TestFile { file })
+    }
+}
+
+impl VfsFile for TestFile {
+    fn read(&mut self, buf: &mut [u8], count: usize, offset: i64) -> ExtResult<i32> {
+        if self.file.seek(SeekFrom::Start(offset as u64)).is_err() {
+            return Err(ResultCode::Error);
+        }
+        self.file
+            .read(&mut buf[..count])
+            .map_err(|_| ResultCode::Error)
+            .map(|n| n as i32)
+    }
+
+    fn write(&mut self, buf: &[u8], count: usize, offset: i64) -> ExtResult<i32> {
+        if self.file.seek(SeekFrom::Start(offset as u64)).is_err() {
+            return Err(ResultCode::Error);
+        }
+        self.file
+            .write(&buf[..count])
+            .map_err(|_| ResultCode::Error)
+            .map(|n| n as i32)
+    }
+
+    fn sync(&self) -> ExtResult<()> {
+        self.file.sync_all().map_err(|_| ResultCode::Error)
+    }
+
+    fn size(&self) -> i64 {
+        self.file.metadata().map(|m| m.len() as i64).unwrap_or(-1)
+    }
+}