From 270d08ccfb57c8dd4c4f2b83cef24f34205bb817 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Thu, 12 Jan 2023 15:48:33 +0100
Subject: [PATCH] Introduce the RememberedPoly/Database and OpenedDatabases
 types

---
 heed/src/db/mod.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/heed/src/db/mod.rs b/heed/src/db/mod.rs
index 26c56881..b5d818d7 100644
--- a/heed/src/db/mod.rs
+++ b/heed/src/db/mod.rs
@@ -1,5 +1,67 @@
+use std::marker;
+
+use crate::mdb::ffi;
+#[allow(unused_imports)]
+use crate::{RoTxn, RwTxn};
+
 mod database;
 mod poly_database;
 
 pub use self::database::Database;
 pub use self::poly_database::PolyDatabase;
+
+/// A dummy handle to a [`PolyDatabase`] that can be used with an [`OpenedDatabases`]
+/// to retrieve a global handle to a database and operate on them.
+///
+/// An `OpenedDatabases` type can be retrieve via a call to [`RoTxn::commit`] or [`RwTxn::commit`].
+pub struct RememberedPolyDatabase {
+    pub(crate) unique_txn_id: u64,
+    pub(crate) env_ident: usize,
+    pub(crate) dbi: ffi::MDB_dbi,
+}
+
+/// A dummy handle to a [`Database`] that can be used with an [`OpenedDatabases`]
+/// to retrieve a global handle to a database and operate on them.
+///
+/// An `OpenedDatabases` type can be retrieve via a call to [`RoTxn::commit`] or [`RwTxn::commit`].
+pub struct RememberedDatabase<KC, DC> {
+    pub(crate) unique_txn_id: u64,
+    pub(crate) env_ident: usize,
+    pub(crate) dbi: ffi::MDB_dbi,
+    pub(crate) marker: marker::PhantomData<(KC, DC)>,
+}
+
+/// This is used to retrieve global handles to [`Database`]s and [`PolyDatabase`]s
+/// once a [`RoTxn::commit`]/[`RwTxn::commit`] has been executed.
+pub struct OpenedDatabases {
+    pub(crate) unique_txn_id: u64,
+}
+
+impl OpenedDatabases {
+    /// Give it a [`RememberedPolyDatabase`] and it will return a global handle to this very [`PolyDatabase`].
+    pub fn retrieve_poly_database(
+        &self,
+        remembered: RememberedPolyDatabase,
+    ) -> PolyDatabase<'static> {
+        let RememberedPolyDatabase { unique_txn_id, env_ident, dbi } = remembered;
+        // create a macro like the `assert_eq_env_txn` or `assert_eq_env_db_txn` ones
+        assert_eq!(
+            self.unique_txn_id, unique_txn_id,
+            "A PolyDatabase must be retrieved from the same transaction as the one it was opened with"
+        );
+        PolyDatabase::new(env_ident, dbi)
+    }
+
+    /// Give it a [`RememberedDatabase`] and it will return a global handle to this very [`Database`].
+    pub fn retrieve_database<KC, DC>(
+        &self,
+        remembered: RememberedDatabase<KC, DC>,
+    ) -> Database<'static, KC, DC> {
+        let RememberedDatabase { unique_txn_id, env_ident, dbi, marker: _ } = remembered;
+        assert_eq!(
+            self.unique_txn_id, unique_txn_id,
+            "A Database must be retrieved from the same transaction as the one it was opened with"
+        );
+        Database::new(env_ident, dbi)
+    }
+}