From 0fad41d1e2ce6aece709559032796432fca48d57 Mon Sep 17 00:00:00 2001
From: Imbris <imbrisf@gmail.com>
Date: Mon, 13 Feb 2023 00:34:32 -0500
Subject: [PATCH] Uncomment `drain` module, update safety comments and
 implement `LendJoin` for `Drain`.

---
 src/join/lend_join.rs |  4 ++--
 src/join/mod.rs       |  4 +---
 src/storage/drain.rs  | 43 +++++++++++++++++++++++++++++++++++++------
 src/storage/mod.rs    | 18 +++++++++---------
 4 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/src/join/lend_join.rs b/src/join/lend_join.rs
index b25358953..4d0fac6c3 100644
--- a/src/join/lend_join.rs
+++ b/src/join/lend_join.rs
@@ -139,8 +139,8 @@ pub unsafe trait LendJoin {
 
 /// # Safety
 ///
-/// Implementing this trait guarantees that `<Self as LendJoin>::get` can soundly be called
-/// multiple times with the same ID.
+/// Implementing this trait guarantees that `<Self as LendJoin>::get` can
+/// soundly be called multiple times with the same ID.
 pub unsafe trait RepeatableLendGet: LendJoin {}
 
 /// Type alias to refer to the `<J as LendJoin>::Type<'next>` (except this
diff --git a/src/join/mod.rs b/src/join/mod.rs
index 40ee5a001..7b41d6ca2 100644
--- a/src/join/mod.rs
+++ b/src/join/mod.rs
@@ -95,8 +95,7 @@ pub use par_join::{JoinParIter, ParJoin};
 /// The `Self::Mask` value returned with the `Self::Value` must correspond such
 /// that it is safe to retrieve items from `Self::Value` whose presence is
 /// indicated in the mask. As part of this, `BitSetLike::iter` must not produce
-/// an iterator that repeats an `Index` value. (S-TODO update impls: probably
-/// drain)
+/// an iterator that repeats an `Index` value.
 pub unsafe trait Join {
     /// Type of joined components.
     type Type;
@@ -125,7 +124,6 @@ pub unsafe trait Join {
 
     /// Get a joined component value by a given index.
     ///
-    // S-TODO: evaluate all impls (TODO: probably restrict, entry, and drain)
     ///
     /// # Safety
     ///
diff --git a/src/storage/drain.rs b/src/storage/drain.rs
index 8946fcff4..294b996f5 100644
--- a/src/storage/drain.rs
+++ b/src/storage/drain.rs
@@ -1,7 +1,9 @@
 use hibitset::BitSet;
 
+#[nougat::gat(Type)]
+use crate::join::LendJoin;
 use crate::{
-    join::Join,
+    join::{Join, RepeatableLendGet},
     storage::MaskedStorage,
     world::{Component, Index},
 };
@@ -13,10 +15,41 @@ pub struct Drain<'a, T: Component> {
     pub data: &'a mut MaskedStorage<T>,
 }
 
-// S-TODO implement LendJoin
-// S-TODO implement RepeatableLendGet
+// SAFETY: Calling `get` is always safe! Iterating the mask does not repeat
+// indices.
+#[nougat::gat]
+unsafe impl<'a, T> LendJoin for Drain<'a, T>
+where
+    T: Component,
+{
+    type Mask = BitSet;
+    type Type<'next> = T;
+    type Value = &'a mut MaskedStorage<T>;
+
+    unsafe fn open(self) -> (Self::Mask, Self::Value) {
+        // TODO: Cloning the whole bitset here seems expensive, and it is
+        // hidden from the user, but there is no obvious way to restructure
+        // things to avoid this with the way that bitsets are composed together
+        // for iteration.
+        let mask = self.data.mask.clone();
+
+        (mask, self.data)
+    }
+
+    unsafe fn get<'next>(value: &'next mut Self::Value, id: Index) -> T
+    where
+        Self: 'next,
+    {
+        value.remove(id).expect("Tried to access same index twice")
+    }
+}
+
+// SAFETY: Calling `get` is always safe!
+unsafe impl<'a, T> RepeatableLendGet for Drain<'a, T> where T: Component {}
 
-impl<'a, T> Join for Drain<'a, T>
+// SAFETY: Calling `get` is always safe! Iterating the mask does not repeat
+// indices.
+unsafe impl<'a, T> Join for Drain<'a, T>
 where
     T: Component,
 {
@@ -24,7 +57,6 @@ where
     type Type = T;
     type Value = &'a mut MaskedStorage<T>;
 
-    // SAFETY: No invariants to meet and no unsafe code.
     unsafe fn open(self) -> (Self::Mask, Self::Value) {
         // TODO: Cloning the whole bitset here seems expensive, and it is
         // hidden from the user, but there is no obvious way to restructure
@@ -35,7 +67,6 @@ where
         (mask, self.data)
     }
 
-    // SAFETY: No invariants to meet and no unsafe code.
     unsafe fn get(value: &mut Self::Value, id: Index) -> T {
         value.remove(id).expect("Tried to access same index twice")
     }
diff --git a/src/storage/mod.rs b/src/storage/mod.rs
index 9c72f7184..0fccdbedb 100644
--- a/src/storage/mod.rs
+++ b/src/storage/mod.rs
@@ -37,12 +37,12 @@ use crate::{
     world::{Component, EntitiesRes, Entity, Index},
 };
 
-// D-TODO use self::drain::Drain;
+use self::drain::Drain;
 use self::sync_unsafe_cell::SyncUnsafeCell;
 
 mod data;
 mod deref_flagged;
-// D-TODO mod drain;
+mod drain;
 mod entry;
 mod flagged;
 mod generic;
@@ -449,13 +449,13 @@ where
         self.data.clear();
     }
 
-    // /// Creates a draining storage wrapper which can be `.join`ed
-    // /// to get a draining iterator.
-    // D-TODO pub fn drain(&mut self) -> Drain<T> {
-    //    Drain {
-    //        data: &mut self.data,
-    //    }
-    //}
+    /// Creates a draining storage wrapper which can be `.join`ed
+    /// to get a draining iterator.
+    pub fn drain(&mut self) -> Drain<T> {
+        Drain {
+            data: &mut self.data,
+        }
+    }
 }
 
 impl<'a, T, D: Clone> Clone for Storage<'a, T, D> {