#![expect(
clippy::cast_possible_truncation,
reason = "level_mod (usize->u8) — always <= MAX_CELL_LEVEL"
)]
use crate::s2::coords::Level;
use crate::s2::region_coverer::RegionCoverer;
use crate::s2::{CellId, CellUnion, Point, Region};
#[derive(Clone, Debug, PartialEq)]
pub struct Options {
pub max_cells: usize,
pub min_level: Level,
pub max_level: Level,
pub level_mod: usize,
pub index_contains_points_only: bool,
pub optimize_for_space: bool,
pub marker_character: char,
}
impl Default for Options {
fn default() -> Self {
Options {
max_cells: 8,
min_level: Level::new(4),
max_level: Level::new(16),
level_mod: 1,
index_contains_points_only: false,
optimize_for_space: false,
marker_character: '$',
}
}
}
#[derive(Debug)]
pub struct RegionTermIndexer {
options: Options,
}
impl RegionTermIndexer {
pub fn new() -> Self {
RegionTermIndexer {
options: Options::default(),
}
}
pub fn with_options(options: Options) -> Self {
RegionTermIndexer { options }
}
pub fn options(&self) -> &Options {
&self.options
}
pub fn options_mut(&mut self) -> &mut Options {
&mut self.options
}
fn true_max_level(&self) -> Level {
let opts = &self.options;
if opts.level_mod == 1 {
opts.max_level
} else {
let range = (opts.max_level - opts.min_level) as usize;
let adjusted = range - (range % opts.level_mod);
opts.min_level + adjusted as u8
}
}
pub fn get_index_terms(&self, region: &dyn Region) -> Vec<String> {
let covering = self.get_covering(region);
self.get_index_terms_for_covering(&covering)
}
pub fn get_index_terms_for_point(&self, point: Point) -> Vec<String> {
let cell_id = CellId::from_point(&point);
let level = self.true_max_level();
let id = cell_id.parent_at_level(level);
let mut terms = Vec::new();
let mut current = id;
while current.level() >= self.options.min_level {
terms.push(self.ancestor_term(current));
if current.level() == 0 {
break;
}
current = current.parent();
}
if !self.options.index_contains_points_only {
terms.push(self.covering_term(id));
}
terms
}
pub fn get_query_terms(&self, region: &dyn Region) -> Vec<String> {
let covering = self.get_covering(region);
self.get_query_terms_for_covering(&covering)
}
pub fn get_query_terms_for_point(&self, point: Point) -> Vec<String> {
let cell_id = CellId::from_point(&point);
let level = self.true_max_level();
let id = cell_id.parent_at_level(level);
let mut terms = Vec::new();
terms.push(self.ancestor_term(id));
let mut current = id;
while current.level() >= self.options.min_level {
terms.push(self.covering_term(current));
if current.level() == 0 {
break;
}
current = current.parent();
}
terms
}
pub fn get_index_terms_for_covering(&self, covering: &CellUnion) -> Vec<String> {
let true_max = self.true_max_level();
let mut terms = Vec::new();
for &id in covering.cell_ids() {
let level = id.level();
debug_assert!(level >= self.options.min_level);
debug_assert!(level <= self.options.max_level);
debug_assert_eq!(
0,
(level - self.options.min_level) % self.options.level_mod as u8
);
let mut current = id;
while current.level() >= self.options.min_level {
if !self.options.optimize_for_space || current == id {
terms.push(self.ancestor_term(current));
}
if current.level() == 0 {
break;
}
current = current.parent();
}
if !self.options.index_contains_points_only && id.level() < true_max {
terms.push(self.covering_term(id));
}
}
terms.sort_unstable();
terms.dedup();
terms
}
pub fn get_query_terms_for_covering(&self, covering: &CellUnion) -> Vec<String> {
let _true_max = self.true_max_level();
let mut terms = Vec::new();
for &id in covering.cell_ids() {
let level = id.level();
debug_assert!(level >= self.options.min_level);
debug_assert!(level <= self.options.max_level);
debug_assert_eq!(
0,
(level - self.options.min_level) % self.options.level_mod as u8
);
terms.push(self.ancestor_term(id));
if !self.options.index_contains_points_only {
let mut current = id;
while current.level() >= self.options.min_level {
terms.push(self.covering_term(current));
if current.level() == 0 {
break;
}
current = current.parent();
}
}
if self.options.optimize_for_space {
let mut current = id;
while current.level() > self.options.min_level {
current = current.parent();
terms.push(self.ancestor_term(current));
}
}
}
terms.sort_unstable();
terms.dedup();
terms
}
fn get_covering(&self, region: &dyn Region) -> CellUnion {
let coverer = RegionCoverer::new()
.max_cells(self.options.max_cells)
.min_level(self.options.min_level)
.max_level(self.options.max_level)
.level_mod(self.options.level_mod as u8);
coverer.covering(region)
}
#[expect(clippy::unused_self, reason = "matches C++ method signature")]
fn ancestor_term(&self, id: CellId) -> String {
id.to_token()
}
fn covering_term(&self, id: CellId) -> String {
format!("{}{}", self.options.marker_character, id.to_token())
}
}
impl Default for RegionTermIndexer {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::s2::cap::Cap;
use crate::s2::text_format;
#[test]
fn test_default_options() {
let indexer = RegionTermIndexer::new();
assert_eq!(indexer.options().max_cells, 8);
assert_eq!(indexer.options().min_level, 4);
assert_eq!(indexer.options().max_level, 16);
assert_eq!(indexer.options().marker_character, '$');
}
#[test]
fn test_point_index_terms() {
let indexer = RegionTermIndexer::new();
let p = text_format::parse_point("0:0");
let terms = indexer.get_index_terms_for_point(p);
assert!(!terms.is_empty());
let covering_count = terms.iter().filter(|t| t.starts_with('$')).count();
assert!(covering_count >= 1, "should have covering terms");
}
#[test]
fn test_point_only_index() {
let opts = Options {
index_contains_points_only: true,
..Options::default()
};
let indexer = RegionTermIndexer::with_options(opts);
let p = text_format::parse_point("0:0");
let terms = indexer.get_index_terms_for_point(p);
let covering_count = terms.iter().filter(|t| t.starts_with('$')).count();
assert_eq!(
covering_count, 0,
"points-only should have no covering terms"
);
}
#[test]
fn test_point_query_terms() {
let indexer = RegionTermIndexer::new();
let p = text_format::parse_point("0:0");
let terms = indexer.get_query_terms_for_point(p);
assert!(!terms.is_empty());
}
#[test]
fn test_region_index_and_query() {
let indexer = RegionTermIndexer::new();
let cap = Cap::from_center_angle(
text_format::parse_point("0:0"),
crate::s1::Angle::from_degrees(1.0),
);
let index_terms = indexer.get_index_terms(&cap);
let query_terms = indexer.get_query_terms(&cap);
assert!(!index_terms.is_empty());
assert!(!query_terms.is_empty());
let has_overlap = index_terms.iter().any(|t| query_terms.contains(t));
assert!(has_overlap, "index and query terms should overlap");
}
#[test]
fn test_covering_terms_contain_marker() {
let indexer = RegionTermIndexer::new();
let p = text_format::parse_point("45:45");
let terms = indexer.get_index_terms_for_point(p);
for term in &terms {
if term.starts_with('$') {
assert!(term.len() > 1, "covering term too short: {term}");
} else {
assert!(!term.is_empty(), "ancestor term should not be empty");
}
}
}
#[test]
fn test_custom_marker() {
let opts = Options {
marker_character: '#',
..Options::default()
};
let indexer = RegionTermIndexer::with_options(opts);
let p = text_format::parse_point("0:0");
let terms = indexer.get_index_terms_for_point(p);
let has_hash = terms.iter().any(|t| t.starts_with('#'));
assert!(has_hash, "should use custom marker '#'");
let has_dollar = terms.iter().any(|t| t.starts_with('$'));
assert!(!has_dollar, "should not use default marker '$'");
}
#[test]
fn test_point_containment() {
let indexer = RegionTermIndexer::new();
let p = text_format::parse_point("10:20");
let index_terms = indexer.get_index_terms_for_point(p);
let cap = Cap::from_center_angle(p, crate::s1::Angle::from_degrees(5.0));
let query_terms = indexer.get_query_terms(&cap);
let intersection: Vec<_> = index_terms
.iter()
.filter(|t| query_terms.contains(t))
.collect();
assert!(
!intersection.is_empty(),
"point inside cap should have matching terms"
);
}
#[test]
fn test_marker_character() {
let opts = Options {
min_level: Level::new(20),
max_level: Level::new(20),
..Options::default()
};
let mut indexer = RegionTermIndexer::with_options(opts);
let point = crate::s2::LatLng::from_degrees(10.0, 20.0).to_point();
assert_eq!(indexer.options().marker_character, '$');
let terms = indexer.get_query_terms_for_point(point);
assert_eq!(2, terms.len());
assert!(terms.iter().any(|t| t.starts_with('$')));
indexer.options_mut().marker_character = ':';
assert_eq!(indexer.options().marker_character, ':');
let terms2 = indexer.get_query_terms_for_point(point);
assert_eq!(2, terms2.len());
assert!(terms2.iter().any(|t| t.starts_with(':')));
}
#[test]
fn test_max_level_set_loosely() {
let opts1 = Options {
min_level: Level::new(1),
level_mod: 2,
max_level: Level::new(19),
..Options::default()
};
let indexer1 = RegionTermIndexer::with_options(opts1);
let opts2 = Options {
min_level: Level::new(1),
level_mod: 2,
max_level: Level::new(20),
..Options::default()
};
let indexer2 = RegionTermIndexer::with_options(opts2);
let point = crate::s2::LatLng::from_degrees(37.7749, -122.4194).to_point();
assert_eq!(
indexer1.get_index_terms_for_point(point),
indexer2.get_index_terms_for_point(point),
);
assert_eq!(
indexer1.get_query_terms_for_point(point),
indexer2.get_query_terms_for_point(point),
);
}
}