Add support for borrowed and owned storage in Negentropy

Signed-off-by: Yuki Kishimoto <yukikishimoto@protonmail.com>
rust-nostr · Jan 25, 2025 · 7d120e8 · 7d120e8
1 parent cc3973b
commit 7d120e8
Show file tree

Hide file tree

Showing 6 changed files with 98 additions and 64 deletions.
diff --git a/negentropy-ffi/src/lib.rs b/negentropy-ffi/src/lib.rs
@@ -25,7 +25,7 @@ pub struct ReconcileWithIds {
 
 #[derive(Object)]
 pub struct Negentropy {
-    inner: Mutex<negentropy::Negentropy<negentropy::NegentropyStorageVector>>,
+    inner: Mutex<negentropy::Negentropy<'static, negentropy::NegentropyStorageVector>>,
 }
 
 #[uniffi::export]
@@ -36,7 +36,7 @@ impl Negentropy {
     #[uniffi::constructor]
     pub fn new(storage: &NegentropyStorageVector, frame_size_limit: Option<u64>) -> Result<Self> {
         Ok(Self {
-            inner: Mutex::new(negentropy::Negentropy::new(
+            inner: Mutex::new(negentropy::Negentropy::owned(
                 storage.to_inner()?,
                 frame_size_limit.unwrap_or_default(),
             )?),

diff --git a/negentropy/examples/negentropy.rs b/negentropy/examples/negentropy.rs
@@ -19,7 +19,7 @@ fn main() {
         )
         .unwrap();
     storage_client.seal().unwrap();
-    let mut client = Negentropy::new(storage_client, 0).unwrap();
+    let mut client = Negentropy::borrowed(&storage_client, 0).unwrap();
     let init_output = client.initiate().unwrap();
     println!("Initiator Output: {:x?}", init_output.clone());
 
@@ -56,7 +56,7 @@ fn main() {
         )
         .unwrap();
     storage_relay.seal().unwrap();
-    let mut relay = Negentropy::new(storage_relay, 0).unwrap();
+    let mut relay = Negentropy::borrowed(&storage_relay, 0).unwrap();
     let reconcile_output = relay.reconcile(&init_output).unwrap();
     println!("Reconcile Output: {:x?}", reconcile_output.clone());
 

diff --git a/negentropy/fuzz/harness/src/main.rs b/negentropy/fuzz/harness/src/main.rs
@@ -38,7 +38,7 @@ fn main() {
         }
     }
 
-    let mut ne = Negentropy::new(storage, frame_size_limit as u64).unwrap();
+    let mut ne = Negentropy::borrowed(&storage, frame_size_limit as u64).unwrap();
 
     for line in io::stdin().lock().lines() {
         let line_unwrapped = line.unwrap();

diff --git a/negentropy/fuzz/perf/src/main.rs b/negentropy/fuzz/perf/src/main.rs
@@ -33,7 +33,7 @@ fn main() {
         )
         .unwrap();
     storage_client.seal().unwrap();
-    let mut client = Negentropy::new(storage_client, 0).unwrap();
+    let mut client = Negentropy::borrowed(&storage_client, 0).unwrap();
     let now = Instant::now();
     let init_output = client.initiate().unwrap();
     println!("Client init took {} ms", now.elapsed().as_millis());
@@ -47,7 +47,7 @@ fn main() {
             .unwrap();
     }
     storage_relay.seal().unwrap();
-    let mut relay = Negentropy::new(storage_relay, 0).unwrap();
+    let mut relay = Negentropy::borrowed(&storage_relay, 0).unwrap();
     let now = Instant::now();
     let reconcile_output = relay.reconcile(&init_output).unwrap();
     println!("Relay reconcile took {} ms", now.elapsed().as_millis());

diff --git a/negentropy/src/lib.rs b/negentropy/src/lib.rs
@@ -36,7 +36,7 @@ pub use self::constants::{FINGERPRINT_SIZE, ID_SIZE, PROTOCOL_VERSION};
 use self::encoding::{decode_var_int, encode_var_int, get_byte_array, get_bytes};
 pub use self::error::Error;
 pub use self::id::Id;
-pub use self::storage::{NegentropyStorageBase, NegentropyStorageVector};
+pub use self::storage::{NegentropyStorageBase, NegentropyStorageVector, Storage};
 use self::types::Mode;
 pub use self::types::{Bound, Item};
 
@@ -45,22 +45,22 @@ const BUCKETS: usize = 16;
 const DOUBLE_BUCKETS: usize = BUCKETS * 2;
 
 /// Negentropy
-pub struct Negentropy<T> {
-    storage: T,
+pub struct Negentropy<'a, T> {
+    storage: Storage<'a, T>,
     frame_size_limit: u64,
     is_initiator: bool,
     last_timestamp_in: u64,
     last_timestamp_out: u64,
 }
 
-impl<T> Negentropy<T>
+impl<'a, T> Negentropy<'a, T>
 where
     T: NegentropyStorageBase,
 {
     /// Create new [`Negentropy`] instance
     ///
     /// Frame size limit must be `equal to 0` or `greater than 4096`
-    pub fn new(storage: T, frame_size_limit: u64) -> Result<Self, Error> {
+    pub fn new(storage: Storage<'a, T>, frame_size_limit: u64) -> Result<Self, Error> {
         if frame_size_limit != 0 && frame_size_limit < 4096 {
             return Err(Error::FrameSizeLimitTooSmall);
         }
@@ -74,6 +74,20 @@ where
         })
     }
 
+    /// Create new [`Negentropy`] instance from owned storage
+    ///
+    /// Frame size limit must be `equal to 0` or `greater than 4096`
+    pub fn owned(storage: T, frame_size_limit: u64) -> Result<Self, Error> {
+        Self::new(Storage::Owned(storage), frame_size_limit)
+    }
+
+    /// Create new [`Negentropy`] instance from owned storage
+    ///
+    /// Frame size limit must be `equal to 0` or `greater than 4096`
+    pub fn borrowed(storage: &'a T, frame_size_limit: u64) -> Result<Self, Error> {
+        Self::new(Storage::Borrowed(storage), frame_size_limit)
+    }
+
     /// Initiate reconciliation set
     pub fn initiate(&mut self) -> Result<Vec<u8>, Error> {
         if self.is_initiator {
@@ -457,7 +471,7 @@ mod tests {
             .unwrap();
         storage_client.seal().unwrap();
 
-        let mut client = Negentropy::new(storage_client, 0).unwrap();
+        let mut client = Negentropy::new(Storage::Borrowed(&storage_client), 0).unwrap();
         let init_output = client.initiate().unwrap();
 
         // Relay
@@ -513,7 +527,7 @@ mod tests {
             )
             .unwrap();
         storage_relay.seal().unwrap();
-        let mut relay = Negentropy::new(storage_relay, 0).unwrap();
+        let mut relay = Negentropy::new(Storage::Borrowed(&storage_relay), 0).unwrap();
         let reconcile_output = relay.reconcile(&init_output).unwrap();
 
         // Client

diff --git a/negentropy/src/storage.rs b/negentropy/src/storage.rs
@@ -5,10 +5,30 @@
 //! Module that contains the various storage implementations
 
 use alloc::vec::Vec;
+use core::ops::Deref;
 
 use crate::types::{Accumulator, Bound, Fingerprint, Item};
 use crate::{Error, Id};
 
+/// Storage
+pub enum Storage<'a, T: 'a> {
+    /// Borrowed
+    Borrowed(&'a T),
+    /// Owned
+    Owned(T),
+}
+
+impl<T> Deref for Storage<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Borrowed(b) => b,
+            Self::Owned(b) => b,
+        }
+    }
+}
+
 /// NegentropyStorageBase
 pub trait NegentropyStorageBase {
     /// Size
@@ -48,56 +68,6 @@ pub struct NegentropyStorageVector {
     sealed: bool,
 }
 
-impl NegentropyStorageBase for NegentropyStorageVector {
-    fn size(&self) -> Result<usize, Error> {
-        self.check_sealed()?;
-        Ok(self.items.len())
-    }
-
-    fn get_item(&self, i: usize) -> Result<Option<Item>, Error> {
-        self.check_sealed()?;
-        Ok(self.items.get(i).copied())
-    }
-
-    fn iterate(
-        &self,
-        begin: usize,
-        end: usize,
-        cb: &mut dyn FnMut(Item, usize) -> Result<bool, Error>,
-    ) -> Result<(), Error> {
-        self.check_sealed()?;
-        self.check_bounds(begin, end)?;
-
-        for i in begin..end {
-            if !cb(self.items[i], i)? {
-                break;
-            }
-        }
-
-        Ok(())
-    }
-
-    fn find_lower_bound(&self, mut first: usize, last: usize, value: &Bound) -> usize {
-        let mut count: usize = last - first;
-
-        while count > 0 {
-            let mut it: usize = first;
-            let step: usize = count / 2;
-            it += step;
-
-            if self.items[it] < value.item {
-                it += 1;
-                first = it;
-                count -= step + 1;
-            } else {
-                count = step;
-            }
-        }
-
-        first
-    }
-}
-
 impl NegentropyStorageVector {
     /// Create new storage
     #[inline]
@@ -159,3 +129,53 @@ impl NegentropyStorageVector {
         Ok(())
     }
 }
+
+impl NegentropyStorageBase for NegentropyStorageVector {
+    fn size(&self) -> Result<usize, Error> {
+        self.check_sealed()?;
+        Ok(self.items.len())
+    }
+
+    fn get_item(&self, i: usize) -> Result<Option<Item>, Error> {
+        self.check_sealed()?;
+        Ok(self.items.get(i).copied())
+    }
+
+    fn iterate(
+        &self,
+        begin: usize,
+        end: usize,
+        cb: &mut dyn FnMut(Item, usize) -> Result<bool, Error>,
+    ) -> Result<(), Error> {
+        self.check_sealed()?;
+        self.check_bounds(begin, end)?;
+
+        for i in begin..end {
+            if !cb(self.items[i], i)? {
+                break;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn find_lower_bound(&self, mut first: usize, last: usize, value: &Bound) -> usize {
+        let mut count: usize = last - first;
+
+        while count > 0 {
+            let mut it: usize = first;
+            let step: usize = count / 2;
+            it += step;
+
+            if self.items[it] < value.item {
+                it += 1;
+                first = it;
+                count -= step + 1;
+            } else {
+                count = step;
+            }
+        }
+
+        first
+    }
+}