Merge pull request #29 from aloucks/iter_mut

Add mutable iteration over arena contents

src/lib.rs

@@ -328,6 +328,75 @@ impl<T> Arena<T> {
         result.append(&mut chunks.current);
         result
     }
+
+    /// Returns an iterator that allows modifying each value.
+    ///
+    /// Items are yielded in the order that they were allocated.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// use typed_arena::Arena;
+    ///
+    /// #[derive(Debug, PartialEq, Eq)]
+    /// struct Point { x: i32, y: i32 };
+    ///
+    /// let mut arena = Arena::new();
+    ///
+    /// arena.alloc(Point { x: 0, y: 0 });
+    /// arena.alloc(Point { x: 1, y: 1 });
+    ///
+    /// for point in arena.iter_mut() {
+    ///     point.x += 10;
+    /// }
+    ///
+    /// let points = arena.into_vec();
+    ///
+    /// assert_eq!(points, vec![Point { x: 10, y: 0 }, Point { x: 11, y: 1 }]);
+    ///
+    /// ```
+    ///
+    /// ## Immutable Iteration
+    ///
+    /// Note that there is no corresponding `iter` method. Access to the arena's contents
+    /// requries mutable access to the arena itself.
+    ///
+    /// ```compile_fail
+    /// use typed_arena::Arena;
+    ///
+    /// let mut arena = Arena::new();
+    /// let x = arena.alloc(1);
+    ///
+    /// // borrow error!
+    /// for i in arena.iter_mut() {
+    ///     println!("i: {}", i);
+    /// }
+    ///
+    /// // borrow error!
+    /// *x = 2;
+    /// ```
+    #[inline]
+    pub fn iter_mut(&mut self) -> IterMut<T> {
+        let chunks = self.chunks.get_mut();
+        let position = if chunks.rest.len() > 0 {
+            let index = 0;
+            let inner_iter = chunks.rest[index].iter_mut();
+            // Extend the lifetime of the individual elements to that of the arena.
+            // This is OK because we borrow the arena mutably to prevent new allocations
+            // and we take care here to never move items inside the arena while the
+            // iterator is alive.
+            let inner_iter = unsafe { mem::transmute(inner_iter) };
+            IterMutState::ChunkListRest { index, inner_iter }
+        } else {
+            // Extend the lifetime of the individual elements to that of the arena.
+            let iter = unsafe { mem::transmute(chunks.current.iter_mut()) };
+            IterMutState::ChunkListCurrent { iter }
+        };
+        IterMut {
+            chunks,
+            state: position,
+        }
+    }
 }
 
 impl<T> Default for Arena<T> {
@@ -353,3 +422,75 @@ impl<T> ChunkList<T> {
         self.rest.push(chunk);
     }
 }
+
+enum IterMutState<'a, T> {
+    ChunkListRest {
+        index: usize,
+        inner_iter: slice::IterMut<'a, T>,
+    },
+    ChunkListCurrent {
+        iter: slice::IterMut<'a, T>,
+    },
+}
+
+/// Mutable arena iterator.
+///
+/// This struct is created by the [`iter_mut`](struct.Arena.html#method.iter_mut) method on [Arenas](struct.Arena.html).
+pub struct IterMut<'a, T: 'a> {
+    chunks: &'a mut ChunkList<T>,
+    state: IterMutState<'a, T>,
+}
+
+impl<'a, T> Iterator for IterMut<'a, T> {
+    type Item = &'a mut T;
+    fn next(&mut self) -> Option<&'a mut T> {
+        match self.state {
+            IterMutState::ChunkListRest {
+                mut index,
+                ref mut inner_iter,
+            } => {
+                match inner_iter.next() {
+                    Some(item) => Some(item),
+                    None => {
+                        index += 1;
+                        if index < self.chunks.rest.len() {
+                            let inner_iter = self.chunks.rest[index].iter_mut();
+                            // Extend the lifetime of the individual elements to that of the arena.
+                            let inner_iter = unsafe { mem::transmute(inner_iter) };
+                            self.state = IterMutState::ChunkListRest { index, inner_iter };
+                            self.next()
+                        } else {
+                            let iter = self.chunks.current.iter_mut();
+                            // Extend the lifetime of the individual elements to that of the arena.
+                            let iter = unsafe { mem::transmute(iter) };
+                            self.state = IterMutState::ChunkListCurrent { iter };
+                            self.next()
+                        }
+                    }
+                }
+            }
+            IterMutState::ChunkListCurrent { ref mut iter } => iter.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let current_len = self.chunks.current.len();
+        let current_cap = self.chunks.current.capacity();
+        if self.chunks.rest.is_empty() {
+            (current_len, Some(current_len))
+        } else {
+            let rest_len = self.chunks.rest.len();
+            let last_chunk_len = self
+                .chunks
+                .rest
+                .last()
+                .map(|chunk| chunk.len())
+                .unwrap_or(0);
+
+            let min = current_len + last_chunk_len;
+            let max = min + (rest_len * current_cap / rest_len);
+
+            (min, Some(max))
+        }
+    }
+}

