Implementing an iterator over mutable references is hard in general. It becomes unsound if the iterator ever returns references to the same element twice. That means that if you want to write one in purely safe code, you have to somehow convince the compiler that each element is only visited once. That rules out simply using an index: you could always forget to increment the index or set it somewhere and the compiler wouldn't be able to reason about it.
One possible way around is chaining together several std::iter::once
s (one for each reference you want to iterate over).
For example,
impl StatusRegister {
fn iter_mut(&mut self) -> impl Iterator<Item = &mut bool> {
use std::iter::once;
once(&mut self.CarryFlag)
.chain(once(&mut self.ZeroFlag))
.chain(once(&mut self.OverflowFlag))
}
}
(playground)
Upsides:
- Fairly simple to implement.
- No allocations.
- No external dependencies.
Downsides:
- The iterator has a very complicated type:
std::iter::Chain<std::iter::Chain<std::iter::Once<&mut bool>, std::iter::Once<&mut bool>>, std::iter::Once<&mut bool>>
.
So you if don't want to use impl Iterator<Item = &mut bool>
, you'll have to have that in your code. That includes implementing IntoIterator
for &mut StatusRegister
, since you'd have to explicitly indicate what the IntoIter
type is.
Another approach is using an array or Vec
to hold all the mutable references (with the correct lifetime) and then delegate to its iterator implementation to get the values. For example,
impl StatusRegister {
fn iter_mut(&mut self) -> std::vec::IntoIter<&mut bool> {
vec![
&mut self.CarryFlag,
&mut self.ZeroFlag,
&mut self.OverflowFlag,
]
.into_iter()
}
}
(playground)
Upsides:
- The type is the much more manageable
std::vec::IntoIter<&mut bool>
.
- Still fairly simple to implement.
- No external dependencies.
Downsides:
- Requires an allocation every time
iter_mut
is called.
I also mentioned using an array. That would avoid the allocation, but it turns out that arrays don't yet implement an iterator over their values, so the above code with a [&mut bool; 3]
instead of a Vec<&mut bool>
won't work. However, there exist crates that implement this functionality for fixed-length arrays with limited size, e.g. arrayvec
(or array_vec
).
Upsides:
- No allocation.
- Simple iterator type.
- Simple to implement.
Downsides:
The last approach I'll talk about is using unsafe
. Since this doesn't have many good upsides over the other approaches, I wouldn't recommend it in general. This is mainly to show you how you could implement this.
Like your original code, we'll implement Iterator
on our own struct.
impl<'a> IntoIterator for &'a mut StatusRegister {
type IntoIter = StatusRegisterIterMut<'a>;
type Item = &'a mut bool;
fn into_iter(self) -> Self::IntoIter {
StatusRegisterIterMut {
status: self,
index: 0,
}
}
}
pub struct StatusRegisterIterMut<'a> {
status: &'a mut StatusRegister,
index: usize,
}
The unsafety comes from the next
method, where we'll have to (essentially) convert something of type &mut &mut T
to &mut T
, which is generally unsafe. However, as long as we ensure that next
isn't allowed to alias these mutable references, we should be fine. There may be some other subtle issues, so I won't guarantee that this is sound. For what it's worth, MIRI doesn't find any problems with this.
impl<'a> Iterator for StatusRegisterIterMut<'a> {
type Item = &'a mut bool;
// Invariant to keep: index is 0, 1, 2 or 3
// Every call, this increments by one, capped at 3
// index should never be 0 on two different calls
// and similarly for 1 and 2.
fn next(&mut self) -> Option<Self::Item> {
let result = unsafe {
match self.index {
// Safety: Since each of these three branches are
// executed exactly once, we hand out no more than one mutable reference
// to each part of self.status
// Since self.status is valid for 'a
// Each partial borrow is also valid for 'a
0 => &mut *(&mut self.status.CarryFlag as *mut _),
1 => &mut *(&mut self.status.ZeroFlag as *mut _),
2 => &mut *(&mut self.status.OverflowFlag as *mut _),
_ => return None
}
};
// If self.index isn't 0, 1 or 2, we'll have already returned
// So this bumps us up to 1, 2 or 3.
self.index += 1;
Some(result)
}
}
(playground)
Upsides:
- No allocations.
- Simple iterator type name.
- No external dependencies.
Downsides:
- Complicated to implement. To successfully use
unsafe
, you need to be very familiar with what is and isn't allowed. This part of the answer took me the longest by far to make sure I wasn't doing something wrong.
- Unsafety infects the module. Within the module defining this iterator, I could "safely" cause unsoundness by messing with the
status
or index
fields of StatusRegisterIterMut
. The only thing allowing encapsulation is that outside of this module, those fields aren't visible.