vox_geometry_rust/

kdtree2.rs

/*
 * // Copyright (c) 2021 Feng Yang
 * //
 * // I am making my contributions/submissions to this project solely in my
 * // personal capacity and am not conveying any rights to any intellectual
 * // property of any third parties.
 */

use crate::vector2::Vector2D;
use crate::bounding_box2::BoundingBox2D;
use std::cmp::Ordering;

/// Simple K-d tree node.
#[derive(Clone)]
pub struct Node {
    /// Split axis if flags < K, leaf indicator if flags == K.
    flags: usize,

    /// \brief Right child index.
    /// Note that left child index is this node index + 1.
    child: usize,

    /// Item index.
    item: usize,

    /// Point stored in the node.
    point: Vector2D,
}

impl Node {
    /// Default constructor.
    pub fn new() -> Node {
        return Node {
            flags: 0,
            child: usize::MAX,
            item: usize::MAX,
            point: Vector2D::new_default(),
        };
    }

    /// Initializes leaf node.
    pub fn init_leaf(&mut self, it: usize, pt: &Vector2D) {
        self.flags = 2;
        self.item = it;
        self.child = usize::MAX;
        self.point = pt.clone();
    }

    /// Initializes internal node.
    pub fn init_internal(&mut self, axis: usize, it: usize, c: usize, pt: &Vector2D) {
        self.flags = axis;
        self.item = it;
        self.child = c;
        self.point = pt.clone();
    }

    /// Returns true if leaf.
    pub fn is_leaf(&self) -> bool {
        return self.flags == 2;
    }
}

/// Generic k-d tree structure.
#[derive(Clone)]
pub struct KdTree2 {
    _points: Vec<Vector2D>,
    _nodes: Vec<Node>,
}

impl KdTree2 {
    /// Constructs an empty kD-tree instance.
    pub fn new() -> KdTree2 {
        return KdTree2 {
            _points: vec![],
            _nodes: vec![],
        };
    }

    /// Builds internal acceleration structure for given points list.
    pub fn build(&mut self, points: &Vec<Vector2D>) {
        self._points = points.clone();

        if self._points.is_empty() {
            return;
        }

        self._nodes.clear();
        let mut item_indices: Vec<usize> = (0..self._points.len()).collect();

        self.build_internal(0, &mut item_indices, self._points.len(), 0);
    }

    ///
    /// Invokes the callback function for each nearby point around the origin
    /// within given radius.
    ///
    /// - parameter:  origin   The origin position.
    /// - parameter:  radius   The search radius.
    /// - parameter:  callback The callback function.
    ///
    pub fn for_each_nearby_point<Callback>(&self, origin: &Vector2D, radius: f64,
                                           callback: &mut Callback) where Callback: FnMut(usize, &Vector2D) {
        let r2 = radius * radius;

        // prepare to traverse the tree for sphere
        const K_MAX_TREE_DEPTH: usize = 8 * 32;
        let mut todo: [Option<usize>; K_MAX_TREE_DEPTH] = [None; K_MAX_TREE_DEPTH];
        let mut todo_pos = 0;

        // traverse the tree nodes for sphere
        let mut node: usize = 0;
        while node < self._nodes.len() {
            if self._nodes[node].item != usize::MAX &&
                (self._nodes[node].point - *origin).length_squared() <= r2 {
                callback(self._nodes[node].item, &self._nodes[node].point);
            }

            if self._nodes[node].is_leaf() {
                // grab next node to process from todo stack
                if todo_pos > 0 {
                    // Dequeue
                    todo_pos -= 1;
                    node = todo[todo_pos].unwrap();
                } else {
                    break;
                }
            } else {
                // get node children pointers for sphere
                let first_child = node + 1;
                let second_child = self._nodes[node].child;

                // advance to next child node, possibly enqueue other child
                let axis = self._nodes[node].flags;
                let plane = self._nodes[node].point[axis];
                if plane - origin[axis] > radius {
                    node = first_child;
                } else if origin[axis] - plane > radius {
                    node = second_child;
                } else {
                    // enqueue second_child in todo stack
                    todo[todo_pos] = Some(second_child);
                    todo_pos += 1;
                    node = first_child;
                }
            }
        }
    }

