kd_tree simplification 2

mod.rs

use na::*;
use na::allocator::Allocator;
/// A kd tree
pub struct KdTree<N: DimName>(Box<kd_tree::KdTreeImpl<N>>)
where
    N: DimName,
    DefaultAllocator: Allocator<f64, N>;

impl<N: DimName> KdTree<N>
where
    N: DimName,
    DefaultAllocator: Allocator<f64, N>,
{
    // /// Get all the points from the kd-tree
    // fn get_points(self) -> Vec<VectorN<f64, N>>  {
    //     let mut pts = Vec::new();
    //     let KdTree(tree) = self;
    //     kd_tree::add_points_to_vector(tree, &mut pts);
    //     pts
    // }

    // fn build(pts : &Vec<VectorN<f64, N>>) -> Self {
    //     KdTree(Box::new(kd_tree::build(pts, 0)))
    // }
}

mod kd_tree {
    use na::*;
    use na::allocator::Allocator;

    /// kd-Tree implementation structure
    pub(super) enum KdTreeImpl<N>
    where
        N: DimName,
        DefaultAllocator: Allocator<f64, N>,
    {
        Node(usize, Box<KdTreeImpl<N>>, VectorN<f64, N>, Box<KdTreeImpl<N>>),
        Empty(),
    }

    impl<N> KdTreeImpl<N>
    where
        N: DimName,
        DefaultAllocator: Allocator<f64, N>,
    {
        /// mutably enunerate points from a kd tree into a vector - TODO: expose as iterator instead?
        pub fn add_points_to_vector(&self, pts: &mut Vec<VectorN<f64, N>>) {
            if let KdTreeImpl::Node(_, ref l, ref pt, ref r) = *self {
                l.add_points_to_vector(pts);
                pts.push(pt.clone());
                r.add_points_to_vector(pts);
            }
        }

        pub fn build(pts: &Vec<VectorN<f64, N>>, dim_idx: usize) -> KdTreeImpl<N>
        where
            N: DimName,
            DefaultAllocator: Allocator<f64, N>,
        {
            match pts.first() {
                Some(first_pt) => {
                    let next_dim_idx = if dim_idx + 1 >= first_pt.ncols() {
                        0
                    } else {
                        dim_idx + 1
                    };
                    match split_vec_with_median(pts, dim_idx) {
                        (Some(vec1), Some(pt), Some(vec2)) => {
                            KdTreeImpl::Node(
                                dim_idx,
                                Box::new(KdTreeImpl::build(&vec1, next_dim_idx)),
                                pt,
                                Box::new(KdTreeImpl::build(&vec2, next_dim_idx)),
                            )
                        }
                        (Some(vec1), Some(pt), None) => {
                            KdTreeImpl::Node(
                                dim_idx,
                                Box::new(KdTreeImpl::build(&vec1, next_dim_idx)),
                                pt,
                                Box::new(KdTreeImpl::Empty()),
                            )
                        }
                        (None, Some(pt), Some(vec2)) => {
                            KdTreeImpl::Node(
                                dim_idx,
                                Box::new(KdTreeImpl::Empty()),
                                pt,
                                Box::new(KdTreeImpl::build(&vec2, next_dim_idx)),
                            )
                        }
                        (None, Some(pt), None) => {
                            KdTreeImpl::Node(
                                dim_idx,
                                Box::new(KdTreeImpl::Empty()),
                                pt,
                                Box::new(KdTreeImpl::Empty()),
                            )
                        }
                        (None, None, None) => {
                            panic!(
                                "Should be impossible to have a first point but still split into (None, None, None)."
                            )
                        }
                        (_, None, _) => {
                            panic!("Meaningless result: point must be contained in a kd-tree node.")
                        }
                    }
                }
                None => KdTreeImpl::Empty(),
            }
        }
    }

    fn split_vec_with_median<N>(
        pts: &Vec<VectorN<f64, N>>,
        dim_idx: usize,
    ) -> (Option<Vec<VectorN<f64, N>>>, Option<VectorN<f64, N>>, Option<Vec<VectorN<f64, N>>>)
    where
        N: DimName,
        DefaultAllocator: Allocator<f64, N>,
    {
        unsafe {
            let mut pts = pts.to_vec();
            pts.sort_by(|a, b| {
                a.get_unchecked(0, dim_idx)
                    .partial_cmp(b.get_unchecked(0, dim_idx))
                    .unwrap()
            });
            if pts.len() >= 3 {
                let length = pts.len();
                let mut middle = pts.split_off(length / 2);
                let right = middle.split_off(1);
                // (Some(pts), Some(middle[0]), Some(right))
                panic!()
            } else if pts.len() == 2 {
                let middle = pts.split_off(1);
                // (Some(pts), Some(middle[0]), None)
                panic!()
            } else if pts.len() == 1 {
                // (None, Some(pts[0]), None)
                panic!()
            } else {
                (None, None, None)
            }
        }
    }
}