ds: Fix ref error in pop_front() of linked list

XuShaohua · Dec 15, 2023 · 45a477d · 45a477d
1 parent 7a07bd1
commit 45a477d
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 30 deletions.
diff --git a/data_structures/src/list/double.rs b/data_structures/src/list/double.rs
@@ -86,24 +86,22 @@ impl<T> DoublyLinkedList<T> {
     }
 
     pub fn pop_front(&mut self) -> Option<T> {
-        if self.length == 1 {
-            // Reset tail to None if both head and tail points to the same node.
-            self.tail.take();
-        }
+        self.head.take().and_then(|old_head| {
+            if let Some(new_head) = old_head.borrow_mut().next.take() {
+                // Reset previous pointer.
+                new_head.borrow_mut().previous = None;
+                self.head = Some(new_head);
+            } else {
+                // Reset tail to None if both head and tail points to the same node.
+                self.tail.take();
+            }
+            self.length -= 1;
 
-        self.head
-            .take()
             // Extract value from head if it has only one strong reference.
-            .and_then(|head: Rc<RefCell<Node<T>>>| Rc::try_unwrap(head).ok())
-            .map(|head: RefCell<Node<T>>| {
-                if let Some(next) = head.borrow_mut().next.take() {
-                    // Reset previous pointer.
-                    next.borrow_mut().previous = None;
-                    self.head = Some(next);
-                }
-                self.length -= 1;
-                head.into_inner().value
-            })
+            Rc::try_unwrap(old_head)
+                .ok()
+                .map(|head| head.into_inner().value)
+        })
     }
 }
 
@@ -158,4 +156,38 @@ mod tests {
         let list = DoublyLinkedList::<i32>::new();
         assert!(list.is_empty());
     }
+
+    #[test]
+    fn test_push() {
+        let mut list = DoublyLinkedList::new();
+        list.push_front(2);
+        list.push_front(3);
+        list.push_front(5);
+        list.push_front(7);
+        list.push_front(11);
+        assert_eq!(list.len(), 5);
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut list = DoublyLinkedList::new();
+        list.push_front(3);
+        list.push_front(5);
+        list.push_front(7);
+        assert_eq!(list.pop_front(), Some(7));
+        assert_eq!(list.len(), 2);
+        assert_eq!(list.pop_front(), Some(5));
+        assert_eq!(list.pop_front(), Some(3));
+        println!("len of list: {}", list.len());
+        assert!(list.is_empty());
+    }
+
+    #[test]
+    fn test_drop() {
+        let mut list = DoublyLinkedList::new();
+        for i in 0..(128 * 200) {
+            list.push_front(i);
+        }
+        drop(list);
+    }
 }
diff --git a/data_structures/src/list/single_v2.rs b/data_structures/src/list/single_v2.rs
@@ -111,21 +111,20 @@ impl<T> LinkedListV2<T> {
 
     /// Remove a node from head of list.
     pub fn pop_front(&mut self) -> Option<T> {
-        if self.length == 1 {
-            // Reset tail to None if both head and tail points to the same node.
-            self.tail.take();
-        }
-        self.head
-            .take()
+        self.head.take().and_then(|old_head| {
+            if let Some(new_head) = old_head.borrow_mut().next.take() {
+                self.head = Some(new_head);
+            } else {
+                // Reset tail to None if both head and tail points to the same node.
+                self.tail.take();
+            }
+            self.length -= 1;
+
             // Extract value from head if it has only one strong reference.
-            .and_then(|head: Rc<RefCell<ListNode<T>>>| Rc::try_unwrap(head).ok())
-            .map(|head| {
-                if let Some(next) = head.borrow_mut().next.take() {
-                    self.head = Some(next);
-                }
-                self.length -= 1;
-                head.into_inner().value
-            })
+            Rc::try_unwrap(old_head)
+                .ok()
+                .map(|head| head.into_inner().value)
+        })
     }
 
     /// Remove a node from tail of list.