Skip to content

Commit

Permalink
Merge pull request #34 from JayKickliter/jsk/subtree-iter
Browse files Browse the repository at this point in the history 
Add subtree iteration
  • Loading branch information
@JayKickliter
JayKickliter authored Oct 31, 2023
2 parents a60a603 + ac67942 commit e6923d6
Show file tree
Hide file tree
Showing 7 changed files with 354 additions and 8 deletions.
4 changes: 3 additions & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,10 @@ serde = { version = "1", optional = true, features = ["derive"] }
[dev-dependencies]
byteorder = "1"
criterion = { version = "0.3", features = ["html_reports"] }
h3ron = "0.15.1"
geo = "0.26.0"
geo-types = "0.7"
h3o = { version = "0.4.0", features = ["geo"] }
h3ron = "0.15.1"

[dev-dependencies.h3-lorawan-regions]
git = "https://github.com/JayKickliter/h3-lorawan-regions.git"
Expand Down
59 changes: 59 additions & 0 deletions benches/benches.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,13 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
use geo::polygon;
use geo_types::coord;
use h3_lorawan_regions::{
compact::US915 as COMPACT_US915_INDICES, nocompact::US915 as PLAIN_US915_INDICES,
};
use h3o::{
geom::{ContainmentMode, PolyfillConfig, Polygon, ToCells},
CellIndex, Resolution,
};
use h3ron::H3Cell;
use hextree::{compaction::EqCompactor, Cell, HexTreeMap, HexTreeSet};
use std::convert::TryFrom;
Expand Down Expand Up @@ -196,8 +201,62 @@ fn set_iteration(c: &mut Criterion) {
group.bench_function("count", |b| b.iter(|| set_precompacted.iter().count()));
}

fn subtree_iter(c: &mut Criterion) {
let mut group = c.benchmark_group("Subtree iteration");

let eiffel_tower_cells = {
let eiffel_tower_poly: geo::Polygon<f64> = polygon![
(x: 2.2918576408729336, y: 48.85772170856845),
(x: 2.295281693366718, y: 48.86007711794011),
(x: 2.2968743826623665, y: 48.859023236935656),
(x: 2.293404431342765, y: 48.85672213596601),
(x: 2.2918484611075485, y: 48.85772774822141),
(x: 2.2918576408729336, y: 48.85772170856845),
];
let eiffel_tower_poly = Polygon::from_degrees(eiffel_tower_poly).unwrap();
let mut eiffel_tower_cells: Vec<CellIndex> = eiffel_tower_poly
.to_cells(
PolyfillConfig::new(Resolution::Twelve)
.containment_mode(ContainmentMode::IntersectsBoundary),
)
.collect();
eiffel_tower_cells.sort();
eiffel_tower_cells.dedup();
eiffel_tower_cells
.into_iter()
.map(|cell| Cell::try_from(u64::from(cell)).unwrap())
.collect::<Vec<Cell>>()
};
let hex_map: HexTreeMap<i32> = eiffel_tower_cells
.iter()
.enumerate()
.map(|(i, &cell)| (cell, i as i32))
.collect();

let eiffel_tower_res1_parent = eiffel_tower_cells[0].to_parent(1).unwrap();
group.bench_function("Eiffel Tower Sum - Res1", |b| {
b.iter(|| {
hex_map
.subtree_iter(eiffel_tower_res1_parent)
.map(|(_cell, val)| val)
.sum::<i32>()
})
});

let eiffel_tower_res6_parent = eiffel_tower_cells[0].to_parent(7).unwrap();
group.bench_function("Eiffel Tower Sum - Res7", |b| {
b.iter(|| {
hex_map
.subtree_iter(eiffel_tower_res6_parent)
.map(|(_cell, val)| val)
.sum::<i32>()
})
});
}

criterion_group!(
benches,
subtree_iter,
set_lookup,
map_lookup,
set_iteration,
Expand Down
6 changes: 6 additions & 0 deletions src/cell.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,12 @@ impl CellStack {
}
}

impl From<Cell> for CellStack {
fn from(cell: Cell) -> CellStack {
CellStack(Some(cell))
}
}

