1 files changed, 109 insertions, 0 deletions
diff --git a/src/iterator/iterator.rs b/src/iterator/iterator.rs
new file mode 100644
index 0000000..1e389a4
--- /dev/null
+++ b/src/iterator/iterator.rs
@@ -0,0 +1,109 @@
+use std::iter::{Enumerate, Fuse, Skip, Take};
+
+use super::{ConsumedIteratorGuard, IteratorGuard, LockingIterator};
+
+use crate::{
+	lockable::{Lockable, RawLock, Sharable},
+	Keyable,
+};
+
+impl<L, Key> LockingIterator<L, Key> {
+	pub(crate) const fn new(key: Key, lockable: L) -> Self {
+		Self { key, lockable }
+	}
+
+	fn with_iterator<M>(self, f: impl FnOnce(L) -> M) -> LockingIterator<M, Key> {
+		LockingIterator {
+			key: self.key,
+			lockable: f(self.lockable),
+		}
+	}
+}
+
+impl<'a, I: Iterator<Item = &'a L>, L: 'a + Lockable + RawLock, Key: Keyable>
+	LockingIterator<I, Key>
+{
+	pub fn lock_next(&'a mut self) -> Option<IteratorGuard<'a, <L as Lockable>::Guard<'a>, Key>> {
+		if let Some(lock) = self.lockable.next() {
+			unsafe {
+				lock.raw_write();
+				let guard = lock.guard();
+
+				Some(IteratorGuard {
+					_key: &self.key,
+					guard,
+				})
+			}
+		} else {
+			None
+		}
+	}
+
+	pub fn lock_last(self) -> Option<ConsumedIteratorGuard<<L as Lockable>::Guard<'a>, Key>> {
+		self.lockable.last().map(|lock| unsafe {
+			lock.raw_write();
+			let guard = lock.guard();
+
+			ConsumedIteratorGuard {
+				key: self.key,
+				guard,
+			}
+		})
+	}
+}
+
+impl<'a, I: Iterator<Item = &'a L>, L: 'a + Sharable + RawLock, Key: Keyable>
+	LockingIterator<I, Key>
+{
+	pub fn read_next(
+		&'a mut self,
+	) -> Option<IteratorGuard<'a, <L as Sharable>::ReadGuard<'a>, Key>> {
+		if let Some(lock) = self.lockable.next() {
+			unsafe {
+				lock.raw_read();
+				let guard = lock.read_guard();
+
+				Some(IteratorGuard {
+					_key: &self.key,
+					guard,
+				})
+			}
+		} else {
+			None
+		}
+	}
+
+	pub fn read_last(self) -> Option<ConsumedIteratorGuard<<L as Sharable>::ReadGuard<'a>, Key>> {
+		self.lockable.last().map(|lock| unsafe {
+			lock.raw_read();
+			let guard = lock.read_guard();
+
+			ConsumedIteratorGuard {
+				key: self.key,
+				guard,
+			}
+		})
+	}
+}
+
+impl<L: Iterator, Key> LockingIterator<L, Key> {
+	pub fn size_hint(&self) -> (usize, Option<usize>) {
+		self.lockable.size_hint()
+	}
+
+	pub fn enumerate(self) -> LockingIterator<Enumerate<L>, Key> {
+		self.with_iterator(Iterator::enumerate)
+	}
+
+	pub fn skip(self, n: usize) -> LockingIterator<Skip<L>, Key> {
+		self.with_iterator(|i| i.skip(n))
+	}
+
+	pub fn take(self, n: usize) -> LockingIterator<Take<L>, Key> {
+		self.with_iterator(|i| i.take(n))
+	}
+
+	pub fn fuse(self) -> LockingIterator<Fuse<L>, Key> {
+		self.with_iterator(Iterator::fuse)
+	}
+}