From ff994c34913e92e0c43241e8233975e2cb5aae5b Mon Sep 17 00:00:00 2001
From: Micah <micah@uplift.games>
Date: Mon, 6 May 2024 15:24:12 -0700
Subject: [PATCH] Add `clear` and `drain` implementations to `Attributes`
 (#409)

---
 rbx_types/CHANGELOG.md          |  2 ++
 rbx_types/src/attributes/mod.rs | 57 +++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/rbx_types/CHANGELOG.md b/rbx_types/CHANGELOG.md
index c4e5d85b6..aa139fe2a 100644
--- a/rbx_types/CHANGELOG.md
+++ b/rbx_types/CHANGELOG.md
@@ -3,8 +3,10 @@
 ## Unreleased Changes
 * Implement `IntoIterator` for `&Attributes`. ([#386])
 * Implement `Extend<(String, Variant)>` for `Attributes`. ([#386])
+* Implement `clear` and `drain` for `Attributes`. ([#409])
 
 [#386]: https://github.com/rojo-rbx/rbx-dom/pull/386
+[#409]: https://github.com/rojo-rbx/rbx-dom/pull/409
 
 ## 1.8.0 (2024-01-16)
 * Add `len` and `is_empty` methods to `Attributes` struct. ([#377])
diff --git a/rbx_types/src/attributes/mod.rs b/rbx_types/src/attributes/mod.rs
index e858d0d05..6432f575a 100644
--- a/rbx_types/src/attributes/mod.rs
+++ b/rbx_types/src/attributes/mod.rs
@@ -74,6 +74,12 @@ impl Attributes {
         self.data.remove(key.borrow())
     }
 
+    /// Removes all attributes.
+    #[inline]
+    pub fn clear(&mut self) {
+        self.data.clear()
+    }
+
     /// Returns an iterator of borrowed attributes.
     #[inline]
     pub fn iter(&self) -> AttributesIter<'_> {
@@ -82,6 +88,14 @@ impl Attributes {
         }
     }
 
+    /// Removes all elements from the `Attributes`, returning them as an
+    /// iterator. If the iterator is dropped before being fully consumed,
+    /// it drops the remaining removed elements.
+    #[inline]
+    pub fn drain(&mut self) -> AttributesDrain<'_> {
+        AttributesDrain { inner: self }
+    }
+
     /// Returns the number of attributes.
     #[inline]
     pub fn len(&self) -> usize {
@@ -159,6 +173,28 @@ impl<'a> Iterator for AttributesIter<'a> {
     }
 }
 
+/// A draining iterator for `Attributes`.
+/// This is created by [`Attributes::drain`].
+///
+/// If dropped before fully used, all remaining values will be dropped.
+pub struct AttributesDrain<'a> {
+    inner: &'a mut Attributes,
+}
+
+impl<'a> Iterator for AttributesDrain<'a> {
+    type Item = (String, Variant);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.data.pop_first()
+    }
+}
+
+impl<'a> Drop for AttributesDrain<'a> {
+    fn drop(&mut self) {
+        self.inner.clear()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -231,4 +267,25 @@ mod tests {
         attributes.insert("key".to_owned(), Variant::Bool(true));
         assert_eq!(attributes.remove("key"), Some(Variant::Bool(true)));
     }
+
+    #[test]
+    fn attribute_drain() {
+        let mut attributes = Attributes::new();
+        attributes.extend([
+            ("string".into(), "value".into()),
+            ("float64".into(), 10.0_f64.into()),
+            ("bool".into(), true.into()),
+        ]);
+        assert!(!attributes.is_empty());
+
+        let mut map = BTreeMap::new();
+        for (key, value) in attributes.drain() {
+            map.insert(key, value);
+        }
+
+        assert_eq!(map.get("string"), Some(&Variant::String("value".into())));
+        assert_eq!(map.get("float64"), Some(&Variant::Float64(10.0)));
+        assert_eq!(map.get("bool"), Some(&Variant::Bool(true)));
+        assert!(attributes.is_empty());
+    }
 }