impl fmt::Debug for Cell {
/// [H3 Index](https://h3geo.org/docs/core-library/h3Indexing/):
/// > The canonical string representation of an H3Index is the
Expand Down
47 changes: 45 additions & 2 deletions src/hex_tree_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

pub use crate::entry::{Entry, OccupiedEntry, VacantEntry};
use crate::{
cell::CellStack,
compaction::{Compactor, NullCompactor},
digits::Digits,
node::Node,
Expand Down Expand Up @@ -214,13 +215,55 @@ impl<V, C> HexTreeMap<V, C> {

/// An iterator visiting all cell-value pairs in arbitrary order.
pub fn iter(&self) -> impl Iterator<Item = (Cell, &V)> {
crate::iteration::Iter::new(&self.nodes)
crate::iteration::Iter::new(&self.nodes, CellStack::new())
}

/// An iterator visiting all cell-value pairs in arbitrary order
/// with mutable references to the values.
pub fn iter_mut(&mut self) -> impl Iterator<Item = (Cell, &mut V)> {
crate::iteration::IterMut::new(&mut self.nodes)
crate::iteration::IterMut::new(&mut self.nodes, CellStack::new())
}

/// An iterator visiting the specified cell or its children
/// references to the values.
pub fn subtree_iter(&self, cell: Cell) -> impl Iterator<Item = (Cell, &V)> {
let base_cell = cell.base();
match self.nodes[base_cell as usize].as_ref() {
Some(node) => {
let digits = Digits::new(cell);
match node.get_raw(0, cell, digits) {
Some((cell, Node::Leaf(val))) => Some((cell, val))
.into_iter()
.chain(crate::iteration::Iter::empty()),
Some((cell, Node::Parent(children))) => None
.into_iter()
.chain(crate::iteration::Iter::new(children, CellStack::from(cell))),
None => None.into_iter().chain(crate::iteration::Iter::empty()),
}
}
None => None.into_iter().chain(crate::iteration::Iter::empty()),
}
}

/// An iterator visiting the specified cell or its children with
/// mutable references to the values.
pub fn subtree_iter_mut(&mut self, cell: Cell) -> impl Iterator<Item = (Cell, &mut V)> {
let base_cell = cell.base();
match self.nodes[base_cell as usize].as_mut() {
Some(node) => {
let digits = Digits::new(cell);
match node.get_raw_mut(0, cell, digits) {
Some((cell, Node::Leaf(val))) => Some((cell, val))
.into_iter()
.chain(crate::iteration::IterMut::empty()),
Some((cell, Node::Parent(children))) => None.into_iter().chain(
crate::iteration::IterMut::new(children, CellStack::from(cell)),
),
None => None.into_iter().chain(crate::iteration::IterMut::empty()),
}
}
None => None.into_iter().chain(crate::iteration::IterMut::empty()),
}
}
}

Expand Down
160 changes: 156 additions & 4 deletions src/iteration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,11 @@ pub(crate) struct Iter<'a, V> {
}

impl<'a, V> Iter<'a, V> {
pub(crate) fn new(base: &'a [Option<Box<Node<V>>>]) -> Self {
pub(crate) fn new(base: &'a [Option<Box<Node<V>>>], mut cell_stack: CellStack) -> Self {
let mut iter = make_node_stack_iter(base);
let curr = iter.next();
let mut stack = Vec::with_capacity(16);
stack.push(iter);
let mut cell_stack = CellStack::new();
if let Some((digit, _)) = curr {
cell_stack.push(digit as u8)
}
Expand All @@ -44,6 +43,17 @@ impl<'a, V> Iter<'a, V> {
cell_stack,
}
}

pub(crate) fn empty() -> Self {
let stack = Vec::new();
let curr = None;
let cell_stack = CellStack::new();
Self {
stack,
curr,
cell_stack,
}
}
}

impl<'a, V> Iterator for Iter<'a, V> {
Expand Down Expand Up @@ -119,12 +129,11 @@ pub(crate) struct IterMut<'a, V> {
}

impl<'a, V> IterMut<'a, V> {
pub(crate) fn new(base: &'a mut [Option<Box<Node<V>>>]) -> Self {
pub(crate) fn new(base: &'a mut [Option<Box<Node<V>>>], mut cell_stack: CellStack) -> Self {
let mut iter = make_node_stack_iter_mut(base);
let curr = iter.next();
let mut stack = Vec::with_capacity(16);
stack.push(iter);
let mut cell_stack = CellStack::new();
if let Some((digit, _)) = curr {
cell_stack.push(digit as u8)
}
Expand All @@ -134,6 +143,17 @@ impl<'a, V> IterMut<'a, V> {
cell_stack,
}
}

pub(crate) fn empty() -> Self {
let stack = Vec::new();
let curr = None;
let cell_stack = CellStack::new();
Self {
stack,
curr,
cell_stack,
}
}
}

impl<'a, V> Iterator for IterMut<'a, V> {
Expand Down Expand Up @@ -184,9 +204,38 @@ impl<'a, V> Iterator for IterMut<'a, V> {
mod tests {
use crate::{Cell, HexTreeMap};
use byteorder::{LittleEndian as LE, ReadBytesExt};
use geo::polygon;
use h3_lorawan_regions::compact::US915 as COMPACT_US915_INDICES;
use h3o::{
geom::{ContainmentMode, PolyfillConfig, Polygon, ToCells},
CellIndex, Resolution,
};
use std::convert::TryFrom;

#[test]
fn test_visit() {
let parent = Cell::try_from(0x825997fffffffff).unwrap();
let children = [
Cell::try_from(0x835990fffffffff).unwrap(),
Cell::try_from(0x835991fffffffff).unwrap(),
Cell::try_from(0x835992fffffffff).unwrap(),
Cell::try_from(0x835993fffffffff).unwrap(),
Cell::try_from(0x835994fffffffff).unwrap(),
Cell::try_from(0x835995fffffffff).unwrap(),
Cell::try_from(0x835996fffffffff).unwrap(),
];

let hexmap: HexTreeMap<Cell> = children.iter().map(|cell| (cell, cell)).collect();
let visited = hexmap.subtree_iter(parent).collect::<Vec<_>>();

for (expected, (actual_k, actual_v)) in children.iter().zip(visited.iter()) {
assert_eq!(expected, *actual_v);
assert_eq!(expected.res(), actual_k.res());
assert_eq!(expected, actual_k);
}
assert_eq!(children.len(), visited.len());
}

#[test]
fn test_kv_iter_derives_key_cells() {
// Create a map where the key==value
Expand Down Expand Up @@ -260,4 +309,107 @@ mod tests {
.iter()
.all(|(cell, value)| map_plus_one[cell] == value + 1));
}

#[test]
fn test_subtree_iter_sum() {
// {
// "type": "FeatureCollection",
// "features": [
// {
// "type": "Feature",
// "properties": {},
// "geometry": {
// "coordinates": [
// [
// 2.2918576408729336,
// 48.85772170856845
// ],
// [
// 2.295281693366718,
// 48.86007711794011
// ],
// [
// 2.2968743826623665,
// 48.859023236935656
// ],
// [
// 2.293404431342765,
// 48.85672213596601
// ],
// [
// 2.2918484611075485,
// 48.85772774822141
// ]
// ],
// "type": "LineString"
// }
// }
// ],
// "bbox": null
// }
let eiffel_tower_cells = {
let eiffel_tower_poly: geo::Polygon<f64> = polygon![
(x: 2.2918576408729336, y: 48.85772170856845),
(x: 2.295281693366718, y: 48.86007711794011),
(x: 2.2968743826623665, y: 48.859023236935656),
(x: 2.293404431342765, y: 48.85672213596601),
(x: 2.2918484611075485, y: 48.85772774822141),
(x: 2.2918576408729336, y: 48.85772170856845),
];
let eiffel_tower_poly = Polygon::from_degrees(eiffel_tower_poly).unwrap();
let mut eiffel_tower_cells: Vec<CellIndex> = eiffel_tower_poly
.to_cells(
PolyfillConfig::new(Resolution::Twelve)
.containment_mode(ContainmentMode::ContainsCentroid),
)
.collect();
eiffel_tower_cells.sort();
eiffel_tower_cells.dedup();
eiffel_tower_cells
.into_iter()
.map(|cell| Cell::try_from(u64::from(cell)).unwrap())
.collect::<Vec<Cell>>()
};
let mut hex_map: HexTreeMap<i32> = eiffel_tower_cells
.iter()
.enumerate()
.map(|(i, &cell)| (cell, i as i32))
.collect();
let eiffel_tower_res1_parent = Cell::try_from(0x811fbffffffffff).unwrap();
let value_sum: i32 = hex_map
.subtree_iter(eiffel_tower_res1_parent)
.map(|(_cell, val)| val)
.sum();
// Establish the sum of map values in the eiffel tower block.
assert_eq!(value_sum, 22578);

let west_mask_res9 = Cell::try_from(0x891fb46741bffff).unwrap();
let east_mask_res9 = Cell::try_from(0x891fb467413ffff).unwrap();
let west_value_sum: i32 = hex_map
.subtree_iter(west_mask_res9)
.map(|(_cell, val)| val)
.sum();
let east_value_sum: i32 = hex_map
.subtree_iter(east_mask_res9)
.map(|(_cell, val)| val)
.sum();
// Now we have the sum of two difference subtrees, both of
// which should cover the entire eiffel tower block. Therefore
// they their individual sums should be the same as the
// overall sum.
assert_eq!(value_sum, west_value_sum + east_value_sum);

let expected_sum = hex_map.len() as i32 + value_sum;

for (_cell, val) in hex_map.subtree_iter_mut(eiffel_tower_res1_parent) {
*val += 1;
}

let value_sum: i32 = hex_map
.subtree_iter(eiffel_tower_res1_parent)
.map(|(_cell, val)| val)
.sum();

assert_eq!(value_sum, expected_sum);
}
}
Loading

0 comments on commit e6923d6

Please sign in to comment.