Added allocation of points to leaf indices. (#5)

* Added allocation of points to leaf indices.

* Fixed style

examples/mpi_complete_tree.rs

@@ -89,7 +89,23 @@ pub fn main() {
 
     assert_eq!(nglobals.iter().unique().count(), 1);
 
+    // Check that the points are associated with the correct leaf keys.
+    let mut npoints = 0;
+    let leaf_point_map = tree.leafs_to_point_indices();
+
+    for (key, point_indices) in leaf_point_map {
+        for &index in point_indices {
+            assert!(key.is_ancestor(tree.point_keys()[index]));
+        }
+        npoints += point_indices.len();
+    }
+
+    // Make sure that the number of points and point keys lines up
+    // with the points stored for each leaf key.
+    assert_eq!(npoints, tree.points().len());
+    assert_eq!(npoints, tree.point_keys().len());
+
     if comm.rank() == 0 {
-        println!("Distributed tree is complete and linear.");
+        println!("No errors were found in setting up tree.");
     }
 }

src/octree.rs

@@ -34,6 +34,7 @@ pub struct Octree<'o, C> {
     leaf_tree: Vec<MortonKey>,
     coarse_tree_bounds: Vec<MortonKey>,
     all_keys: HashMap<MortonKey, KeyType>,
+    leaf_keys_to_points: HashMap<MortonKey, Vec<usize>>,
     bounding_box: PhysicalBox,
     comm: &'o C,
 }
@@ -107,13 +108,16 @@ impl<'o, C: CommunicatorCollectives> Octree<'o, C> {
 
         let all_keys = generate_all_keys(&leaf_tree, &coarse_tree, &coarse_tree_bounds, comm);
 
+        let leaf_keys_to_points = assign_points_to_leaf_keys(&point_keys, &leaf_tree);
+
         Self {
             points: points.to_vec(),
             point_keys,
             coarse_tree,
             leaf_tree,
             coarse_tree_bounds,
             all_keys,
+            leaf_keys_to_points,
             bounding_box,
             comm,
         }
@@ -147,6 +151,11 @@ impl<'o, C: CommunicatorCollectives> Octree<'o, C> {
         &self.leaf_tree
     }
 
+    /// Return the map from leaf keys to point indices
+    pub fn leafs_to_point_indices(&self) -> &HashMap<MortonKey, Vec<usize>> {
+        &self.leaf_keys_to_points
+    }
+
     /// Get the coarse tree bounds.
     ///
     /// This returns an array of size the number of ranks,

src/octree/parallel.rs

@@ -851,6 +851,25 @@ pub fn get_key_index(arr: &[MortonKey], key: MortonKey) -> usize {
     }
 }
 
+/// Generate a map that associates each leaf with the corresponding point indices.
+pub fn assign_points_to_leaf_keys(
+    point_keys: &[MortonKey],
+    leaf_keys: &[MortonKey],
+) -> HashMap<MortonKey, Vec<usize>> {
+    let mut point_map = HashMap::<MortonKey, Vec<usize>>::new();
+
+    for (index, point_key) in point_keys.iter().enumerate() {
+        let leaf_key_index = get_key_index(leaf_keys, *point_key);
+
+        let leaf_key = leaf_keys[leaf_key_index];
+        debug_assert!(leaf_key.is_ancestor(*point_key));
+
+        point_map.entry(leaf_key).or_default().push(index);
+    }
+
+    point_map
+}
+
 /// Check if a key is associated with the current rank.
 ///
 /// Note that the key does not need to exist as leaf. It just needs