Implementing Ord trait for an Iterator Wrapper

Question

Recently, I needed to sort iterators by their first value and/or use them in BinaryHeap. There are things I don't like about this implementation:

Using RefCell.
Mutating an underlying object while calling eq/cmp/etc methods.
Using unsafe in Deref.

It seems to be doing exactly what I want but I feel there must be a better way of achieving the same result.

use std::cell::RefCell;
use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;
use std::iter::Peekable;
use std::ops::{Deref, DerefMut};

struct IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    iter: RefCell<Peekable<I>>,
}

impl<I, T> Eq for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
}

impl<I, T> PartialEq<Self> for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn eq(&self, other: &Self) -> bool {
        self.iter.borrow_mut().peek().eq(&other.iter.borrow_mut().peek())
    }
}

impl<I, T> PartialOrd<Self> for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.iter.borrow_mut().peek().partial_cmp(&other.iter.borrow_mut().peek())
    }
}

impl<I, T> Ord for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn cmp(&self, other: &Self) -> Ordering {
        self.iter.borrow_mut().peek().cmp(&other.iter.borrow_mut().peek())
    }
}

impl<I, T> From<I> for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn from(iter: I) -> Self {
        Self { iter: RefCell::new(iter.peekable()) }
    }
}

impl<I, T> Deref for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    type Target = Peekable<I>;

    fn deref(&self) -> &Self::Target {
        let ptr = self.iter.as_ptr();
        unsafe { &*ptr }
    }
}

impl<I, T> DerefMut for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        let ptr = self.iter.as_ptr();
        unsafe { &mut *ptr }
    }
}

Usage example:

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test() {
        let a = [1, 2, 3];
        let b = [3, 2, 1];
        let mut heap = BinaryHeap::new();
        heap.push(IterHeapWrapper::from(a.iter()));
        heap.push(IterHeapWrapper::from(b.iter()));
        assert_eq!(heap.pop().unwrap().next().unwrap(), &3);
    }
}

Winston Ewert · Accepted Answer · 2022-06-20 15:28:55Z

The fundamental problem is a mismatch between Eq/PartialCmp/Cmp and Peekable. The comparison functions assume that comparing objects does not require modifying them. But the peekable() method is a mutable method, so you can't call it.

But why isn't peekable an immutable function? After all, it doesn't logically modify the iterator, it just peeks. The reason is that it its lazy, it doesn't progress the iterator until you request to take a peek. If instead, it was eager, pulling the next item from the iterator right away, you could peek at it without modifying the iterator.

So you could do something like this:

struct IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>
{
    iter: I,
    peeked: Option<T>
}

impl<I, T> From<I> for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    fn from(mut iter: I) -> Self {
        let peeked = iter.next();
        Self { iter, peeked }
    }
}


impl<I, T> Iterator for IterHeapWrapper<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    type Item = T;
    
    fn next(&mut self) -> Option<Self::Item> {
        let mut next = self.iter.next();
        std::mem::swap(&mut next, &mut self.peeked);
        next
    }
}

This implements the logic of peekable, but in such a way that the next item is stored right there on the struct, so you can use it to do comparisons.

Your unsafe code is not a good idea. I think implementing Iterator on IterHeapWrapper probably does what you want better. As it stands, by using unsafe code you've completely defeated the protection of RefCell. RefCell tracks at runtime the references created out of it to make sure it remains safe, but you skip that completely.

Well, I agree with everything you said. I didn't feel right about reimplementing Peekable functionality or abusing unsafe. Now it seems that lesser evil is on the side of reimplementing a part of Peekable and making the Wrapper an iterator itself. — night-crawler
– night-crawler, Commented Jun 20, 2022 at 18:51
@night-crawler, there is a crate which provides a version of peekable that has the properties you'd need: crates.io/crates/peepable — Winston Ewert
– Winston Ewert, Commented Jun 21, 2022 at 14:54
Alternatively, it might make sense isntead of implementing Iterator to provide a method which decomposes the wrapper into a first and rest. I'm guessing that you probably don't actually need the power of iterators, but probably always just fetch the first item and stick the iterator back in the BinaryHeap. — Winston Ewert
– Winston Ewert, Commented Jun 21, 2022 at 15:02

Stack Exchange Network

Implementing Ord trait for an Iterator Wrapper

1 Answer 1

You must log in to answer this question.

Hot Network Questions

Implementing Ord trait for an Iterator Wrapper

1 Answer 1

You must log in to answer this question.

Related

Hot Network Questions