From 3e8cb239ca816e2cd4a1bc1ef453585e29ff45cc Mon Sep 17 00:00:00 2001
From: Marc Nijdam <marc_nijdam@mac.com>
Date: Mon, 3 Apr 2023 16:37:18 -0400
Subject: [PATCH] Add index check command

---
 Cargo.lock             | 17 +++++++++++++++++
 generator/Cargo.toml   |  1 +
 generator/src/index.rs | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index e085bb3..4b60c85 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -108,6 +108,12 @@ dependencies = [
  "rustc-demangle",
 ]
 
+[[package]]
+name = "bitfield"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d7e60934ceec538daadb9d8432424ed043a904d8e0243f3c6446bce549a46ac"
+
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -476,6 +482,16 @@ version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286"
 
+[[package]]
+name = "hextree"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a6ffc35fc0fb2af599742de1938d44cd18b0ae7cf7a7faf4b067f43cc863902"
+dependencies = [
+ "bitfield",
+ "thiserror",
+]
+
 [[package]]
 name = "indexmap"
 version = "1.9.2"
@@ -587,6 +603,7 @@ dependencies = [
  "geojson",
  "h3ron",
  "helium-proto",
+ "hextree",
  "log",
  "rayon",
  "serde",
diff --git a/generator/Cargo.toml b/generator/Cargo.toml
index dd052f3..cc45837 100644
--- a/generator/Cargo.toml
+++ b/generator/Cargo.toml
@@ -18,6 +18,7 @@ serde_json = "1"
 flate2 = "1"
 h3ron = {version = "0.16"}
 geo-types = "*"
+hextree = "0"
 byteorder = "1.4"
 rayon = "1.6"
 log = "0"
diff --git a/generator/src/index.rs b/generator/src/index.rs
index 8a1f876..463fb3c 100644
--- a/generator/src/index.rs
+++ b/generator/src/index.rs
@@ -23,6 +23,7 @@ impl Cmd {
 pub enum IndexCmd {
     Generate(Generate),
     Export(Export),
+    Check(Check),
 }
 
 impl IndexCmd {
@@ -30,6 +31,7 @@ impl IndexCmd {
         match self {
             Self::Generate(cmd) => cmd.run(),
             Self::Export(cmd) => cmd.run(),
+            Self::Check(cmd) => cmd.run(),
         }
     }
 }
@@ -101,6 +103,18 @@ fn read_cells<P: AsRef<path::Path>>(file: P) -> Result<Vec<H3Cell>> {
     Ok(vec)
 }
 
+fn read_hexset<P: AsRef<path::Path>>(file: P) -> Result<hextree::HexTreeSet> {
+    let file = fs::File::open(file.as_ref())?;
+    let mut reader = GzDecoder::new(file);
+    let mut vec = Vec::new();
+
+    while let Ok(entry) = reader.read_u64::<byteorder::LittleEndian>() {
+        vec.push(hextree::Cell::from_raw(entry)?);
+    }
+
+    Ok(vec.iter().collect())
+}
+
 fn write_cells<P: AsRef<path::Path>>(cells: Vec<H3Cell>, output: P) -> Result<()> {
     let file = fs::File::create(output.as_ref())?;
     let mut writer = GzEncoder::new(file, Compression::default());
@@ -148,3 +162,26 @@ impl Export {
         Ok(())
     }
 }
+
+/// Check membership of a given h3 index in the given binary file
+#[derive(Debug, clap::Args)]
+pub struct Check {
+    input: path::PathBuf,
+    cell: h3ron::H3Cell,
+}
+
+impl Check {
+    pub fn run(&self) -> Result<()> {
+        let hex_set = read_hexset(&self.input)?;
+        if hex_set.contains(hextree::Cell::from_raw(*self.cell)?) {
+            println!("Cell {} in {}", self.cell.to_string(), self.input.display())
+        } else {
+            anyhow::bail!(
+                "Cell {} not in {}",
+                self.cell.to_string(),
+                self.input.display()
+            )
+        }
+        Ok(())
+    }
+}