chore: add iterator trait for linked list

shard-org · Aug 2, 2024 · 8a8212b · 8a8212b
1 parent d5beee1
commit 8a8212b
Showing 1 changed file with 76 additions and 55 deletions.
diff --git a/src/linked_list.rs b/src/linked_list.rs
@@ -1,11 +1,17 @@
-use std::marker::PhantomData;
 use std::fmt::Debug;
+use std::marker::PhantomData;
 
+/// A stack-based linked list.
 pub struct LinkedList<T> {
-   pub front: Link<T>,
-   back: Link<T>,
-   len: usize,
-   _boo: PhantomData<T>,
+    /// To the left.
+    pub front: Link<T>,
+    /// To the right.
+    back: Link<T>,
+    len: usize,
+    _boo: PhantomData<T>,
+
+    next: Option<Link<T>>,
+    current: Option<Link<T>>,
 }
 
 type Link<T> = Option<*mut Node<T>>;
@@ -19,64 +25,60 @@ pub struct Node<T> {
 
 impl<T> LinkedList<T> {
     pub fn new() -> Self {
-        LinkedList {
-            front: None,
-            back: None,
-            len: 0,
-            _boo: PhantomData,
-        }
+        LinkedList { front: None, back: None, len: 0, _boo: PhantomData, next: None, current: None }
     }
 
-    pub fn push(&mut self, elem: T) { unsafe {
-        let new = Box::into_raw(Box::new(Node {
-            back: None,
-            front: None,
-            elem,
-        }));
-
-        match self.back {
-            Some(old) => {
-                (*old).back = Some(new);
-                (*new).front = Some(old);
-            }
-            None => self.front = Some(new),
-        }
+    pub fn push(&mut self, elem: T) {
+        unsafe {
+            let new = Box::into_raw(Box::new(Node { back: None, front: None, elem }));
 
-        self.back = Some(new);
-        self.len += 1;
-    } }
+            match self.back {
+                Some(old) => {
+                    (*old).back = Some(new);
+                    (*new).front = Some(old);
+                },
+                None => self.front = Some(new),
+            }
 
-    pub fn pop_front(&mut self) -> Option<T> { unsafe {
-        self.front.map(|node| {
-            let boxed_node = Box::from_raw(node);
-            let result = boxed_node.elem;
+            self.back = Some(new);
+            self.len += 1;
+        }
+    }
 
-            self.front = boxed_node.back;
-            match self.front {
-                Some(new) => (*new).front = None,
-                None => self.back = None,
-            }
+    pub fn pop_front(&mut self) -> Option<T> {
+        unsafe {
+            self.front.map(|node| {
+                let boxed_node = Box::from_raw(node);
+                let result = boxed_node.elem;
 
-            self.len -= 1;
-            result
-        })
-    } }
+                self.front = boxed_node.back;
+                match self.front {
+                    Some(new) => (*new).front = None,
+                    None => self.back = None,
+                }
 
-    pub fn pop_back(&mut self) -> Option<T> { unsafe {
-        self.back.map(|node| {
-            let boxed_node = Box::from_raw(node);
-            let result = boxed_node.elem;
+                self.len -= 1;
+                result
+            })
+        }
+    }
 
-            self.back = boxed_node.front;
-            match self.back {
-                Some(new) => (*new).back = None,
-                None => self.front = None,
-            }
-            self.len -= 1;
-            result
-        })
-    } }
+    pub fn pop_back(&mut self) -> Option<T> {
+        unsafe {
+            self.back.map(|node| {
+                let boxed_node = Box::from_raw(node);
+                let result = boxed_node.elem;
 
+                self.back = boxed_node.front;
+                match self.back {
+                    Some(new) => (*new).back = None,
+                    None => self.front = None,
+                }
+                self.len -= 1;
+                result
+            })
+        }
+    }
 
     pub fn front_mut(&self) -> Option<*mut Node<T>> {
         self.front
@@ -118,12 +120,31 @@ impl<T> Node<T> {
     pub fn into_elem(self) -> T {
         self.elem
     }
-    
+
     pub fn next(&self) -> Option<*mut Node<T>> {
         self.back.map(|node| unsafe { node })
     }
 }
 
+impl<T> Iterator for LinkedList<T> {
+    type Item = Link<T>;
+
+    fn next(&mut self) -> Option<Link<T>> {
+        match self.current {
+            Some(current) => {
+                self.current = Some(self.next?);
+                self.next = self.current.unwrap().map(|x| unsafe { &(*x) }.front);
+            },
+            None => {
+                self.current = Some(self.back);
+                self.next = self.current.unwrap().map(|x| unsafe { &(*x) }.front);
+            },
+        }
+
+        self.current
+    }
+}
+
 impl<T> From<Vec<T>> for LinkedList<T> {
     fn from(vec: Vec<T>) -> Self {
         vec.into_iter().fold(LinkedList::new(), |mut list, elem| {