src/test.rs

@@ -179,3 +179,143 @@ fn arena_is_send() {
     let arena: Arena<u32> = Arena::new();
     assert_is_send(arena);
 }
+
+#[test]
+fn iter_mut_low_capacity() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 1_000;
+    const CAP: usize = 16;
+
+    let mut arena = Arena::with_capacity(CAP);
+    for i in 1..MAX {
+        arena.alloc(NonCopy(i));
+    }
+
+    assert!(
+        arena.chunks.borrow().rest.len() > 1,
+        "expected multiple chunks"
+    );
+
+    let mut iter = arena.iter_mut();
+    for i in 1..MAX {
+        assert_eq!(Some(&mut NonCopy(i)), iter.next());
+    }
+
+    assert_eq!(None, iter.next());
+}
+
+#[test]
+fn iter_mut_high_capacity() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 1_000;
+    const CAP: usize = 8192;
+
+    let mut arena = Arena::with_capacity(CAP);
+    for i in 1..MAX {
+        arena.alloc(NonCopy(i));
+    }
+
+    assert!(
+        arena.chunks.borrow().rest.is_empty(),
+        "expected single chunk"
+    );
+
+    let mut iter = arena.iter_mut();
+    for i in 1..MAX {
+        assert_eq!(Some(&mut NonCopy(i)), iter.next());
+    }
+
+    assert_eq!(None, iter.next());
+}
+
+fn assert_size_hint<T>(arena_len: usize, iter: IterMut<'_, T>) {
+    let (min, max) = iter.size_hint();
+
+    assert!(max.is_some());
+    let max = max.unwrap();
+
+    // Check that the actual arena length lies between the estimated min and max
+    assert!(min <= arena_len);
+    assert!(max >= arena_len);
+
+    // Check that the min and max estimates are within a factor of 3
+    assert!(min >= arena_len / 3);
+    assert!(max <= arena_len * 3);
+}
+
+#[test]
+fn size_hint() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 32;
+    const CAP: usize = 0;
+
+    for cap in CAP..(CAP + 16/* check some non-power-of-two capacities */) {
+        let mut arena = Arena::with_capacity(cap);
+        for i in 1..MAX {
+            arena.alloc(NonCopy(i));
+            let iter = arena.iter_mut();
+            assert_size_hint(i, iter);
+        }
+    }
+}
+
+#[test]
+#[cfg(not(miri))]
+fn size_hint_low_initial_capacities() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 25_000;
+    const CAP: usize = 0;
+
+    for cap in CAP..(CAP + 128/* check some non-power-of-two capacities */) {
+        let mut arena = Arena::with_capacity(cap);
+        for i in 1..MAX {
+            arena.alloc(NonCopy(i));
+            let iter = arena.iter_mut();
+            assert_size_hint(i, iter);
+        }
+    }
+}
+
+#[test]
+#[cfg(not(miri))]
+fn size_hint_high_initial_capacities() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 25_000;
+    const CAP: usize = 8164;
+
+    for cap in CAP..(CAP + 128/* check some non-power-of-two capacities */) {
+        let mut arena = Arena::with_capacity(cap);
+        for i in 1..MAX {
+            arena.alloc(NonCopy(i));
+            let iter = arena.iter_mut();
+            assert_size_hint(i, iter);
+        }
+    }
+}
+
+#[test]
+#[cfg(not(miri))]
+fn size_hint_many_items() {
+    #[derive(Debug, PartialEq, Eq)]
+    struct NonCopy(usize);
+
+    const MAX: usize = 5_000_000;
+    const CAP: usize = 16;
+
+    let mut arena = Arena::with_capacity(CAP);
+    for i in 1..MAX {
+        arena.alloc(NonCopy(i));
+        let iter = arena.iter_mut();
+        assert_size_hint(i, iter);
+    }
+}