CODEWITHSUNDEEP
rustiterator
Blue7
Blue7

Reputation: 2044

How to return an iterator for a tuple of slices that iterates the first slice then the second slice?

I have a function that splits a slice into three parts, a leading and trailing slice, and a reference to the middle element.

/// The leading and trailing parts of a slice.
struct LeadingTrailing<'a, T>(&'a mut [T], &'a mut [T]);

/// Divides one mutable slice into three parts, a leading and trailing slice,
/// and a reference to the middle element.
pub fn split_at_rest_mut<T>(x: &mut [T], index: usize) -> (&mut T, LeadingTrailing<T>) {
    debug_assert!(index < x.len());
    let (leading, trailing) = x.split_at_mut(index);
    let (val, trailing) = trailing.split_first_mut().unwrap();
    (val, LeadingTrailing(leading, trailing))
}

I would like to implement Iterator for LeadingTrailing<'a, T> so that it first iterates over the first slice, and then over the second. i.e., it will behave like:

let mut foo = [0,1,2,3,4,5];
let (item, lt) = split_at_rest_mut(&foo, 2);
for num in lt.0 {
    ...
}
for num in lt.1 {
    ...
}

I have tried converting to a Chain:

struct LeadingTrailing<'a, T>(&'a mut [T], &'a mut [T]);
impl <'a, T> LeadingTrailing<'a, T> {
    fn to_chain(&mut self) -> std::iter::Chain<&'a mut [T], &'a mut [T]> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }
}

But I get the error:

89 |         self.0.iter_mut().chain(self.1.iter_mut())
   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `&mut [T]`, found struct `std::slice::IterMut`

I have also tried creating a custom Iterator

/// The leading and trailing parts of a slice.
struct LeadingTrailing<'a, T>(&'a mut [T], &'a mut [T]);

struct LTOthersIterator<'a, T> {
    data: LeadingTrailing<'a, T>,
    index: usize,
}

/// Iterates over the first slice, then the second slice.
impl<'a, T> Iterator for LTOthersIterator<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        let leading_len = self.data.0.len();
        let trailing_len = self.data.1.len();
        let total_len = leading_len + trailing_len;
        match self.index {
            0..=leading_len => {
                self.index += 1;
                self.data.0.get(self.index - 1)
            }
            leading_len..=total_len => {
                self.index += 1;
                self.data.1.get(self.index - leading_len - 1)
            }
        }
    }
}

But I get the error:

error[E0495]: cannot infer an appropriate lifetime for autoref due to conflicting requirements
   --> src\main.rs:104:29
    |
104 |                 self.data.0.get(self.index - 1)
                                  ^^^

What is the correct way to do this?

Upvotes: 0

Views: 314

Answers (2)

cafce25
cafce25

Reputation: 27537

You either let the compiler do the work:

impl <'a, T> LeadingTrailing<'a, T> {
    fn to_chain(&mut self) -> impl Iterator<Item = &mut T> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }
}

Or perscribe the correct type, Chain takes the iterators, not the thing they got created from.

impl <'a, T> LeadingTrailing<'a, T> {
    fn to_chain(&'a mut self) -> std::iter::Chain<std::slice::IterMut<'a, T>, std::slice::IterMut<'a, T>> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }
}

Upvotes: 3

Finomnis
Finomnis

Reputation: 22738

The return value of to_chain() is incorrect. For simplicity, just use impl Iterator.

/// The leading and trailing parts of a slice.
#[derive(Debug)]
pub struct LeadingTrailing<'a, T>(&'a mut [T], &'a mut [T]);

/// Divides one mutable slice into three parts, a leading and trailing slice,
/// and a reference to the middle element.
pub fn split_at_rest_mut<T>(x: &mut [T], index: usize) -> (&mut T, LeadingTrailing<T>) {
    debug_assert!(index < x.len());
    let (leading, trailing) = x.split_at_mut(index);
    let (val, trailing) = trailing.split_first_mut().unwrap();
    (val, LeadingTrailing(leading, trailing))
}

impl<T> LeadingTrailing<'_, T> {
    fn to_chain(&mut self) -> impl Iterator<Item = &mut T> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }
}

fn main() {
    let mut arr = [0, 1, 2, 3, 4, 5, 6, 7, 8];
    let (x, mut leadtrail) = split_at_rest_mut(&mut arr, 5);

    println!("x: {}", x);
    println!("leadtrail: {:?}", leadtrail);

    for el in leadtrail.to_chain() {
        *el *= 2;
    }

    println!("leadtrail: {:?}", leadtrail);
}

x: 5
leadtrail: LeadingTrailing([0, 1, 2, 3, 4], [6, 7, 8])
leadtrail: LeadingTrailing([0, 2, 4, 6, 8], [12, 14, 16])

The fully written out version would be:

impl<T> LeadingTrailing<'_, T> {
    fn to_chain(&mut self) -> std::iter::Chain<std::slice::IterMut<T>, std::slice::IterMut<T>> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }
}

Upvotes: 2

Related Questions