add Database::get_or_put_reserved* variations

heed/src/database.rs

@@ -2093,14 +2093,163 @@ impl<KC, DC, C> Database<KC, DC, C> {
  lmdb_master_sys::MDB_KEYEXIST => {
  let bytes = crate::from_val(data_val);
  let data = DC::bytes_decode(bytes).map_err(Error::Decoding)?;
- Ok(Some(data))
+ return Ok(Some(data));
  }
- err_code => {
- mdb_result(err_code)?;
- Ok(None)
+ err_code => mdb_result(err_code)?,
+ }
+ }
+
+ Ok(None)
+ }
+
+ /// Attempt to insert a key-value pair in this database, where the value can be directly
+ /// written to disk, or if a value already exists for the key, returns the previous value.
+ ///
+ /// The entry is written with no specific flags.
+ ///
+ /// ```
+ /// # use heed::EnvOpenOptions;
+ /// use std::io::Write;
+ /// use heed::{Database, PutFlags};
+ /// use heed::types::*;
+ /// use heed::byteorder::BigEndian;
+ ///
+ /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+ /// # let dir = tempfile::tempdir()?;
+ /// # let env = unsafe { EnvOpenOptions::new()
+ /// # .map_size(10 * 1024 * 1024) // 10MB
+ /// # .max_dbs(3000)
+ /// # .open(dir.path())?
+ /// # };
+ /// type BEI32 = I32<BigEndian>;
+ ///
+ /// let mut wtxn = env.write_txn()?;
+ /// let db = env.create_database::<BEI32, Str>(&mut wtxn, Some("number-string"))?;
+ ///
+ /// # db.clear(&mut wtxn)?;
+ /// let long = "I am a long long long value";
+ /// assert_eq!(
+ /// db.get_or_put_reserved(&mut wtxn, &42, long.len(), |reserved| {
+ /// reserved.write_all(long.as_bytes())
+ /// })?,
+ /// None
+ /// );
+ ///
+ /// let longer = "I am an even longer long long long value";
+ /// assert_eq!(
+ /// db.get_or_put_reserved(&mut wtxn, &42, longer.len(), |reserved| {
+ /// unreachable!()
+ /// })?,
+ /// Some(long)
+ /// );
+ ///
+ /// let ret = db.get(&mut wtxn, &42)?;
+ /// assert_eq!(ret, Some(long));
+ ///
+ /// wtxn.commit()?;
+ /// # Ok(()) }
+ /// ```
+ pub fn get_or_put_reserved<'a, 'txn, F>(
+ &'txn self,
+ txn: &mut RwTxn,
+ key: &'a KC::EItem,
+ data_size: usize,
+ write_func: F,
+ ) -> Result<Option<DC::DItem>>
+ where
+ KC: BytesEncode<'a>,
+ F: FnOnce(&mut ReservedSpace) -> io::Result<()>,
+ DC: BytesDecode<'a>,
+ {
+ self.get_or_put_reserved_with_flags(txn, PutFlags::empty(), key, data_size, write_func)
+ }
+
+ /// Attempt to insert a key-value pair in this database, where the value can be directly
+ /// written to disk, or if a value already exists for the key, returns the previous value.
+ ///
+ /// The entry is written with the specified flags.
+ ///
+ /// ```
+ /// # use heed::EnvOpenOptions;
+ /// use std::io::Write;
+ /// use heed::{Database, PutFlags};
+ /// use heed::types::*;
+ /// use heed::byteorder::BigEndian;
+ ///
+ /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+ /// # let dir = tempfile::tempdir()?;
+ /// # let env = unsafe { EnvOpenOptions::new()
+ /// # .map_size(10 * 1024 * 1024) // 10MB
+ /// # .max_dbs(3000)
+ /// # .open(dir.path())?
+ /// # };
+ /// type BEI32 = I32<BigEndian>;
+ ///
+ /// let mut wtxn = env.write_txn()?;
+ /// let db = env.create_database::<BEI32, Str>(&mut wtxn, Some("number-string"))?;
+ ///
+ /// # db.clear(&mut wtxn)?;
+ /// let long = "I am a long long long value";
+ /// assert_eq!(
+ /// db.get_or_put_reserved_with_flags(&mut wtxn, PutFlags::empty(), &42, long.len(), |reserved| {
+ /// reserved.write_all(long.as_bytes())
+ /// })?,
+ /// None
+ /// );
+ ///
+ /// let longer = "I am an even longer long long long value";
+ /// assert_eq!(
+ /// db.get_or_put_reserved_with_flags(&mut wtxn, PutFlags::empty(), &42, longer.len(), |reserved| {
+ /// unreachable!()
+ /// })?,
+ /// Some(long)
+ /// );
+ ///
+ /// let ret = db.get(&mut wtxn, &42)?;
+ /// assert_eq!(ret, Some(long));
+ ///
+ /// wtxn.commit()?;
+ /// # Ok(()) }
+ /// ```
+ pub fn get_or_put_reserved_with_flags<'a, 'txn, F>(
+ &'txn self,
+ txn: &mut RwTxn,
+ flags: PutFlags,
+ key: &'a KC::EItem,
+ data_size: usize,
+ write_func: F,
+ ) -> Result<Option<DC::DItem>>
+ where
+ KC: BytesEncode<'a>,
+ F: FnOnce(&mut ReservedSpace) -> io::Result<()>,
+ DC: BytesDecode<'a>,
+ {
+ assert_eq_env_db_txn!(self, txn);
+
+ 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 reserved = ffi::reserve_size_val(data_size);
+ let flags = (flags | PutFlags::NO_OVERWRITE).bits() | lmdb_master_sys::MDB_RESERVE;
+
+ unsafe {
+ match ffi::mdb_put(txn.txn.txn, self.dbi, &mut key_val, &mut reserved, flags) {
+ lmdb_master_sys::MDB_KEYEXIST => {
+ let bytes = crate::from_val(reserved);
+ let data = DC::bytes_decode(bytes).map_err(Error::Decoding)?;
+ return Ok(Some(data));
  }
+ err_code => mdb_result(err_code)?,
  }
  }
+
+ let mut reserved = unsafe { ReservedSpace::from_val(reserved) };
+ write_func(&mut reserved)?;
+ if reserved.remaining() == 0 {
+ Ok(None)
+ } else {
+ Err(io::Error::from(io::ErrorKind::UnexpectedEof).into())
+ }
  }
 
  /// Deletes an entry or every duplicate data items of a key