Как создать итератор для QuadTree, который выдает каждую точку, лист, в котором она находится, и соседние листья?

У меня простой QuadTree. Для удобочитаемости и чтобы помочь мне понять, что происходит, он избегает рекурсии и имеет статическую глубину. QuadTree хранит ссылки на точки, принадлежащие другому контейнеру.

struct QuadLeaf<'a> {
    vec: Vec<&'a (f32,f32)>,
    rect: (f32, f32, f32, f32)
}
struct QuadTwig<'a> {
    cells: [QuadLeaf<'a>; 4],
}
struct QuadBranch<'a> {
    cells: [QuadTwig<'a>; 4],
}
struct QuadTree<'a> {
    cells: [QuadBranch<'a>; 4],
}

Построение и вставка в это дерево относительно просты. QuadLeaf состоит из ограничивающего rect и пустого vec и имеет метод, который пытается вставить точку. Он возвращает true, если точка находится внутри rect и была вставлена.

impl<'a> QuadLeaf<'a> {
    fn new(rect: (f32, f32, f32, f32)) -> Self {
        QuadLeaf {vec: Vec::new(),rect}
    }
    fn insert(&mut self, point: &'a (f32, f32)) -> bool {
        if is_point_in_rect(point.0, point.1, self.rect) {
            self.vec.push(point);
            true
        } else {
            false
        }
    }
}

Функция QuadTwignew разбивает ограничивающий rect на 4 и создает 4 новых QuadLeafs. Это insert функция пытается вставить в каждый лист, замыкаясь в случае успеха и возвращая false в случае неудачи.

impl<'a> QuadTwig<'a> {
    fn new(rect: (f32, f32, f32, f32)) -> Self {
        let rects = divide_into_4(rect);
        QuadTwig {
            cells: [
                QuadLeaf::new(rects[0]),
                QuadLeaf::new(rects[1]),
                QuadLeaf::new(rects[2]),
                QuadLeaf::new(rects[3])
            ]
        }
    }
    fn insert(&mut self, point: &'a (f32, f32)) -> bool {
        for cell in self.cells.iter_mut() {
            if cell.insert(point) {
                return true;
            }
        }
        false
    }
}

Реализации для QuadBranch и QuadTree абсолютно одинаковы, но функция new просто строит следующий уровень в дереве. Это можно было бы отрефакторить позже для уменьшения дублирования кода, но в целях демонстрации я оставлю это. Я также думаю, что это не имеет значения для контекста этого вопроса.

Вопрос:

Я хочу создать итератор, который выдает каждую точку дерева и 9 листьев, к которым он близок (или внутри).

Мне удалось создать более простую версию, которая просто дает каждую точку и лист, в котором она находится:

/// An Iterator that yields each point and the leaf it is in
struct PointAndLeafIterator<'a> {
    ptr: &'a QuadTree<'a>,
    index: (usize, usize, usize, usize)
}

/// An Iterator that yields each point and the leaf it is in
impl<'a> Iterator for PointAndLeafIterator<'a> {
    /// Returns (point, leaf)
    type Item = (&'a (f32, f32), Vec<&'a (f32, f32)>);

    /// Starts at (0,0,0,0) and ends at (3, 3, 3, num_points_in_leaf)
    /// It increases the index by 1 each time, and if it reaches the end of the cell, it moves to the next cell
    fn next(&mut self) -> Option<Self::Item> {
        let (branch_index, twig_index, leaf_index, point_index) = &mut self.index;
        let branch = &self.ptr.cells[*branch_index];
        let twig = &branch.cells[*twig_index];
        let leaf = &twig.cells[*leaf_index];
        let point = leaf.vec.get(*point_index);
        //go through all the points in the leaf
        if let Some(point) = point {
            *point_index += 1;
            return Some((point, leaf.vec.clone()));
        }

        //if we reach the end of the leaf, go to the next leaf
        *point_index = 0;
        *leaf_index += 1;
        if *leaf_index < 4 {
            return self.next();
        }

        //if we reach the end of the twig, go to the next twig
        *leaf_index = 0;
        *twig_index += 1;
        if *twig_index < 4 {
            return self.next();
        }

        //if we reach the end of the branch, go to the next branch
        *twig_index = 0;
        *branch_index += 1;
        if *branch_index < 4 {
            return self.next();
        }

        //if we reach the end of the tree, we are done
        None

    }
}

Это можно использовать так:

fn main() {
    let points: Vec<(f32, f32)> = vec![
        (0.0, 0.0),
        (1.0, 1.0),
        (31.0,31.0),
        (2.0, 2.0),
        (3.0, 3.0),
        (32.0,32.0),
    ];
    let mut quadtree = QuadTree::new((0.0, 0.0, 40.0, 40.0));
    for point in points.iter() {
        quadtree.insert(point);
    }
    for (point, leaf) in quadtree.into_point_and_leaf_iter() {
        println!("Point: {:?}", point);
        println!("Leaf: {:?}", leaf);
    }
}

Однако соседний вариант оказывается намного сложнее. Как я могу написать этот алгоритм?

/// An Iterator that yields each point, the leaf it is in, and the neighbouring leaves
struct PointAndLeafAndNeighboursIterator<'a> {
    ptr: &'a QuadTree<'a>,
    index: (usize, usize, usize, usize)
}

impl<'a> Iterator for PointAndLeafAndNeighboursIterator<'a> {
    ///Return the 9 leaves that surround the point
    ///If there is no leaf in a direction, it will return an empty leaf
    type Item = (&'a (f32, f32), [Vec<&'a (f32, f32)>; 9]);

    /// Starts at (0,0,0,0) and ends at (3, 3, 3, num_points_in_leaf)
    /// It increases the index by 1 each time, and if it reaches the end of the cell, it moves to the next cell
    fn next(&mut self) -> Option<Self::Item> {
        unimplemented!()
    }
}

Вот ссылка на игровую площадку Rust на весь код в этом вопросе.