    ///
    /// Returns true if there are any nearby points for given origin within
    /// radius.
    ///
    /// - parameter:  origin The origin.
    /// - parameter:  radius The radius.
    ///
    /// \return     True if has nearby point, false otherwise.
    ///
    pub fn has_nearby_point(&self, origin: &Vector2D, radius: f64) -> bool {
        let r2 = radius * radius;

        // prepare to traverse the tree for sphere
        const K_MAX_TREE_DEPTH: usize = 8 * 32;
        let mut todo: [Option<usize>; K_MAX_TREE_DEPTH] = [None; K_MAX_TREE_DEPTH];
        let mut todo_pos = 0;

        // traverse the tree nodes for sphere
        let mut node: usize = 0;
        while node < self._nodes.len() {
            if self._nodes[node].item != usize::MAX &&
                (self._nodes[node].point - *origin).length_squared() <= r2 {
                return true;
            }

            if self._nodes[node].is_leaf() {
                // grab next node to process from todo stack
                if todo_pos > 0 {
                    // Dequeue
                    todo_pos -= 1;
                    node = todo[todo_pos].unwrap();
                } else {
                    break;
                }
            } else {
                // get node children pointers for sphere
                let first_child = node + 1;
                let second_child = self._nodes[node].child;

                // advance to next child node, possibly enqueue other child
                let axis = self._nodes[node].flags;
                let plane = self._nodes[node].point[axis];
                if origin[axis] < plane && plane - origin[axis] > radius {
                    node = first_child;
                } else if origin[axis] > plane && origin[axis] - plane > radius {
                    node = second_child;
                } else {
                    // enqueue second_child in todo stack
                    todo[todo_pos] = Some(second_child);
                    todo_pos += 1;
                    node = first_child;
                }
            }
        }

        return false;
    }

    /// Returns index of the nearest point.
    pub fn nearest_point(&self, origin: &Vector2D) -> usize {
        // prepare to traverse the tree for sphere
        const K_MAX_TREE_DEPTH: usize = 8 * 32;
        let mut todo: [Option<usize>; K_MAX_TREE_DEPTH] = [None; K_MAX_TREE_DEPTH];
        let mut todo_pos = 0;

        // traverse the tree nodes for sphere
        let mut node: usize = 0;
        let mut nearest = 0;
        let mut min_dist2 = (self._nodes[node].point - *origin).length_squared();

        while node < self._nodes.len() {
            let new_dist2 = (self._nodes[node].point - *origin).length_squared();
            if new_dist2 <= min_dist2 {
                nearest = self._nodes[node].item;
                min_dist2 = new_dist2;
            }

            if self._nodes[node].is_leaf() {
                // grab next node to process from todo stack
                if todo_pos > 0 {
                    // Dequeue
                    todo_pos -= 1;
                    node = todo[todo_pos].unwrap();
                } else {
                    break;
                }
            } else {
                // get node children pointers for sphere
                let first_child = node + 1;
                let second_child = self._nodes[node].child;

                // advance to next child node, possibly enqueue other child
                let axis = self._nodes[node].flags;
                let plane = self._nodes[node].point[axis];
                let min_dist = f64::sqrt(min_dist2);
                if plane - origin[axis] > min_dist {
                    node = first_child;
                } else if origin[axis] - plane > min_dist {
                    node = second_child;
                } else {
                    // enqueue second_child in todo stack
                    todo[todo_pos] = Some(second_child);
                    todo_pos += 1;
                    node = first_child;
                }
            }
        }

        return nearest;
    }

    /// Reserves memory space for this tree.
    pub fn reserve(&mut self, num_points: usize, num_nodes: usize) {
        self._points.resize(num_points, Vector2D::new_default());
        self._nodes.resize(num_nodes, Node::new());
    }

    pub fn build_internal(&mut self, node_index: usize, item_indices: &mut [usize], n_items: usize,
                          current_depth: usize) -> usize {
        // add a node
        self._nodes.push(Node::new());

        // initialize leaf node if termination criteria met
        if n_items == 0 {
            self._nodes[node_index].init_leaf(usize::MAX, &Vector2D::new_default());
            return current_depth + 1;
        }
        if n_items == 1 {
            self._nodes[node_index].init_leaf(item_indices[0], &self._points[item_indices[0]]);
            return current_depth + 1;
        }

        // choose which axis to split along
        let mut node_bound = BoundingBox2D::new_default();
        for i in 0..n_items {
            node_bound.merge_vec(&self._points[item_indices[i]]);
        }
        let d = node_bound.upper_corner - node_bound.lower_corner;
        let axis = d.dominant_axis();

        // pick mid point
        let (left, mid, right) = item_indices.select_nth_unstable_by(n_items, |a: &usize, b: &usize| {
            return match self._points[*a][axis] < self._points[*b][axis] {
                true => Ordering::Less,
                false => Ordering::Greater
            };
        });

        let mid_point = n_items / 2;
        let node_len = self._nodes.len();
        // recursively initialize children nodes
        let d0 = self.build_internal(node_index + 1, left, mid_point, current_depth + 1);
        self._nodes[node_index].init_internal(axis, *mid, node_len,
                                              &self._points[*mid]);
        let d1 = self.build_internal(self._nodes[node_index].child, right,
                                     n_items - mid_point - 1, current_depth + 1);

        return usize::max(d0, d1);
    }
}