feat(ios): nbstore swift native binding (#9211)

packages/frontend/mobile-native/src/lib.rs

@@ -1,4 +1,15 @@
+use std::fmt::Display;
+use std::str::FromStr;
+use std::time::SystemTime;
+
 use affine_common::hashcash::Stamp;
+use affine_nbstore::storage;
+use dashmap::{mapref::one::RefMut, DashMap, Entry};
+
+use crate::error::UniffiError;
+
+mod error;
+mod utils;
 
 uniffi::setup_scaffolding!("affine_mobile_native");
 
@@ -6,3 +17,680 @@ uniffi::setup_scaffolding!("affine_mobile_native");
 pub fn hashcash_mint(resource: String, bits: u32) -> String {
   Stamp::mint(resource, Some(bits)).format()
 }
+
+#[derive(uniffi::Record)]
+pub struct DocRecord {
+  pub doc_id: String,
+  // base64 encoded data
+  pub data: String,
+  pub timestamp: SystemTime,
+}
+
+impl From<affine_nbstore::DocRecord> for DocRecord {
+  fn from(record: affine_nbstore::DocRecord) -> Self {
+    Self {
+      doc_id: record.doc_id,
+      data: base64_simd::STANDARD.encode_to_string(&record.data),
+      timestamp: record.timestamp.and_utc().into(),
+    }
+  }
+}
+
+impl TryFrom<DocRecord> for affine_nbstore::DocRecord {
+  type Error = UniffiError;
+
+  fn try_from(record: DocRecord) -> Result<Self, Self::Error> {
+    Ok(Self {
+      doc_id: record.doc_id,
+      data: base64_simd::STANDARD
+        .decode_to_vec(record.data)
+        .map_err(|e| UniffiError::Base64DecodingError(e.to_string()))?,
+      timestamp: chrono::DateTime::<chrono::Utc>::from(record.timestamp).naive_utc(),
+    })
+  }
+}
+
+#[derive(uniffi::Record)]
+pub struct DocUpdate {
+  pub doc_id: String,
+  pub created_at: SystemTime,
+  // base64 encoded data
+  pub data: String,
+}
+
+impl From<affine_nbstore::DocUpdate> for DocUpdate {
+  fn from(update: affine_nbstore::DocUpdate) -> Self {
+    Self {
+      doc_id: update.doc_id,
+      created_at: update.created_at.and_utc().into(),
+      data: base64_simd::STANDARD.encode_to_string(&update.data),
+    }
+  }
+}
+
+impl From<DocUpdate> for affine_nbstore::DocUpdate {
+  fn from(update: DocUpdate) -> Self {
+    Self {
+      doc_id: update.doc_id,
+      created_at: chrono::DateTime::<chrono::Utc>::from(update.created_at).naive_utc(),
+      data: update.data.into(),
+    }
+  }
+}
+
+#[derive(uniffi::Record)]
+pub struct DocClock {
+  pub doc_id: String,
+  pub timestamp: SystemTime,
+}
+
+impl From<affine_nbstore::DocClock> for DocClock {
+  fn from(clock: affine_nbstore::DocClock) -> Self {
+    Self {
+      doc_id: clock.doc_id,
+      timestamp: clock.timestamp.and_utc().into(),
+    }
+  }
+}
+
+impl From<DocClock> for affine_nbstore::DocClock {
+  fn from(clock: DocClock) -> Self {
+    Self {
+      doc_id: clock.doc_id,
+      timestamp: chrono::DateTime::<chrono::Utc>::from(clock.timestamp).naive_utc(),
+    }
+  }
+}
+
+#[derive(uniffi::Record)]
+pub struct Blob {
+  pub key: String,
+  // base64 encoded data
+  pub data: String,
+  pub mime: String,
+  pub size: i64,
+  pub created_at: SystemTime,
+}
+
+impl From<affine_nbstore::Blob> for Blob {
+  fn from(blob: affine_nbstore::Blob) -> Self {
+    Self {
+      key: blob.key,
+      data: base64_simd::STANDARD.encode_to_string(&blob.data),
+      mime: blob.mime,
+      size: blob.size,
+      created_at: blob.created_at.and_utc().into(),
+    }
+  }
+}
+
+#[derive(uniffi::Record)]
+pub struct SetBlob {
+  pub key: String,
+  // base64 encoded data
+  pub data: String,
+  pub mime: String,
+}
+
+impl TryFrom<SetBlob> for affine_nbstore::SetBlob {
+  type Error = UniffiError;
+
+  fn try_from(blob: SetBlob) -> Result<Self, Self::Error> {
+    Ok(Self {
+      key: blob.key,
+      data: base64_simd::STANDARD
+        .decode_to_vec(blob.data)
+        .map_err(|e| UniffiError::Base64DecodingError(e.to_string()))?,
+      mime: blob.mime,
+    })
+  }
+}
+
+#[derive(uniffi::Record)]
+pub struct ListedBlob {
+  pub key: String,
+  pub size: i64,
+  pub mime: String,
+  pub created_at: SystemTime,
+}
+
+impl From<affine_nbstore::ListedBlob> for ListedBlob {
+  fn from(blob: affine_nbstore::ListedBlob) -> Self {
+    Self {
+      key: blob.key,
+      size: blob.size,
+      mime: blob.mime,
+      created_at: blob.created_at.and_utc().into(),
+    }
+  }
+}
+
+#[derive(uniffi::Object)]
+pub struct DocStoragePool {
+  inner: DashMap<String, storage::SqliteDocStorage>,
+}
+
+impl DocStoragePool {
+  fn ensure_storage<'a>(
+    &'a self,
+    universal_id: &str,
+  ) -> Result<RefMut<'a, String, storage::SqliteDocStorage>, UniffiError> {
+    let entry = self.inner.entry(universal_id.to_string());
+
+    if let Entry::Occupied(storage) = entry {
+      return Ok(storage.into_ref());
+    }
+    let options = parse_universal_id(entry.key())?;
+    let db_path = utils::get_db_path(&options)?;
+    if db_path.is_empty() {
+      return Err(UniffiError::EmptyDocStoragePath);
+    }
+    let storage = storage::SqliteDocStorage::new(db_path);
+    Ok(entry.or_insert(storage))
+  }
+}
+
+#[uniffi::export]
+impl DocStoragePool {
+  /// Initialize the database and run migrations.
+  pub async fn connect(&self, universal_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.connect().await?)
+  }
+
+  pub async fn close(&self, universal_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    storage.close().await;
+    self.inner.remove(&universal_id);
+    Ok(())
+  }
+
+  pub fn is_closed(&self, universal_id: String) -> bool {
+    let storage = self.ensure_storage(&universal_id).unwrap();
+    storage.is_closed()
+  }
+
+  pub async fn checkpoint(&self, universal_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.checkpoint().await?)
+  }
+
+  pub async fn validate(&self, universal_id: String) -> Result<bool, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.validate().await?)
+  }
+
+  pub async fn set_space_id(
+    &self,
+    universal_id: String,
+    space_id: String,
+  ) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    if space_id.is_empty() {
+      return Err(UniffiError::EmptySpaceId);
+    }
+    Ok(storage.set_space_id(space_id).await?)
+  }
+
+  pub async fn push_update(
+    &self,
+    universal_id: String,
+    doc_id: String,
+    update: String,
+  ) -> Result<SystemTime, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .push_update(
+          doc_id,
+          base64_simd::STANDARD
+            .decode_to_vec(update)
+            .map_err(|e| UniffiError::Base64DecodingError(e.to_string()))?,
+        )
+        .await?
+        .and_utc()
+        .into(),
+    )
+  }
+
+  pub async fn get_doc_snapshot(
+    &self,
+    universal_id: String,
+    doc_id: String,
+  ) -> Result<Option<DocRecord>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.get_doc_snapshot(doc_id).await?.map(Into::into))
+  }
+
+  pub async fn set_doc_snapshot(
+    &self,
+    universal_id: String,
+    snapshot: DocRecord,
+  ) -> Result<bool, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.set_doc_snapshot(snapshot.try_into()?).await?)
+  }
+
+  pub async fn get_doc_updates(
+    &self,
+    universal_id: String,
+    doc_id: String,
+  ) -> Result<Vec<DocUpdate>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_doc_updates(doc_id)
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn mark_updates_merged(
+    &self,
+    universal_id: String,
+    doc_id: String,
+    updates: Vec<SystemTime>,
+  ) -> Result<u32, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .mark_updates_merged(
+          doc_id,
+          updates
+            .into_iter()
+            .map(|t| chrono::DateTime::<chrono::Utc>::from(t).naive_utc())
+            .collect(),
+        )
+        .await?,
+    )
+  }
+
+  pub async fn delete_doc(&self, universal_id: String, doc_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.delete_doc(doc_id).await?)
+  }
+
+  pub async fn get_doc_clocks(
+    &self,
+    universal_id: String,
+    after: Option<SystemTime>,
+  ) -> Result<Vec<DocClock>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_doc_clocks(after.map(|t| chrono::DateTime::<chrono::Utc>::from(t).naive_utc()))
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn get_doc_clock(
+    &self,
+    universal_id: String,
+    doc_id: String,
+  ) -> Result<Option<DocClock>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.get_doc_clock(doc_id).await?.map(Into::into))
+  }
+
+  pub async fn get_blob(
+    &self,
+    universal_id: String,
+    key: String,
+  ) -> Result<Option<Blob>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.get_blob(key).await?.map(Into::into))
+  }
+
+  pub async fn set_blob(&self, universal_id: String, blob: SetBlob) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.set_blob(blob.try_into()?).await?)
+  }
+
+  pub async fn delete_blob(
+    &self,
+    universal_id: String,
+    key: String,
+    permanently: bool,
+  ) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.delete_blob(key, permanently).await?)
+  }
+
+  pub async fn release_blobs(&self, universal_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.release_blobs().await?)
+  }
+
+  pub async fn list_blobs(&self, universal_id: String) -> Result<Vec<ListedBlob>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .list_blobs()
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn get_peer_remote_clocks(
+    &self,
+    universal_id: String,
+    peer: String,
+  ) -> Result<Vec<DocClock>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_peer_remote_clocks(peer)
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn get_peer_remote_clock(
+    &self,
+    universal_id: String,
+    peer: String,
+    doc_id: String,
+  ) -> Result<DocClock, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.get_peer_remote_clock(peer, doc_id).await?.into())
+  }
+
+  pub async fn set_peer_remote_clock(
+    &self,
+    universal_id: String,
+    peer: String,
+    doc_id: String,
+    clock: SystemTime,
+  ) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .set_peer_remote_clock(
+          peer,
+          doc_id,
+          chrono::DateTime::<chrono::Utc>::from(clock).naive_utc(),
+        )
+        .await?,
+    )
+  }
+
+  pub async fn get_peer_pulled_remote_clocks(
+    &self,
+    universal_id: String,
+    peer: String,
+  ) -> Result<Vec<DocClock>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_peer_pulled_remote_clocks(peer)
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn get_peer_pulled_remote_clock(
+    &self,
+    universal_id: String,
+    peer: String,
+    doc_id: String,
+  ) -> Result<DocClock, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_peer_pulled_remote_clock(peer, doc_id)
+        .await?
+        .into(),
+    )
+  }
+
+  pub async fn set_peer_pulled_remote_clock(
+    &self,
+    universal_id: String,
+    peer: String,
+    doc_id: String,
+    clock: SystemTime,
+  ) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .set_peer_pulled_remote_clock(
+          peer,
+          doc_id,
+          chrono::DateTime::<chrono::Utc>::from(clock).naive_utc(),
+        )
+        .await?,
+    )
+  }
+
+  pub async fn get_peer_pushed_clocks(
+    &self,
+    universal_id: String,
+    peer: String,
+  ) -> Result<Vec<DocClock>, UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .get_peer_pushed_clocks(peer)
+        .await?
+        .into_iter()
+        .map(Into::into)
+        .collect(),
+    )
+  }
+
+  pub async fn set_peer_pushed_clock(
+    &self,
+    universal_id: String,
+    peer: String,
+    doc_id: String,
+    clock: SystemTime,
+  ) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(
+      storage
+        .set_peer_pushed_clock(
+          peer,
+          doc_id,
+          chrono::DateTime::<chrono::Utc>::from(clock).naive_utc(),
+        )
+        .await?,
+    )
+  }
+
+  pub async fn clear_clocks(&self, universal_id: String) -> Result<(), UniffiError> {
+    let storage = self.ensure_storage(&universal_id)?;
+    Ok(storage.clear_clocks().await?)
+  }
+}
+
+#[uniffi::export]
+pub fn get_db_path(peer: String, space_type: String, id: String) -> Result<String, UniffiError> {
+  let options = StorageOptions {
+    peer,
+    space_type: SpaceType::from_str(&space_type)?,
+    id,
+  };
+  utils::get_db_path(&options)
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Default)]
+pub enum SpaceType {
+  #[default]
+  Userspace,
+  Workspace,
+}
+
+impl Display for SpaceType {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    match self {
+      SpaceType::Userspace => write!(f, "userspace"),
+      SpaceType::Workspace => write!(f, "workspace"),
+    }
+  }
+}
+
+impl FromStr for SpaceType {
+  type Err = UniffiError;
+
+  fn from_str(s: &str) -> Result<Self, Self::Err> {
+    Ok(match s {
+      "userspace" => Self::Userspace,
+      "workspace" => Self::Workspace,
+      _ => return Err(UniffiError::InvalidSpaceType(s.to_string())),
+    })
+  }
+}
+
+pub struct StorageOptions {
+  pub peer: String,
+  pub space_type: SpaceType,
+  pub id: String,
+}
+
+pub fn parse_universal_id(id: &str) -> Result<StorageOptions, UniffiError> {
+  let mut result = StorageOptions {
+    peer: String::new(),
+    space_type: SpaceType::default(),
+    id: String::new(),
+  };
+
+  let mut key = String::new();
+  let mut value = String::new();
+  let mut is_in_value = false;
+  let mut chars = id.chars().peekable();
+
+  while let Some(ch) = chars.next() {
+    if is_in_value {
+      if ch == ')' && chars.peek() == Some(&';') {
+        // Store the collected value in the appropriate field
+        match key.as_str() {
+          "peer" => result.peer = value.clone(),
+          "type" => result.space_type = SpaceType::from_str(&value)?,
+          "id" => result.id = value.clone(),
+          _ => return Err(UniffiError::InvalidUniversalId(id.to_string())),
+        }
+        key.clear();
+        value.clear();
+        is_in_value = false;
+        chars.next(); // Skip the semicolon
+        continue;
+      }
+      value.push(ch);
+      continue;
+    }
+
+    if ch == '@' {
+      // Find the position of next '('
+      let mut temp_chars = chars.clone();
+      let mut found_paren = false;
+      let mut key_chars = Vec::new();
+
+      while let Some(next_ch) = temp_chars.next() {
+        if next_ch == '(' {
+          found_paren = true;
+          break;
+        }
+        key_chars.push(next_ch);
+      }
+
+      // Invalid format if no '(' found or it's immediately after '@'
+      if !found_paren || key_chars.is_empty() {
+        return Err(UniffiError::InvalidUniversalId(id.to_string()));
+      }
+
+      key = key_chars.into_iter().collect();
+      // Advance the original iterator to the position after the key
+      for _ in 0..key.len() + 1 {
+        chars.next();
+      }
+      is_in_value = true;
+    } else {
+      return Err(UniffiError::InvalidUniversalId(id.to_string()));
+    }
+  }
+
+  // Validate the parsed results
+  if result.peer.is_empty() || result.id.is_empty() {
+    return Err(UniffiError::InvalidUniversalId(id.to_string()));
+  }
+
+  Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  // ... existing test functions ...
+
+  #[test]
+  fn test_universal_id() {
+    let options = StorageOptions {
+      peer: "123".to_string(),
+      space_type: SpaceType::Workspace,
+      id: "456".to_string(),
+    };
+
+    let id = format!(
+      "@peer({});@type({});@id({});",
+      options.peer, options.space_type, options.id
+    );
+    let result = parse_universal_id(&id).unwrap();
+
+    assert_eq!(result.peer, "123");
+    assert_eq!(result.space_type, SpaceType::Workspace);
+    assert_eq!(result.id, "456");
+  }
+
+  #[test]
+  fn test_parse_universal_id_valid_cases() {
+    let testcases = vec![
+      "@peer(123);@type(userspace);@id(456);",
+      "@peer(123);@type(workspace);@id(456);",
+      "@peer(https://app.affine.pro);@type(userspace);@id(hello:world);",
+      "@peer(@name);@type(userspace);@id(@id);",
+      "@peer(@peer(name);@type(userspace);@id(@id);",
+    ];
+
+    for id in testcases {
+      let result = parse_universal_id(id);
+      assert!(result.is_ok(), "Failed to parse: {}", id);
+
+      let parsed = result.unwrap();
+      assert!(!parsed.peer.is_empty());
+      assert!(!parsed.id.is_empty());
+    }
+  }
+
+  #[test]
+  fn test_parse_universal_id_invalid_cases() {
+    let testcases = vec![
+      // invalid space type
+      "@peer(123);@type(anyspace);@id(456);",
+      // invalid peer
+      "@peer(@peer(name););@type(userspace);@id(@id);",
+    ];
+
+    for id in testcases {
+      let result = parse_universal_id(id);
+      assert!(result.is_err(), "Should have failed to parse: {}", id);
+
+      match result {
+        Err(UniffiError::InvalidUniversalId(_)) => (),
+        Err(UniffiError::InvalidSpaceType(_)) => (),
+        _ => panic!("Expected InvalidUniversalId error for: {}", id),
+      }
+    }
+  }
+}