Reintroduce #289

examples/heed3-encryption.rs → examples/heed3-encrypted.rs

@@ -4,8 +4,8 @@ use std::path::Path;
 
 use argon2::Argon2;
 use chacha20poly1305::{ChaCha20Poly1305, Key};
-use heed3_encryption::types::*;
-use heed3_encryption::{Database, EnvOpenOptions};
+use heed3::types::*;
+use heed3::{Database, EnvOpenOptions};
 
 fn main() -> Result<(), Box<dyn Error>> {
     let env_path = Path::new("target").join("encrypt.mdb");
@@ -21,12 +21,12 @@ fn main() -> Result<(), Box<dyn Error>> {
     Argon2::default().hash_password_into(password.as_bytes(), salt.as_bytes(), &mut key)?;
 
     // We open the environment
-    let mut options = EnvOpenOptions::<ChaCha20Poly1305>::new_encrypted_with(key);
+    let mut options = EnvOpenOptions::new();
     let env = unsafe {
         options
             .map_size(10 * 1024 * 1024) // 10MB
             .max_dbs(3)
-            .open(&env_path)?
+            .open_encrypted::<ChaCha20Poly1305, _>(key, &env_path)?
     };
 
     let key1 = "first-key";
@@ -36,11 +36,11 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // We create database and write secret values in it
     let mut wtxn = env.write_txn()?;
-    let db: Database<Str, Str> = env.create_database(&mut wtxn, Some("first"))?;
+    let db = env.create_database::<Str, Str>(&mut wtxn, Some("first"))?;
     db.put(&mut wtxn, key1, val1)?;
     db.put(&mut wtxn, key2, val2)?;
     wtxn.commit()?;
-    env.prepare_for_closing().wait();
+    // env.prepare_for_closing().wait();
 
     // We reopen the environment now
     let env = unsafe { options.open(&env_path)? };

heed/src/databases/database.rs

@@ -139,8 +139,7 @@ impl<'e, 'n, KC, DC, C> DatabaseOpenOptions<'e, 'n, KC, DC, C> {
     {
         assert_eq_env_txn!(self.env, rtxn);
 
-        let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() };
-        match self.env.raw_init_database::<C>(raw_txn, self.name, self.flags) {
+        match self.env.raw_init_database::<C>(rtxn.txn.unwrap(), self.name, self.flags) {
             Ok(dbi) => Ok(Some(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi))),
             Err(Error::Mdb(e)) if e.not_found() => Ok(None),
             Err(e) => Err(e),
@@ -165,8 +164,7 @@ impl<'e, 'n, KC, DC, C> DatabaseOpenOptions<'e, 'n, KC, DC, C> {
         assert_eq_env_txn!(self.env, wtxn);
 
         let flags = self.flags | AllDatabaseFlags::CREATE;
-        let raw_txn = unsafe { wtxn.txn.txn.unwrap().as_mut() };
-        match self.env.raw_init_database::<C>(raw_txn, self.name, flags) {
+        match self.env.raw_init_database::<C>(wtxn.txn.txn.unwrap(), self.name, flags) {
             Ok(dbi) => Ok(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi)),
             Err(e) => Err(e),
         }
@@ -352,7 +350,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
 
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut data_val = mem::MaybeUninit::uninit();
-        let mut txn = txn.txn.unwrap();
 
         let result = unsafe {
             mdb_result(ffi::mdb_get(
@@ -1129,7 +1126,7 @@ impl<KC, DC, C> Database<KC, DC, C> {
         RoCursor::new(txn, self.dbi).map(|cursor| RoRevIter::new(cursor))
     }
 
-    /// Return a mutable reversed lexicographically ordered iterator of all key-value\
+    /// Return a mutable reverse ordered iterator of all key-value\
     /// pairs in this database.
     ///
     /// ```
@@ -1842,7 +1839,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut data_val = unsafe { crate::into_val(&data_bytes) };
         let flags = 0;
-        let mut txn = txn.txn.txn.unwrap();
 
         unsafe {
             mdb_result(ffi::mdb_put(
@@ -1910,7 +1906,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut reserved = ffi::reserve_size_val(data_size);
         let flags = ffi::MDB_RESERVE;
-        let mut txn = txn.txn.txn.unwrap();
 
         unsafe {
             mdb_result(ffi::mdb_put(
@@ -2009,7 +2004,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut data_val = unsafe { crate::into_val(&data_bytes) };
         let flags = flags.bits();
-        let mut txn = txn.txn.txn.unwrap();
 
         unsafe {
             mdb_result(ffi::mdb_put(
@@ -2123,7 +2117,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut data_val = unsafe { crate::into_val(&data_bytes) };
         let flags = (flags | PutFlags::NO_OVERWRITE).bits();
-        let mut txn = txn.txn.txn.unwrap();
 
         let result = unsafe {
             mdb_result(ffi::mdb_put(
@@ -2280,7 +2273,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut reserved = ffi::reserve_size_val(data_size);
         let flags = (flags | PutFlags::NO_OVERWRITE).bits() | ffi::MDB_RESERVE;
-        let mut txn = txn.txn.txn.unwrap();
 
         let result = unsafe {
             mdb_result(ffi::mdb_put(
@@ -2364,7 +2356,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
 
         let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?;
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
-        let mut txn = txn.txn.txn.unwrap();
 
         let result = unsafe {
             mdb_result(ffi::mdb_del(
@@ -2457,7 +2448,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
         let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Encoding)?;
         let mut key_val = unsafe { crate::into_val(&key_bytes) };
         let mut data_val = unsafe { crate::into_val(&data_bytes) };
-        let mut txn = txn.txn.txn.unwrap();
 
         let result = unsafe {
             mdb_result(ffi::mdb_del(
@@ -2589,7 +2579,6 @@ impl<KC, DC, C> Database<KC, DC, C> {
     /// ```
     pub fn clear(&self, txn: &mut RwTxn) -> Result<()> {
         assert_eq_env_db_txn!(self, txn);
-        let mut txn = txn.txn.txn.unwrap();
 
         unsafe {
             mdb_result(ffi::mdb_drop(txn.txn.txn.unwrap().as_mut(), self.dbi, 0))
@@ -2682,7 +2671,7 @@ mod tests {
     use heed_types::*;
 
     use super::*;
-    use crate::env::IntegerComparator;
+    use crate::IntegerComparator;
 
     #[test]
     fn put_overwrite() -> Result<()> {

heed/src/databases/encrypted_database.rs

@@ -1015,9 +1015,9 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
         self.inner.rev_iter_mut(txn)
     }
 
-    /// Return a lexicographically ordered iterator of a range of key-value pairs in this database.
+    /// Return an ordered iterator of a range of key-value pairs in this database.
     ///
-    /// Comparisons are made by using the bytes representation of the key.
+    /// Comparisons are made by using the comparator `C`.
     ///
     /// You can make this iterator `Send`able between threads by
     /// using the `read-txn-no-tls` crate feature.
@@ -1062,18 +1062,18 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
         &self,
         txn: &'txn mut RoTxn,
         range: &'a R,
-    ) -> Result<RoRange<'txn, KC, DC>>
+    ) -> Result<RoRange<'txn, KC, DC, C>>
     where
         KC: BytesEncode<'a>,
         R: RangeBounds<KC::EItem>,
     {
         self.inner.range(txn, range)
     }
 
-    /// Return a mutable lexicographically ordered iterator of a range of
+    /// Return a mutable ordered iterator of a range of
     /// key-value pairs in this database.
     ///
-    /// Comparisons are made by using the bytes representation of the key.
+    /// Comparisons are made by using the comparator `C`.
     ///
     /// ```
     /// # use std::fs;
@@ -1128,18 +1128,18 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
         &self,
         txn: &'txn mut RwTxn,
         range: &'a R,
-    ) -> Result<RwRange<'txn, KC, DC>>
+    ) -> Result<RwRange<'txn, KC, DC, C>>
     where
         KC: BytesEncode<'a>,
         R: RangeBounds<KC::EItem>,
     {
         self.inner.range_mut(txn, range)
     }
 
-    /// Return a reversed lexicographically ordered iterator of a range of key-value
+    /// Return a reverse ordered iterator of a range of key-value
     /// pairs in this database.
     ///
-    /// Comparisons are made by using the bytes representation of the key.
+    /// Comparisons are made by using the comparator `C`.
     ///
     /// You can make this iterator `Send`able between threads by
     /// using the `read-txn-no-tls` crate feature.
@@ -1184,18 +1184,18 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
         &self,
         txn: &'txn mut RoTxn,
         range: &'a R,
-    ) -> Result<RoRevRange<'txn, KC, DC>>
+    ) -> Result<RoRevRange<'txn, KC, DC, C>>
     where
         KC: BytesEncode<'a>,
         R: RangeBounds<KC::EItem>,
     {
         self.inner.rev_range(txn, range)
     }
 
-    /// Return a mutable reversed lexicographically ordered iterator of a range of
+    /// Return a mutable reverse ordered iterator of a range of
     /// key-value pairs in this database.
     ///
-    /// Comparisons are made by using the bytes representation of the key.
+    /// Comparisons are made by using the comparator `C`.
     ///
     /// ```
     /// # use std::fs;
@@ -1250,7 +1250,7 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
         &self,
         txn: &'txn mut RwTxn,
         range: &'a R,
-    ) -> Result<RwRevRange<'txn, KC, DC>>
+    ) -> Result<RwRevRange<'txn, KC, DC, C>>
     where
         KC: BytesEncode<'a>,
         R: RangeBounds<KC::EItem>,
@@ -2020,7 +2020,7 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
     ///
     /// Prefer using [`clear`] instead of a call to this method with a full range ([`..`]).
     ///
-    /// Comparisons are made by using the bytes representation of the key.
+    /// Comparisons are made by using the comparator `C`.
     ///
     /// [`clear`]: crate::Database::clear
     /// [`..`]: std::ops::RangeFull
@@ -2068,6 +2068,7 @@ impl<KC, DC, C> EncryptedDatabase<KC, DC, C> {
     pub fn delete_range<'a, 'txn, R>(&self, txn: &'txn mut RwTxn, range: &'a R) -> Result<usize>
     where
         KC: BytesEncode<'a> + BytesDecode<'txn>,
+        C: Comparator,
         R: RangeBounds<KC::EItem>,
     {
         self.inner.delete_range(txn, range)

heed/src/envs/env.rs

@@ -9,7 +9,7 @@ use heed_traits::Comparator;
 
 use super::{
     custom_key_cmp_wrapper, get_file_fd, metadata_from_fd, DefaultComparator, EnvInfo, FlagSetMode,
-    OPENED_ENV,
+    IntegerComparator, OPENED_ENV,
 };
 use crate::cursor::{MoveOperation, RoCursor};
 use crate::mdb::ffi::{self, MDB_env};
@@ -158,8 +158,7 @@ impl Env {
 
         let rtxn = self.read_txn()?;
         // Open the main database
-        let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() };
-        let dbi = self.raw_open_dbi::<DefaultComparator>(raw_txn, None, 0)?;
+        let dbi = self.raw_open_dbi::<DefaultComparator>(rtxn.txn.unwrap(), None, 0)?;
 
         // We're going to iterate on the unnamed database
         let mut cursor = RoCursor::new(&rtxn, dbi)?;
@@ -172,8 +171,9 @@ impl Env {
             let key = String::from_utf8(key.to_vec()).unwrap();
             // Calling `ffi::db_stat` on a database instance does not involve key comparison
             // in LMDB, so it's safe to specify a noop key compare function for it.
-            let raw_txn = unsafe { rtxn.txn.unwrap().as_mut() };
-            if let Ok(dbi) = self.raw_open_dbi::<DefaultComparator>(raw_txn, Some(&key), 0) {
+            if let Ok(dbi) =
+                self.raw_open_dbi::<DefaultComparator>(rtxn.txn.unwrap(), Some(&key), 0)
+            {
                 let mut stat = mem::MaybeUninit::uninit();
                 unsafe {
                     mdb_result(ffi::mdb_stat(rtxn.txn.unwrap().as_mut(), dbi, stat.as_mut_ptr()))?
@@ -268,7 +268,7 @@ impl Env {
 
     fn raw_open_dbi<C: Comparator + 'static>(
         &self,
-        raw_txn: NonNull<ffi::MDB_txn>,
+        mut raw_txn: NonNull<ffi::MDB_txn>,
         name: Option<&str>,
         flags: u32,
     ) -> std::result::Result<u32, crate::mdb::lmdb_error::Error> {
@@ -282,13 +282,17 @@ impl Env {
         // safety: The name cstring is cloned by LMDB, we can drop it after.
         //         If a read-only is used with the MDB_CREATE flag, LMDB will throw an error.
         unsafe {
-            mdb_result(ffi::mdb_dbi_open(raw_txn, name_ptr, flags, &mut dbi))?;
+            mdb_result(ffi::mdb_dbi_open(raw_txn.as_mut(), name_ptr, flags, &mut dbi))?;
             let cmp_type_id = TypeId::of::<C>();
 
             if cmp_type_id != TypeId::of::<DefaultComparator>()
                 && cmp_type_id != TypeId::of::<IntegerComparator>()
             {
-                mdb_result(ffi::mdb_set_compare(raw_txn, dbi, Some(custom_key_cmp_wrapper::<C>)))?;
+                mdb_result(ffi::mdb_set_compare(
+                    raw_txn.as_mut(),
+                    dbi,
+                    Some(custom_key_cmp_wrapper::<C>),
+                ))?;
             }
         };
 

heed/src/envs/mod.rs

@@ -8,9 +8,7 @@ use std::os::unix::io::{AsRawFd, BorrowedFd, RawFd};
 use std::panic::catch_unwind;
 use std::path::{Path, PathBuf};
 use std::process::abort;
-use std::ptr::NonNull;
 use std::sync::{LazyLock, RwLock};
-use std::time::Duration;
 #[cfg(windows)]
 use std::{
     ffi::OsStr,
@@ -181,6 +179,39 @@ impl LexicographicComparator for DefaultComparator {
     }
 }
 
+/// A representation of LMDB's `MDB_INTEGERKEY` comparator behavior.
+///
+/// This enum is used to indicate a table should be sorted by the keys numeric
+/// value in native byte order. When a [`Database`] is created or opened with
+/// [`IntegerComparator`], it signifies that the comparator should not be explicitly
+/// set via [`ffi::mdb_set_compare`], instead the flag [`AllDatabaseFlags::INTEGER_KEY`]
+/// is set on the table.
+///
+/// This can only be used on certain types: either `u32` or `usize`.
+/// The keys must all be of the same size.
+pub enum IntegerComparator {}
+
+impl Comparator for IntegerComparator {
+    fn compare(a: &[u8], b: &[u8]) -> Ordering {
+        #[cfg(target_endian = "big")]
+        return a.cmp(b);
+
+        #[cfg(target_endian = "little")]
+        {
+            let len = a.len();
+
+            for i in (0..len).rev() {
+                match a[i].cmp(&b[i]) {
+                    Ordering::Equal => continue,
+                    other => return other,
+                }
+            }
+
+            Ordering::Equal
+        }
+    }
+}
+
 /// Whether to perform compaction while copying an environment.
 #[derive(Debug, Copy, Clone)]
 pub enum CompactionOption {

heed/src/txn.rs

@@ -47,7 +47,7 @@ use crate::Result;
 pub struct RoTxn<'e> {
     /// Makes the struct covariant and !Sync
     pub(crate) txn: Option<NonNull<ffi::MDB_txn>>,
-    env: Cow<'e, Env>,
+    env: &'e Env,
 }
 
 impl<'e> RoTxn<'e> {

heed3/Cargo.toml

@@ -118,5 +118,5 @@ longer-keys = ["lmdb-master3-sys/longer-keys"]
 # Examples are located outside the standard heed/examples directory to prevent
 # conflicts between heed3 and heed examples when working on both crates.
 [[example]]
-name = "heed3-encryption"
-path = "../examples/heed3-encryption.rs"
+name = "heed3-encrypted"
+path = "../examples/heed3-encrypted.rs"