Rust `group_by` (Reader discretion is advised)

group_by.rs

use core::cell::RefCell;
use std::collections::{BTreeMap, VecDeque};
use std::rc::Rc;

pub trait GroupBy<Item> {
    fn group(self) -> impl Iterator<Item = (Item, impl Iterator<Item = Item>)>
    where
        Item: PartialEq + Copy;

    fn group_by<Key, KeyFn>(
        self,
        key: KeyFn,
    ) -> impl Iterator<Item = (Key, impl Iterator<Item = Item>)>
    where
        Key: PartialEq + Copy,
        KeyFn: Fn(&Item) -> Key;
}

impl<Iter, Item> GroupBy<Item> for Iter
where
    Iter: Iterator<Item = Item>,
{
    fn group(self) -> impl Iterator<Item = (Item, impl Iterator<Item = Item>)>
    where
        Item: PartialEq + Copy,
    {
        Groups {
            internal: Rc::new(RefCell::new(GroupsInternal {
                iter: self,
                key: |item| *item,
                buffers: BTreeMap::new(),
                current_id: 0,
                current_key: None,
                pending: false,
            })),
        }
    }

    fn group_by<Key, KeyFn>(
        self,
        key: KeyFn,
    ) -> impl Iterator<Item = (Key, impl Iterator<Item = Item>)>
    where
        Key: PartialEq + Copy,
        KeyFn: Fn(&Item) -> Key,
    {
        Groups {
            internal: Rc::new(RefCell::new(GroupsInternal {
                iter: self,
                key: key,
                buffers: BTreeMap::new(),
                current_id: 0,
                current_key: None,
                pending: false,
            })),
        }
    }
}

struct GroupsInternal<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    iter: Iter,
    key: KeyFn,
    buffers: BTreeMap<u64, VecDeque<Item>>,
    current_id: u64,
    current_key: Option<Key>,
    pending: bool,
}

struct Groups<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    internal: Rc<RefCell<GroupsInternal<Iter, Item, Key, KeyFn>>>,
}

impl<Iter, Item, Key, KeyFn> Iterator for Groups<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    type Item = (Key, Group<Iter, Item, Key, KeyFn>);

    fn next(&mut self) -> Option<Self::Item> {
        let mut internal = self.internal.borrow_mut();
        let current_id = internal.current_id;

        if internal.pending {
            internal.pending = false;
            return Some((
                internal.current_key.unwrap(),
                Group {
                    internal: self.internal.clone(),
                    id: current_id,
                },
            ));
        }

        loop {
            if let Some(next) = internal.iter.next() {
                let key = (internal.key)(&next);

                if let Some(current_key) = internal.current_key {
                    if key == current_key {
                        if let Some(buffer) = internal.buffers.get_mut(&current_id) {
                            buffer.push_back(next);
                        }
                    } else {
                        let next_id = current_id + 1;
                        internal.buffers.insert(next_id, VecDeque::from([next]));
                        internal.current_id = next_id;
                        internal.current_key = Some(key);
                        return Some((
                            key,
                            Group {
                                internal: self.internal.clone(),
                                id: next_id,
                            },
                        ));
                    }
                } else {
                    internal.buffers.insert(current_id, VecDeque::from([next]));
                    internal.current_key = Some(key);
                    return Some((
                        key,
                        Group {
                            internal: self.internal.clone(),
                            id: current_id,
                        },
                    ));
                }
            } else {
                return None;
            }
        }
    }
}

struct Group<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    internal: Rc<RefCell<GroupsInternal<Iter, Item, Key, KeyFn>>>,
    id: u64,
}

impl<Iter, Item, Key, KeyFn> Iterator for Group<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    type Item = Item;

    fn next(&mut self) -> Option<Self::Item> {
        let mut internal = self.internal.borrow_mut();

        if let Some(next) = internal.buffers.get_mut(&self.id).unwrap().pop_front() {
            return Some(next);
        }

        if internal.current_id == self.id {
            if let Some(next) = internal.iter.next() {
                let key = (internal.key)(&next);

                if key == internal.current_key.unwrap() {
                    return Some(next);
                } else {
                    let next_id = internal.current_id + 1;
                    internal.buffers.insert(next_id, VecDeque::from([next]));
                    internal.current_id = next_id;
                    internal.current_key = Some(key);
                    internal.pending = true;
                }
            }
        }

        None
    }
}

impl<Iter, Item, Key, KeyFn> Drop for Group<Iter, Item, Key, KeyFn>
where
    Iter: Iterator<Item = Item>,
    Key: PartialEq + Copy,
    KeyFn: Fn(&Item) -> Key,
{
    fn drop(&mut self) {
        self.internal.borrow_mut().buffers.remove(&self.id);
    }
}