From 7eb37a4c5e4272d015b9a98ebf97ca1f6d5dafab Mon Sep 17 00:00:00 2001
From: Xu Shaohua <shaohua@biofan.org>
Date: Thu, 14 Dec 2023 11:37:44 +0800
Subject: [PATCH] ds: Add IterMut in linked list

---
 data_structures/src/list/single_v1.rs | 38 +++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/data_structures/src/list/single_v1.rs b/data_structures/src/list/single_v1.rs
index fc99753b..7219bc3c 100644
--- a/data_structures/src/list/single_v1.rs
+++ b/data_structures/src/list/single_v1.rs
@@ -22,6 +22,10 @@ pub struct Iter<'a, T> {
     next: Option<&'a ListNode<T>>,
 }
 
+pub struct IterMut<'a, T> {
+    next: Option<&'a mut ListNode<T>>,
+}
+
 impl<T> ListNode<T> {
     #[must_use]
     pub fn new(value: T) -> Box<Self> {
@@ -114,6 +118,13 @@ impl<T> LinkedListV1<T> {
             next: self.head.as_deref(),
         }
     }
+
+    #[must_use]
+    pub fn iter_mut(&mut self) -> IterMut<T> {
+        IterMut {
+            next: self.head.as_deref_mut(),
+        }
+    }
 }
 
 impl<T> Drop for LinkedListV1<T> {
@@ -184,6 +195,17 @@ impl<'a, T> Iterator for Iter<'a, T> {
     }
 }
 
+impl<'a, T> Iterator for IterMut<'a, T> {
+    type Item = &'a mut T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next.take().map(|node| {
+            self.next = node.next.as_deref_mut();
+            &mut node.value
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::LinkedListV1;
@@ -273,4 +295,20 @@ mod tests {
             assert_eq!(val, num);
         }
     }
+
+    #[test]
+    fn test_iter_mut() {
+        let mut list = LinkedListV1::new();
+        list.push(2);
+        list.push(3);
+        list.push(5);
+        list.push(7);
+        for val in list.iter_mut() {
+            *val *= 2;
+        }
+        let nums = &[14, 10, 6, 4];
+        for (val, num) in list.iter().zip(nums) {
+            assert_eq!(val, num);
+        }
+    }
 }