diff --git a/Cargo.toml b/Cargo.toml index 21521b9ac7..6fde595e33 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -91,6 +91,10 @@ datasketches = { version = "0.5.0", features = ["hll"] } futures-util = { version = "0.3.28", optional = true } futures-channel = { version = "0.3.28", optional = true } fnv = "1.0.7" +num-bigint = "0.4" +num-rational = "0.4" +num-traits = "0.2" +robust = "1" typetag = "0.2.21" unwrap-infallible = "1.0.0" serde-tuple-vec-map = "1.0.1" diff --git a/examples/geo_json.rs b/examples/geo_json.rs new file mode 100644 index 0000000000..51b6539639 --- /dev/null +++ b/examples/geo_json.rs @@ -0,0 +1,51 @@ +use tantivy::collector::TopDocs; +use tantivy::query::SpatialQuery; +use tantivy::schema::{Schema, Value, SPATIAL, STORED, TEXT}; +use tantivy::{Index, IndexWriter, TantivyDocument}; +fn main() -> tantivy::Result<()> { + let mut schema_builder = Schema::builder(); + schema_builder.add_json_field("properties", STORED | TEXT); + schema_builder.add_spatial_field("geometry", STORED | SPATIAL); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema.clone()); + let mut index_writer: IndexWriter = index.writer(50_000_000)?; + let doc = TantivyDocument::parse_json( + &schema, + r#"{ + "type":"Feature", + "geometry":{ + "type":"Polygon", + "coordinates":[[[-99.483911,45.577697],[-99.483869,45.571457],[-99.481739,45.571461],[-99.474881,45.571584],[-99.473167,45.571615],[-99.463394,45.57168],[-99.463391,45.57883],[-99.463368,45.586076],[-99.48177,45.585926],[-99.48384,45.585953],[-99.483885,45.57873],[-99.483911,45.577697]]] + }, + "properties":{ + "admin_level":"8", + "border_type":"city", + "boundary":"administrative", + "gnis:feature_id":"1267426", + "name":"Hosmer", + "place":"city", + "source":"TIGER/Line® 2008 Place Shapefiles (http://www.census.gov/geo/www/tiger/)", + "wikidata":"Q2442118", + "wikipedia":"en:Hosmer, South Dakota" + } + }"#, + )?; + index_writer.add_document(doc)?; + index_writer.commit()?; + + let reader = index.reader()?; + let searcher = reader.searcher(); + let field = schema.get_field("geometry").unwrap(); + let query = SpatialQuery::from_bounds(field, [[-99.49, 45.56], [-99.45, 45.59]]); + let hits = searcher.search(&query, &TopDocs::with_limit(10).order_by_score())?; + for (_score, doc_address) in &hits { + let retrieved_doc: TantivyDocument = searcher.doc(*doc_address)?; + if let Some(field_value) = retrieved_doc.get_first(field) { + if let Some(geometry_box) = field_value.as_value().into_geometry() { + println!("Retrieved geometry: {:?}", geometry_box); + } + } + } + assert_eq!(hits.len(), 1); + Ok(()) +} diff --git a/query-grammar/Cargo.toml b/query-grammar/Cargo.toml index 9ea2875f26..4ea0ac6c20 100644 --- a/query-grammar/Cargo.toml +++ b/query-grammar/Cargo.toml @@ -15,5 +15,5 @@ edition = "2024" nom = "7" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.140" -ordered-float = "5.0.0" +ordered-float = { version = "5.0.0", features = ["serde"] } fnv = "1.0.7" diff --git a/query-grammar/src/lib.rs b/query-grammar/src/lib.rs index 46ea65c84f..58d50f02df 100644 --- a/query-grammar/src/lib.rs +++ b/query-grammar/src/lib.rs @@ -11,7 +11,7 @@ pub use crate::infallible::LenientError; pub use crate::occur::Occur; use crate::query_grammar::{parse_to_ast, parse_to_ast_lenient}; pub use crate::user_input_ast::{ - Delimiter, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral, + Delimiter, SpatialPredicateKind, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral, }; #[derive(Debug, Serialize)] @@ -58,6 +58,52 @@ mod tests { assert!(parse_query("title:").is_err()); } + #[test] + fn test_spatial_intersects() { + // Test simple case first. + let ast = parse_query(r#"geo:$intersects(1.0 2.0, 3.0 4.0, 5.0 6.0)"#).unwrap(); + let debug = format!("{ast:?}"); + assert!(debug.contains("$intersects"), "got: {debug}"); + } + + #[test] + fn test_spatial_contains() { + let ast = parse_query(r#"geo:$contains(1.0 2.0, 3.0 4.0, 5.0 6.0)"#).unwrap(); + let debug = format!("{ast:?}"); + assert!(debug.contains("$contains"), "got: {debug}"); + } + + #[test] + fn test_spatial_join_debug() { + for (query, expected) in [ + ( + "geo:$within(1rad, $query(name:point))", + "\"geo\":$within(1rad, $query(name:point))", + ), + ( + "geo:$between(1rad, 2rad, $query(a))", + "\"geo\":$between(1rad, 2rad, $query(a))", + ), + ( + "geo:$intersects($query(name:x))", + "\"geo\":$intersects($query(name:x))", + ), + ( + "geo:$contains($query(name:x))", + "\"geo\":$contains($query(name:x))", + ), + ("geo:$knn(3, 1 2)", "\"geo\":$knn(3, 1 2)"), + ] { + let debug = format!("{:?}", parse_query(query).unwrap()); + assert!(debug.contains(expected), "{query} gives {debug}"); + } + } + + #[test] + fn test_spatial_too_few_coords() { + assert!(parse_query(r#"geo:$intersects(1.0 2.0, 3.0 4.0)"#).is_err()); + } + #[test] fn test_parse_query_lenient_wrong_query() { let (_, errors) = parse_query_lenient("title:"); diff --git a/query-grammar/src/query_grammar.rs b/query-grammar/src/query_grammar.rs index 3d0d58b641..82124d7f59 100644 --- a/query-grammar/src/query_grammar.rs +++ b/query-grammar/src/query_grammar.rs @@ -13,7 +13,9 @@ use nom::error::{Error, ErrorKind}; use nom::multi::{many0, many1, separated_list0}; use nom::sequence::{delimited, preceded, separated_pair, terminated, tuple}; -use super::user_input_ast::{UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral}; +use super::user_input_ast::{ + SpatialPredicateKind, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral, +}; use crate::Occur; use crate::infallible::*; use crate::user_input_ast::Delimiter; @@ -364,6 +366,319 @@ fn exists_infallible(inp: &str) -> JResult<&str, UserInputAst> { Ok((inp, (exists, Vec::new()))) } +/// Parse a float that may be negative. Returns the f64 value. +fn spatial_float(inp: &str) -> IResult<&str, f64> { + map_res( + recognize(tuple(( + opt(char('-')), + digit1, + opt(tuple((char('.'), digit1))), + ))), + |s: &str| s.parse::(), + )(inp) +} + +/// Earth radius in miles. +const EARTH_RADIUS_MI: f64 = 3958.8; +/// Earth radius in kilometers. +const EARTH_RADIUS_KM: f64 = 6371.0; +/// Earth radius in meters. +const EARTH_RADIUS_M: f64 = 6_371_000.0; +/// Earth radius in feet. +const EARTH_RADIUS_FT: f64 = 20_902_464.0; + +/// Parse a distance with unit: 50mi, 75km, 200ft, 1000m, 0.5rad. +fn spatial_distance(inp: &str) -> IResult<&str, f64> { + let (inp, value) = spatial_float(inp)?; + let (inp, unit) = alt((tag("mi"), tag("km"), tag("ft"), tag("m"), tag("rad")))(inp)?; + let radians = match unit { + "mi" => value / EARTH_RADIUS_MI, + "km" => value / EARTH_RADIUS_KM, + "ft" => value / EARTH_RADIUS_FT, + "m" => value / EARTH_RADIUS_M, + "rad" => value, + _ => unreachable!(), + }; + Ok((inp, radians)) +} + +/// Extract the content of a $query(...) argument, handling balanced parentheses. +fn query_arg(inp: &str) -> IResult<&str, String> { + let (inp, _) = tag("$query(")(inp)?; + let mut depth = 1usize; + let mut end = 0; + for (i, c) in inp.char_indices() { + match c { + '(' => depth += 1, + ')' => { + depth -= 1; + if depth == 0 { + end = i; + break; + } + } + _ => {} + } + } + if depth != 0 { + return Err(nom::Err::Failure(Error::new(inp, ErrorKind::Verify))); + } + let query_str = inp[..end].trim().to_string(); + Ok((&inp[end + 1..], query_str)) +} + +/// Parse a spatial predicate: $intersects, $contains, $within, $between, or $knn. +fn spatial(inp: &str) -> IResult<&str, UserInputLeaf> { + alt(( + spatial_polygon, + spatial_within, + spatial_between, + spatial_knn, + ))(inp) +} + +fn spatial_polygon(inp: &str) -> IResult<&str, UserInputLeaf> { + let mut predicate_tag = alt(( + value(SpatialPredicateKind::Intersects, tag("$intersects")), + value(SpatialPredicateKind::Contains, tag("$contains")), + )); + + let (inp, predicate) = predicate_tag(inp)?; + let (inp, _) = tuple((multispace0, char('('), multispace0))(inp)?; + + // Try $query(...) for a join. + if let Ok((inp, inner)) = query_arg(inp) { + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + return Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate, + coordinates: vec![], + inner_query: Some(inner), + }, + )); + } + + let (inp, coordinates) = polygon_coordinates(inp)?; + Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate, + coordinates, + inner_query: None, + }, + )) +} + +/// Lon/lat pairs, as a spatial leaf holds them. +type Coordinates = Vec<( + ordered_float::OrderedFloat, + ordered_float::OrderedFloat, +)>; + +/// Parse the vertices of a polygon, at least 3 lon/lat pairs, and the closing parenthesis. +fn polygon_coordinates(inp: &str) -> IResult<&str, Coordinates> { + let coord_sep = || delimited(multispace0, char(','), multispace0); + let coord_pair = separated_pair(spatial_float, multispace1, spatial_float); + let (inp, pairs) = separated_list0(coord_sep(), coord_pair)(inp)?; + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + + if pairs.len() < 3 { + return Err(nom::Err::Failure(Error::new(inp, ErrorKind::Verify))); + } + + let coordinates = pairs + .into_iter() + .map(|(lon, lat)| { + ( + ordered_float::OrderedFloat(lon), + ordered_float::OrderedFloat(lat), + ) + }) + .collect(); + Ok((inp, coordinates)) +} + +fn spatial_within(inp: &str) -> IResult<&str, UserInputLeaf> { + let coord_sep = || delimited(multispace0, char(','), multispace0); + let mut coord_pair = separated_pair(spatial_float, multispace1, spatial_float); + + let (inp, _) = tag("$within")(inp)?; + let (inp, _) = tuple((multispace0, char('('), multispace0))(inp)?; + + // With no distance, the coordinates are a polygon to lie inside. + if let Ok((inp, coordinates)) = polygon_coordinates(inp) { + return Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::CoveredBy, + coordinates, + inner_query: None, + }, + )); + } + + let (inp, radius) = spatial_distance(inp)?; + let (inp, _) = coord_sep()(inp)?; + + // Try $query(...) for a join. + if let Ok((inp, inner)) = query_arg(inp) { + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + return Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::Within(ordered_float::OrderedFloat(radius)), + coordinates: vec![], + inner_query: Some(inner), + }, + )); + } + + let (inp, (lon, lat)) = coord_pair(inp)?; + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + + Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::Within(ordered_float::OrderedFloat(radius)), + coordinates: vec![( + ordered_float::OrderedFloat(lon), + ordered_float::OrderedFloat(lat), + )], + inner_query: None, + }, + )) +} + +fn spatial_between(inp: &str) -> IResult<&str, UserInputLeaf> { + let coord_sep = || delimited(multispace0, char(','), multispace0); + let mut coord_pair = separated_pair(spatial_float, multispace1, spatial_float); + + let (inp, _) = tag("$between")(inp)?; + let (inp, _) = tuple((multispace0, char('('), multispace0))(inp)?; + let (inp, inner) = spatial_distance(inp)?; + let (inp, _) = coord_sep()(inp)?; + let (inp, outer) = spatial_distance(inp)?; + let (inp, _) = coord_sep()(inp)?; + + // Try $query(...) for a join. + if let Ok((inp, inner_q)) = query_arg(inp) { + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + return Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::Between( + ordered_float::OrderedFloat(inner), + ordered_float::OrderedFloat(outer), + ), + coordinates: vec![], + inner_query: Some(inner_q), + }, + )); + } + + let (inp, (lon, lat)) = coord_pair(inp)?; + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + + Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::Between( + ordered_float::OrderedFloat(inner), + ordered_float::OrderedFloat(outer), + ), + coordinates: vec![( + ordered_float::OrderedFloat(lon), + ordered_float::OrderedFloat(lat), + )], + inner_query: None, + }, + )) +} + +fn spatial_knn(inp: &str) -> IResult<&str, UserInputLeaf> { + let mut coord_pair = separated_pair(spatial_float, multispace1, spatial_float); + + let (inp, _) = tag("$knn")(inp)?; + let (inp, _) = tuple((multispace0, char('('), multispace0))(inp)?; + let (inp, k) = nom::character::complete::u32(inp)?; + let (inp, _) = delimited(multispace0, char(','), multispace0)(inp)?; + let (inp, (lon, lat)) = coord_pair(inp)?; + let (inp, _) = tuple((multispace0, char(')')))(inp)?; + + Ok(( + inp, + UserInputLeaf::Spatial { + field: None, + predicate: SpatialPredicateKind::Knn(k as usize), + coordinates: vec![( + ordered_float::OrderedFloat(lon), + ordered_float::OrderedFloat(lat), + )], + inner_query: None, + }, + )) +} + +fn spatial_precond(inp: &str) -> IResult<&str, (), ()> { + let predicate: IResult<&str, (&str, &str, char)> = peek(tuple(( + alt(( + tag("$intersects"), + tag("$contains"), + tag("$within"), + tag("$between"), + tag("$knn"), + )), + multispace0, + char('('), + )))(inp); + predicate + .map(|(rest, _)| (rest, ())) + .map_err(|e| e.map(|_| ())) +} + +/// The lenient counterpart of `spatial`, for input that starts with a spatial predicate. A +/// malformed predicate is reported and skipped, through its closing parenthesis, so that the +/// clauses after it still parse. +fn spatial_infallible(inp: &str) -> JResult<&str, Option> { + if let Ok((rest, leaf)) = spatial(inp) { + return Ok((rest, (Some(leaf), Vec::new()))); + } + let open = inp.find('(').unwrap_or(inp.len()); + let mut depth = 0usize; + let mut end = inp.len(); + for (i, c) in inp[open..].char_indices() { + match c { + '(' => depth += 1, + ')' => { + depth -= 1; + if depth == 0 { + end = open + i + 1; + break; + } + } + _ => {} + } + } + Ok(( + &inp[end..], + ( + None, + vec![LenientErrorInternal { + pos: inp.len(), + message: "invalid spatial predicate".to_string(), + }], + ), + )) +} + fn literal(inp: &str) -> IResult<&str, UserInputAst> { // * alone is already parsed by our caller, so if `exists` succeed, we can be confident // something (a field name) got parsed before @@ -371,7 +686,7 @@ fn literal(inp: &str) -> IResult<&str, UserInputAst> { map( tuple(( opt(field_name), - alt((range, set, exists, regex, term_or_phrase)), + alt((range, set, exists, regex, spatial, term_or_phrase)), )), |(field_name, leaf): (Option, UserInputLeaf)| leaf.set_field(field_name).into(), ), @@ -398,6 +713,7 @@ fn literal_no_group_infallible(inp: &str) -> JResult<&str, Option> value((), peek(one_of("/"))), map(regex_infallible, |(regex, errs)| (Some(regex), errs)), ), + (spatial_precond, spatial_infallible), ), delimited_infallible(space0_infallible, term_or_phrase_infallible, nothing), ), @@ -1211,6 +1527,12 @@ mod test { // Ok(()) // } + #[test] + fn test_spatial_parser_direct() { + let result = spatial("$intersects(1.0 2.0, 3.0 4.0, 5.0 6.0)"); + assert!(result.is_ok(), "spatial parser failed: {:?}", result); + } + #[test] fn test_positive_float_number() { fn valid_parse(float_str: &str, expected_val: f64, expected_remaining: &str) { @@ -1279,6 +1601,40 @@ mod test { test_parse_query_to_ast_helper("", ""); } + #[test] + fn test_parse_malformed_spatial_lenient() { + let (ast, errors) = parse_to_ast_lenient("geometry:$within(oops) name:box"); + assert_eq!(format!("{ast:?}"), "\"name\":box"); + assert_eq!(errors.len(), 1); + let (ast, errors) = parse_to_ast_lenient("geometry:$knn(1, 10) name:box"); + assert_eq!(format!("{ast:?}"), "\"name\":box"); + assert_eq!(errors.len(), 1); + } + + #[test] + fn test_parse_spatial_lenient() { + test_parse_query_to_ast_helper( + "geometry:$within(1rad, 1 2)", + "\"geometry\":$within(1rad, 1 2)", + ); + test_parse_query_to_ast_helper( + "geometry:$intersects(0 0, 1 0, 1 1)", + "\"geometry\":$intersects(0 0, 1 0, 1 1)", + ); + test_parse_query_to_ast_helper( + "+name:cafe +geometry:$knn(5, 1 2)", + "(+\"name\":cafe +\"geometry\":$knn(5, 1 2))", + ); + test_parse_query_to_ast_helper( + "geometry:$contains($query(name:small))", + "\"geometry\":$contains($query(name:small))", + ); + test_parse_query_to_ast_helper( + "geometry:$within(0 0, 1 0, 1 1)", + "\"geometry\":$within(0 0, 1 0, 1 1)", + ); + } + #[test] fn test_parse_query_to_ast_hyphen() { test_parse_query_to_ast_helper("\"www-form-encoded\"", "\"www-form-encoded\""); diff --git a/query-grammar/src/user_input_ast.rs b/query-grammar/src/user_input_ast.rs index 0451b244dc..2e48ff6055 100644 --- a/query-grammar/src/user_input_ast.rs +++ b/query-grammar/src/user_input_ast.rs @@ -5,6 +5,24 @@ use serde::Serialize; use crate::Occur; +/// The spatial predicate kind, shared between parser and formatter. +#[derive(PartialEq, Eq, Hash, Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum SpatialPredicateKind { + Intersects, + /// The field's geometry contains the query's: `$contains`. + Contains, + /// The field's geometry lies inside the query polygon: `$within` with a polygon and no + /// distance. + CoveredBy, + Within(ordered_float::OrderedFloat), + Between( + ordered_float::OrderedFloat, + ordered_float::OrderedFloat, + ), + Knn(usize), +} + #[derive(PartialEq, Eq, Hash, Clone, Serialize)] #[serde(tag = "type")] #[serde(rename_all = "snake_case")] @@ -27,6 +45,16 @@ pub enum UserInputLeaf { field: Option, pattern: String, }, + Spatial { + field: Option, + predicate: SpatialPredicateKind, + coordinates: Vec<( + ordered_float::OrderedFloat, + ordered_float::OrderedFloat, + )>, + /// If set, this spatial predicate is a join. The string is the inner query text. + inner_query: Option, + }, } impl UserInputLeaf { @@ -52,6 +80,17 @@ impl UserInputLeaf { None => UserInputLeaf::All, }, UserInputLeaf::Regex { field: _, pattern } => UserInputLeaf::Regex { field, pattern }, + UserInputLeaf::Spatial { + field: _, + predicate, + coordinates, + inner_query, + } => UserInputLeaf::Spatial { + field, + predicate, + coordinates, + inner_query, + }, } } @@ -68,7 +107,8 @@ impl UserInputLeaf { UserInputLeaf::Range { field, .. } if field.is_none() => *field = Some(default_field), UserInputLeaf::Set { field, .. } if field.is_none() => *field = Some(default_field), UserInputLeaf::Regex { field, .. } if field.is_none() => *field = Some(default_field), - _ => (), // field was already set, do nothing + UserInputLeaf::Spatial { .. } => (), + _ => (), } } } @@ -118,6 +158,39 @@ impl Debug for UserInputLeaf { // TODO properly escape pattern (in case of \") write!(formatter, "/{pattern}/") } + UserInputLeaf::Spatial { + field, + predicate, + coordinates, + inner_query, + } => { + if let Some(field) = field { + write!(formatter, "\"{field}\":")?; + } + match predicate { + SpatialPredicateKind::Intersects => write!(formatter, "$intersects(")?, + SpatialPredicateKind::Contains => write!(formatter, "$contains(")?, + SpatialPredicateKind::CoveredBy => write!(formatter, "$within(")?, + SpatialPredicateKind::Within(radius) => { + write!(formatter, "$within({}rad, ", radius.0)? + } + SpatialPredicateKind::Between(inner, outer) => { + write!(formatter, "$between({}rad, {}rad, ", inner.0, outer.0)? + } + SpatialPredicateKind::Knn(k) => write!(formatter, "$knn({k}, ")?, + } + // A join has an inner query in place of coordinates. + if let Some(inner_query) = inner_query { + write!(formatter, "$query({inner_query})")?; + } + for (i, (lon, lat)) in coordinates.iter().enumerate() { + if i > 0 { + write!(formatter, ", ")?; + } + write!(formatter, "{} {}", lon.0, lat.0)?; + } + write!(formatter, ")") + } } } } diff --git a/src/core/json_utils.rs b/src/core/json_utils.rs index c2f2f17655..e2f0b9aac9 100644 --- a/src/core/json_utils.rs +++ b/src/core/json_utils.rs @@ -227,6 +227,8 @@ pub(crate) fn index_json_value<'a, V: Value<'a>>( ReferenceValueLeaf::IpAddr(_) => { unimplemented!("IP address support in dynamic fields is not yet implemented") } + // Only a spatial field indexes geometry. + ReferenceValueLeaf::Geometry(_) => {} ReferenceValueLeaf::Custom(_) => { unimplemented!("the JSON field does not support custom field types") } diff --git a/src/fastfield/mod.rs b/src/fastfield/mod.rs index dbaa68a52e..66d8aaaad2 100644 --- a/src/fastfield/mod.rs +++ b/src/fastfield/mod.rs @@ -718,7 +718,7 @@ mod tests { } #[test] - fn test_datefastfield() -> crate::Result<()> { + fn test_datefastfield() { let mut schema_builder = Schema::builder(); let date_field = schema_builder.add_date_field( "date", @@ -732,22 +732,28 @@ mod tests { ); let schema = schema_builder.build(); let index = Index::create_in_ram(schema); - let mut index_writer = index.writer_for_tests()?; + let mut index_writer = index.writer_for_tests().unwrap(); index_writer.set_merge_policy(Box::new(NoMergePolicy)); - index_writer.add_document(doc!( - date_field => DateTime::from_u64(1i64.to_u64()), - multi_date_field => DateTime::from_u64(2i64.to_u64()), - multi_date_field => DateTime::from_u64(3i64.to_u64()) - ))?; - index_writer.add_document(doc!( - date_field => DateTime::from_u64(4i64.to_u64()) - ))?; - index_writer.add_document(doc!( - multi_date_field => DateTime::from_u64(5i64.to_u64()), - multi_date_field => DateTime::from_u64(6i64.to_u64()) - ))?; - index_writer.commit()?; - let reader = index.reader()?; + index_writer + .add_document(doc!( + date_field => DateTime::from_u64(1i64.to_u64()), + multi_date_field => DateTime::from_u64(2i64.to_u64()), + multi_date_field => DateTime::from_u64(3i64.to_u64()) + )) + .unwrap(); + index_writer + .add_document(doc!( + date_field => DateTime::from_u64(4i64.to_u64()) + )) + .unwrap(); + index_writer + .add_document(doc!( + multi_date_field => DateTime::from_u64(5i64.to_u64()), + multi_date_field => DateTime::from_u64(6i64.to_u64()) + )) + .unwrap(); + index_writer.commit().unwrap(); + let reader = index.reader().unwrap(); let searcher = reader.searcher(); assert_eq!(searcher.segment_readers().len(), 1); let segment_reader = searcher.segment_reader(0); @@ -781,7 +787,6 @@ mod tests { assert_eq!(dates[0].into_timestamp_nanos(), 5i64); assert_eq!(dates[1].into_timestamp_nanos(), 6i64); } - Ok(()) } #[test] diff --git a/src/fastfield/writer.rs b/src/fastfield/writer.rs index e387d9541c..87ee37d2c7 100644 --- a/src/fastfield/writer.rs +++ b/src/fastfield/writer.rs @@ -211,6 +211,8 @@ impl FastFieldsWriter { .record_str(doc_id, field_name, &token.text); } } + // Only a spatial field indexes geometry, and it has no fast field. + ReferenceValueLeaf::Geometry(_) => {} // Custom fields are never fast (`is_fast() == false`), so they are not registered // in `fast_field_names` and `add_doc_value` returns early above. ReferenceValueLeaf::Custom(_) => { @@ -354,6 +356,8 @@ fn record_json_value_to_columnar_writer<'a, V: Value<'a>>( "Pre-tokenized string support in dynamic fields is not yet implemented" ) } + // Only a spatial field indexes geometry. + ReferenceValueLeaf::Geometry(_) => {} ReferenceValueLeaf::Custom(_) => { unimplemented!("the JSON field does not support custom field types") } diff --git a/src/index/index.rs b/src/index/index.rs index 2506747755..63620c1743 100644 --- a/src/index/index.rs +++ b/src/index/index.rs @@ -31,6 +31,8 @@ use crate::plugin::SegmentPlugin; use crate::reader::{IndexReader, IndexReaderBuilder}; use crate::schema::document::Document; use crate::schema::{Field, FieldType, Schema, Type}; +use crate::spatial::plugin::SpatialPlugin; +use crate::spatial::spatial_index_manager::SpatialIndexManager; use crate::store::StorePlugin; use crate::tokenizer::{TextAnalyzer, TokenizerManager}; use crate::SegmentReader; @@ -63,6 +65,7 @@ pub(crate) fn builtin_plugins() -> Vec> { Arc::new(InvertedIndexPlugin), Arc::new(FastFieldsPlugin), Arc::new(StorePlugin), + Arc::new(SpatialPlugin), ] } @@ -165,6 +168,7 @@ pub struct IndexBuilder { index_settings: IndexSettings, tokenizer_manager: TokenizerManager, fast_field_tokenizer_manager: TokenizerManager, + spatial_index_manager: SpatialIndexManager, custom_plugins: Vec>, } impl Default for IndexBuilder { @@ -180,6 +184,7 @@ impl IndexBuilder { index_settings: IndexSettings::default(), tokenizer_manager: TokenizerManager::default(), fast_field_tokenizer_manager: TokenizerManager::default(), + spatial_index_manager: SpatialIndexManager::default(), custom_plugins: Vec::new(), } } @@ -292,6 +297,7 @@ impl IndexBuilder { } let mut index = Index::open(dir)?; index.set_tokenizers(self.tokenizer_manager.clone()); + index.set_spatial_indices(self.spatial_index_manager.clone()); if index.schema() == self.get_expect_schema()? { index.custom_plugins.extend(self.custom_plugins); Ok(index) @@ -375,6 +381,7 @@ impl IndexBuilder { let mut index = Index::open_from_metas(directory, &metas, SegmentMetaInventory::default()); index.set_tokenizers(self.tokenizer_manager); index.set_fast_field_tokenizers(self.fast_field_tokenizer_manager); + index.set_spatial_indices(self.spatial_index_manager); index.custom_plugins.extend(self.custom_plugins); Ok(index) } @@ -389,6 +396,7 @@ pub struct Index { executor: Executor, tokenizers: TokenizerManager, fast_field_tokenizers: TokenizerManager, + spatial_indices: SpatialIndexManager, inventory: SegmentMetaInventory, custom_plugins: Vec>, #[cfg(feature = "jitexpr")] @@ -510,6 +518,7 @@ impl Index { schema, tokenizers: TokenizerManager::default(), fast_field_tokenizers: TokenizerManager::default(), + spatial_indices: SpatialIndexManager::default(), executor: Executor::single_thread(), inventory, custom_plugins: Vec::new(), @@ -554,6 +563,16 @@ impl Index { &self.fast_field_tokenizers } + /// Setter for the spatial index manager. + pub fn set_spatial_indices(&mut self, spatial_indices: SpatialIndexManager) { + self.spatial_indices = spatial_indices; + } + + /// Accessor for the spatial index manager. + pub fn spatial_indices(&self) -> &SpatialIndexManager { + &self.spatial_indices + } + /// Get the tokenizer associated with a specific field. pub fn tokenizer_for_field(&self, field: Field) -> crate::Result { let field_entry = self.schema.get_field_entry(field); diff --git a/src/index/index_meta.rs b/src/index/index_meta.rs index ad40f1c612..d72947f748 100644 --- a/src/index/index_meta.rs +++ b/src/index/index_meta.rs @@ -434,8 +434,17 @@ mod tests { fn expected_builtin_files(segment_id: SegmentId, delete_opstamp: u64) -> HashSet { let segment_uuid = segment_id.uuid_string(); - let mut files = HashSet::with_capacity(7); - for extension in ["idx", "pos", "term", "store", "fast", "fieldnorm"] { + let mut files = HashSet::with_capacity(9); + for extension in [ + "idx", + "pos", + "term", + "store", + "fast", + "fieldnorm", + "cells", + "edges", + ] { files.insert(PathBuf::from(format!("{segment_uuid}.{extension}"))); } files.insert(PathBuf::from(format!( diff --git a/src/index/inverted_index_plugin.rs b/src/index/inverted_index_plugin.rs index 19273c9260..7cb6bac2b8 100644 --- a/src/index/inverted_index_plugin.rs +++ b/src/index/inverted_index_plugin.rs @@ -392,7 +392,7 @@ impl InvertedIndexPluginWriter { } } // These fields are not indexed; the `is_indexed()` guard above skips them. - FieldType::TieBreaker | FieldType::Custom(_) => { + FieldType::TieBreaker | FieldType::Custom(_) | FieldType::Spatial(_) => { unreachable!("the inverted index does not support this field type") } } diff --git a/src/index/segment_reader.rs b/src/index/segment_reader.rs index 0aa0c4a96b..1926b0fe72 100644 --- a/src/index/segment_reader.rs +++ b/src/index/segment_reader.rs @@ -16,6 +16,7 @@ use crate::index::{Index, InvertedIndexReader, Segment, SegmentComponent, Segmen use crate::json_utils::json_path_sep_to_dot; use crate::schema::{Field, IndexRecordOption, Schema, Type}; use crate::space_usage::{ComponentSpaceUsage, SegmentSpaceUsage}; +use crate::spatial::reader::SpatialReaders; use crate::store::StoreReader; use crate::termdict::TermDictionary; use crate::{DocId, Opstamp}; @@ -46,6 +47,7 @@ pub struct SegmentReader { positions_composite: CompositeFile, fast_fields_readers: FastFieldReaders, fieldnorm_readers: FieldNormReaders, + spatial_readers: SpatialReaders, store_file: FileSlice, alive_bitset_opt: Option, @@ -100,6 +102,11 @@ impl SegmentReader { &self.fast_fields_readers } + /// HUSH + pub fn spatial_fields(&self) -> &SpatialReaders { + &self.spatial_readers + } + /// Accessor to the `FacetReader` associated with a given `Field`. pub fn facet_reader(&self, field_name: &str) -> crate::Result { let schema = self.schema(); @@ -179,6 +186,13 @@ impl SegmentReader { let fast_fields_readers = FastFieldReaders::open(fast_fields_data, schema.clone())?; let fieldnorm_data = segment.open_read(SegmentComponent::FieldNorms)?; let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?; + let spatial_readers = if schema.contains_spatial_field() { + let cells_data = segment.open_read(crate::spatial::plugin::cells_component())?; + let edges_data = segment.open_read(crate::spatial::plugin::edges_component())?; + SpatialReaders::open(cells_data, edges_data)? + } else { + SpatialReaders::empty() + }; let original_bitset = if segment.meta().has_deletes() { let alive_doc_file_slice = segment.open_read(SegmentComponent::Delete)?; @@ -205,6 +219,7 @@ impl SegmentReader { postings_composite, fast_fields_readers, fieldnorm_readers, + spatial_readers, segment_id: segment.id(), delete_opstamp: segment.meta().delete_opstamp(), store_file, diff --git a/src/indexer/merger.rs b/src/indexer/merger.rs index 4d0d9d2deb..bf0044303a 100644 --- a/src/indexer/merger.rs +++ b/src/indexer/merger.rs @@ -542,6 +542,8 @@ impl IndexMerger { #[cfg(test)] mod tests { + use std::collections::HashSet; + use columnar::Column; use proptest::prop_oneof; use proptest::strategy::Strategy; @@ -550,7 +552,7 @@ mod tests { use crate::collector::tests::{ BytesFastFieldTestCollector, FastFieldTestCollector, TEST_COLLECTOR_WITH_SCORE, }; - use crate::collector::{Count, FacetCollector}; + use crate::collector::{Count, DocSetCollector, FacetCollector}; use crate::index::{Index, SegmentId}; use crate::indexer::NoMergePolicy; use crate::query::{AllQuery, BooleanQuery, EnableScoring, Scorer, TermQuery}; @@ -564,6 +566,162 @@ mod tests { IndexSortByField, IndexWriter, Order, Searcher, }; + /// Commits each JSON document as its own segment, then merges the segments in commit order, + /// so that each document's id in the merged segment is its position in `docs`. + fn merge_one_segment_per_doc( + docs: &[&str], + ) -> crate::Result<(Index, schema::Field, schema::Field)> { + let mut schema_builder = schema::Schema::builder(); + let text = schema_builder.add_text_field("text", TEXT); + let geometry = schema_builder.add_spatial_field("geometry", schema::SPATIAL); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.set_merge_policy(Box::new(NoMergePolicy)); + let mut segment_ids: Vec = Vec::new(); + for json in docs { + writer.add_document(TantivyDocument::parse_json(&index.schema(), json)?)?; + writer.commit()?; + let committed = index.searchable_segment_ids()?; + segment_ids.extend(committed.into_iter().find(|id| !segment_ids.contains(id))); + } + writer.merge(&segment_ids).wait()?; + Ok((index, text, geometry)) + } + + #[test] + fn test_merge_spatial_segments_without_geometry() -> crate::Result<()> { + let (index, text, geometry) = merge_one_segment_per_doc(&[ + r#"{"text": "no place"}"#, + r#"{"text": "nowhere"}"#, + r#"{"text": "here", "geometry": {"type": "Point", "coordinates": [19.94, 50.06]}}"#, + ])?; + + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.segment_readers().len(), 1); + assert_eq!(searcher.num_docs(), 3); + // The point moved from doc 0 of its own segment to doc 2 of the merged one. + let krakow = + crate::query::SpatialQuery::from_bounds(geometry, [[19.9, 50.0], [20.0, 50.1]]); + let here = TermQuery::new( + Term::from_field_text(text, "here"), + IndexRecordOption::Basic, + ); + let found = searcher.search(&krakow, &DocSetCollector)?; + assert_eq!(found, searcher.search(&here, &DocSetCollector)?); + assert_eq!(found, HashSet::from([DocAddress::new(0, 2)])); + Ok(()) + } + + #[test] + fn test_merge_never_puts_a_line_around_a_cell_center() -> crate::Result<()> { + use crate::spatial::cell_index_reader::CellIndexReader; + use crate::spatial::edge_reader::EdgeReader; + use crate::spatial::sphere::Sphere; + + // A zigzag line among points from other segments, so that the merge subdivides the cells + // the line crosses. + let mut docs = vec![r#"{"text": "line", "geometry": {"type": "LineString", + "coordinates": [[0, 0], [1, 2], [2, 0], [3, 2], [4, 0], [5, 2], [6, 0]]}}"# + .to_string()]; + for i in 0..48 { + let (lon, lat) = (i as f64 / 8.0, (i % 5) as f64 / 2.5); + docs.push(format!( + r#"{{"text": "point", "geometry": {{"type": "Point", "coordinates": [{lon}, {lat}]}}}}"# + )); + } + let docs: Vec<&str> = docs.iter().map(String::as_str).collect(); + let (index, _, geometry) = merge_one_segment_per_doc(&docs)?; + + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.segment_readers().len(), 1); + let spatial = searcher + .segment_reader(0) + .spatial_fields() + .get_field(geometry)? + .unwrap(); + let cells = CellIndexReader::open(spatial.cells_bytes())?; + let edges = EdgeReader::::open(spatial.edges_bytes())?; + let mut line_cells = 0; + for cell in cells.iter() { + for shape in &cell.shapes { + let (_, set) = edges.read_geometry_set(shape.geometry_id.1); + if set.members[0].closed { + continue; + } + // A line or point is in a cell only where its edges are. + assert!(!shape.contains_center, "{:?}", cell.cell_id); + assert!(!shape.edge_indices.is_empty(), "{:?}", cell.cell_id); + line_cells += usize::from(set.members[0].vertices.len() > 1); + } + } + assert!(line_cells > 1); + Ok(()) + } + + #[test] + fn test_merge_never_puts_an_absorbed_line_around_a_cell_center() -> crate::Result<()> { + use crate::spatial::cell_index_reader::CellIndexReader; + use crate::spatial::edge_reader::EdgeReader; + use crate::spatial::sphere::Sphere; + + // Short lines inside the spread of a few points: the points' coarse cell absorbs the + // lines' fine cells when merged, and the cells are flattened into one. + let mut docs = Vec::new(); + for i in 0..24 { + let angle = i as f64 * 0.7; + let (lon, lat) = ((i % 5) as f64 - 2.0, (i % 7) as f64 * 0.6 - 1.8); + docs.push(format!( + r#"{{"text": "line", "geometry": {{"type": "LineString", + "coordinates": [[{lon}, {lat}], [{}, {}]]}}}}"#, + lon + angle.cos(), + lat + angle.sin() + )); + } + for (lon, lat) in [(-4.0, -4.0), (5.0, -4.5), (4.5, 5.0), (-4.5, 4.0)] { + docs.push(format!( + r#"{{"text": "point", "geometry": {{"type": "Point", "coordinates": [{lon}, {lat}]}}}}"# + )); + } + let docs: Vec<&str> = docs.iter().map(String::as_str).collect(); + let (index, _, geometry) = merge_one_segment_per_doc(&docs)?; + + let searcher = index.reader()?.searcher(); + let spatial = searcher + .segment_reader(0) + .spatial_fields() + .get_field(geometry)? + .unwrap(); + let cells = CellIndexReader::open(spatial.cells_bytes())?; + let edges = EdgeReader::::open(spatial.edges_bytes())?; + for cell in cells.iter() { + for shape in &cell.shapes { + let (_, set) = edges.read_geometry_set(shape.geometry_id.1); + assert!( + set.members[0].closed || !shape.contains_center, + "{:?}", + cell.cell_id + ); + } + } + Ok(()) + } + + #[test] + fn test_merge_spatial_field_never_written() -> crate::Result<()> { + let (index, _, geometry) = + merge_one_segment_per_doc(&[r#"{"text": "no place"}"#, r#"{"text": "nowhere"}"#])?; + + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.num_docs(), 2); + // Like the segments merged into it, the merged segment has no entry for the field. + let spatial_fields = searcher.segment_reader(0).spatial_fields(); + assert!(spatial_fields.get_field(geometry)?.is_none()); + let anywhere = + crate::query::SpatialQuery::from_bounds(geometry, [[-10.0, -10.0], [10.0, 10.0]]); + assert_eq!(searcher.search(&anywhere, &Count)?, 0); + Ok(()) + } + #[test] fn test_index_merger_no_deletes() -> crate::Result<()> { let mut schema_builder = schema::Schema::builder(); diff --git a/src/indexer/segment_writer.rs b/src/indexer/segment_writer.rs index dfb08806f2..43625c1bb9 100644 --- a/src/indexer/segment_writer.rs +++ b/src/indexer/segment_writer.rs @@ -5,6 +5,7 @@ use crate::index::{InvertedIndexPluginWriter, Segment}; use crate::plugin::{PluginWriter, PluginWriterContext}; use crate::schema::document::{Document, ErasedDocument}; use crate::schema::Schema; +use crate::spatial::plugin::SpatialPluginWriter; use crate::store::StorePluginWriter; use crate::{DocId, Opstamp, TantivyError}; @@ -33,6 +34,7 @@ pub struct SegmentWriter { pub(crate) inverted_index: InvertedIndexPluginWriter, pub(crate) fast_fields: FastFieldsPluginWriter, pub(crate) store_writer: StorePluginWriter, + pub(crate) spatial: SpatialPluginWriter, custom_plugins: Vec>, pub(crate) doc_opstamps: Vec, schema: Schema, @@ -57,6 +59,7 @@ impl SegmentWriter { let inverted_index = InvertedIndexPluginWriter::new(&ctx)?; let fast_fields = FastFieldsPluginWriter::new(&ctx)?; let store_writer = StorePluginWriter::new(&ctx)?; + let spatial = SpatialPluginWriter::new(&ctx); let custom_plugins = segment .index() .custom_plugins() @@ -69,6 +72,7 @@ impl SegmentWriter { inverted_index, fast_fields, store_writer, + spatial, custom_plugins, doc_opstamps: Vec::with_capacity(1_000), schema, @@ -144,11 +148,13 @@ impl SegmentWriter { } fn finalize_inner(self, mapping: Option<&DocIdMapping>) -> crate::Result> { - // Serialize in write order: inverted index, fast fields, store, then custom plugins. + // Serialize in write order: inverted index, fast fields, store, spatial, then custom + // plugins. let segment = &self.segment; Box::new(self.inverted_index).serialize(segment, mapping)?; Box::new(self.fast_fields).serialize(segment, mapping)?; Box::new(self.store_writer).serialize(segment, mapping)?; + Box::new(self.spatial).serialize(segment, mapping)?; for writer in self.custom_plugins { writer.serialize(segment, mapping)?; } @@ -163,6 +169,7 @@ impl SegmentWriter { self.inverted_index.mem_usage() + self.fast_fields.mem_usage() + self.store_writer.mem_usage() + + self.spatial.mem_usage() + self .custom_plugins .iter() @@ -185,6 +192,8 @@ impl SegmentWriter { self.inverted_index.index_document(doc_id, &document)?; self.fast_fields.index_document(&document)?; self.store_writer.store(&document, &self.schema)?; + self.spatial + .index_document(doc_id, &document, &self.schema)?; // Custom plugins are trait objects, so they take the document type-erased. Leaves stay // borrowed; a plugin that knows the concrete type recovers it for free via `as_any`. diff --git a/src/lib.rs b/src/lib.rs index 8ffcab1f6a..642c30395a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -197,6 +197,7 @@ pub mod fieldnorm; pub mod index; pub mod positions; pub mod postings; +pub mod spatial; pub mod plugin; /// Module containing the different query implementations. diff --git a/src/postings/per_field_postings_writer.rs b/src/postings/per_field_postings_writer.rs index 7147c356e4..32c7174b9e 100644 --- a/src/postings/per_field_postings_writer.rs +++ b/src/postings/per_field_postings_writer.rs @@ -51,6 +51,7 @@ fn posting_writer_from_field_entry(field_entry: &FieldEntry) -> Box Box::>::default(), FieldType::JsonObject(ref json_object_options) => { if let Some(text_indexing_option) = json_object_options.get_text_indexing_options() { diff --git a/src/query/boolean_query/boolean_query.rs b/src/query/boolean_query/boolean_query.rs index 8c09678175..1486bb7d67 100644 --- a/src/query/boolean_query/boolean_query.rs +++ b/src/query/boolean_query/boolean_query.rs @@ -257,6 +257,11 @@ impl BooleanQuery { pub fn clauses(&self) -> &[(Occur, Box)] { &self.subqueries[..] } + + /// The clauses making up this query, to change in place. + pub(crate) fn clauses_mut(&mut self) -> &mut [(Occur, Box)] { + &mut self.subqueries[..] + } } #[cfg(test)] diff --git a/src/query/boost_query.rs b/src/query/boost_query.rs index eec2d1e0b1..b7906b0f8b 100644 --- a/src/query/boost_query.rs +++ b/src/query/boost_query.rs @@ -20,6 +20,16 @@ impl BoostQuery { pub fn new(query: Box, boost: Score) -> BoostQuery { BoostQuery { query, boost } } + + /// The boosted query. + pub(crate) fn query(&self) -> &dyn Query { + self.query.as_ref() + } + + /// The boosted query, to change in place. + pub(crate) fn query_mut(&mut self) -> &mut dyn Query { + self.query.as_mut() + } } impl Clone for BoostQuery { diff --git a/src/query/mod.rs b/src/query/mod.rs index dd168e03b1..d1a4948528 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -27,6 +27,7 @@ mod reqopt_scorer; mod scorer; mod set_query; mod size_hint; +mod spatial_query; mod term_query; mod union; mod weight; @@ -50,6 +51,7 @@ pub use self::disjunction_max_query::DisjunctionMaxQuery; pub use self::empty_query::{EmptyQuery, EmptyScorer, EmptyWeight}; pub use self::exclude::{Exclude, ExclusionSet}; pub use self::exist_query::ExistsQuery; +pub(crate) use self::explanation::does_not_match; pub use self::explanation::Explanation; pub use self::fuzzy_query::{DfaWrapper, FuzzyTermQuery}; pub use self::intersection::{intersect_scorers, Intersection}; @@ -65,6 +67,7 @@ pub use self::reqopt_scorer::RequiredOptionalScorer; pub use self::score_combiner::{DisjunctionMaxCombiner, ScoreCombiner, SumCombiner}; pub use self::scorer::Scorer; pub use self::set_query::TermSetQuery; +pub use self::spatial_query::{SpatialPredicate, SpatialQuery}; pub use self::term_query::TermQuery; pub use self::union::BufferedUnionScorer; #[cfg(test)] diff --git a/src/query/query_parser/logical_ast.rs b/src/query/query_parser/logical_ast.rs index 1bcb864bd7..548954359f 100644 --- a/src/query/query_parser/logical_ast.rs +++ b/src/query/query_parser/logical_ast.rs @@ -4,8 +4,9 @@ use std::sync::Arc; use tantivy_fst::Regex; -use crate::query::Occur; +use crate::query::{Occur, SpatialPredicate}; use crate::schema::{Field, Term}; +use crate::spatial::executor::SpatialRelation; use crate::Score; #[derive(Clone)] @@ -28,11 +29,24 @@ pub enum LogicalLiteral { pattern: Arc, field: Field, }, + Spatial { + field: Field, + predicate: SpatialPredicate, + coordinates: Vec<[f64; 2]>, + }, + /// Matches the documents whose geometry relates to the geometry of a document `inner` + /// matches. + SpatialJoin { + field: Field, + relation: SpatialRelation, + inner: Box, + }, Exists { field: String, }, } +#[derive(Clone)] pub enum LogicalAst { Clause(Vec<(Occur, LogicalAst)>), Leaf(Box), @@ -162,6 +176,20 @@ impl fmt::Debug for LogicalLiteral { ref pattern, ref field, } => write!(formatter, "Regex({field:?}, {pattern:?})"), + LogicalLiteral::Spatial { + ref field, + ref predicate, + ref coordinates, + } => write!( + formatter, + "Spatial({field:?}, {predicate:?}, {} coords)", + coordinates.len() + ), + LogicalLiteral::SpatialJoin { + ref field, + ref relation, + ref inner, + } => write!(formatter, "SpatialJoin({field:?}, {relation:?}, {inner:?})"), LogicalLiteral::Exists { ref field } => write!(formatter, "Exists({field})"), } } diff --git a/src/query/query_parser/query_parser.rs b/src/query/query_parser/query_parser.rs index ba9f1a25b6..f46af66070 100644 --- a/src/query/query_parser/query_parser.rs +++ b/src/query/query_parser/query_parser.rs @@ -7,7 +7,9 @@ use std::sync::Arc; use base64::engine::general_purpose::STANDARD as BASE64; use base64::Engine; use itertools::Itertools; -use query_grammar::{UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral}; +use query_grammar::{ + SpatialPredicateKind, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral, +}; use rustc_hash::FxHashMap; use tantivy_fst::Regex; @@ -17,12 +19,14 @@ use crate::json_utils::convert_to_fast_value_and_append_to_json_term; use crate::query::range_query::{is_type_valid_for_fastfield_range_query, RangeQuery}; use crate::query::{ AllQuery, BooleanQuery, BoostQuery, EmptyQuery, ExistsQuery, FuzzyTermQuery, Occur, - PhrasePrefixQuery, PhraseQuery, Query, RegexQuery, TermQuery, TermSetQuery, + PhrasePrefixQuery, PhraseQuery, Query, RegexQuery, SpatialPredicate, SpatialQuery, TermQuery, + TermSetQuery, }; use crate::schema::{ Facet, FacetParseError, Field, FieldType, IndexRecordOption, IntoIpv6Addr, JsonObjectOptions, Schema, Term, TextFieldIndexing, Type, }; +use crate::spatial::executor::{PlanNode, SpatialExecutor, SpatialRelation}; use crate::time::format_description::well_known::Rfc3339; use crate::time::OffsetDateTime; use crate::tokenizer::{TextAnalyzer, TokenizerManager}; @@ -187,6 +191,24 @@ fn trim_ast(logical_ast: LogicalAst) -> Option { /// /// * all docs query: A plain `*` will match all documents in the index. /// +/// * spatial predicates on a spatial field, with lon/lat coordinates and distances in `mi`, +/// `km`, `m`, `ft` or `rad`: +/// - `geo:$intersects(0 0, 10 0, 10 10, 0 10)`: geometries that meet the polygon. +/// - `geo:$contains(0 0, 10 0, 10 10, 0 10)`: geometries that contain the polygon. +/// - `geo:$within(0 0, 10 0, 10 10, 0 10)`: geometries inside the polygon. +/// - `geo:$within(5km, 2.35 48.85)`: geometries within a distance of a point. +/// - `geo:$between(1km, 5km, 2.35 48.85)`: geometries within the outer distance of a point, +/// but not the inner one. +/// - `geo:$knn(10, 2.35 48.85)`: the 10 nearest geometries to a point, scored by rank. The +/// required and excluded clauses beside it narrow its candidates: `+type:cafe +/// +geo:$knn(10, ..)` finds the 10 nearest cafes. +/// +/// Polygon rings are closed if they are not, and bound the smaller region whichever way they +/// run, unless the field is oriented (see [`crate::schema::SpatialOptions::set_oriented`]). A +/// join replaces the coordinates with `$query(...)`, a query for the documents to relate to: +/// `geo:$within(1km, $query(type:station))` finds geometries within 1 km of a station's. +/// Joins take `$intersects`, `$contains`, `$within` with a distance and `$between`. +/// /// Parts of the queries can be boosted by appending `^boostfactor`. /// For instance, `"SRE"^2.0 OR devops^0.4` will boost documents containing `SRE` instead of /// devops. Negative boosts are not allowed. @@ -528,6 +550,9 @@ impl QueryParser { let ip_v6 = IpAddr::from_str(phrase)?.into_ipv6_addr(); Ok(Term::from_field_ip_addr(field, ip_v6)) } + FieldType::Spatial(_) => Err(QueryParserError::UnsupportedQuery( + "Spatial queries are not yet supported in text query parser".to_string(), + )), // Custom fields are neither indexed nor fast, so the guard above returns // `FieldNotIndexed` before this match. FieldType::Custom(_) => { @@ -635,6 +660,10 @@ impl QueryParser { let term = Term::from_field_ip_addr(field, ip_v6); Ok(vec![LogicalLiteral::Term(term)]) } + FieldType::Spatial(_) => Err(QueryParserError::UnsupportedQuery(format!( + "Spatial queries are not yet supported for field '{}'", + field_name + ))), // These fields are not indexed, so the `is_indexed()` guard above returns // `FieldNotIndexed` before this match. FieldType::TieBreaker | FieldType::Custom(_) => { @@ -876,6 +905,105 @@ impl QueryParser { Some(LogicalAst::Leaf(Box::new(LogicalLiteral::Exists { field }))), Vec::new(), ), + UserInputLeaf::Spatial { + field, + predicate, + coordinates, + inner_query, + } => { + let full_path = try_tuple!(field.ok_or_else(|| { + QueryParserError::UnsupportedQuery( + "Spatial query requires an explicit field name.".to_string(), + ) + })); + let (field, json_path) = try_tuple!(self + .split_full_path(&full_path) + .ok_or_else(|| QueryParserError::FieldDoesNotExist(full_path.clone()))); + if !json_path.is_empty() { + return ( + None, + vec![QueryParserError::UnsupportedQuery( + "Spatial query does not support json paths.".to_string(), + )], + ); + } + if !matches!( + self.schema.get_field_entry(field).field_type(), + FieldType::Spatial(_) + ) { + return ( + None, + vec![QueryParserError::UnsupportedQuery( + "Spatial predicates require a spatial field.".to_string(), + )], + ); + } + if let Some(inner_query) = inner_query { + let relation = match predicate { + SpatialPredicateKind::Intersects => SpatialRelation::Intersects, + SpatialPredicateKind::Contains => SpatialRelation::Contains, + SpatialPredicateKind::Within(r) => SpatialRelation::Near(r.0), + SpatialPredicateKind::Between(inner, outer) => { + SpatialRelation::Between(inner.0, outer.0) + } + SpatialPredicateKind::Knn(_) | SpatialPredicateKind::CoveredBy => { + return ( + None, + vec![QueryParserError::UnsupportedQuery(format!( + "A {predicate:?} predicate cannot join a $query." + ))], + ); + } + }; + // The inner query's errors are the query's errors. The lenient grammar + // accepts some queries the strict one rejects, so check with both. + let (inner, errors) = match query_grammar::parse_query(&inner_query) { + Ok(user_input_ast) => self.compute_logical_ast_lenient(user_input_ast), + Err(_) => { + let (inner, mut errors) = + self.parse_query_to_logical_ast_lenient(&inner_query); + errors.insert(0, QueryParserError::SyntaxError(inner_query)); + (inner, errors) + } + }; + let logical_ast = LogicalAst::Leaf(Box::new(LogicalLiteral::SpatialJoin { + field, + relation, + inner: Box::new(inner.simplify()), + })); + return (Some(logical_ast), errors); + } + let (predicate, min_coordinates) = match predicate { + SpatialPredicateKind::Intersects => (SpatialPredicate::Intersects, 3), + // `$contains` reads from the field's side, as a join's does: the field's + // geometry contains the polygon. `$within` a polygon is the other way around. + SpatialPredicateKind::Contains => (SpatialPredicate::Contains, 3), + SpatialPredicateKind::CoveredBy => (SpatialPredicate::CoveredBy, 3), + SpatialPredicateKind::Within(r) => (SpatialPredicate::Within(r.0), 1), + SpatialPredicateKind::Between(inner, outer) => { + (SpatialPredicate::Between(inner.0, outer.0), 1) + } + SpatialPredicateKind::Knn(k) => (SpatialPredicate::Knn(k), 1), + }; + if coordinates.len() < min_coordinates { + return ( + None, + vec![QueryParserError::UnsupportedQuery(format!( + "Spatial predicate {predicate:?} needs at least {min_coordinates} \ + coordinates." + ))], + ); + } + let logical_ast = LogicalAst::Leaf(Box::new(LogicalLiteral::Spatial { + field, + predicate, + coordinates: coordinates + .into_iter() + .map(|(lon, lat)| [lon.0, lat.0]) + .collect(), + })); + (Some(logical_ast), Vec::new()) + } UserInputLeaf::Regex { field, pattern } => { if !self.regexes_allowed { return ( @@ -966,6 +1094,27 @@ fn convert_literal_to_query( LogicalLiteral::Regex { pattern, field } => { Box::new(RegexQuery::from_regex(pattern, field)) } + // A kNN query is a plan of its own, so that its siblings can narrow its candidates. + LogicalLiteral::Spatial { + field, + predicate: SpatialPredicate::Knn(k), + coordinates, + } => Box::new(SpatialExecutor::knn(field, coordinates[0], k)), + LogicalLiteral::Spatial { + field, + predicate, + coordinates, + } => Box::new(SpatialQuery::with_predicate(field, coordinates, predicate)), + LogicalLiteral::SpatialJoin { + field, + relation, + inner, + } => Box::new(SpatialExecutor::new(PlanNode::Join { + field, + outer: Box::new(PlanNode::Query(Box::new(AllQuery))), + inner: Box::new(PlanNode::Query(convert_to_query(fuzzy, *inner))), + relation, + })), LogicalLiteral::Exists { field } => Box::new(ExistsQuery::new(field, true)), } } @@ -1078,10 +1227,138 @@ fn generate_literals_for_json_object( Ok(logical_literals) } +/// A spatial join probes the geometry of every one of its outer documents, so it pays to narrow +/// them down. Every document a clause matches satisfies its required siblings and none of its +/// excluded ones, so restricting the outer documents of the clause's joins to those leaves the +/// clause's matches and scores unchanged, whatever the occur of the join and of its siblings. That +/// holds as well for a join that is boosted, or nested in a sibling boolean query, which the +/// clause's own restriction then adds to. +/// +/// A kNN query is restricted the same way, but to change its matches: `+name:cafe +/// +geometry:$knn(5, ..)` asks for the 5 nearest cafes, not for the cafes among the 5 nearest +/// documents, however the kNN query is boosted or nested. Siblings that hold a join or a kNN +/// query of their own do not restrict. +fn restrict_join_outers(occur_subqueries: &mut [(Occur, Box)]) { + if !occur_subqueries + .iter() + .any(|(_, query)| holds_spatial_plan(query.as_ref())) + { + return; + } + // With no required clause, a document must match one of the optional ones. + let has_must = occur_subqueries + .iter() + .any(|(occur, _)| *occur == Occur::Must); + for i in 0..occur_subqueries.len() { + if !holds_spatial_plan(occur_subqueries[i].1.as_ref()) { + continue; + } + let mut restriction: Vec<(Occur, Box)> = Vec::new(); + let siblings = occur_subqueries + .iter() + .enumerate() + .filter(|(j, _)| *j != i) + .map(|(_, sibling)| sibling); + for (occur, query) in siblings.clone() { + collect_restriction(*occur, query.as_ref(), &mut restriction); + } + let optional: Vec<&Box> = siblings + .filter(|(occur, _)| *occur == Occur::Should) + .map(|(_, query)| query) + .collect(); + let plan_is_optional = occur_subqueries[i].0 == Occur::Should; + if !has_must + && !plan_is_optional + && !optional.is_empty() + && !optional + .iter() + .any(|query| holds_spatial_plan(query.as_ref())) + { + let any_of = optional + .iter() + .map(|query| (Occur::Should, query.box_clone())) + .collect(); + restriction.push((Occur::Must, Box::new(BooleanQuery::new(any_of)))); + } + if restriction.is_empty() { + continue; + } + if !restriction.iter().any(|(occur, _)| *occur == Occur::Must) { + // A boolean query of exclusions alone matches nothing. + restriction.push((Occur::Must, Box::new(AllQuery))); + } + let outer: Box = if restriction.len() == 1 { + restriction.pop().unwrap().1 + } else { + Box::new(BooleanQuery::new(restriction)) + }; + restrict_spatial_plans(occur_subqueries[i].1.as_mut(), outer.as_ref()); + } +} + +/// Collect what every document a clause matches satisfies, through a sibling with the given +/// occur: the sibling itself, required or excluded, if it holds no spatial plan of its own, and +/// otherwise the required and excluded clauses of a required group, however nested or boosted. +fn collect_restriction(occur: Occur, query: &dyn Query, out: &mut Vec<(Occur, Box)>) { + if occur == Occur::Should { + return; + } + if !holds_spatial_plan(query) { + out.push((occur, query.box_clone())); + return; + } + if occur != Occur::Must { + return; + } + if let Some(boost) = query.downcast_ref::() { + collect_restriction(occur, boost.query(), out); + } else if let Some(boolean) = query.downcast_ref::() { + let has_must = boolean + .clauses() + .iter() + .any(|(occur, _)| *occur == Occur::Must); + for (occur, clause) in boolean.clauses() { + if *occur == Occur::MustNot || has_must { + collect_restriction(*occur, clause.as_ref(), out); + } + } + } +} + +/// Whether the query is a join or kNN query, or holds one in a boost or a boolean query. +fn holds_spatial_plan(query: &dyn Query) -> bool { + if query.is::() { + return true; + } + if let Some(boost) = query.downcast_ref::() { + return holds_spatial_plan(boost.query()); + } + if let Some(boolean) = query.downcast_ref::() { + return boolean + .clauses() + .iter() + .any(|(_, clause)| holds_spatial_plan(clause.as_ref())); + } + false +} + +/// Restrict every join and kNN query the query is or holds. +fn restrict_spatial_plans(query: &mut dyn Query, restriction: &dyn Query) { + if let Some(executor) = query.downcast_mut::() { + executor.restrict(restriction.box_clone()); + } else if let Some(boost) = query.downcast_mut::() { + restrict_spatial_plans(boost.query_mut(), restriction); + } else if let Some(boolean) = query.downcast_mut::() { + for (_, clause) in boolean.clauses_mut() { + restrict_spatial_plans(clause.as_mut(), restriction); + } + } +} + fn convert_to_query(fuzzy: &FxHashMap, logical_ast: LogicalAst) -> Box { match trim_ast(logical_ast) { Some(LogicalAst::Clause(trimmed_clause)) => { - let occur_subqueries = trimmed_clause + let mut occur_subqueries = trimmed_clause .into_iter() .map(|(occur, subquery)| (occur, convert_to_query(fuzzy, subquery))) .collect::>(); @@ -1089,6 +1366,7 @@ fn convert_to_query(fuzzy: &FxHashMap, logical_ast: LogicalAst) -> !occur_subqueries.is_empty(), "Should not be empty after trimming" ); + restrict_join_outers(&mut occur_subqueries); Box::new(BooleanQuery::new(occur_subqueries)) } Some(LogicalAst::Leaf(trimmed_logical_literal)) => { diff --git a/src/query/range_query/mod.rs b/src/query/range_query/mod.rs index ce625d8381..7845ce9ff2 100644 --- a/src/query/range_query/mod.rs +++ b/src/query/range_query/mod.rs @@ -21,7 +21,7 @@ pub(crate) fn is_type_valid_for_fastfield_range_query(typ: Type) -> bool { | Type::Json | Type::IpAddr | Type::Bytes => true, - Type::Facet => false, + Type::Facet | Type::Spatial => false, Type::Custom => false, } } diff --git a/src/query/range_query/range_query_fastfield.rs b/src/query/range_query/range_query_fastfield.rs index 23944d23b4..98154d1fc6 100644 --- a/src/query/range_query/range_query_fastfield.rs +++ b/src/query/range_query/range_query_fastfield.rs @@ -133,6 +133,7 @@ impl Weight for FastFieldRangeWeight { | Type::Bytes | Type::Json | Type::IpAddr + | Type::Spatial | Type::Custom => Err(crate::TantivyError::InvalidArgument(format!( "unsupported value bytes type in json term value_bytes {:?}", term_value.typ() @@ -459,7 +460,7 @@ pub(crate) fn maps_to_u64_fastfield(typ: Type) -> bool { match typ { Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date => true, Type::IpAddr => false, - Type::Str | Type::Facet | Type::Bytes | Type::Json | Type::Custom => false, + Type::Str | Type::Facet | Type::Bytes | Type::Json | Type::Spatial | Type::Custom => false, } } diff --git a/src/query/spatial_query.rs b/src/query/spatial_query.rs new file mode 100644 index 0000000000..3b2eda0651 --- /dev/null +++ b/src/query/spatial_query.rs @@ -0,0 +1,2598 @@ +//! Spatial polygon query. +//! +//! Finds indexed geometries that match a query polygon via contains or intersects predicates. The +//! query polygon is specified as lon/lat vertices, converted to unit sphere coordinates, and +//! searched against each segment's cell index and edge index. + +use common::BitSet; + +use crate::query::explanation::does_not_match; +use crate::query::{BitSetDocSet, Explanation, Query, Scorer, Weight}; +use crate::schema::{Field, FieldType}; +use crate::spatial::cell_index_reader::CellIndexReader; +use crate::spatial::closest_edge_query::ClosestEdgeQuery; +use crate::spatial::contains_query::{ContainingQuery, ContainsQuery}; +use crate::spatial::edge_cache::EdgeCache; +use crate::spatial::edge_reader::EdgeReader; +use crate::spatial::executor::{SpatialExecutor, QUERY_EDGE_CACHE_BYTES}; +use crate::spatial::geometry::Geometry; +use crate::spatial::geometry_set::{ + to_geometry_set, to_geometry_set_with_orientation, RingOrientation, +}; +use crate::spatial::intersects_query::IntersectsQuery; +use crate::spatial::plane::Plane; +use crate::spatial::rect_query::{RectQuery, RectRelation}; +use crate::spatial::region_coverer::CovererOptions; +use crate::spatial::s1chord_angle::S1ChordAngle; +use crate::spatial::sphere::Sphere; +use crate::spatial::surface::Surface; +use crate::spatial::within_query::WithinQuery; +use crate::{DocId, DocSet, Score, TERMINATED}; + +/// The spatial predicate to apply. +#[derive(Clone, Debug)] +pub enum SpatialPredicate { + /// Return geometries inside the query polygon, its boundary included (ST_CoveredBy). + CoveredBy, + /// Return geometries that contain the query polygon. + Contains, + /// Return geometries that intersect the query polygon. + Intersects, + /// Return geometries within a distance of a point. Distance in radians. + Within(f64), + /// Return geometries between two distances of a point. Distances in radians. + Between(f64, f64), + /// Return the K nearest geometries to a point, across all segments. A hit scores by its + /// rank by distance, from 1 for the nearest down to 1/K, so that the nearest score highest + /// and a kNN clause lifts the nearest documents of a boolean query. Without a searcher, as + /// when deleting by query, each segment keeps its own K nearest. + Knn(usize), +} + +/// Spatial query. +#[derive(Clone, Debug)] +pub struct SpatialQuery { + field: Field, + predicate: SpatialPredicate, + /// Polygon vertices (lon/lat) for Contains and Intersects, a single point for distance, or + /// the south-west and north-east corners of a rectangle. + coordinates: Vec<[f64; 2]>, + /// How the polygon's ring picks the region it bounds, if not as the field's rings do. + orientation: Option, + /// Whether the coordinates are the corners of a rectangle of longitude and latitude. + rectangle: bool, +} + +impl SpatialQuery { + /// Create a query for the geometries inside a polygon specified as lon/lat vertices. + pub fn new(field: Field, polygon: Vec<[f64; 2]>) -> Self { + SpatialQuery { + field, + coordinates: polygon, + predicate: SpatialPredicate::CoveredBy, + orientation: None, + rectangle: false, + } + } + + /// Create a spatial query with the given predicate. + pub fn with_predicate( + field: Field, + coordinates: Vec<[f64; 2]>, + predicate: SpatialPredicate, + ) -> Self { + SpatialQuery { + field, + coordinates, + predicate, + orientation: None, + rectangle: false, + } + } + + /// Create a query for the geometries within a rectangle of longitude and latitude, its + /// boundary included. The bounds are the rectangle's west and south, then its east and + /// north, as a GeoJSON bounding box has them: its sides are the meridians and the parallels + /// through them. A west greater than the east is a rectangle across the antimeridian, and a + /// west of -180 with an east of 180 is every longitude. + pub fn covered_by_rect(field: Field, bounds: [[f64; 2]; 2]) -> Self { + SpatialQuery { + field, + coordinates: bounds.to_vec(), + predicate: SpatialPredicate::CoveredBy, + orientation: None, + rectangle: true, + } + } + + /// Create a query for the geometries that have a point in common with a rectangle of + /// longitude and latitude. See [`SpatialQuery::covered_by_rect`] for its bounds. + pub fn intersects_rect(field: Field, bounds: [[f64; 2]; 2]) -> Self { + SpatialQuery { + field, + coordinates: bounds.to_vec(), + predicate: SpatialPredicate::Intersects, + orientation: None, + rectangle: true, + } + } + + /// Create a spatial query from a bounding box. The box is converted to a 4-vertex polygon, + /// whose edges are geodesics: its north and south edges bow toward the nearer pole, away + /// from the parallels through its corners, by more the wider the box is, and an edge between + /// corners more than 180 degrees apart runs the short way around. For the rectangle between + /// those parallels, of any width, see [`SpatialQuery::covered_by_rect`]. + pub fn from_bounds(field: Field, bounds: [[f64; 2]; 2]) -> Self { + let [lo, hi] = bounds; + let polygon = vec![ + [lo[0], lo[1]], + [hi[0], lo[1]], + [hi[0], hi[1]], + [lo[0], hi[1]], + ]; + SpatialQuery { + field, + coordinates: polygon, + predicate: SpatialPredicate::CoveredBy, + orientation: None, + rectangle: false, + } + } + + /// Create an intersects query from a bounding box, converted to a 4-vertex polygon as + /// [`SpatialQuery::from_bounds`] has it. For the rectangle between the parallels through its + /// corners, see [`SpatialQuery::intersects_rect`]. + pub fn intersects_bounds(field: Field, bounds: [[f64; 2]; 2]) -> Self { + let [lo, hi] = bounds; + let polygon = vec![ + [lo[0], lo[1]], + [hi[0], lo[1]], + [hi[0], hi[1]], + [lo[0], hi[1]], + ]; + SpatialQuery { + field, + coordinates: polygon, + predicate: SpatialPredicate::Intersects, + orientation: None, + rectangle: false, + } + } + + /// Have the query polygon's ring pick the region it bounds as given, rather than as the + /// field's rings do. With [`RingOrientation::RightHand`], a query polygon may be larger than + /// a hemisphere on any field. + #[must_use] + pub fn with_ring_orientation(mut self, orientation: RingOrientation) -> Self { + self.orientation = Some(orientation); + self + } +} + +impl SpatialQuery { + /// The weight of a rectangle query, once its corners are checked. + fn rectangle_weight(&self) -> crate::Result> { + let invalid = |reason: &str| { + crate::TantivyError::InvalidArgument(format!( + "Invalid rectangle {:?}: {reason}", + self.coordinates + )) + }; + let &[[west, south], [east, north]] = self.coordinates.as_slice() else { + return Err(invalid("a rectangle has two corners")); + }; + if ![west, east] + .iter() + .all(|lon| (-180.0..=180.0).contains(lon)) + || ![south, north] + .iter() + .all(|lat| (-90.0..=90.0).contains(lat)) + { + return Err(invalid( + "longitude must be within [-180, 180] and latitude within [-90, 90]", + )); + } + if south > north { + return Err(invalid("its south is above its north")); + } + let relation = match self.predicate { + SpatialPredicate::CoveredBy => RectRelation::CoveredBy, + SpatialPredicate::Intersects => RectRelation::Intersects, + _ => return Err(invalid("a rectangle covers or intersects")), + }; + Ok(Box::new(SpatialWeight { + field: self.field, + query: Box::new(RectQuery::new([west, south], [east, north], relation)), + })) + } +} + +impl Query for SpatialQuery { + fn weight( + &self, + enable_scoring: super::EnableScoring<'_>, + ) -> crate::Result> { + if self.rectangle { + return self.rectangle_weight(); + } + let min_coordinates = match self.predicate { + SpatialPredicate::CoveredBy + | SpatialPredicate::Contains + | SpatialPredicate::Intersects => 3, + _ => 1, + }; + if self.coordinates.len() < min_coordinates { + return Err(crate::TantivyError::InvalidArgument(format!( + "Spatial predicate {:?} needs at least {min_coordinates} coordinates", + self.predicate + ))); + } + // GeoJSON values are checked as they are parsed; query coordinates are checked here. + if let Some([lon, lat]) = self + .coordinates + .iter() + .find(|[lon, lat]| !(-180.0..=180.0).contains(lon) || !(-90.0..=90.0).contains(lat)) + { + return Err(crate::TantivyError::InvalidArgument(format!( + "Invalid coordinate {lon} {lat}: longitude must be within [-180, 180] and \ + latitude within [-90, 90]" + ))); + } + let distances = match self.predicate { + SpatialPredicate::Within(radius) => vec![radius], + SpatialPredicate::Between(inner, outer) => vec![inner, outer], + _ => Vec::new(), + }; + if distances.iter().any(|distance| distance.is_nan()) { + return Err(crate::TantivyError::InvalidArgument(format!( + "Spatial predicate {:?} needs distances that are numbers", + self.predicate + ))); + } + let prepared: Box = match &self.predicate { + SpatialPredicate::CoveredBy + | SpatialPredicate::Contains + | SpatialPredicate::Intersects => { + let mut ring: Vec<[f64; 2]> = self.coordinates.clone(); + if ring.first() != ring.last() { + ring.push(ring[0]); + } + let orientation = self.orientation.unwrap_or_else(|| { + // A field the schema does not have matches nothing, however it is read. + let schema = enable_scoring.schema(); + if self.field.field_id() as usize >= schema.num_fields() { + return RingOrientation::default(); + } + match schema.get_field_entry(self.field).field_type() { + FieldType::Spatial(options) => options.ring_orientation(), + _ => RingOrientation::default(), + } + }); + let plane_geometry = Geometry::::Polygon(vec![ring]); + let projected = plane_geometry.project::(); + let set = to_geometry_set_with_orientation(&projected, 0, orientation); + // Smashing keeps a polygon with no area as the line or point it is, so that the + // bounding box of a point queries the point, as it would be indexed. + match &self.predicate { + SpatialPredicate::CoveredBy => { + Box::new(ContainsQuery::new(set, CovererOptions::default())) + } + SpatialPredicate::Contains => { + Box::new(ContainingQuery::new(set, CovererOptions::default())?) + } + SpatialPredicate::Intersects => { + Box::new(IntersectsQuery::new(set, CovererOptions::default())) + } + _ => unreachable!(), + } + } + SpatialPredicate::Within(radius_radians) => { + let [lon, lat] = self.coordinates[0]; + let radius = S1ChordAngle::from_radians(*radius_radians); + Box::new(WithinQuery::new(Sphere::project(lon, lat), radius)) + } + SpatialPredicate::Between(inner_radians, outer_radians) => { + let query_geometry = Geometry::::Point(self.coordinates[0]); + let projected = query_geometry.project::(); + let set = to_geometry_set(&projected, 0); + let inner = S1ChordAngle::from_radians(*inner_radians); + let outer = S1ChordAngle::from_radians(*outer_radians); + Box::new(ClosestEdgeQuery::between(set, inner, outer)) + } + SpatialPredicate::Knn(k) => { + // The nearest K are ranked across segments, so the executor runs the query on + // every segment, keeps the global K, and scores each by its distance. + return SpatialExecutor::knn(self.field, self.coordinates[0], *k) + .weight(enable_scoring); + } + }; + Ok(Box::new(SpatialWeight { + field: self.field, + query: prepared, + })) + } +} + +/// Shared interface for prepared spatial queries. +trait PreparedSpatialQuery: Send + Sync { + /// Add the documents of one segment that match to `docs`. + fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ); +} + +/// Prepared queries that list the documents of a segment that match. +macro_rules! prepared_from_doc_ids { + ($($query:ty),*) => {$( + impl PreparedSpatialQuery for $query { + fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + for doc_id in <$query>::search_segment(self, cell_reader, edge_cache) { + docs.insert(doc_id); + } + } + } + )*}; +} + +prepared_from_doc_ids!(ContainsQuery, ContainingQuery, IntersectsQuery); + +impl PreparedSpatialQuery for ClosestEdgeQuery { + fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + for result in ClosestEdgeQuery::search_segment(self, cell_reader, edge_cache) { + docs.insert(result.doc_id); + } + } +} + +impl PreparedSpatialQuery for WithinQuery { + fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + WithinQuery::search_segment(self, cell_reader, edge_cache, docs) + } +} + +impl PreparedSpatialQuery for RectQuery { + fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + RectQuery::search_segment(self, cell_reader, edge_cache, docs) + } +} + +struct SpatialWeight { + field: Field, + query: Box, +} + +impl Weight for SpatialWeight { + fn scorer( + &self, + reader: &crate::SegmentReader, + boost: crate::Score, + ) -> crate::Result> { + let spatial_reader = match reader.spatial_fields().get_field(self.field)? { + Some(reader) => reader, + None => { + let empty_bitset = BitSet::with_max_value(reader.max_doc()); + return Ok(Box::new(SpatialScorer::new(boost, empty_bitset))); + } + }; + + let cell_reader = CellIndexReader::open(spatial_reader.cells_bytes())?; + let edge_reader = EdgeReader::::open(spatial_reader.edges_bytes())?; + let mut edge_cache = EdgeCache::new(vec![edge_reader], QUERY_EDGE_CACHE_BYTES); + + let mut include = BitSet::with_max_value(reader.max_doc()); + self.query + .search_segment(&cell_reader, &mut edge_cache, &mut include); + + Ok(Box::new(SpatialScorer::new(boost, include))) + } + + fn explain( + &self, + reader: &crate::SegmentReader, + doc: DocId, + ) -> crate::Result { + let mut scorer = self.scorer(reader, 1.0)?; + // Seeking backward is not allowed: a scorer may start past the document. + if scorer.doc() > doc || scorer.seek(doc) != doc { + return Err(does_not_match(doc)); + } + let score = scorer.score(); + let explanation = Explanation::new("SpatialQuery", score); + Ok(explanation) + } +} + +struct SpatialScorer { + include: BitSetDocSet, + doc_id: DocId, + score: Score, +} + +impl SpatialScorer { + pub fn new(score: Score, include: BitSet) -> Self { + let mut scorer = SpatialScorer { + include: BitSetDocSet::from(include), + doc_id: 0, + score, + }; + scorer.doc_id = scorer.include.doc(); + scorer + } +} + +impl Scorer for SpatialScorer { + fn score(&mut self) -> Score { + self.score + } +} + +impl DocSet for SpatialScorer { + fn advance(&mut self) -> DocId { + if self.doc_id == TERMINATED { + return TERMINATED; + } + self.doc_id = self.include.advance(); + self.doc_id + } + + fn seek(&mut self, target: DocId) -> DocId { + if self.doc_id >= target { + return self.doc_id; + } + self.doc_id = self.include.seek(target); + self.doc_id + } + + fn size_hint(&self) -> u32 { + self.include.size_hint() + } + + fn doc(&self) -> DocId { + self.doc_id + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + use crate::collector::{DocSetCollector, TopDocs}; + use crate::query::{AllQuery, BoostQuery, TermQuery}; + use crate::schema::{IndexRecordOption, Schema, Value, SPATIAL, STORED, STRING}; + use crate::spatial::containment::{brute_force_contains, compute_origin_inside}; + use crate::spatial::executor::{PlanNode, SpatialRelation}; + use crate::spatial::geometry_set::{EdgeSet, GeometrySet}; + use crate::spatial::s2edge_distances::{update_edge_pair_min_distance, update_min_distance}; + use crate::{DocAddress, Index, IndexWriter, Searcher, TantivyDocument, Term}; + + /// Indexes each inner slice of `(name, GeoJSON geometry)` pairs as one segment. + fn index_segments(segments: &[&[(&str, &str)]]) -> crate::Result<(Index, Field, Field)> { + let mut schema_builder = Schema::builder(); + let name = schema_builder.add_text_field("name", STRING | STORED); + let geometry = schema_builder.add_spatial_field("geometry", SPATIAL); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.set_merge_policy(Box::new(crate::indexer::NoMergePolicy)); + for docs in segments { + for (doc_name, doc_geometry) in *docs { + let json = format!(r#"{{"name": "{doc_name}", "geometry": {doc_geometry}}}"#); + writer.add_document(TantivyDocument::parse_json(&index.schema(), &json)?)?; + } + writer.commit()?; + } + Ok((index, name, geometry)) + } + + fn name_of(searcher: &Searcher, name: Field, doc_address: DocAddress) -> crate::Result { + let doc: TantivyDocument = searcher.doc(doc_address)?; + Ok(doc + .get_first(name) + .and_then(|v| v.as_str()) + .unwrap() + .to_string()) + } + + /// Names of the documents the query matches, sorted. + fn matching(index: &Index, name: Field, query: &dyn Query) -> crate::Result> { + let searcher = index.reader()?.searcher(); + let mut names = searcher + .search(query, &DocSetCollector)? + .into_iter() + .map(|doc_address| name_of(&searcher, name, doc_address)) + .collect::>>()?; + names.sort(); + Ok(names) + } + + fn name_query(name: Field, value: &str) -> Box { + Box::new(TermQuery::new( + Term::from_field_text(name, value), + IndexRecordOption::Basic, + )) + } + + const BOX: &str = + r#"{"type": "Polygon", "coordinates": [[[0, 0], [10, 0], [10, 10], [0, 10], [0, 0]]]}"#; + + #[test] + fn test_within_matches_a_polygon_containing_the_point() -> crate::Result<()> { + let (index, name, geometry) = index_segments(&[&[("box", BOX)]])?; + let query = SpatialQuery::with_predicate( + geometry, + vec![[5.0, 5.0]], + SpatialPredicate::Within(1e-6), + ); + assert_eq!(matching(&index, name, &query)?, ["box"]); + Ok(()) + } + + #[test] + fn test_within_matches_a_line_near_the_point() -> crate::Result<()> { + let line = r#"{"type": "LineString", "coordinates": [[0, 0], [10, 0]]}"#; + let (index, name, geometry) = index_segments(&[&[("equator", line)]])?; + let near = SpatialQuery::with_predicate( + geometry, + vec![[5.0, 0.001]], + SpatialPredicate::Within(0.01f64.to_radians()), + ); + assert_eq!(matching(&index, name, &near)?, ["equator"]); + let far = SpatialQuery::with_predicate( + geometry, + vec![[5.0, 1.0]], + SpatialPredicate::Within(0.01f64.to_radians()), + ); + assert!(matching(&index, name, &far)?.is_empty()); + Ok(()) + } + + #[test] + fn test_knn_returns_the_nearest_across_segments() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (a, b, c, d) = (point(1.0), point(2.0), point(3.0), point(4.0)); + let (index, name, geometry) = + index_segments(&[&[("a", &a), ("c", &c)], &[("b", &b), ("d", &d)]])?; + let searcher = index.reader()?.searcher(); + let knn = + SpatialQuery::with_predicate(geometry, vec![[0.0, 0.0]], SpatialPredicate::Knn(3)); + let hits = searcher.search(&knn, &TopDocs::with_limit(10).order_by_score())?; + let names = hits + .iter() + .map(|&(_, doc_address)| name_of(&searcher, name, doc_address)) + .collect::>>()?; + assert_eq!(names, ["a", "b", "c"]); + + // Explaining a hit gives its score; explaining the fourth nearest fails. + let (score, nearest) = hits[0]; + assert_eq!(knn.explain(&searcher, nearest)?.value(), score); + let d = searcher.search(&name_query(name, "d"), &DocSetCollector)?; + assert!(knn.explain(&searcher, *d.iter().next().unwrap()).is_err()); + + let none = + SpatialQuery::with_predicate(geometry, vec![[0.0, 0.0]], SpatialPredicate::Knn(0)); + assert!(matching(&index, name, &none)?.is_empty()); + Ok(()) + } + + #[test] + fn test_knn_skips_deleted_documents_before_ranking() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (deleted, near, middle, far) = (point(1.0), point(2.0), point(3.0), point(4.0)); + let (index, name, geometry) = index_segments(&[ + &[("deleted", &deleted), ("far", &far)], + &[("near", &near), ("middle", &middle)], + ])?; + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.set_merge_policy(Box::new(crate::indexer::NoMergePolicy)); + writer.delete_term(Term::from_field_text(name, "deleted")); + writer.commit()?; + + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.segment_readers().len(), 2); + assert_eq!(searcher.num_docs(), 3); + for k in 1..=4 { + let unfiltered = + SpatialQuery::with_predicate(geometry, vec![[0.0, 0.0]], SpatialPredicate::Knn(k)); + let filtered = SpatialExecutor::new(PlanNode::Knn { + field: geometry, + geometry: geometry_set(r#"{"type": "Point", "coordinates": [0, 0]}"#), + k, + filter: Some(Box::new(PlanNode::Query(Box::new(AllQuery)))), + }); + for query in [&unfiltered as &dyn Query, &filtered] { + let hits = searcher.search(query, &TopDocs::with_limit(10).order_by_score())?; + let names = hits + .iter() + .map(|&(_, address)| name_of(&searcher, name, address)) + .collect::>>()?; + assert_eq!(names, ["near", "middle", "far"][..k.min(3)]); + } + } + // The deletion mask must intersect the caller's filter rather than replace it. + for only in ["deleted", "far"] { + let filtered = SpatialExecutor::new(PlanNode::Knn { + field: geometry, + geometry: geometry_set(r#"{"type": "Point", "coordinates": [0, 0]}"#), + k: 1, + filter: Some(Box::new(PlanNode::Query(name_query(name, only)))), + }); + let expected = if only == "deleted" { + vec![] + } else { + vec!["far"] + }; + assert_eq!(matching(&index, name, &filtered)?, expected); + } + Ok(()) + } + + #[test] + fn test_knn_applies_boosts_and_explains_the_same_score() -> crate::Result<()> { + let a = r#"{"type": "Point", "coordinates": [0, 1]}"#; + let b = r#"{"type": "Point", "coordinates": [0, 2]}"#; + let (index, _, geometry) = index_segments(&[&[("a", a)], &[("b", b)]])?; + let searcher = index.reader()?.searcher(); + let knn = + SpatialQuery::with_predicate(geometry, vec![[0.0, 0.0]], SpatialPredicate::Knn(2)); + let plain = searcher.search(&knn, &TopDocs::with_limit(10).order_by_score())?; + assert_eq!(plain.len(), 2); + for boost in [0.0, 3.0, -1.0] { + let boosted = BoostQuery::new(Box::new(knn.clone()), boost); + let hits = searcher.search(&boosted, &TopDocs::with_limit(10).order_by_score())?; + assert_eq!(hits.len(), plain.len()); + for (score, address) in hits { + let unboosted = plain.iter().find(|(_, doc)| *doc == address).unwrap().0; + assert_eq!(score, unboosted * boost); + assert_eq!(boosted.explain(&searcher, address)?.value(), score); + } + } + Ok(()) + } + + #[test] + fn test_empty_distance_ranges_match_no_geometries() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let (index, name, geometry) = index_segments(&[&[("box", BOX), ("point", point)]])?; + for predicate in [ + SpatialPredicate::Within(0.0), + SpatialPredicate::Within(-1.0), + SpatialPredicate::Between(0.0, 0.0), + SpatialPredicate::Between(-1.0, 0.0), + SpatialPredicate::Between(1.0, 1.0), + SpatialPredicate::Between(1.0, 0.5), + ] { + let query = SpatialQuery::with_predicate(geometry, vec![[5.0, 5.0]], predicate); + assert!(matching(&index, name, &query)?.is_empty(), "{query:?}"); + } + let nonempty = SpatialQuery::with_predicate( + geometry, + vec![[5.0, 5.0]], + SpatialPredicate::Between(0.0, 1e-6), + ); + assert_eq!(matching(&index, name, &nonempty)?, ["box", "point"]); + Ok(()) + } + + #[test] + fn test_overlapping_polygons_match_knn_and_within() -> crate::Result<()> { + let docs: Vec<(String, &str)> = (0..256).map(|i| (format!("box{i}"), BOX)).collect(); + let docs: Vec<(&str, &str)> = docs + .iter() + .map(|(name, shape)| (name.as_str(), *shape)) + .collect(); + let (index, _, geometry) = index_segments(&[&docs])?; + let searcher = index.reader()?.searcher(); + for k in [1, 10, docs.len(), docs.len() + 1] { + let knn = + SpatialQuery::with_predicate(geometry, vec![[5.0, 5.0]], SpatialPredicate::Knn(k)); + let hits = searcher.search(&knn, &TopDocs::with_limit(docs.len()).order_by_score())?; + assert_eq!(hits.len(), k.min(docs.len())); + // All at distance zero, the hits share the top rank. + assert!(hits.iter().all(|(score, _)| *score == 1.0)); + } + let within = SpatialQuery::with_predicate( + geometry, + vec![[5.0, 5.0]], + SpatialPredicate::Within(1e-6), + ); + assert_eq!( + searcher.search(&within, &DocSetCollector)?.len(), + docs.len() + ); + Ok(()) + } + + #[test] + fn test_contains_needs_every_member_of_a_multi_polygon() -> crate::Result<()> { + let split = r#"{"type": "MultiPolygon", "coordinates": [ + [[[1, 1], [2, 1], [2, 2], [1, 2], [1, 1]]], + [[[20, 20], [21, 20], [21, 21], [20, 21], [20, 20]]]]}"#; + let inside = r#"{"type": "MultiPolygon", "coordinates": [ + [[[1, 1], [2, 1], [2, 2], [1, 2], [1, 1]]], + [[[3, 3], [4, 3], [4, 4], [3, 4], [3, 3]]]]}"#; + let (index, name, geometry) = index_segments(&[&[("split", split), ("inside", inside)]])?; + let query = SpatialQuery::from_bounds(geometry, [[0.0, 0.0], [5.0, 5.0]]); + assert_eq!(matching(&index, name, &query)?, ["inside"]); + Ok(()) + } + + #[test] + fn test_near_join_matches_a_polygon_containing_the_point() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let (index, name, geometry) = index_segments(&[&[("box", BOX), ("point", point)]])?; + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "box"))), + inner: Box::new(PlanNode::Query(name_query(name, "point"))), + relation: SpatialRelation::Near(1e-6), + }); + assert_eq!(matching(&index, name, &join)?, ["box"]); + Ok(()) + } + + fn parse(index: &Index, name: Field, query: &str) -> crate::Result> { + let parser = crate::query::QueryParser::for_index(index, vec![name]); + matching(index, name, &*parser.parse_query(query)?) + } + + #[test] + fn test_join_keeps_the_occur_of_the_join_and_its_siblings() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let far = r#"{"type": "Point", "coordinates": [50, 50]}"#; + let (index, name, _) = index_segments(&[&[("box", BOX), ("point", point), ("far", far)]])?; + let near_point = "geometry:$within(1km, $query(name:point))"; + assert_eq!(parse(&index, name, near_point)?, ["box"]); + assert_eq!( + parse(&index, name, &format!("+name:box +{near_point}"))?, + ["box"] + ); + assert!(parse(&index, name, &format!("-name:box {near_point}"))?.is_empty()); + assert_eq!( + parse(&index, name, &format!("name:far OR {near_point}"))?, + ["box", "far"] + ); + assert_eq!( + parse(&index, name, &format!("name:far -{near_point}"))?, + ["far"] + ); + assert_eq!( + parse(&index, name, &format!("name:box name:far -{near_point}"))?, + ["far"] + ); + let any_near = "geometry:$within(1km, $query(*))"; + assert_eq!( + parse(&index, name, &format!("+(name:box name:point) +{any_near}"))?, + ["box", "point"] + ); + // A join with a boost is narrowed too, and still matches. + assert_eq!( + parse(&index, name, &format!("+name:box +{near_point}^2"))?, + ["box"] + ); + Ok(()) + } + + #[test] + fn test_join_rejects_an_invalid_inner_query() -> crate::Result<()> { + let (index, name, _) = index_segments(&[&[("box", BOX)]])?; + let parser = crate::query::QueryParser::for_index(&index, vec![name]); + let query = "geometry:$within(1km, $query(name:point AND))"; + assert!(parser.parse_query(query).is_err()); + assert!(!parser.parse_query_lenient(query).1.is_empty()); + assert!(parser + .parse_query("geometry:$within(1km, $query(nofield:x))") + .is_err()); + + let knn_join = query_grammar::UserInputAst::from(query_grammar::UserInputLeaf::Spatial { + field: Some("geometry".to_string()), + predicate: query_grammar::SpatialPredicateKind::Knn(3), + coordinates: vec![], + inner_query: Some("name:box".to_string()), + }); + assert!(parser.build_query_from_user_input_ast(knn_join).is_err()); + let no_coordinates = + query_grammar::UserInputAst::from(query_grammar::UserInputLeaf::Spatial { + field: Some("geometry".to_string()), + predicate: query_grammar::SpatialPredicateKind::Intersects, + coordinates: vec![], + inner_query: None, + }); + assert!(parser + .build_query_from_user_input_ast(no_coordinates) + .is_err()); + Ok(()) + } + + #[test] + fn test_a_query_with_too_few_coordinates_is_an_error() -> crate::Result<()> { + let (index, _, geometry) = index_segments(&[&[("box", BOX)]])?; + let searcher = index.reader()?.searcher(); + for (coordinates, predicate) in [ + (vec![], SpatialPredicate::Within(1.0)), + (vec![], SpatialPredicate::Knn(1)), + (vec![[0.0, 0.0], [1.0, 1.0]], SpatialPredicate::CoveredBy), + ] { + let query = SpatialQuery::with_predicate(geometry, coordinates, predicate); + assert!(searcher.search(&query, &DocSetCollector).is_err()); + } + Ok(()) + } + + #[test] + fn test_joins_with_point_and_line_outer_geometries() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let line = r#"{"type": "LineString", "coordinates": [[-5, 5.001], [15, 5.001]]}"#; + let (index, name, _) = + index_segments(&[&[("box", BOX), ("point", point), ("line", line)]])?; + assert_eq!( + parse(&index, name, "geometry:$intersects($query(name:box))")?, + ["line", "point"] + ); + assert_eq!( + parse(&index, name, "geometry:$intersects($query(name:line))")?, + ["box"] + ); + assert_eq!( + parse(&index, name, "geometry:$intersects($query(name:point))")?, + ["box"] + ); + assert!(parse(&index, name, "geometry:$contains($query(name:line))")?.is_empty()); + assert_eq!( + parse(&index, name, "geometry:$contains($query(name:point))")?, + ["box"] + ); + Ok(()) + } + + #[test] + fn test_queries_use_every_member_of_a_multi_polygon() -> crate::Result<()> { + let multi = r#"{"type": "MultiPolygon", "coordinates": [ + [[[0, 0], [1, 0], [1, 1], [0, 1], [0, 0]]], + [[[20, 20], [30, 20], [30, 30], [20, 30], [20, 20]]]]}"#; + let point = r#"{"type": "Point", "coordinates": [25, 25]}"#; + let square = r#"{"type": "Polygon", "coordinates": [ + [[24, 24], [26, 24], [26, 26], [24, 26], [24, 24]]]}"#; + let (index, name, geometry) = + index_segments(&[&[("multi", multi), ("point", point), ("square", square)]])?; + let intersects = SpatialExecutor::new(PlanNode::Intersects { + field: geometry, + geometry: geometry_set(multi), + filter: Box::new(PlanNode::Query(Box::new(AllQuery))), + }); + assert_eq!( + matching(&index, name, &intersects)?, + ["multi", "point", "square"] + ); + for relation in [SpatialRelation::Intersects, SpatialRelation::Contains] { + for inner in ["point", "square"] { + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "multi"))), + inner: Box::new(PlanNode::Query(name_query(name, inner))), + relation: relation.clone(), + }); + assert_eq!( + matching(&index, name, &join)?, + ["multi"], + "{relation:?} {inner}" + ); + } + } + Ok(()) + } + + #[test] + fn test_a_line_query_has_no_interior() -> crate::Result<()> { + let point = + |lon: f64, lat: f64| format!(r#"{{"type": "Point", "coordinates": [{lon}, {lat}]}}"#); + let (east, far_east) = (point(10.0, 0.0), point(20.0, 1.0)); + let line = r#"{"type": "LineString", "coordinates": [[5, -5], [5, 5]]}"#; + let crossing = r#"{"type": "LineString", "coordinates": [[0, 0], [10, 0]]}"#; + let (index, name, geometry) = index_segments(&[&[ + ("east", &east), + ("far_east", &far_east), + ("box", BOX), + ("crossing", crossing), + ("line", line), + ]])?; + let intersects = SpatialExecutor::new(PlanNode::Intersects { + field: geometry, + geometry: geometry_set(line), + filter: Box::new(PlanNode::Query(Box::new(AllQuery))), + }); + // The line has no interior, but it meets the lines it crosses, and itself. + assert_eq!( + matching(&index, name, &intersects)?, + ["box", "crossing", "line"] + ); + for relation in [SpatialRelation::Intersects, SpatialRelation::Contains] { + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "line"))), + inner: Box::new(PlanNode::Query(Box::new(AllQuery))), + relation: relation.clone(), + }); + let expected: &[&str] = match relation { + SpatialRelation::Intersects => &["line"], + _ => &[], + }; + assert_eq!(matching(&index, name, &join)?, expected, "{relation:?}"); + } + Ok(()) + } + + #[test] + fn test_join_ignores_deleted_inner_documents() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let (index, name, geometry) = index_segments(&[&[("box", BOX), ("point", point)]])?; + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "box"))), + inner: Box::new(PlanNode::Query(Box::new(AllQuery))), + relation: SpatialRelation::Near(1e-6), + }); + assert_eq!(matching(&index, name, &join)?, ["box"]); + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.set_merge_policy(Box::new(crate::indexer::NoMergePolicy)); + writer.delete_term(Term::from_field_text(name, "point")); + writer.commit()?; + assert!(matching(&index, name, &join)?.is_empty()); + Ok(()) + } + + #[test] + fn test_join_does_not_match_a_document_to_itself() -> crate::Result<()> { + let point = |lon: f64| format!(r#"{{"type": "Point", "coordinates": [{lon}, 0]}}"#); + let (a, b, c) = (point(0.0), point(10.0), point(10.000001)); + let (index, name, geometry) = index_segments(&[&[("a", &a), ("b", &b)], &[("c", &c)]])?; + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(Box::new(AllQuery))), + inner: Box::new(PlanNode::Query(Box::new(AllQuery))), + relation: SpatialRelation::Near(1e-6), + }); + assert_eq!(matching(&index, name, &join)?, ["b", "c"]); + Ok(()) + } + + #[test] + fn test_contains_rejects_a_polygon_leaving_through_a_shared_vertex() -> crate::Result<()> { + let polygon = |ring: &str| format!(r#"{{"type": "Polygon", "coordinates": [{ring}]}}"#); + let diamond = polygon("[[5, 5], [10, 0], [15, 5], [10, 10], [5, 5]]"); + let corner = polygon("[[5, 5], [10, 0], [10, 10], [5, 5]]"); + let inside_corner = polygon("[[10, 0], [8, 5], [5, 5], [10, 0]]"); + let outside_corner = polygon("[[10, 0], [15, 0], [15, 5], [10, 0]]"); + let (index, name, geometry) = index_segments(&[&[ + ("diamond", &diamond), + ("corner", &corner), + ("inside_corner", &inside_corner), + ("outside_corner", &outside_corner), + ("same", BOX), + ]])?; + let query = SpatialQuery::from_bounds(geometry, [[0.0, 0.0], [10.0, 10.0]]); + assert_eq!( + matching(&index, name, &query)?, + ["corner", "inside_corner", "same"] + ); + Ok(()) + } + + #[test] + fn test_polygons_sharing_boundary_intersect_only_where_interiors_overlap() -> crate::Result<()> + { + let polygon = |ring: &str| format!(r#"{{"type": "Polygon", "coordinates": [{ring}]}}"#); + let east = polygon("[[10, 0], [20, 0], [20, 10], [10, 10], [10, 0]]"); + let corner_touch = polygon("[[10, 10], [20, 10], [20, 20], [10, 20], [10, 10]]"); + let inner_touch = polygon("[[10, 0], [8, 5], [5, 5], [10, 0]]"); + let surround = polygon("[[-5, -5], [15, -5], [15, 15], [-5, 15], [-5, -5]]"); + let surround_touch = polygon("[[0, 0], [15, -5], [15, 15], [-5, 15], [0, 0]]"); + let (index, name, geometry) = index_segments(&[&[ + ("same", BOX), + ("east", &east), + ("corner_touch", &corner_touch), + ("inner_touch", &inner_touch), + ("surround", &surround), + ("surround_touch", &surround_touch), + ]])?; + let bounds = [[0.0, 0.0], [10.0, 10.0]]; + let intersects = SpatialQuery::intersects_bounds(geometry, bounds); + assert_eq!( + matching(&index, name, &intersects)?, + ["inner_touch", "same", "surround", "surround_touch"] + ); + let contains = SpatialQuery::from_bounds(geometry, bounds); + assert_eq!(matching(&index, name, &contains)?, ["inner_touch", "same"]); + Ok(()) + } + + #[test] + fn test_thin_polygons_sharing_edges_use_their_interior_sides() -> crate::Result<()> { + for height in [1e-6, 1e-8, 1e-12] { + let triangle = format!( + r#"{{"type": "Polygon", "coordinates": [[[0, 0], [10, 0], [5, {height}], [0, 0]]]}}"# + ); + let below = format!( + r#"{{"type": "Polygon", "coordinates": [[[10, 0], [0, 0], [5, {}], [10, 0]]]}}"#, + -height + ); + let (index, name, geometry) = + index_segments(&[&[("same", &triangle), ("below", &below)]])?; + for predicate in [SpatialPredicate::Intersects, SpatialPredicate::CoveredBy] { + let query = SpatialQuery::with_predicate( + geometry, + vec![[0.0, 0.0], [10.0, 0.0], [5.0, height]], + predicate, + ); + assert_eq!(matching(&index, name, &query)?, ["same"], "height {height}"); + } + } + Ok(()) + } + + #[test] + fn test_contains_rejects_a_polygon_around_a_hole_of_the_query() -> crate::Result<()> { + let holed = format!( + r#"{{"type": "Polygon", "coordinates": [{}, {}]}}"#, + "[[0, 0], [20, 0], [20, 20], [0, 20], [0, 0]]", + "[[8, 8], [8, 12], [12, 12], [12, 8], [8, 8]]" + ); + let square = |lo: f64, hi: f64| { + format!( + r#"{{"type": "Polygon", "coordinates": [[[{lo}, {lo}], [{hi}, {lo}], [{hi}, {hi}], [{lo}, {hi}], [{lo}, {lo}]]]}}"# + ) + }; + let (around, beside, in_hole) = (square(5.0, 15.0), square(1.0, 3.0), square(9.0, 11.0)); + let hole = square(8.0, 12.0); + let (index, name, geometry) = index_segments(&[&[ + ("holed", &holed), + ("around", &around), + ("beside", &beside), + ("in_hole", &in_hole), + ("hole", &hole), + ("copy", &holed), + ]])?; + let contains = |inner: &str| { + SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "holed"))), + inner: Box::new(PlanNode::Query(name_query(name, inner))), + relation: SpatialRelation::Contains, + }) + }; + assert!(matching(&index, name, &contains("around"))?.is_empty()); + assert!(matching(&index, name, &contains("in_hole"))?.is_empty()); + assert!(matching(&index, name, &contains("hole"))?.is_empty()); + assert_eq!(matching(&index, name, &contains("beside"))?, ["holed"]); + assert_eq!(matching(&index, name, &contains("copy"))?, ["holed"]); + Ok(()) + } + + #[test] + fn test_contains_joins_use_the_union_of_polygon_values() -> crate::Result<()> { + let square = |lo: f64, hi: f64| { + format!( + r#"{{"type": "Polygon", "coordinates": [[[{lo}, {lo}], [{hi}, {lo}], [{hi}, {hi}], [{lo}, {hi}], [{lo}, {lo}]]]}}"# + ) + }; + let inner = + r#"{"type": "Polygon", "coordinates": [[[2, 2], [8, 2], [8, 8], [2, 8], [2, 2]]]}"#; + for (outer, expected) in [ + // The candidate crosses boundaries of both members but stays inside their union. + ( + r#"[ + {"type": "Polygon", "coordinates": [[[0, 0], [6, 0], [6, 10], [0, 10], [0, 0]]]}, + {"type": "Polygon", "coordinates": [[[4, 0], [10, 0], [10, 10], [4, 10], [4, 0]]]} + ]"#, + true, + ), + // The oppositely directed shared edge is internal to the union. + ( + r#"[ + {"type": "Polygon", "coordinates": [[[0, 0], [5, 0], [5, 10], [0, 10], [0, 0]]]}, + {"type": "Polygon", "coordinates": [[[5, 0], [10, 0], [10, 10], [5, 10], [5, 0]]]} + ]"#, + true, + ), + // A gap between the members must still reject containment. + ( + r#"[ + {"type": "Polygon", "coordinates": [[[0, 0], [3, 0], [3, 10], [0, 10], [0, 0]]]}, + {"type": "Polygon", "coordinates": [[[7, 0], [10, 0], [10, 10], [7, 10], [7, 0]]]} + ]"#, + false, + ), + // No individual member has a hole, but the union encloses an uncovered gap. + ( + r#"[ + {"type": "Polygon", "coordinates": [[[0, 0], [10, 0], [10, 3], [0, 3], [0, 0]]]}, + {"type": "Polygon", "coordinates": [[[0, 7], [10, 7], [10, 10], [0, 10], [0, 7]]]}, + {"type": "Polygon", "coordinates": [[[0, 2], [3, 2], [3, 8], [0, 8], [0, 2]]]}, + {"type": "Polygon", "coordinates": [[[7, 2], [10, 2], [10, 8], [7, 8], [7, 2]]]} + ]"#, + false, + ), + ] { + let (index, name, _) = + index_segments(&[&[("outer", outer)], &[("inner", inner), ("copy", outer)]])?; + let query = "+name:outer +geometry:$contains($query(name:inner))"; + let wanted = if expected { vec!["outer"] } else { vec![] }; + assert_eq!(parse(&index, name, query)?, wanted, "{outer}"); + assert_eq!( + parse( + &index, + name, + "+name:outer +geometry:$contains($query(name:copy))" + )?, + ["outer"], + "contains a copy: {outer}" + ); + + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.merge(&index.searchable_segment_ids()?).wait()?; + assert_eq!(parse(&index, name, query)?, wanted, "merged {outer}"); + assert_eq!( + parse( + &index, + name, + "+name:outer +geometry:$contains($query(name:copy))" + )?, + ["outer"], + "merged contains a copy: {outer}" + ); + } + + let donut = r#"{"type": "Polygon", "coordinates": [ + [[0, 0], [20, 0], [20, 20], [0, 20], [0, 0]], + [[8, 8], [8, 12], [12, 12], [12, 8], [8, 8]]]}"#; + for (filling, expected) in [(square(7.0, 13.0), true), (square(7.0, 10.0), false)] { + let outer = format!("[{donut}, {filling}]"); + let inner = square(5.0, 15.0); + let (index, name, _) = index_segments(&[&[("outer", &outer), ("inner", &inner)]])?; + let wanted = if expected { vec!["outer"] } else { vec![] }; + assert_eq!( + parse( + &index, + name, + "+name:outer +geometry:$contains($query(name:inner))" + )?, + wanted, + "{outer}" + ); + } + Ok(()) + } + + #[test] + fn test_queries_with_many_vertices_find_crossings() -> crate::Result<()> { + let polygon = |ring: &str| format!(r#"{{"type": "Polygon", "coordinates": [{ring}]}}"#); + let straddling = polygon("[[4.5, -0.5], [5.5, -0.5], [5.5, 0.5], [4.5, 0.5], [4.5, -0.5]]"); + let inside = polygon("[[-1, -1], [1, -1], [1, 1], [-1, 1], [-1, -1]]"); + let outside = polygon("[[5.5, -0.5], [6.5, -0.5], [6.5, 0.5], [5.5, 0.5], [5.5, -0.5]]"); + let (index, name, geometry) = index_segments(&[&[ + ("straddling", &straddling), + ("inside", &inside), + ("outside", &outside), + ]])?; + let circle: Vec<[f64; 2]> = (0..200) + .map(|i| { + let (sin, cos) = (std::f64::consts::TAU * i as f64 / 200.0).sin_cos(); + [5.0 * cos, 5.0 * sin] + }) + .collect(); + let contains = SpatialQuery::new(geometry, circle.clone()); + assert_eq!(matching(&index, name, &contains)?, ["inside"]); + let intersects = + SpatialQuery::with_predicate(geometry, circle, SpatialPredicate::Intersects); + assert_eq!( + matching(&index, name, &intersects)?, + ["inside", "straddling"] + ); + Ok(()) + } + + #[test] + fn test_contains_needs_every_value_of_a_document() -> crate::Result<()> { + let point = |lon: f64| format!(r#"{{"type": "Point", "coordinates": [{lon}, 5]}}"#); + let split = format!("[{}, {}]", point(5.0), point(50.0)); + let both = format!("[{}, {}]", point(4.0), point(6.0)); + let (index, name, geometry) = index_segments(&[&[("split", &split), ("both", &both)]])?; + let query = SpatialQuery::from_bounds(geometry, [[0.0, 0.0], [10.0, 10.0]]); + assert_eq!(matching(&index, name, &query)?, ["both"]); + let intersects = SpatialQuery::intersects_bounds(geometry, [[0.0, 0.0], [10.0, 10.0]]); + assert_eq!(matching(&index, name, &intersects)?, ["both", "split"]); + Ok(()) + } + + #[test] + fn test_geometry_collection_is_indexed_and_round_trips() -> crate::Result<()> { + let collection = r#"{"type": "GeometryCollection", "geometries": [ + {"type": "Point", "coordinates": [50, 50]}, + {"type": "LineString", "coordinates": [[1, 1], [2, 2]]}]}"#; + let (index, name, geometry) = index_segments(&[&[("collection", collection)]])?; + let query = SpatialQuery::intersects_bounds(geometry, [[0.0, 0.0], [10.0, 10.0]]); + assert_eq!(matching(&index, name, &query)?, ["collection"]); + + let object: serde_json::Map = + serde_json::from_str(collection).unwrap(); + let parsed = Geometry::::from_geojson(&object).unwrap(); + assert_eq!( + Geometry::::from_geojson(&parsed.to_geojson()).unwrap(), + parsed + ); + Ok(()) + } + + #[test] + fn test_index_a_polygon_with_more_vertices_than_a_u16() -> crate::Result<()> { + let ring = regular_ring(0.0, 0.0, 70_000, 5.0, false); + let big = format!(r#"{{"type": "Polygon", "coordinates": [{ring}]}}"#); + let point = |lon: f64| format!(r#"{{"type": "Point", "coordinates": [{lon}, 0]}}"#); + let (center, edge, out) = (point(0.0), point(4.99999), point(5.01)); + let (index, name, geometry) = index_segments(&[&[ + ("big", &big), + ("center", ¢er), + ("edge", &edge), + ("out", &out), + ]])?; + let at = |lon: f64| { + SpatialQuery::with_predicate(geometry, vec![[lon, 0.0]], SpatialPredicate::Within(1e-9)) + }; + assert_eq!(matching(&index, name, &at(0.0))?, ["big", "center"]); + // The vertices near (5, 0) have the highest edge ids. + assert_eq!(matching(&index, name, &at(4.999995))?, ["big"]); + assert!(matching(&index, name, &at(5.005))?.is_empty()); + let near_edge = SpatialQuery::intersects_bounds(geometry, [[4.9, -0.1], [5.1, 0.1]]); + assert_eq!(matching(&index, name, &near_edge)?, ["big", "edge", "out"]); + Ok(()) + } + + /// A closed regular ring around (lon, lat): counterclockwise, or clockwise for a hole. + fn regular_ring(lon: f64, lat: f64, sides: usize, radius: f64, clockwise: bool) -> String { + let mut positions: Vec = (0..=sides) + .map(|i| { + let angle = std::f64::consts::TAU * (i % sides) as f64 / sides as f64; + let (sin, cos) = angle.sin_cos(); + format!("[{}, {}]", lon + radius * cos, lat + radius * sin) + }) + .collect(); + if clockwise { + positions.reverse(); + } + format!("[{}]", positions.join(", ")) + } + + /// A random GeoJSON point, line string, polygon (half of them with a hole) or two-polygon + /// multi-polygon, a few degrees from the origin. + fn random_geojson(rng: &mut StdRng) -> String { + let (lon, lat): (f64, f64) = (rng.random_range(0.0..3.0), rng.random_range(0.0..3.0)); + let sides = rng.random_range(3..9); + match rng.random_range(0..4) { + 0 => format!(r#"{{"type": "Point", "coordinates": [{lon}, {lat}]}}"#), + 1 => { + let positions: Vec = (0..rng.random_range(2..6)) + .map(|_| { + let (dlon, dlat): (f64, f64) = + (rng.random_range(-0.3..0.3), rng.random_range(-0.3..0.3)); + format!("[{}, {}]", lon + dlon, lat + dlat) + }) + .collect(); + format!( + r#"{{"type": "LineString", "coordinates": [{}]}}"#, + positions.join(", ") + ) + } + 2 => { + let radius = rng.random_range(0.05..0.5); + let mut rings = vec![regular_ring(lon, lat, sides, radius, false)]; + if rng.random_bool(0.5) { + rings.push(regular_ring(lon, lat, sides, radius / 3.0, true)); + } + format!( + r#"{{"type": "Polygon", "coordinates": [{}]}}"#, + rings.join(", ") + ) + } + _ => { + let radius = rng.random_range(0.05..0.2); + format!( + r#"{{"type": "MultiPolygon", "coordinates": [[{}], [{}]]}}"#, + regular_ring(lon, lat, sides, radius, false), + regular_ring(lon + 0.6, lat - 0.4, sides, radius, false) + ) + } + } + } + + fn geometry_set(geojson: &str) -> GeometrySet { + let object: serde_json::Map = + serde_json::from_str(geojson).unwrap(); + let geometry = Geometry::::from_geojson(&object).unwrap(); + to_geometry_set(&geometry.project::(), 0) + } + + /// The edges of a member, ring by ring. A point is an edge from itself to itself. + fn member_edges(member: &EdgeSet) -> Vec<([f64; 3], [f64; 3])> { + if let [point] = member.vertices[..] { + return vec![(point, point)]; + } + member + .ring_offsets + .windows(2) + .flat_map(|ring| ring[0]..ring[1] - 1) + .map(|i| (member.vertices[i], member.vertices[i + 1])) + .collect() + } + + fn member_contains(member: &EdgeSet, point: &[f64; 3]) -> bool { + member.closed + && member.ring_offsets.windows(2).fold(false, |inside, ring| { + let ring = &member.vertices[ring[0]..ring[1]]; + inside ^ brute_force_contains(point, ring, compute_origin_inside(ring)) + }) + } + + /// The distance between two geometry sets, interiors included, from every pair of edges. + fn brute_force_distance(a: &GeometrySet, b: &GeometrySet) -> f64 { + let mut min_dist = S1ChordAngle::infinity(); + for a_member in &a.members { + for b_member in &b.members { + if member_contains(a_member, &b_member.vertices[0]) + || member_contains(b_member, &a_member.vertices[0]) + { + return 0.0; + } + for (a0, a1) in member_edges(a_member) { + for (b0, b1) in member_edges(b_member) { + if a0 == a1 { + update_min_distance(&a0, &b0, &b1, &mut min_dist); + } else if b0 == b1 { + update_min_distance(&b0, &a0, &a1, &mut min_dist); + } else { + update_edge_pair_min_distance(&a0, &a1, &b0, &b1, &mut min_dist); + } + } + } + } + } + min_dist.to_radians() + } + + /// Asserts that the found names hold every geometry whose distance is in [min, max) and none + /// outside it, leaving out distances within rounding of either bound. + fn assert_found_in_range(found: &[String], distances: &[(f64, String)], min: f64, max: f64) { + const ROUNDING: f64 = 1e-12; + for (distance, name) in distances { + let is_found = found.contains(name); + if *distance >= min + ROUNDING && *distance < max - ROUNDING { + assert!(is_found, "{name} at {distance} is in [{min}, {max})"); + } + if *distance < min - ROUNDING || *distance >= max + ROUNDING { + assert!(!is_found, "{name} at {distance} is not in [{min}, {max})"); + } + } + } + + /// Asserts that the hits are the k nearest geometries, nearest first. Ties may be broken + /// either way, so the distances are compared rather than the names. + fn assert_nearest( + searcher: &Searcher, + name: Field, + hits: &[(Score, DocAddress)], + distances: &[(f64, String)], + k: usize, + ) -> crate::Result<()> { + assert_eq!(hits.len(), k.min(distances.len())); + for (rank, (_, doc_address)) in hits.iter().enumerate() { + let hit = name_of(searcher, name, *doc_address)?; + let distance = distances.iter().find(|(_, n)| *n == hit).unwrap().0; + assert!( + (distance - distances[rank].0).abs() <= 1e-12, + "hit {rank} is {hit} at {distance}, the nearest are {:?}", + &distances[..k] + ); + } + Ok(()) + } + + #[test] + fn test_distance_queries_agree_with_brute_force() -> crate::Result<()> { + check_distance_queries_against_brute_force(7) + } + + /// Indexes random geometries and checks distance queries from random points and random + /// geometries against brute force. + fn check_distance_queries_against_brute_force(seed: u64) -> crate::Result<()> { + let mut rng = StdRng::seed_from_u64(seed); + let docs: Vec<(String, String)> = (0..120) + .map(|i| (format!("g{i}"), random_geojson(&mut rng))) + .collect(); + let segments: Vec> = docs + .chunks(40) + .map(|chunk| { + chunk + .iter() + .map(|(n, g)| (n.as_str(), g.as_str())) + .collect() + }) + .collect(); + let segments: Vec<&[(&str, &str)]> = segments.iter().map(Vec::as_slice).collect(); + let (index, name, geometry) = index_segments(&segments)?; + let searcher = index.reader()?.searcher(); + let stored: Vec<(String, GeometrySet)> = docs + .iter() + .map(|(n, g)| (n.clone(), geometry_set(g))) + .collect(); + let distances_from = |query: &GeometrySet| { + let mut distances: Vec<(f64, String)> = stored + .iter() + .map(|(n, set)| (brute_force_distance(query, set), n.clone())) + .collect(); + distances.sort_by(|a, b| a.0.total_cmp(&b.0)); + distances + }; + + for _ in 0..40 { + let point = [rng.random_range(-0.5..3.5), rng.random_range(-0.5..3.5)]; + let query = format!( + r#"{{"type": "Point", "coordinates": [{}, {}]}}"#, + point[0], point[1] + ); + let distances = distances_from(&geometry_set(&query)); + let radius = rng.random_range(0.0..0.5f64).to_radians(); + + let within = SpatialQuery::with_predicate( + geometry, + vec![point], + SpatialPredicate::Within(radius), + ); + let found = matching(&index, name, &within)?; + assert_found_in_range(&found, &distances, 0.0, radius); + + let between = SpatialQuery::with_predicate( + geometry, + vec![point], + SpatialPredicate::Between(radius / 2.0, radius), + ); + let found = matching(&index, name, &between)?; + assert_found_in_range(&found, &distances, radius / 2.0, radius); + + let k = rng.random_range(1..6); + let knn = SpatialQuery::with_predicate(geometry, vec![point], SpatialPredicate::Knn(k)); + let hits = searcher.search(&knn, &TopDocs::with_limit(10).order_by_score())?; + assert_nearest(&searcher, name, &hits, &distances, k)?; + } + + for _ in 0..40 { + let query = geometry_set(&random_geojson(&mut rng)); + let distances = distances_from(&query); + let radius = rng.random_range(0.0..0.5f64).to_radians(); + + let within = SpatialExecutor::new(PlanNode::Within { + field: geometry, + geometry: query.clone(), + radius: S1ChordAngle::from_radians(radius).length2(), + filter: Box::new(PlanNode::Query(Box::new(AllQuery))), + }); + let found = matching(&index, name, &within)?; + assert_found_in_range(&found, &distances, 0.0, radius); + + let k = rng.random_range(1..6); + let knn = SpatialExecutor::new(PlanNode::Knn { + field: geometry, + geometry: query, + k, + filter: None, + }); + let hits = searcher.search(&knn, &TopDocs::with_limit(10).order_by_score())?; + assert_nearest(&searcher, name, &hits, &distances, k)?; + } + Ok(()) + } + + #[test] + fn test_between_join_skips_geometries_inside_the_inner_radius() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (origin, near, far) = (point(0.0), point(0.5), point(2.0)); + let (index, name, geometry) = + index_segments(&[&[("origin", &origin), ("near", &near), ("far", &far)]])?; + let join = SpatialExecutor::new(PlanNode::Join { + field: geometry, + outer: Box::new(PlanNode::Query(name_query(name, "origin"))), + inner: Box::new(PlanNode::Query(Box::new(AllQuery))), + relation: SpatialRelation::Between(1f64.to_radians(), 3f64.to_radians()), + }); + assert_eq!(matching(&index, name, &join)?, ["origin"]); + Ok(()) + } + + /// Merges every segment of the index into one. + fn merge_all(index: &Index) -> crate::Result<()> { + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.merge(&index.searchable_segment_ids()?).wait()?; + writer.wait_merging_threads() + } + + /// A query for the geometries within a few meters of a point. + fn near(geometry: Field, [lon, lat]: [f64; 2]) -> SpatialQuery { + SpatialQuery::with_predicate(geometry, vec![[lon, lat]], SpatialPredicate::Within(1e-6)) + } + + /// A query for the geometries that intersect a small box around a point. + fn box_around(geometry: Field, [lon, lat]: [f64; 2]) -> SpatialQuery { + let d = 0.01; + SpatialQuery::intersects_bounds(geometry, [[lon - d, lat - d], [lon + d, lat + d]]) + } + + #[test] + fn test_a_polygon_around_a_whole_cube_face_is_indexed_on_it() -> crate::Result<()> { + // Face 0 lies within 45 degrees of (0, 0), so the polygon has no edges on it. The point + // is in a segment of its own, so that it adds none either. + let ring = [[-50.0, -50.0], [50.0, -50.0], [50.0, 50.0], [-50.0, 50.0]]; + let big = r#"{"type": "Polygon", "coordinates": + [[[-50, -50], [50, -50], [50, 50], [-50, 50], [-50, -50]]]}"#; + let origin = r#"{"type": "Point", "coordinates": [0, 0]}"#; + let (index, name, geometry) = index_segments(&[&[("big", big)], &[("origin", origin)]])?; + assert_eq!( + matching(&index, name, &near(geometry, [0.0, 0.0]))?, + ["big", "origin"] + ); + assert_eq!( + matching(&index, name, &box_around(geometry, [0.0, 0.0]))?, + ["big", "origin"] + ); + // As a query, the polygon covers the face too. + let intersects = + SpatialQuery::with_predicate(geometry, ring.to_vec(), SpatialPredicate::Intersects); + assert_eq!(matching(&index, name, &intersects)?, ["big", "origin"]); + let contains = SpatialQuery::new(geometry, ring.to_vec()); + assert!(matching(&index, name, &contains)?.contains(&"origin".to_string())); + Ok(()) + } + + #[test] + fn test_rings_of_either_orientation_bound_the_smaller_region() -> crate::Result<()> { + let clockwise = r#"{"type": "Polygon", "coordinates": + [[[0, 0], [0, 10], [10, 10], [10, 0], [0, 0]]]}"#; + let (index, name, geometry) = index_segments(&[&[("clockwise", clockwise)]])?; + assert_eq!( + matching(&index, name, &box_around(geometry, [4.5, 4.5]))?, + ["clockwise"] + ); + assert!(matching(&index, name, &box_around(geometry, [50.0, 50.0]))?.is_empty()); + + let inside = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let far = r#"{"type": "Point", "coordinates": [50, 50]}"#; + let (index, name, _) = index_segments(&[&[("inside", inside), ("far", far)]])?; + let query = "geometry:$intersects(0 0, 0 10, 10 10, 10 0)"; + assert_eq!(parse(&index, name, query)?, ["inside"]); + Ok(()) + } + + #[test] + fn test_a_ring_with_a_repeated_vertex_is_indexed() -> crate::Result<()> { + let repeated = r#"{"type": "Polygon", "coordinates": + [[[0, 0], [0, 0], [10, 0], [10, 10], [0, 10], [0, 0]]]}"#; + let (index, name, geometry) = index_segments(&[&[("repeated", repeated)]])?; + assert_eq!( + matching(&index, name, &box_around(geometry, [4.5, 4.5]))?, + ["repeated"] + ); + assert!(matching(&index, name, &box_around(geometry, [50.0, 50.0]))?.is_empty()); + Ok(()) + } + + #[test] + fn test_a_query_polygon_with_no_area_is_what_is_left_of_it() -> crate::Result<()> { + let (index, name, geometry) = index_segments(&[&[("box", BOX)]])?; + // A line inside the box, and the bounding box of a point inside it. + let line = SpatialQuery::with_predicate( + geometry, + vec![[2.0, 2.0], [3.0, 3.0], [2.0, 2.0]], + SpatialPredicate::Intersects, + ); + assert_eq!(matching(&index, name, &line)?, ["box"]); + let point = SpatialQuery::intersects_bounds(geometry, [[5.0, 5.0], [5.0, 5.0]]); + assert_eq!(matching(&index, name, &point)?, ["box"]); + Ok(()) + } + + #[test] + fn test_query_coordinates_and_distances_are_checked() -> crate::Result<()> { + let (index, name, geometry) = index_segments(&[&[("box", BOX)]])?; + let searcher = index.reader()?.searcher(); + for (coordinates, predicate) in [ + ( + vec![[0.0, 0.0], [1.0, 0.0], [f64::INFINITY, 1.0]], + SpatialPredicate::Intersects, + ), + (vec![[0.0, 95.0]], SpatialPredicate::Within(1e-6)), + (vec![[f64::NAN, 0.0]], SpatialPredicate::Knn(1)), + (vec![[0.0, 0.0]], SpatialPredicate::Within(f64::NAN)), + ] { + let query = SpatialQuery::with_predicate(geometry, coordinates, predicate); + assert!(searcher.search(&query, &DocSetCollector).is_err()); + } + let huge = format!("geometry:$intersects(0 0, 1 0, 1{} 1)", "0".repeat(400)); + let parser = crate::query::QueryParser::for_index(&index, vec![name]); + assert!(searcher + .search(&*parser.parse_query(&huge)?, &DocSetCollector) + .is_err()); + Ok(()) + } + + #[test] + fn test_a_line_beside_a_polygon_never_contains_a_cell_center() -> crate::Result<()> { + use crate::spatial::cell_index_reader::CellIndexReader; + + let zigzag: Vec = (0..200) + .map(|i| { + format!( + "[{}, {}]", + 0.5 + i as f64 / 25.0, + 1.0 + (i % 2) as f64 / 10.0 + ) + }) + .collect(); + let line = format!( + r#"{{"type": "LineString", "coordinates": [{}]}}"#, + zigzag.join(", ") + ); + let (index, _, geometry) = index_segments(&[&[("box", BOX), ("line", &line)]])?; + let searcher = index.reader()?.searcher(); + let spatial = searcher + .segment_reader(0) + .spatial_fields() + .get_field(geometry)? + .unwrap(); + let cells = CellIndexReader::open(spatial.cells_bytes())?; + let edges = EdgeReader::::open(spatial.edges_bytes())?; + let mut line_cells = 0; + for cell in cells.iter() { + for shape in &cell.shapes { + let (_, set) = edges.read_geometry_set(shape.geometry_id.1); + if !set.members[0].closed { + assert!(!shape.contains_center, "{:?}", cell.cell_id); + line_cells += 1; + } + } + } + assert!(line_cells > 1); + Ok(()) + } + + #[test] + fn test_merge_keeps_a_long_line_in_every_cell_it_crosses() -> crate::Result<()> { + use crate::spatial::math::normalize; + use crate::spatial::surface::Surface; + + let line = r#"{"type": "LineString", "coordinates": [[1, -1], [100, -1]]}"#; + let points: Vec = (0..80) + .map(|i| { + let lon = 30.0 + i as f64 / 8.0; + format!(r#"{{"type": "Point", "coordinates": [{lon}, -0.3]}}"#) + }) + .collect(); + let points: Vec<(&str, &str)> = points.iter().map(|p| ("point", p.as_str())).collect(); + let (index, name, geometry) = index_segments(&[&[("line", line)], &points])?; + + // Points on the line, from 10 to 42 degrees of longitude. + let (a, b) = (Sphere::project(1.0, -1.0), Sphere::project(100.0, -1.0)); + let on_line: Vec<[f64; 2]> = [0.09, 0.19, 0.29, 0.34, 0.37, 0.41] + .iter() + .map(|&t| { + let p = normalize(&[ + a[0] + t * (b[0] - a[0]), + a[1] + t * (b[1] - a[1]), + a[2] + t * (b[2] - a[2]), + ]); + [p[1].atan2(p[0]).to_degrees(), p[2].asin().to_degrees()] + }) + .collect(); + let finds_line = |index: &Index| -> crate::Result { + for &point in &on_line { + if !matching(index, name, &near(geometry, point))?.contains(&"line".to_string()) { + return Ok(false); + } + } + Ok(true) + }; + assert!(finds_line(&index)?); + merge_all(&index)?; + assert_eq!(index.searchable_segment_ids()?.len(), 1); + assert!(finds_line(&index)?); + Ok(()) + } + + #[test] + fn test_points_and_lines_intersect_where_they_meet() -> crate::Result<()> { + let point = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let line = r#"{"type": "LineString", "coordinates": [[5, 5], [6, 6]]}"#; + let next = r#"{"type": "LineString", "coordinates": [[6, 6], [7, 5]]}"#; + let on_equator = r#"{"type": "Point", "coordinates": [20.5, 0]}"#; + let equator = r#"{"type": "LineString", "coordinates": [[20, 0], [21, 0]]}"#; + let far = r#"{"type": "Point", "coordinates": [50, 50]}"#; + let (index, name, _) = index_segments(&[ + &[("a", point), ("b", point), ("line", line)], + &[ + ("next", next), + ("on_equator", on_equator), + ("equator", equator), + ], + &[("far", far)], + ])?; + let intersecting = |inner: &str| { + parse( + &index, + name, + &format!("geometry:$intersects($query(name:{inner}))"), + ) + }; + assert_eq!(intersecting("a")?, ["b", "line"]); + assert_eq!(intersecting("line")?, ["a", "b", "next"]); + assert_eq!(intersecting("next")?, ["line"]); + // A point on the interior of an edge. + assert_eq!(intersecting("on_equator")?, ["equator"]); + assert_eq!(intersecting("equator")?, ["on_equator"]); + assert!(intersecting("far")?.is_empty()); + Ok(()) + } + + #[test] + fn test_geometry_outside_a_spatial_field_is_not_indexed() -> crate::Result<()> { + use crate::schema::{OwnedValue, FAST, TEXT}; + + // Fast fields, and the fields of a JSON object, skip geometry. + let mut schema_builder = Schema::builder(); + let number = schema_builder.add_u64_field("number", FAST); + let text = schema_builder.add_text_field("text", STRING | FAST); + let json = schema_builder.add_json_field("json", TEXT | FAST); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + let geometry = OwnedValue::Geometry(Geometry::Point([1.0, 2.0])); + let mut doc = TantivyDocument::new(); + doc.add_field_value(number, &geometry); + doc.add_field_value(text, &geometry); + doc.add_field_value( + json, + &OwnedValue::Object(vec![("place".to_string(), geometry.clone())]), + ); + writer.add_document(doc)?; + writer.commit()?; + assert_eq!(index.reader()?.searcher().num_docs(), 1); + Ok(()) + } + + #[test] + fn test_knn_finds_the_nearest_of_the_documents_its_siblings_allow() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (a, b, cafe1, cafe2, cafe3) = + (point(1.0), point(2.0), point(3.0), point(4.0), point(5.0)); + let (index, name, _) = index_segments(&[ + &[("a", &a), ("cafe2", &cafe2)], + &[("b", &b), ("cafe1", &cafe1), ("cafe3", &cafe3)], + ])?; + assert_eq!(parse(&index, name, "geometry:$knn(2, 0 0)")?, ["a", "b"]); + let cafes = "+(name:cafe1 name:cafe2 name:cafe3)"; + assert_eq!( + parse(&index, name, &format!("{cafes} +geometry:$knn(2, 0 0)"))?, + ["cafe1", "cafe2"] + ); + assert_eq!( + parse(&index, name, "+name:cafe3 +geometry:$knn(1, 0 0)")?, + ["cafe3"] + ); + assert_eq!(parse(&index, name, "-name:a geometry:$knn(1, 0 0)")?, ["b"]); + Ok(()) + } + + #[test] + fn test_contains_reads_from_the_field_side_with_coordinates_and_joins() -> crate::Result<()> { + let small = r#"{"type": "Polygon", "coordinates": + [[[4, 4], [5, 4], [5, 5], [4, 5], [4, 4]]]}"#; + let halves = r#"{"type": "MultiPolygon", "coordinates": [ + [[[20, 0], [25, 0], [25, 5], [20, 5], [20, 0]]], + [[[25, 0], [30, 0], [30, 5], [25, 5], [25, 0]]]]}"#; + let point = r#"{"type": "Point", "coordinates": [4.5, 4.5]}"#; + let (index, name, _) = index_segments(&[ + &[("big", BOX), ("small", small)], + &[("halves", halves), ("point", point)], + ])?; + let found = |query: &str| parse(&index, name, query); + assert_eq!( + found("geometry:$contains(4.2 4.2, 4.8 4.2, 4.8 4.8, 4.2 4.8)")?, + ["big", "small"] + ); + assert_eq!(found("geometry:$contains(1 1, 6 1, 6 6, 1 6)")?, ["big"]); + // A polygon contains itself. + assert_eq!( + found("geometry:$contains(4 4, 5 4, 5 5, 4 5)")?, + ["big", "small"] + ); + assert_eq!(found("geometry:$contains($query(name:small))")?, ["big"]); + // The members of a set may contain the polygon jointly. + assert_eq!( + found("geometry:$contains(24 1, 26 1, 26 2, 24 2)")?, + ["halves"] + ); + assert!(found("geometry:$contains(9 9, 11 9, 11 11, 9 11)")?.is_empty()); + // `$within` a polygon is the other way around. + assert_eq!( + found("geometry:$within(3 3, 6 3, 6 6, 3 6)")?, + ["point", "small"] + ); + assert_eq!( + found("geometry:$within(-1 -1, 11 -1, 11 11, -1 11)")?, + ["big", "point", "small"] + ); + Ok(()) + } + + #[test] + fn test_oriented_rings_may_bound_more_than_a_hemisphere() -> crate::Result<()> { + // On an oriented field, a clockwise ring around the box bounds the rest of the sphere. + let mut schema_builder = Schema::builder(); + let name = schema_builder.add_text_field("name", STRING | STORED); + let geometry = schema_builder.add_spatial_field("geometry", SPATIAL.set_oriented()); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + let everything_else = r#"{"type": "Polygon", "coordinates": + [[[0, 0], [0, 10], [10, 10], [10, 0], [0, 0]]]}"#; + let point = + |lon: f64, lat: f64| format!(r#"{{"type": "Point", "coordinates": [{lon}, {lat}]}}"#); + for (doc_name, doc_geometry) in [ + ("everything_else", everything_else.to_string()), + ("inside", point(5.0, 5.0)), + ("far", point(50.0, 50.0)), + ("antipode", point(-175.0, -5.0)), + ] { + let json = format!(r#"{{"name": "{doc_name}", "geometry": {doc_geometry}}}"#); + writer.add_document(TantivyDocument::parse_json(&index.schema(), &json)?)?; + } + writer.commit()?; + + let near_everything_else = |at: [f64; 2]| -> crate::Result { + Ok(matching(&index, name, &near(geometry, at))? + .contains(&"everything_else".to_string())) + }; + assert!(!near_everything_else([5.0, 5.0])?); + assert!(near_everything_else([50.0, 50.0])?); + assert!(near_everything_else([-175.0, -5.0])?); + + // Query polygons on the field are read the same way. + let clockwise = vec![[0.0, 0.0], [0.0, 10.0], [10.0, 10.0], [10.0, 0.0]]; + let contained = SpatialQuery::new(geometry, clockwise.clone()); + assert_eq!( + matching(&index, name, &contained)?, + ["antipode", "everything_else", "far"] + ); + let parsed = parse(&index, name, "geometry:$within(0 0, 0 10, 10 10, 10 0)")?; + assert_eq!(parsed, ["antipode", "everything_else", "far"]); + // A query may read its polygon otherwise. + let smaller = contained.with_ring_orientation(RingOrientation::Smaller); + assert_eq!(matching(&index, name, &smaller)?, ["inside"]); + Ok(()) + } + + #[test] + fn test_a_query_polygon_may_be_oriented_on_any_field() -> crate::Result<()> { + let inside = r#"{"type": "Point", "coordinates": [5, 5]}"#; + let far = r#"{"type": "Point", "coordinates": [50, 50]}"#; + let (index, name, geometry) = index_segments(&[&[("inside", inside), ("far", far)]])?; + let clockwise = vec![[0.0, 0.0], [0.0, 10.0], [10.0, 10.0], [10.0, 0.0]]; + let query = SpatialQuery::with_predicate(geometry, clockwise, SpatialPredicate::Intersects); + assert_eq!(matching(&index, name, &query)?, ["inside"]); + let oriented = query.with_ring_orientation(RingOrientation::RightHand); + assert_eq!(matching(&index, name, &oriented)?, ["far"]); + Ok(()) + } + + #[test] + fn test_knn_lifts_the_nearest_documents_of_a_boolean_query() -> crate::Result<()> { + // Without the lift, the tie on the text clauses goes to the first document, b. + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (b, a) = (point(2.0), point(1.0)); + let (index, name, _) = index_segments(&[&[("b", &b), ("a", &a)]])?; + let searcher = index.reader()?.searcher(); + let parser = crate::query::QueryParser::for_index(&index, vec![name]); + for query in [ + "name:a name:b geometry:$knn(1, 0 0)", + "name:a name:b geometry:$knn(1, 0 0)^2", + ] { + let query = parser.parse_query(query)?; + let hits = searcher.search(&*query, &TopDocs::with_limit(2).order_by_score())?; + assert_eq!(name_of(&searcher, name, hits[0].1)?, "a"); + assert!(hits[0].0 > hits[1].0); + } + Ok(()) + } + + #[test] + fn test_knn_is_restricted_however_it_is_boosted_or_nested() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (a, b, cafe1, cafe2, cafe3) = + (point(1.0), point(2.0), point(3.0), point(4.0), point(5.0)); + let (index, name, _) = index_segments(&[ + &[("a", &a), ("cafe2", &cafe2)], + &[("b", &b), ("cafe1", &cafe1), ("cafe3", &cafe3)], + ])?; + assert_eq!( + parse(&index, name, "+name:cafe3 +geometry:$knn(1, 0 0)^2")?, + ["cafe3"] + ); + assert_eq!( + parse(&index, name, "+name:cafe3 +(geometry:$knn(1, 0 0) name:b)")?, + ["cafe3"] + ); + // The nearest of cafe1 and cafe3 is cafe1, so cafe3 is left. + let cafes = "(name:cafe1 OR name:cafe3)"; + for query in [ + format!("+{cafes} -geometry:$knn(1, 0 0)"), + format!("+{cafes} -geometry:$knn(1, 0 0)^2"), + ] { + assert_eq!(parse(&index, name, &query)?, ["cafe3"], "{query}"); + } + Ok(()) + } + + #[test] + fn test_knn_without_a_searcher_keeps_the_nearest_of_each_segment() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (a, b, c, d) = (point(1.0), point(2.0), point(3.0), point(4.0)); + let (index, name, _) = index_segments(&[&[("a", &a), ("c", &c)], &[("b", &b), ("d", &d)]])?; + let parser = crate::query::QueryParser::for_index(&index, vec![name]); + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.delete_query(parser.parse_query("geometry:$knn(1, 0 0)")?)?; + writer.commit()?; + assert_eq!(parse(&index, name, "*")?, ["c", "d"]); + Ok(()) + } + + #[test] + fn test_polygons_with_no_area_and_spikes_are_found() -> crate::Result<()> { + let dot = r#"{"type": "Polygon", "coordinates": + [[[10, 20], [10, 20], [10, 20], [10, 20]]]}"#; + let sliver = + r#"{"type": "Polygon", "coordinates": [[[30, 0], [31, 1], [30, 0], [30, 0]]]}"#; + let spiked = r#"{"type": "Polygon", "coordinates": + [[[0, 0], [10, 0], [14, 0], [10, 0], [10, 10], [0, 10], [0, 0]]]}"#; + let (index, name, geometry) = + index_segments(&[&[("dot", dot), ("sliver", sliver), ("spiked", spiked)]])?; + assert_eq!( + matching(&index, name, &near(geometry, [10.0, 20.0]))?, + ["dot"] + ); + assert_eq!( + matching(&index, name, &near(geometry, [31.0, 1.0]))?, + ["sliver"] + ); + assert_eq!( + matching(&index, name, &near(geometry, [14.0, 0.0]))?, + ["spiked"] + ); + // The spike bounds no area: the polygon is the square. + assert_eq!( + matching(&index, name, &box_around(geometry, [5.0, 5.0]))?, + ["spiked"] + ); + assert!(matching(&index, name, &box_around(geometry, [12.0, 5.0]))?.is_empty()); + Ok(()) + } + + #[test] + fn test_a_polygon_with_no_rings_is_empty() -> crate::Result<()> { + let values = r#"[{"type": "Point", "coordinates": [1, 1]}, + {"type": "Polygon", "coordinates": []}]"#; + let (index, name, geometry) = index_segments(&[&[("point", values)]])?; + let around = SpatialQuery::from_bounds(geometry, [[0.0, 0.0], [2.0, 2.0]]); + assert_eq!(matching(&index, name, &around)?, ["point"]); + Ok(()) + } + + #[test] + fn test_a_self_crossing_ring_bounds_its_lobes() -> crate::Result<()> { + let bow_tie = r#"{"type": "Polygon", "coordinates": + [[[0, 0], [2, 2], [2, 0], [0, 2], [0, 0]]]}"#; + let (index, name, geometry) = index_segments(&[&[("bow_tie", bow_tie)]])?; + assert_eq!( + matching(&index, name, &near(geometry, [1.8, 1.0]))?, + ["bow_tie"] + ); + assert!(matching(&index, name, &near(geometry, [100.0, 50.0]))?.is_empty()); + assert!(matching(&index, name, &near(geometry, [-120.0, -60.0]))?.is_empty()); + Ok(()) + } + + #[test] + fn test_boundaries_meeting_inside_edges_are_contact() -> crate::Result<()> { + // a and b touch the box's boundary along parts of its edges and at T-junctions. + let a = r#"{"type": "Polygon", "coordinates": [[[0, 0], [5, 0], [5, 5], [0, 5], [0, 0]]]}"#; + let b = + r#"{"type": "Polygon", "coordinates": [[[5, 0], [10, 0], [10, 5], [5, 5], [5, 0]]]}"#; + let (index, name, _) = index_segments(&[&[("box", BOX), ("a", a), ("b", b)]])?; + assert_eq!( + parse(&index, name, "geometry:$within(0 0, 10 0, 10 10, 0 10)")?, + ["a", "b", "box"] + ); + assert_eq!( + parse( + &index, + name, + "geometry:$within(0 0, 5 0, 10 0, 10 10, 0 10, 0 5)" + )?, + ["a", "b", "box"] + ); + assert_eq!( + parse(&index, name, "geometry:$contains($query(name:a))")?, + ["box"] + ); + + let line = |coordinates: &str| { + format!(r#"{{"type": "LineString", "coordinates": {coordinates}}}"#) + }; + let (equator, south, overlap, bottom) = ( + line("[[0, 0], [10, 0]]"), + line("[[5, 0], [5, -5]]"), + line("[[2, 0], [5, 0]]"), + line("[[0, 0], [10, 0]]"), + ); + let (index, name, _) = index_segments(&[ + &[ + ("equator", &equator), + ("south", &south), + ("overlap", &overlap), + ], + &[("box", BOX), ("bottom", &bottom)], + ])?; + assert_eq!( + parse(&index, name, "geometry:$intersects($query(name:equator))")?, + ["bottom", "box", "overlap", "south"] + ); + // A line along a polygon's edge is within it, and intersects it. + assert_eq!( + parse(&index, name, "geometry:$intersects(0 0, 10 0, 10 10, 0 10)")?, + ["bottom", "box", "equator", "overlap", "south"] + ); + Ok(()) + } + + #[test] + fn test_oriented_query_rings_bound_more_than_a_hemisphere_in_contains() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let name = schema_builder.add_text_field("name", STRING | STORED); + let geometry = schema_builder.add_spatial_field("geometry", SPATIAL.set_oriented()); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + // Everything but a small box, and a box that covers the small one. + let all_but = + r#"{"type": "Polygon", "coordinates": [[[4, 4], [4, 6], [6, 6], [6, 4], [4, 4]]]}"#; + let big = r#"{"type": "Polygon", "coordinates": [[[-5, -5], [15, -5], [15, 15], [-5, 15], [-5, -5]]]}"#; + for (doc_name, doc_geometry) in [("all_but", all_but), ("big", big)] { + let json = format!(r#"{{"name": "{doc_name}", "geometry": {doc_geometry}}}"#); + writer.add_document(TantivyDocument::parse_json(&index.schema(), &json)?)?; + } + writer.commit()?; + let inside_box = SpatialQuery::new( + geometry, + vec![[0.0, 0.0], [10.0, 0.0], [10.0, 10.0], [0.0, 10.0]], + ); + assert!(matching(&index, name, &inside_box)?.is_empty()); + // The world but the box [0, 10]²: big covers part of the box. + assert!(parse(&index, name, "geometry:$within(0 0, 0 10, 10 10, 10 0)")?.is_empty()); + Ok(()) + } + + #[test] + fn test_knn_restriction_does_not_depend_on_how_the_query_is_written() -> crate::Result<()> { + let point = |lat: f64| format!(r#"{{"type": "Point", "coordinates": [0, {lat}]}}"#); + let (a, b, cafe3) = (point(1.0), point(2.0), point(5.0)); + let (index, name, _) = index_segments(&[&[("a", &a), ("b", &b), ("cafe3", &cafe3)]])?; + // The nearest cafe3 is cafe3, however the clause that requires it is written. + for query in [ + "+name:cafe3 -geometry:$knn(1, 0 0)", + "name:cafe3 -geometry:$knn(1, 0 0)", + ] { + assert!(parse(&index, name, query)?.is_empty(), "{query}"); + } + let nested = + "+(+name:cafe3 +geometry:$within(1000km, $query(name:a))) +geometry:$knn(1, 0 0)"; + assert_eq!(parse(&index, name, nested)?, ["cafe3"]); + let parser = crate::query::QueryParser::for_index(&index, vec![name]); + let (lenient, errors) = parser.parse_query_lenient(nested); + assert!(errors.is_empty()); + assert_eq!(matching(&index, name, &*lenient)?, ["cafe3"]); + Ok(()) + } + + fn point(lon: f64, lat: f64) -> String { + format!(r#"{{"type": "Point", "coordinates": [{lon}, {lat}]}}"#) + } + + fn line(positions: &[[f64; 2]]) -> String { + format!(r#"{{"type": "LineString", "coordinates": {positions:?}}}"#) + } + + #[test] + fn test_a_rectangle_follows_its_parallels() -> crate::Result<()> { + // At their middle, the geodesics through the corners run 7 degrees north of the + // parallels: the polygon through the corners holds what lies between the geodesics. + let (south, north) = (point(15.0, 62.0), point(15.0, 69.0)); + let (index, name, geometry) = index_segments(&[&[("south", &south), ("north", &north)]])?; + let bounds = [[-30.0, 60.0], [60.0, 65.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, bounds); + assert_eq!(matching(&index, name, &covered)?, ["south"]); + let intersects = SpatialQuery::intersects_rect(geometry, bounds); + assert_eq!(matching(&index, name, &intersects)?, ["south"]); + let polygon = SpatialQuery::from_bounds(geometry, bounds); + assert_eq!(matching(&index, name, &polygon)?, ["north"]); + Ok(()) + } + + #[test] + fn test_a_rectangle_is_as_wide_as_its_bounds() -> crate::Result<()> { + let docs = [ + ("greenwich", point(0.0, 0.0)), + ("east", point(100.0, 0.0)), + ("date_line", point(180.0, 0.0)), + ("west", point(-175.0, 0.0)), + ]; + let docs: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (*n, g.as_str())).collect(); + let (index, name, geometry) = index_segments(&[&docs[..2], &docs[2..]])?; + let within = |west: f64, east: f64| { + let query = SpatialQuery::covered_by_rect(geometry, [[west, -1.0], [east, 1.0]]); + matching(&index, name, &query) + }; + assert_eq!(within(-170.0, 170.0)?, ["east", "greenwich"]); + assert_eq!(within(170.0, -170.0)?, ["date_line", "west"]); + assert_eq!( + within(-180.0, 180.0)?, + ["date_line", "east", "greenwich", "west"] + ); + assert_eq!(within(0.0, 180.0)?, ["date_line", "east", "greenwich"]); + assert_eq!(within(100.0, 100.0)?, ["east"]); + // The antimeridian alone, however it is written. + for (west, east) in [(180.0, -180.0), (180.0, 180.0), (-180.0, -180.0)] { + assert_eq!(within(west, east)?, ["date_line"], "{west} to {east}"); + } + Ok(()) + } + + #[test] + fn test_a_rectangle_includes_its_boundary() -> crate::Result<()> { + let docs = [ + ("corner", point(10.0, 10.0)), + ("origin", point(0.0, 0.0)), + ("north", point(5.0, 10.0)), + ("east", point(10.0, 5.0)), + ("west", point(0.0, 5.0)), + ("beyond", point(10.000001, 5.0)), + ("above", point(5.0, 10.000001)), + ]; + let docs: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (*n, g.as_str())).collect(); + let (index, name, geometry) = index_segments(&[&docs])?; + let query = SpatialQuery::covered_by_rect(geometry, [[0.0, 0.0], [10.0, 10.0]]); + assert_eq!( + matching(&index, name, &query)?, + ["corner", "east", "north", "origin", "west"] + ); + // A rectangle with no area is the point or the line it is. + let at = SpatialQuery::intersects_rect(geometry, [[10.0, 5.0], [10.0, 5.0]]); + assert_eq!(matching(&index, name, &at)?, ["east"]); + let along = SpatialQuery::covered_by_rect(geometry, [[0.0, 10.0], [10.0, 10.0]]); + assert_eq!(matching(&index, name, &along)?, ["corner", "north"]); + Ok(()) + } + + #[test] + fn test_a_rectangle_covers_lines_and_polygons_by_their_edges() -> crate::Result<()> { + let polygon = |ring: &str| format!(r#"{{"type": "Polygon", "coordinates": [{ring}]}}"#); + let docs = [ + ("inside", line(&[[1.0, 1.0], [9.0, 9.0]])), + // Rises above the north of the rectangle between its ends, which are within it. + ("bowed", line(&[[2.0, 9.99], [8.0, 9.99]])), + ("leaving", line(&[[5.0, 5.0], [15.0, 5.0]])), + // Neither has an end in the rectangle: one crosses its meridians, one its parallels. + ("through", line(&[[-2.0, 9.9], [12.0, 9.9]])), + ("diagonal", line(&[[2.0, 12.0], [8.0, -2.0]])), + ("over", line(&[[2.0, 10.001], [8.0, 10.001]])), + ("outside", line(&[[11.0, 1.0], [12.0, 2.0]])), + ("small", polygon(®ular_ring(5.0, 5.0, 6, 2.0, false))), + // Its north edge is the geodesic between the rectangle's north corners. + ("same", BOX.to_string()), + ("around", polygon(®ular_ring(5.0, 5.0, 4, 20.0, false))), + ("far", polygon(®ular_ring(50.0, 5.0, 6, 2.0, false))), + ]; + let docs: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (*n, g.as_str())).collect(); + let (index, name, geometry) = index_segments(&[&docs[..5], &docs[5..]])?; + let bounds = [[0.0, 0.0], [10.0, 10.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, bounds); + assert_eq!(matching(&index, name, &covered)?, ["inside", "small"]); + let intersects = SpatialQuery::intersects_rect(geometry, bounds); + assert_eq!( + matching(&index, name, &intersects)?, + ["around", "bowed", "diagonal", "inside", "leaving", "same", "small", "through"] + ); + Ok(()) + } + + #[test] + fn test_a_rectangle_wider_than_a_hemisphere_leaves_out_what_an_edge_may_cross( + ) -> crate::Result<()> { + // Both run the short way between their ends: one across the antimeridian. + let across = line(&[[160.0, 0.0], [-160.0, 0.0]]); + let along = line(&[[-60.0, 0.0], [60.0, 0.0]]); + let (index, name, geometry) = index_segments(&[&[("across", &across), ("along", &along)]])?; + let bounds = [[-170.0, -1.0], [170.0, 1.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, bounds); + assert_eq!(matching(&index, name, &covered)?, ["along"]); + let intersects = SpatialQuery::intersects_rect(geometry, bounds); + assert_eq!(matching(&index, name, &intersects)?, ["across", "along"]); + Ok(()) + } + + #[test] + fn test_a_rectangle_to_a_pole_holds_the_pole() -> crate::Result<()> { + let docs = [ + ("pole", point(0.0, 90.0)), + ("near", point(45.0, 85.0)), + // Runs over the pole. + ("over", line(&[[0.0, 85.0], [180.0, 85.0]])), + ("far", point(45.0, 70.0)), + ]; + let docs: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (*n, g.as_str())).collect(); + let (index, name, geometry) = index_segments(&[&docs])?; + let cap = SpatialQuery::covered_by_rect(geometry, [[-180.0, 80.0], [180.0, 90.0]]); + assert_eq!(matching(&index, name, &cap)?, ["near", "over", "pole"]); + let wedge = [[0.0, 80.0], [90.0, 90.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, wedge); + assert_eq!(matching(&index, name, &covered)?, ["near", "pole"]); + let intersects = SpatialQuery::intersects_rect(geometry, wedge); + assert_eq!( + matching(&index, name, &intersects)?, + ["near", "over", "pole"] + ); + let world = SpatialQuery::covered_by_rect(geometry, [[-180.0, -90.0], [180.0, 90.0]]); + assert_eq!( + matching(&index, name, &world)?, + ["far", "near", "over", "pole"] + ); + Ok(()) + } + + #[test] + fn test_a_rectangle_does_not_cover_a_polygon_around_what_it_leaves_out() -> crate::Result<()> { + // The polygon's edges lie between 80 and 83 degrees north, around the pole. + let cap = r#"{"type": "Polygon", "coordinates": + [[[0, 80], [90, 80], [180, 80], [-90, 80], [0, 80]]]}"#; + let (index, name, geometry) = index_segments(&[&[("cap", cap)]])?; + let band = [[-180.0, 70.0], [180.0, 85.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, band); + assert!(matching(&index, name, &covered)?.is_empty()); + let intersects = SpatialQuery::intersects_rect(geometry, band); + assert_eq!(matching(&index, name, &intersects)?, ["cap"]); + let to_pole = SpatialQuery::covered_by_rect(geometry, [[-180.0, 70.0], [180.0, 90.0]]); + assert_eq!(matching(&index, name, &to_pole)?, ["cap"]); + // Within the polygon, the rectangle meets none of its edges. + let at_pole = SpatialQuery::intersects_rect(geometry, [[-180.0, 88.0], [180.0, 90.0]]); + assert_eq!(matching(&index, name, &at_pole)?, ["cap"]); + let south = SpatialQuery::intersects_rect(geometry, [[-180.0, 0.0], [180.0, 70.0]]); + assert!(matching(&index, name, &south)?.is_empty()); + Ok(()) + } + + #[test] + fn test_a_rectangle_needs_every_value_of_a_document() -> crate::Result<()> { + let split = format!("[{}, {}]", point(5.0, 5.0), point(50.0, 5.0)); + let both = format!("[{}, {}]", point(4.0, 5.0), point(6.0, 5.0)); + let ring = |lon: f64| regular_ring(lon, 5.0, 5, 1.0, false); + let multi = |lon: f64| { + format!( + r#"{{"type": "MultiPolygon", "coordinates": [[{}], [{}]]}}"#, + ring(3.0), + ring(lon) + ) + }; + let (apart, together) = (multi(30.0), multi(7.0)); + let (index, name, geometry) = index_segments(&[ + &[("split", &split), ("both", &both)], + &[("apart", &apart), ("together", &together)], + ])?; + let bounds = [[0.0, 0.0], [10.0, 10.0]]; + let covered = SpatialQuery::covered_by_rect(geometry, bounds); + assert_eq!(matching(&index, name, &covered)?, ["both", "together"]); + let intersects = SpatialQuery::intersects_rect(geometry, bounds); + assert_eq!( + matching(&index, name, &intersects)?, + ["apart", "both", "split", "together"] + ); + Ok(()) + } + + #[test] + fn test_rectangle_bounds_are_checked() -> crate::Result<()> { + let (index, _, geometry) = index_segments(&[&[("box", BOX)]])?; + let searcher = index.reader()?.searcher(); + for bounds in [ + [[0.0, 10.0], [10.0, 0.0]], + [[-190.0, 0.0], [10.0, 10.0]], + [[0.0, 0.0], [10.0, 91.0]], + [[0.0, f64::NAN], [10.0, 10.0]], + ] { + for query in [ + SpatialQuery::covered_by_rect(geometry, bounds), + SpatialQuery::intersects_rect(geometry, bounds), + ] { + let result = searcher.search(&query, &DocSetCollector); + assert!( + matches!(result, Err(crate::TantivyError::InvalidArgument(_))), + "{bounds:?}" + ); + } + } + Ok(()) + } + + /// The longitude and latitude of a unit vector, in degrees. + fn lon_lat(p: &[f64; 3]) -> [f64; 2] { + [ + p[1].atan2(p[0]).to_degrees(), + p[2].atan2(p[0].hypot(p[1])).to_degrees(), + ] + } + + /// Whether a rectangle given as the queries take it holds the longitude and latitude. + fn rect_holds([[west, south], [east, north]]: [[f64; 2]; 2], [lon, lat]: [f64; 2]) -> bool { + let between = if west <= east { + west <= lon && lon <= east + } else { + lon >= west || lon <= east + }; + between && south <= lat && lat <= north + } + + /// Points along the edges of a set's members, close enough together that the edges between + /// them leave a rectangle only by a hair. + fn sampled(set: &GeometrySet) -> Vec> { + let samples = 512; + set.members + .iter() + .map(|member| { + member_edges(member) + .iter() + .flat_map(|(a, b)| { + (0..=samples).map(move |i| { + let t = i as f64 / samples as f64; + let mix = |k: usize| a[k] * (1.0 - t) + b[k] * t; + let p = [mix(0), mix(1), mix(2)]; + let norm = (p[0] * p[0] + p[1] * p[1] + p[2] * p[2]).sqrt(); + lon_lat(&[p[0] / norm, p[1] / norm, p[2] / norm]) + }) + }) + .collect() + }) + .collect() + } + + #[test] + fn test_rectangles_agree_with_the_longitude_and_latitude_of_sampled_edges() -> crate::Result<()> + { + let mut rng = StdRng::seed_from_u64(7); + let mut docs: Vec<(String, String)> = Vec::new(); + for i in 0..240 { + let (lon, lat): (f64, f64) = ( + rng.random_range(-180.0..180.0), + rng.random_range(-80.0..80.0), + ); + let geojson = match i % 4 { + 0 => point(lon, lat), + 1 => { + let mut at = [lon, lat]; + let positions: Vec<[f64; 2]> = (0..rng.random_range(2..5)) + .map(|_| { + let next = at; + let lon: f64 = at[0] + rng.random_range(-50.0..50.0); + let lat: f64 = at[1] + rng.random_range(-20.0..20.0); + at = [ + (lon + 180.0).rem_euclid(360.0) - 180.0, + lat.clamp(-89.0, 89.0), + ]; + next + }) + .collect(); + line(&positions) + } + 2 => format!( + r#"{{"type": "Polygon", "coordinates": [{}]}}"#, + regular_ring( + lon.clamp(-160.0, 160.0), + lat.clamp(-60.0, 60.0), + rng.random_range(3..8), + rng.random_range(0.5..15.0), + false + ) + ), + _ => format!( + "[{}, {}]", + point(lon, lat), + point( + (lon + rng.random_range(0.0..40.0)).min(180.0), + (lat + rng.random_range(-10.0..10.0)).clamp(-89.0, 89.0) + ) + ), + }; + docs.push((format!("{i:03}"), geojson)); + } + // Every value of a document, each a set of members sampled along its edges. + let values: Vec> = docs + .iter() + .map(|(_, geojson)| match geojson.strip_prefix('[') { + Some(list) => list + .trim_end_matches(']') + .split("}, {") + .map(|value| { + let value = value.trim_start_matches('{').trim_end_matches('}'); + geometry_set(&format!("{{{value}}}")) + }) + .collect(), + None => vec![geometry_set(geojson)], + }) + .collect(); + let named: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (n.as_str(), g.as_str())).collect(); + let (index, name, geometry) = + index_segments(&[&named[..90], &named[90..170], &named[170..]])?; + + for round in 0..150 { + let (a, b): (f64, f64) = (rng.random_range(-90.0..90.0), rng.random_range(-90.0..90.0)); + let (mut west, mut east): (f64, f64) = ( + rng.random_range(-180.0..180.0), + rng.random_range(-180.0..180.0), + ); + match round % 5 { + // Every longitude, and rectangles a few degrees wide where geometries are. + 0 => (west, east) = (-180.0, 180.0), + 1 => east = (west + rng.random_range(1.0..60.0)).min(180.0), + _ => {} + } + let bounds = [[west, a.min(b)], [east, a.max(b)]]; + let mut covered = Vec::new(); + let mut intersecting = Vec::new(); + for ((doc_name, _), sets) in docs.iter().zip(&values) { + let members: Vec<(Vec<[f64; 2]>, &EdgeSet)> = sets + .iter() + .flat_map(|set| sampled(set).into_iter().zip(&set.members)) + .collect(); + // A polygon here is smaller than a hemisphere and away from the poles: the + // rectangle covers it when it covers its edges, and meets it when it meets its + // edges or lies within it. + let within = |points: &Vec<[f64; 2]>| points.iter().all(|p| rect_holds(bounds, *p)); + let meets = |(points, member): &(Vec<[f64; 2]>, &EdgeSet)| { + let middle = ::project( + if west <= east { + (west + east) / 2.0 + } else { + west + (east + 360.0 - west) / 2.0 + }, + (bounds[0][1] + bounds[1][1]) / 2.0, + ); + points.iter().any(|p| rect_holds(bounds, *p)) + || member_contains(member, &middle) + }; + if members.iter().all(|(points, _)| within(points)) { + covered.push(doc_name.clone()); + } + if members.iter().any(meets) { + intersecting.push(doc_name.clone()); + } + } + let query = SpatialQuery::covered_by_rect(geometry, bounds); + assert_eq!( + matching(&index, name, &query)?, + covered, + "covered by {bounds:?}" + ); + let query = SpatialQuery::intersects_rect(geometry, bounds); + assert_eq!( + matching(&index, name, &query)?, + intersecting, + "intersects {bounds:?}" + ); + } + Ok(()) + } + + #[test] + fn test_within_decides_as_the_distances_between_do() -> crate::Result<()> { + let mut rng = StdRng::seed_from_u64(11); + let mut docs: Vec<(String, String)> = Vec::new(); + for i in 0..300 { + let (lon, lat): (f64, f64) = ( + rng.random_range(-180.0..180.0), + rng.random_range(-85.0..85.0), + ); + let geojson = match i % 5 { + 0 | 1 => point(lon, lat), + 2 => line(&[ + [lon, lat], + [ + (lon + rng.random_range(-30.0..30.0)).clamp(-180.0, 180.0), + (lat + rng.random_range(-20.0..20.0)).clamp(-89.0, 89.0), + ], + ]), + 3 => format!( + r#"{{"type": "Polygon", "coordinates": [{}]}}"#, + regular_ring( + lon.clamp(-150.0, 150.0), + lat.clamp(-55.0, 55.0), + rng.random_range(3..8), + rng.random_range(0.5..25.0), + false + ) + ), + _ => format!( + "[{}, {}]", + point(lon, lat), + point(-lon, (lat + rng.random_range(-5.0..5.0)).clamp(-89.0, 89.0)) + ), + }; + docs.push((format!("{i:03}"), geojson)); + } + let named: Vec<(&str, &str)> = docs.iter().map(|(n, g)| (n.as_str(), g.as_str())).collect(); + let (index, name, geometry) = + index_segments(&[&named[..120], &named[120..200], &named[200..]])?; + let mut found = 0; + for round in 0..200 { + let center: [f64; 2] = [ + rng.random_range(-180.0..180.0), + rng.random_range(-90.0..90.0), + ]; + // From meters to more than half the way around. + let radius: f64 = 10f64.powf(rng.random_range(-6.0..0.6)); + let radius = if round % 20 == 0 { + f64::INFINITY + } else { + radius + }; + let within = SpatialPredicate::Within(radius); + let within = SpatialQuery::with_predicate(geometry, vec![center], within); + let between = SpatialPredicate::Between(0.0, radius); + let between = SpatialQuery::with_predicate(geometry, vec![center], between); + let expected = matching(&index, name, &between)?; + found += expected.len(); + assert_eq!( + matching(&index, name, &within)?, + expected, + "within {radius} of {center:?}" + ); + } + assert!(found > 5_000, "{found}"); + for radius in [0.0, -1.0] { + let within = SpatialPredicate::Within(radius); + let at = SpatialQuery::with_predicate(geometry, vec![[0.0, 0.0]], within); + assert!(matching(&index, name, &at)?.is_empty()); + } + Ok(()) + } +} diff --git a/src/schema/document/de.rs b/src/schema/document/de.rs index 02ee7a6a0c..f1a3ebee78 100644 --- a/src/schema/document/de.rs +++ b/src/schema/document/de.rs @@ -22,6 +22,8 @@ use super::se::BinaryObjectSerializer; use super::{OwnedValue, Value}; use crate::schema::document::type_codes; use crate::schema::{Facet, Field}; +use crate::spatial::geometry::Geometry; +use crate::spatial::plane::Plane; use crate::store::DocStoreVersion; use crate::tokenizer::PreTokenizedString; @@ -129,6 +131,9 @@ pub trait ValueDeserializer<'de> { /// Attempts to deserialize a pre-tokenized string value from the deserializer. fn deserialize_pre_tokenized_string(self) -> Result; + /// HUSH + fn deserialize_geometry(self) -> Result, DeserializeError>; + /// Attempts to deserialize the value using a given visitor. fn deserialize_any(self, visitor: V) -> Result where V: ValueVisitor; @@ -166,6 +171,8 @@ pub enum ValueType { /// A JSON object value. Deprecated. #[deprecated(note = "We keep this for backwards compatibility, use Object instead")] JSONObject, + /// HUSH + Geometry, } /// A value visitor for deserializing a document value. @@ -246,6 +253,12 @@ pub trait ValueVisitor { Err(DeserializeError::UnsupportedType(ValueType::PreTokStr)) } + #[inline] + /// Called when the deserializer visits a geometry value. + fn visit_geometry(&self, _val: Geometry) -> Result { + Err(DeserializeError::UnsupportedType(ValueType::Geometry)) + } + #[inline] /// Called when the deserializer visits an array. fn visit_array<'de, A>(&self, _access: A) -> Result @@ -380,6 +393,7 @@ where R: Read match ext_type_code { type_codes::TOK_STR_EXT_CODE => ValueType::PreTokStr, + type_codes::GEO_EXT_CODE => ValueType::Geometry, _ => { return Err(DeserializeError::from(io::Error::new( io::ErrorKind::InvalidData, @@ -495,6 +509,12 @@ where R: Read .map_err(DeserializeError::from) } + fn deserialize_geometry(self) -> Result, DeserializeError> { + self.validate_type(ValueType::Geometry)?; + as BinarySerializable>::deserialize(self.reader) + .map_err(DeserializeError::from) + } + fn deserialize_any(self, visitor: V) -> Result where V: ValueVisitor { match self.value_type { @@ -539,6 +559,10 @@ where R: Read let val = self.deserialize_pre_tokenized_string()?; visitor.visit_pre_tokenized_string(val) } + ValueType::Geometry => { + let val = self.deserialize_geometry()?; + visitor.visit_geometry(val) + } ValueType::Array => { let access = BinaryArrayDeserializer::from_reader(self.reader, self.doc_store_version)?; diff --git a/src/schema/document/default_document.rs b/src/schema/document/default_document.rs index de4e0e8244..4e02f3388d 100644 --- a/src/schema/document/default_document.rs +++ b/src/schema/document/default_document.rs @@ -13,6 +13,8 @@ use crate::schema::document::{ }; use crate::schema::field_type::ValueParsingError; use crate::schema::{Facet, Field, FieldType, NamedFieldDocument, OwnedValue, Schema}; +use crate::spatial::geometry::Geometry; +use crate::spatial::plane::Plane; use crate::tokenizer::PreTokenizedString; #[repr(C, packed)] @@ -266,6 +268,7 @@ impl CompactDoc { } ReferenceValueLeaf::IpAddr(num) => write_into(&mut self.node_data, num.to_u128()), ReferenceValueLeaf::PreTokStr(pre_tok) => write_into(&mut self.node_data, *pre_tok), + ReferenceValueLeaf::Geometry(geometry) => write_into(&mut self.node_data, *geometry), }; ValueAddr { type_id, val_addr } } @@ -480,6 +483,12 @@ impl<'a> CompactDocValue<'a> { .map(Into::into) .map(ReferenceValueLeaf::PreTokStr) .map(Into::into), + ValueType::Geometry => self + .container + .read_from::>(addr) + .map(Into::into) + .map(ReferenceValueLeaf::Geometry) + .map(Into::into), ValueType::Object => Ok(ReferenceValue::Object(CompactDocObjectIter::new( self.container, addr, @@ -560,6 +569,8 @@ pub enum ValueType { Array = 12, /// Opaque payload of a plugin-defined custom field. Custom = 13, + /// A GeoJSON geometry, for a spatial field. + Geometry = 14, } impl BinarySerializable for ValueType { @@ -606,6 +617,7 @@ impl<'a> From<&ReferenceValueLeaf<'a>> for ValueType { ReferenceValueLeaf::Facet(_) => ValueType::Facet, ReferenceValueLeaf::Bytes(_) => ValueType::Bytes, ReferenceValueLeaf::Custom(_) => ValueType::Custom, + ReferenceValueLeaf::Geometry(_) => ValueType::Geometry, } } } diff --git a/src/schema/document/mod.rs b/src/schema/document/mod.rs index e98716ba19..2532ebd1d9 100644 --- a/src/schema/document/mod.rs +++ b/src/schema/document/mod.rs @@ -290,4 +290,5 @@ pub(crate) mod type_codes { // Extended type codes pub const TOK_STR_EXT_CODE: u8 = 0; + pub const GEO_EXT_CODE: u8 = 1; } diff --git a/src/schema/document/owned_value.rs b/src/schema/document/owned_value.rs index c14553bffd..e489e205fb 100644 --- a/src/schema/document/owned_value.rs +++ b/src/schema/document/owned_value.rs @@ -15,6 +15,8 @@ use crate::schema::document::{ ValueDeserializer, ValueVisitor, }; use crate::schema::Facet; +use crate::spatial::geometry::Geometry; +use crate::spatial::plane::Plane; use crate::tokenizer::PreTokenizedString; use crate::DateTime; @@ -51,6 +53,8 @@ pub enum OwnedValue { IpAddr(Ipv6Addr), /// Opaque payload of a plugin-defined custom field. Custom(Vec), + /// A GeoJSON geometry, indexed by a spatial field. + Geometry(Geometry), } impl AsRef for OwnedValue { @@ -82,6 +86,7 @@ impl OwnedValue { OwnedValue::Object(_) => 11, OwnedValue::IpAddr(_) => 12, OwnedValue::Custom(_) => 13, + OwnedValue::Geometry(_) => 14, } } } @@ -106,6 +111,9 @@ impl<'a> Value<'a> for &'a OwnedValue { OwnedValue::Custom(val) => ReferenceValueLeaf::Custom(val).into(), OwnedValue::Array(array) => ReferenceValue::Array(array.iter()), OwnedValue::Object(object) => ReferenceValue::Object(ObjectMapIter(object.iter())), + OwnedValue::Geometry(geometry) => { + ReferenceValueLeaf::Geometry(Box::new(geometry.clone())).into() + } } } } @@ -165,6 +173,13 @@ impl ValueDeserialize for OwnedValue { Ok(OwnedValue::PreTokStr(val)) } + fn visit_geometry( + &self, + val: Geometry, + ) -> Result { + Ok(OwnedValue::Geometry(val)) + } + fn visit_array<'de, A>(&self, mut access: A) -> Result where A: ArrayAccess<'de> { let mut elements = Vec::with_capacity(access.size_hint()); @@ -231,6 +246,7 @@ impl serde::Serialize for OwnedValue { // round-trip through serde (they are populated programmatically), so this is lossy // by design. OwnedValue::Custom(ref bytes) => serializer.serialize_str(&BASE64.encode(bytes)), + OwnedValue::Geometry(ref geometry) => geometry.to_geojson().serialize(serializer), } } } @@ -319,6 +335,7 @@ impl<'a, V: Value<'a>> From> for OwnedValue { ReferenceValueLeaf::Bool(val) => OwnedValue::Bool(val), ReferenceValueLeaf::PreTokStr(val) => OwnedValue::PreTokStr(*val.clone()), ReferenceValueLeaf::Custom(val) => OwnedValue::Custom(val.to_vec()), + ReferenceValueLeaf::Geometry(val) => OwnedValue::Geometry(*val.clone()), }, ReferenceValue::Array(val) => { OwnedValue::Array(val.map(|v| v.as_value().into()).collect()) diff --git a/src/schema/document/se.rs b/src/schema/document/se.rs index d56c5c8917..18f3de894b 100644 --- a/src/schema/document/se.rs +++ b/src/schema/document/se.rs @@ -133,6 +133,10 @@ where W: Write self.write_type_code(type_codes::EXT_CODE)?; self.serialize_with_type_code(type_codes::TOK_STR_EXT_CODE, &*val) } + ReferenceValueLeaf::Geometry(val) => { + self.write_type_code(type_codes::EXT_CODE)?; + self.serialize_with_type_code(type_codes::GEO_EXT_CODE, &*val) + } // Custom fields are never stored (`is_stored() == false`), so the document // serializer never reaches this arm. Guard it explicitly rather than inventing an // on-disk type code for an opaque, plugin-owned payload. diff --git a/src/schema/document/value.rs b/src/schema/document/value.rs index 100fa151e7..509ea6e3c3 100644 --- a/src/schema/document/value.rs +++ b/src/schema/document/value.rs @@ -3,6 +3,8 @@ use std::net::Ipv6Addr; use common::DateTime; +use crate::spatial::geometry::Geometry; +use crate::spatial::plane::Plane; use crate::tokenizer::PreTokenizedString; /// A single field value. @@ -114,6 +116,12 @@ pub trait Value<'a>: Send + Sync + Debug { None } } + + #[inline] + /// HUSH + fn as_geometry(&self) -> Option>> { + self.as_leaf().and_then(|leaf| leaf.into_geometry()) + } } /// A enum representing a leaf value for tantivy to index. @@ -145,6 +153,8 @@ pub enum ReferenceValueLeaf<'a> { /// Opaque payload of a plugin-defined custom field. Distinct from [`Bytes`](Self::Bytes) so /// the built-ins leave it alone and a consuming plugin can recognize it. Custom(&'a [u8]), + /// A GeoJSON geometry, for a spatial field. + Geometry(Box>), } impl From for ReferenceValueLeaf<'_> { @@ -232,6 +242,9 @@ impl<'a, T: Value<'a> + ?Sized> From> for ReferenceValue< ReferenceValueLeaf::Custom(val) => { ReferenceValue::Leaf(ReferenceValueLeaf::Custom(val)) } + ReferenceValueLeaf::Geometry(val) => { + ReferenceValue::Leaf(ReferenceValueLeaf::Geometry(val)) + } } } } @@ -344,6 +357,16 @@ impl<'a> ReferenceValueLeaf<'a> { } } + #[inline] + /// If the Value is a geometry, returns it. Returns None otherwise. + pub fn into_geometry(self) -> Option>> { + if let Self::Geometry(val) = self { + Some(val) + } else { + None + } + } + #[inline] /// If the Value is a custom payload, returns its bytes. Returns None otherwise. pub fn as_custom(&self) -> Option<&'a [u8]> { @@ -476,4 +499,10 @@ where V: Value<'a> pub fn is_object(&self) -> bool { matches!(self, Self::Object(_)) } + + #[inline] + /// HUSH + pub fn into_geometry(self) -> Option>> { + self.into_leaf().and_then(|leaf| leaf.into_geometry()) + } } diff --git a/src/schema/field_entry.rs b/src/schema/field_entry.rs index f7816b218f..b910db7537 100644 --- a/src/schema/field_entry.rs +++ b/src/schema/field_entry.rs @@ -1,6 +1,7 @@ use serde::{Deserialize, Serialize}; use super::ip_options::IpAddrOptions; +use super::spatial_options::SpatialOptions; use crate::schema::bytes_options::BytesOptions; use crate::schema::custom_options::CustomOptions; use crate::schema::{ @@ -86,6 +87,11 @@ impl FieldEntry { Self::new(field_name, FieldType::JsonObject(json_object_options)) } + /// Creates a field entry for a spatial field + pub fn new_spatial(field_name: String, spatial_options: SpatialOptions) -> FieldEntry { + Self::new(field_name, FieldType::Spatial(spatial_options)) + } + /// Creates a field entry for a plugin-defined custom field. pub fn new_custom(field_name: String, custom_options: CustomOptions) -> FieldEntry { Self::new(field_name, FieldType::Custom(custom_options)) @@ -140,6 +146,7 @@ impl FieldEntry { FieldType::Bytes(ref options) => options.is_stored(), FieldType::JsonObject(ref options) => options.is_stored(), FieldType::IpAddr(ref options) => options.is_stored(), + FieldType::Spatial(ref options) => options.is_stored(), FieldType::TieBreaker | FieldType::Custom(_) => false, } } diff --git a/src/schema/field_type.rs b/src/schema/field_type.rs index 43b1f3e29e..3eb298da76 100644 --- a/src/schema/field_type.rs +++ b/src/schema/field_type.rs @@ -9,6 +9,7 @@ use serde_json::Value as JsonValue; use thiserror::Error; use super::ip_options::IpAddrOptions; +use super::spatial_options::SpatialOptions; use super::IntoIpv6Addr; use crate::schema::bytes_options::BytesOptions; use crate::schema::custom_options::CustomOptions; @@ -17,6 +18,7 @@ use crate::schema::{ DateOptions, Facet, IndexRecordOption, JsonObjectOptions, NumericOptions, OwnedValue, TextFieldIndexing, TextOptions, }; +use crate::spatial::geometry::Geometry; use crate::time::format_description::well_known::Rfc3339; use crate::time::OffsetDateTime; use crate::tokenizer::PreTokenizedString; @@ -81,6 +83,8 @@ pub enum Type { Json = b'j', /// IpAddr IpAddr = b'p', + /// Spatial + Spatial = b't', /// A plugin-defined custom type. Opaque to the built-ins; see [`CustomOptions`]. /// /// [`CustomOptions`]: crate::schema::CustomOptions @@ -154,6 +158,7 @@ impl Type { Type::Bytes => "Bytes", Type::Json => "Json", Type::IpAddr => "IpAddr", + Type::Spatial => "Spatial", Type::Custom => "Custom", } } @@ -210,6 +215,8 @@ pub enum FieldType { JsonObject(JsonObjectOptions), /// IpAddr field IpAddr(IpAddrOptions), + /// Spatial field + Spatial(SpatialOptions), /// A plugin-defined field type. Opaque to the built-ins; populated via /// [`add_custom`](crate::TantivyDocument::add_custom) and consumed by a /// [`SegmentPlugin`](crate::SegmentPlugin). @@ -230,6 +237,7 @@ impl FieldType { FieldType::Bytes(_) => Type::Bytes, FieldType::JsonObject(_) => Type::Json, FieldType::IpAddr(_) => Type::IpAddr, + FieldType::Spatial(_) => Type::Spatial, FieldType::Custom(_) => Type::Custom, } } @@ -277,6 +285,7 @@ impl FieldType { FieldType::Bytes(ref bytes_options) => bytes_options.is_indexed(), FieldType::JsonObject(ref json_object_options) => json_object_options.is_indexed(), FieldType::IpAddr(ref ip_addr_options) => ip_addr_options.is_indexed(), + FieldType::Spatial(ref _spatial_options) => false, FieldType::TieBreaker | FieldType::Custom(_) => false, } } @@ -315,6 +324,7 @@ impl FieldType { FieldType::IpAddr(ref ip_addr_options) => ip_addr_options.is_fast(), FieldType::Facet(_) => true, FieldType::JsonObject(ref json_object_options) => json_object_options.is_fast(), + FieldType::Spatial(_) => false, FieldType::TieBreaker => true, FieldType::Custom(_) => false, } @@ -336,6 +346,7 @@ impl FieldType { FieldType::Bytes(ref bytes_options) => bytes_options.fieldnorms(), FieldType::JsonObject(ref _json_object_options) => false, FieldType::IpAddr(ref ip_addr_options) => ip_addr_options.fieldnorms(), + FieldType::Spatial(_) => false, FieldType::TieBreaker | FieldType::Custom(_) => false, } } @@ -388,6 +399,8 @@ impl FieldType { None } } + FieldType::Spatial(_) => None, /* Geometry types cannot be indexed in the inverted + * index. */ FieldType::TieBreaker | FieldType::Custom(_) => None, } } @@ -491,6 +504,10 @@ impl FieldType { Ok(OwnedValue::IpAddr(ip_addr.into_ipv6_addr())) } + FieldType::Spatial(_) => Err(ValueParsingError::TypeError { + expected: "spatial field parsing not implemented", + json: JsonValue::String(field_text), + }), FieldType::TieBreaker => Err(ValueParsingError::TypeError { expected: "a generated tie-breaker field", json: JsonValue::String(field_text), @@ -554,6 +571,10 @@ impl FieldType { expected: "a string with an ip addr", json: JsonValue::Number(field_val_num), }), + FieldType::Spatial(_) => Err(ValueParsingError::TypeError { + expected: "spatial field parsing not implemented", + json: JsonValue::Number(field_val_num), + }), FieldType::TieBreaker => Err(ValueParsingError::TypeError { expected: "a generated tie-breaker field", json: JsonValue::Number(field_val_num), @@ -576,6 +597,14 @@ impl FieldType { } } FieldType::JsonObject(_) => Ok(OwnedValue::from(json_map)), + FieldType::Spatial(_) => Ok(OwnedValue::Geometry( + Geometry::from_geojson(&json_map).map_err(|e| { + ValueParsingError::ParseError { + error: format!("{:?}", e), + json: JsonValue::Object(json_map), + } + })?, + )), _ => Err(ValueParsingError::TypeError { expected: self.value_type().name(), json: JsonValue::Object(json_map), diff --git a/src/schema/flags.rs b/src/schema/flags.rs index 4f8caa612c..9b03b214d5 100644 --- a/src/schema/flags.rs +++ b/src/schema/flags.rs @@ -1,6 +1,6 @@ use std::ops::BitOr; -use crate::schema::{DateOptions, NumericOptions, TextOptions}; +use crate::schema::{DateOptions, NumericOptions, SpatialOptions, TextOptions}; #[derive(Clone)] pub struct StoredFlag; @@ -95,6 +95,14 @@ impl> BitOr for SchemaFlagList } } +impl> BitOr for SchemaFlagList { + type Output = SpatialOptions; + + fn bitor(self, rhs: SpatialOptions) -> Self::Output { + self.head.into() | rhs + } +} + #[derive(Clone)] pub struct SchemaFlagList { pub head: Head, diff --git a/src/schema/mod.rs b/src/schema/mod.rs index bfffba3653..e4d4198b65 100644 --- a/src/schema/mod.rs +++ b/src/schema/mod.rs @@ -129,6 +129,7 @@ mod ip_options; mod json_object_options; mod named_field_document; mod numeric_options; +mod spatial_options; mod text_options; use columnar::ColumnType; @@ -150,6 +151,7 @@ pub use self::json_object_options::JsonObjectOptions; pub use self::named_field_document::NamedFieldDocument; pub use self::numeric_options::NumericOptions; pub use self::schema::{Schema, SchemaBuilder}; +pub use self::spatial_options::{SpatialOptions, SPATIAL}; pub use self::term::{Term, ValueBytes}; pub use self::text_options::{TextFieldIndexing, TextOptions, STRING, TEXT}; @@ -174,6 +176,7 @@ pub(crate) fn value_type_to_column_type(typ: Type) -> Option { Type::Bytes => Some(ColumnType::Bytes), Type::IpAddr => Some(ColumnType::IpAddr), Type::Json => None, + Type::Spatial => None, Type::Custom => None, } } diff --git a/src/schema/schema.rs b/src/schema/schema.rs index 1b60d91ee6..794bb8a716 100644 --- a/src/schema/schema.rs +++ b/src/schema/schema.rs @@ -209,6 +209,16 @@ impl SchemaBuilder { self.add_field(field_entry) } + /// Adds a spatial entry to the schema in build. + pub fn add_spatial_field>( + &mut self, + field_name: &str, + field_options: T, + ) -> Field { + let field_entry = FieldEntry::new_spatial(field_name.to_string(), field_options.into()); + self.add_field(field_entry) + } + /// Adds a plugin-defined custom field to the schema. /// /// `type_name` identifies the custom type a plugin will consume (e.g. `"vec"`); `params` @@ -242,9 +252,14 @@ impl SchemaBuilder { /// Finalize the creation of a `Schema` /// This will consume your `SchemaBuilder` pub fn build(self) -> Schema { + let contains_spatial_field = self + .fields + .iter() + .any(|field_entry| field_entry.field_type().value_type() == Type::Spatial); Schema(Arc::new(InnerSchema { fields: self.fields, fields_map: self.fields_map, + contains_spatial_field, })) } } @@ -252,6 +267,7 @@ impl SchemaBuilder { struct InnerSchema { fields: Vec, fields_map: HashMap, // transient + contains_spatial_field: bool, } impl PartialEq for InnerSchema { @@ -402,6 +418,11 @@ impl Schema { } Some((field, json_path)) } + + /// Returns true if the schema contains a spatial field. + pub(crate) fn contains_spatial_field(&self) -> bool { + self.0.contains_spatial_field + } } impl Serialize for Schema { @@ -429,16 +450,16 @@ impl<'de> Deserialize<'de> for Schema { fn visit_seq(self, mut seq: A) -> Result where A: SeqAccess<'de> { - let mut schema = SchemaBuilder { + let mut schema_builder = SchemaBuilder { fields: Vec::with_capacity(seq.size_hint().unwrap_or(0)), fields_map: HashMap::with_capacity(seq.size_hint().unwrap_or(0)), }; while let Some(value) = seq.next_element()? { - schema.add_field(value); + schema_builder.add_field(value); } - Ok(schema.build()) + Ok(schema_builder.build()) } } @@ -1054,4 +1075,33 @@ mod tests { Some((default, "foobar")) ); } + + #[test] + fn test_contains_spatial_field() { + // No spatial field + { + let mut schema_builder = Schema::builder(); + schema_builder.add_text_field("title", TEXT); + let schema = schema_builder.build(); + assert!(!schema.contains_spatial_field()); + + // Serialization check + let schema_json = serde_json::to_string(&schema).unwrap(); + let schema_deserialized: Schema = serde_json::from_str(&schema_json).unwrap(); + assert!(!schema_deserialized.contains_spatial_field()); + } + // With spatial field + { + let mut schema_builder = Schema::builder(); + schema_builder.add_text_field("title", TEXT); + schema_builder.add_spatial_field("location", SPATIAL); + let schema = schema_builder.build(); + assert!(schema.contains_spatial_field()); + + // Serialization check + let schema_json = serde_json::to_string(&schema).unwrap(); + let schema_deserialized: Schema = serde_json::from_str(&schema_json).unwrap(); + assert!(schema_deserialized.contains_spatial_field()); + } + } } diff --git a/src/schema/spatial_options.rs b/src/schema/spatial_options.rs new file mode 100644 index 0000000000..e5c9b32f8c --- /dev/null +++ b/src/schema/spatial_options.rs @@ -0,0 +1,90 @@ +use std::ops::BitOr; + +use serde::{Deserialize, Serialize}; + +use crate::schema::flags::StoredFlag; +use crate::spatial::geometry_set::RingOrientation; + +/// Define how a spatial field should be handled by tantivy. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)] +pub struct SpatialOptions { + #[serde(default)] + stored: bool, + /// Whether polygon rings follow the right-hand rule of RFC 7946, rather than bounding the + /// smaller region. See [`RingOrientation`]. + #[serde(default)] + oriented: bool, +} + +/// The field will be untokenized and indexed. +pub const SPATIAL: SpatialOptions = SpatialOptions { + stored: false, + oriented: false, +}; + +impl SpatialOptions { + /// Returns true if the geometry is to be stored. + #[inline] + pub fn is_stored(&self) -> bool { + self.stored + } + + /// Returns true if polygon rings follow the right-hand rule of RFC 7946. + #[inline] + pub fn is_oriented(&self) -> bool { + self.oriented + } + + /// Have polygon rings follow the right-hand rule of RFC 7946: a polygon lies to the left of + /// its exterior ring and to the right of its holes, so that it may be larger than a + /// hemisphere. By default a ring bounds the smaller of the two regions it separates, + /// whichever way it runs. Query polygons on the field are read the same way. + #[must_use] + pub fn set_oriented(mut self) -> SpatialOptions { + self.oriented = true; + self + } + + /// How the field's polygon rings pick the region they bound. + pub fn ring_orientation(&self) -> RingOrientation { + if self.oriented { + RingOrientation::RightHand + } else { + RingOrientation::Smaller + } + } +} + +impl> BitOr for SpatialOptions { + type Output = SpatialOptions; + + fn bitor(self, other: T) -> SpatialOptions { + let other = other.into(); + SpatialOptions { + stored: self.stored | other.stored, + oriented: self.oriented | other.oriented, + } + } +} + +impl From for SpatialOptions { + fn from(_: StoredFlag) -> SpatialOptions { + SpatialOptions { + stored: true, + oriented: false, + } + } +} + +// #[cfg(test)] +// mod tests { +// use crate::schema::*; +// +// #[test] +// fn test_field_options() { +// let field_options = STORED | SPATIAL; +// assert!(field_options.is_stored()); +// let mut schema_builder = Schema::builder(); +// schema_builder.add_spatial_index("where", SPATIAL | STORED); +// } +// } diff --git a/src/schema/term.rs b/src/schema/term.rs index 407290bcd2..cb4a286576 100644 --- a/src/schema/term.rs +++ b/src/schema/term.rs @@ -514,6 +514,9 @@ where B: AsRef<[u8]> Type::IpAddr => { write_opt(f, self.as_ip_addr())?; } + Type::Spatial => { + write!(f, "")?; + } // Custom fields produce no terms, so no term ever carries this type. Type::Custom => unreachable!("the term does not support custom field types"), } diff --git a/src/space_usage/mod.rs b/src/space_usage/mod.rs index ab98d4824f..c55a5d57ff 100644 --- a/src/space_usage/mod.rs +++ b/src/space_usage/mod.rs @@ -28,6 +28,9 @@ pub(crate) const FAST_FIELDS: &str = "fast_fields"; pub(crate) const FIELDNORMS: &str = "fieldnorms"; pub(crate) const STORE: &str = "store"; pub(crate) const DELETES: &str = "deletes"; +/// Keyed by the extension of the spatial plugin's file, as `component` looks custom ones up. +pub(crate) const SPATIAL_CELLS: &str = "cells"; +pub(crate) const SPATIAL_EDGES: &str = "edges"; static EMPTY_PER_FIELD: LazyLock = LazyLock::new(|| PerFieldSpaceUsage::new(Vec::new())); @@ -179,6 +182,22 @@ impl SegmentSpaceUsage { } } + /// Space usage for spatial cell indexes + pub fn spatial_cells(&self) -> &PerFieldSpaceUsage { + self.per_field(SPATIAL_CELLS) + } + + /// Space usage for spatial cell indexes, as older versions reported it. + #[deprecated(note = "use `spatial_cells`, and `spatial_edges` for the edge indexes")] + pub fn spatial(&self) -> &PerFieldSpaceUsage { + self.spatial_cells() + } + + /// Space usage for spatial edge indexes + pub fn spatial_edges(&self) -> &PerFieldSpaceUsage { + self.per_field(SPATIAL_EDGES) + } + /// Space usage for stored documents pub fn store(&self) -> &StoreSpaceUsage { match self.components.get(STORE) { @@ -383,9 +402,9 @@ fn merge_column_space_usage( #[cfg(test)] mod test { use crate::index::Index; - use crate::schema::{Schema, FAST, INDEXED, STORED, TEXT}; - use crate::space_usage::PerFieldSpaceUsage; - use crate::{IndexWriter, Term}; + use crate::schema::{Schema, FAST, INDEXED, SPATIAL, STORED, TEXT}; + use crate::space_usage::{ComponentSpaceUsage, PerFieldSpaceUsage}; + use crate::{IndexWriter, TantivyDocument, Term}; #[test] fn test_empty() { @@ -453,6 +472,41 @@ mod test { Ok(()) } + #[test] + fn test_spatial() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + schema_builder.add_spatial_field("geometry", SPATIAL); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema.clone()); + + { + let mut index_writer: IndexWriter = index.writer_for_tests()?; + let line = r#"{"geometry": {"type": "LineString", "coordinates": [[0, 0], [10, 0]]}}"#; + index_writer.add_document(TantivyDocument::parse_json(&schema, line)?)?; + index_writer.commit()?; + } + + let reader = index.reader()?; + let searcher = reader.searcher(); + let searcher_space_usage = searcher.space_usage()?; + let segment = &searcher_space_usage.segments()[0]; + + expect_single_field(segment.spatial_cells(), "geometry", 1, 512); + // The edge index holds at least the line's two vertices, three f64 each. + expect_single_field(segment.spatial_edges(), "geometry", 48, 512); + let ComponentSpaceUsage::PerField(edges) = + segment.component(crate::spatial::plugin::edges_component()) + else { + panic!("the edge index is stored per field"); + }; + assert_eq!(edges.total(), segment.spatial_edges().total()); + assert!( + segment.total() >= segment.spatial_cells().total() + segment.spatial_edges().total() + ); + + Ok(()) + } + #[test] fn test_text() -> crate::Result<()> { let mut schema_builder = Schema::builder(); diff --git a/src/spatial/area_query.rs b/src/spatial/area_query.rs new file mode 100644 index 0000000000..340129b36b --- /dev/null +++ b/src/spatial/area_query.rs @@ -0,0 +1,132 @@ +//! Search for the geometries that meet an area of the sphere that is not a polygon: a rectangle +//! of latitude and longitude, or what lies within a distance of a point. +//! +//! A covering of the area finds the candidates. A covering cell within the area decides what it +//! holds without decoding it: whatever lies in the cell has a point in the area. A point, most +//! often the whole of its document, is read from its entry in the edge index and nothing else. +//! Only lines and polygons in the cells that the area's boundary crosses are decoded and tested +//! edge by edge. + +use std::collections::HashSet; + +use common::BitSet; + +use super::cell_index_reader::CellIndexReader; +use super::edge_cache::EdgeCache; +use super::geometry_set::EdgeSet; +use super::region::Region; +use super::region_coverer::{CovererOptions, RegionCoverer}; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; +use super::shape_index_region::{contains_point, SegmentIndex}; +use super::sphere::Sphere; + +/// Cells in the covering of an area. More cells follow its boundary more closely, and leave +/// fewer geometries outside it to read and reject. +const COVERING_CELLS: i32 = 64; + +/// A connected area of the sphere that tests points and edges against itself. +pub(crate) trait Area { + /// Whether the area contains the point. + fn contains(&self, point: &[f64; 3]) -> bool; + + /// Whether the area and the geodesic edge between two points outside it have a point in + /// common. + fn meets_edge(&self, a: &[f64; 3], b: &[f64; 3]) -> bool; + + /// A point of the area. A polygon that contains it, and has no edge that meets the area, + /// contains the area. + fn inside(&self) -> &[f64; 3]; +} + +/// The cells of a covering of an area, and whether the area contains each. +pub(crate) struct Covering(Vec<(S2CellId, bool)>); + +impl Covering { + /// Cover a region. `within` tells whether the area contains a cell: it may say no of a cell + /// that the area does contain, whose geometries are then tested one by one. + pub(crate) fn new(region: &R, within: impl Fn(&S2Cell) -> bool) -> Self { + let coverer = RegionCoverer::new(CovererOptions::default().max_cells(COVERING_CELLS)); + let cells = coverer.get_covering(region).into_cell_ids(); + Covering( + cells + .into_iter() + .map(|id| (id, within(&S2Cell::new(id)))) + .collect(), + ) + } + + /// Calls `visit` with the position of each geometry in the segment's cells that the covering + /// overlaps, and whether the cell lies in a covering cell within the area. A geometry in + /// several cells is visited once for each. + pub(crate) fn visit(&self, cell_reader: &CellIndexReader, mut visit: impl FnMut(u32, bool)) { + for &(covering_cell, within) in &self.0 { + cell_reader.visit_range(covering_cell, |index_cell, position| { + visit(position, within && covering_cell.contains(index_cell)) + }); + } + } +} + +/// The edges of a member: consecutive vertices of each ring that differ. +pub(crate) fn edges(member: &EdgeSet) -> impl Iterator { + member.ring_offsets.windows(2).flat_map(move |ring| { + (ring[0]..ring[1].saturating_sub(1)) + .map(move |i| (&member.vertices[i], &member.vertices[i + 1])) + .filter(|(a, b)| a != b) + }) +} + +/// Add the documents of one segment that have a point in common with the area to `docs`. A +/// document meets the area when any member of its set does. +pub(crate) fn search_meeting( + area: &A, + covering: &Covering, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, +) { + // Lines and polygons tested, each of which may lie in many cells. + let mut tested: HashSet = HashSet::new(); + covering.visit(cell_reader, |position, within| { + let summary = edge_cache.summary((0, position)); + if docs.contains(summary.doc_id) { + return; + } + let meets = if within { + true + } else if let Some(point) = summary.point { + area.contains(&point) + } else { + tested.insert(position) && member_meets(area, position, cell_reader, edge_cache) + }; + if meets { + docs.insert(summary.doc_id); + } + }); +} + +/// Whether the area and a line or polygon have a point in common: a vertex within the area, an +/// edge that meets it, or for a polygon the area within it. +fn member_meets( + area: &A, + position: u32, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, +) -> bool { + let entry = edge_cache.get((0, position)); + let member = entry.edge_set(); + if member.vertices.iter().any(|v| area.contains(v)) + || edges(member).any(|(a, b)| area.meets_edge(a, b)) + { + return true; + } + if !member.closed { + return false; + } + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + contains_point(&mut segment, (0, position), area.inside()) +} diff --git a/src/spatial/cell_index.rs b/src/spatial/cell_index.rs new file mode 100644 index 0000000000..7646badc94 --- /dev/null +++ b/src/spatial/cell_index.rs @@ -0,0 +1,968 @@ +//! Spatial index for geometry on the unit sphere. +//! +//! Maps S2CellIds to the shapes that intersect each cell. Cells are subdivided until no cell +//! contains more than a configurable number of edges. +use std::io; +use std::io::Write; + +use common::{write_u32_vint, CountingWriter}; + +use super::crossings::{ilogb, S2EdgeCrosser}; +use super::geometry_set::GeometrySet; +use super::math::normalize; +use super::r1interval::R1Interval; +use super::r2rect::R2Rect; +use super::s2cell_id::S2CellId; +use super::s2coords::{face_uv_to_xyz, get_face, NUM_FACES}; +use super::s2edge_clipping::{ + clip_to_padded_face, EDGE_CLIP_ERROR_UV_COORD, FACE_CLIP_ERROR_UV_COORD, +}; +use super::s2padded_cell::S2PaddedCell; +use crate::spatial::s2coords::valid_face_xyz_to_uv; + +/// The amount by which cells are "padded" to compensate for numerical errors +/// when clipping line segments to cell boundaries. +pub const CELL_PADDING: f64 = 2.0 * (FACE_CLIP_ERROR_UV_COORD + EDGE_CLIP_ERROR_UV_COORD); + +/// Default maximum number of edges per cell (not counting "long" edges). Reasonable values range +/// from 10 to about 50 or so. +pub const DEFAULT_MAX_EDGES_PER_CELL: usize = 10; + +/// Bytes in one cell index directory entry: the cell id (u64) and the cell's offset (u64). +pub(crate) const DIRECTORY_ENTRY_LEN: usize = 2 * std::mem::size_of::(); + +/// Bytes in the cell index footer: the cell count (u32) and the directory offset (u64). +pub(crate) const FOOTER_LEN: usize = std::mem::size_of::() + std::mem::size_of::(); + +/// Average edge length metric derivative for quadratic projection. +/// From s2metrics.cc: kAvgEdge for S2_QUADRATIC_PROJECTION +pub(crate) const K_AVG_EDGE_DERIV: f64 = 1.459_213_746_386_106_1; + +/// Default ratio of cell size to edge length for "long" edge classification. +/// From FLAGS_s2shape_index_cell_size_to_long_edge_ratio default value. +pub(crate) const CELL_SIZE_TO_LONG_EDGE_RATIO: f64 = 1.0; + +/// Returns the minimum level such that the average edge metric is at most the given value. +/// Port of S2::Metric<1>::GetLevelForMaxValue from s2metrics.h. +pub(crate) fn get_level_for_max_value(value: f64) -> i32 { + // Catches non-positive values, including NaN. + if value <= 0.0 || value.is_nan() { + return S2CellId::MAX_LEVEL; + } + + // This code is equivalent to computing a floating-point "level" value and rounding up. ilogb() + // returns the exponent corresponding to a fraction in the range [1,2), exactly, where + // log2().floor() may round up just below a power of two. + let level = ilogb(value / K_AVG_EDGE_DERIV); + // For dim=1, (level >> (dim - 1)) == (level >> 0) == level + (-level).clamp(0, S2CellId::MAX_LEVEL) +} + +/// Computes the maximum level at which an edge is "short" relative to the cell size. Above this +/// level the edge is "long" and doesn't count toward the subdivision threshold. Port of +/// MutableS2ShapeIndex::GetEdgeMaxLevel. +pub(crate) fn get_edge_max_level(v0: &[f64; 3], v1: &[f64; 3]) -> i32 { + let length = + ((v0[0] - v1[0]).powi(2) + (v0[1] - v1[1]).powi(2) + (v0[2] - v1[2]).powi(2)).sqrt(); + let max_cell_edge = length * CELL_SIZE_TO_LONG_EDGE_RATIO; + get_level_for_max_value(max_cell_edge) +} + +/// Represents the part of a shape that intersects an S2Cell. +/// +/// It consists of the set of edge ids that intersect that cell, and a boolean indicating whether +/// the center of the cell is inside the shape (for shapes that have an interior). +/// +/// Identifies a geometry member within a segment or across segments during merge. The first +/// element is the segment index (0 for single-segment operations). The second is the geometry +/// position within that segment's edge index. +pub type GeometryId = (u16, u32); + +/// Note that the edges themselves are not clipped; we always use the original edges for +/// intersection tests so that the results will be the same as the original shape. +#[derive(Clone, Debug, PartialEq)] +pub struct ClippedShape { + /// Identifies the geometry member. (segment, position) during merge, (0, position) otherwise. + pub geometry_id: GeometryId, + /// True if the geometry contains the center of the cell. + pub contains_center: bool, + /// Indicies of the edges stored per geometry id in the edge index. + pub edge_indices: Vec, +} + +/// Stores the index contents for a particular S2CellId. +/// +/// It consists of a set of clipped shapes. +#[derive(Clone, Debug, PartialEq)] +pub struct IndexCell { + /// Cell id. + pub cell_id: S2CellId, + /// Shapes stored in the cell. + pub shapes: Vec, +} + +impl IndexCell { + /// Creates a new empty IndexCell for the given cell ID. + pub fn new(cell_id: S2CellId) -> Self { + Self { + cell_id, + shapes: Vec::new(), + } + } + + /// Adds a clipped shape to this cell. + pub fn add_shape(&mut self, shape: ClippedShape) { + self.shapes.push(shape); + } + + /// Returns the clipped shape corresponding to the given geometry ID, or None if the geometry + /// does not intersect this cell. + pub fn find_shape(&self, geometry_id: GeometryId) -> Option<&ClippedShape> { + self.shapes.iter().find(|s| s.geometry_id == geometry_id) + } + + /// Returns the number of clipped shapes in this cell. + #[inline] + pub fn num_shapes(&self) -> usize { + self.shapes.len() + } + + /// Returns true if this cell contains no clipped shapes. + #[inline] + pub fn is_empty(&self) -> bool { + self.shapes.is_empty() + } +} + +impl ClippedShape { + /// Creates a new ClippedShape with the given geometry ID and contains_center flag. + pub fn new(geometry_id: GeometryId, contains_center: bool) -> Self { + Self { + geometry_id, + contains_center, + edge_indices: Vec::new(), + } + } + + /// Adds an edge index to this clipped shape. + pub fn add_edge(&mut self, edge_id: u32) { + self.edge_indices.push(edge_id); + } + + /// Returns the number of edges that intersect the S2CellId. + #[inline] + pub fn num_edges(&self) -> usize { + self.edge_indices.len() + } +} + +/// A spatial index mapping S2CellIds to the shapes that intersect each cell. +#[derive(Clone, Debug, Default)] +pub struct CellIndex { + /// Cells in the index. + pub cells: Vec, +} + +/// An edge projected onto a cube face. +#[derive(Clone, Debug)] +struct FaceEdge { + /// Geometry the edge belongs to. + geometry_id: u32, + /// Edge index in the original polygon. + edge_id: u32, + /// Maximum level for subdivision (based on edge length). + max_level: i32, + /// Determine if we are indexing a polygon. + has_interior: bool, + /// Edge endpoints in (u,v) coordinates. + a: [f64; 2], + b: [f64; 2], + /// Original edge endpoints in 3D. + v0: [f64; 3], + v1: [f64; 3], +} + +/// A clipped portion of an edge within a cell. +#[derive(Clone, Copy, Debug)] +struct ClippedEdge { + /// Original unclipped edge. + face_edge_index: usize, + /// Bounding box of the clipped edge in (u,v) coordinates. + bound: R2Rect, +} + +/// Tracks which cells' centers are inside the polygon as we traverse in S2CellId order. +/// +/// `is_active` is used for incremental updates in S2, so it can go from true to false. Here it +/// only ever goes from false to true if any of the indexed geometries is a polygon. If there are +/// no polygons to index, it will disable the tracker. It is more for purity of the port than for +/// optimization. +/// +/// The C++ MutableS2ShapeIndex uses a sorted vector rather than a set for tracking contained +/// shapes. The S2 authors benchmarked this and found sorted vector faster for the typical case of +/// 0-2 shapes. This set of contained shapes does not stack deep at any single point, it is the set +/// of overlapping polygons at a certian point on the face. This holds for bulk load as well as +/// incremental update. +/// +/// The C++ uses linear search. We use binary_search because the function is there and reads better +/// than a loop. For 0-2 elements the difference is negligible. +struct InteriorTracker { + /// Determine if any shape with an interior is being tracked. + is_active: bool, + /// Current focus point. + focus: [f64; 3], + /// Edge crosser for testing edge crossings. + crosser: S2EdgeCrosser, + /// Shapes whose interiors contain the current focus point. + shape_ids: Vec, + /// The next cell ID we expect to process. + next_cellid: S2CellId, +} + +impl InteriorTracker { + fn new() -> Self { + let origin = normalize(&face_uv_to_xyz(0, -1.0, -1.0)); + Self { + focus: origin, + crosser: S2EdgeCrosser::new(&origin, &origin), + shape_ids: Vec::new(), + is_active: false, + next_cellid: S2CellId::begin(S2CellId::MAX_LEVEL), + } + } + + fn toggle_shape(&mut self, geometry_id: u32) { + match self.shape_ids.binary_search(&geometry_id) { + Ok(i) => { + self.shape_ids.remove(i); + } + Err(i) => { + self.shape_ids.insert(i, geometry_id); + } + } + } + + fn test_edge(&mut self, geometry_id: u32, v0: &[f64; 3], v1: &[f64; 3]) { + if self.crosser.edge_or_vertex_crossing_two(v0, v1) { + self.toggle_shape(geometry_id); + } + } + + fn shape_ids(&self) -> &[u32] { + &self.shape_ids + } + + #[inline] + fn at_cellid(&self, id: S2CellId) -> bool { + id.range_min() == self.next_cellid + } + + #[inline] + fn set_next_cellid(&mut self, id: S2CellId) { + self.next_cellid = id.range_min(); + } + + #[inline] + fn move_to(&mut self, p: &[f64; 3]) { + self.focus = *p; + } + + fn draw_to(&mut self, p: &[f64; 3]) { + self.crosser = S2EdgeCrosser::new(&self.focus, p); + self.focus = *p; + } +} + +/// Options that affect construction of the CellIndex. +#[derive(Clone, Debug)] +pub struct BuildOptions { + /// Maximum edges per cell before subdivision. + pub max_edges_per_cell: usize, + /// Minimum fraction of 'short' edges required for subdivision. If this parameter is non-zero + /// then the total index size and construction time are guaranteed to be linear in the number + /// of input edges. Default: 0.2 (from FLAGS_s2shape_index_min_short_edge_fraction) + pub min_short_edge_fraction: f64, +} + +impl Default for BuildOptions { + fn default() -> Self { + Self { + max_edges_per_cell: DEFAULT_MAX_EDGES_PER_CELL, + min_short_edge_fraction: 0.2, + } + } +} + +impl BuildOptions { + /// Creates new build options with default values. + pub fn new() -> Self { + Self::default() + } + + /// Sets the maximum number of edges per cell. + /// + /// If a cell has more than this many edges that are not considered "long" relative to the cell + /// size, then it is subdivided. + pub fn with_max_edges_per_cell(mut self, max: usize) -> Self { + self.max_edges_per_cell = max; + self + } + + /// Sets the minimum short edge fraction. + pub fn with_min_short_edge_fraction(mut self, fraction: f64) -> Self { + self.min_short_edge_fraction = fraction; + self + } +} + +impl CellIndex { + /// Creates a new empty CellIndex. + pub fn new() -> Self { + Self { cells: Vec::new() } + } + + /// Returns true if the index contains no cells. + #[inline] + pub fn is_empty(&self) -> bool { + self.cells.is_empty() + } + + /// Returns the number of cells in the index. + #[inline] + pub fn num_cells(&self) -> usize { + self.cells.len() + } + + /// Returns the index cell containing the given cell ID, or None if not found. + pub fn find_cell(&self, target: S2CellId) -> Option<&IndexCell> { + let mut lo = 0; + let mut hi = self.cells.len(); + + while lo < hi { + let mid = lo + (hi - lo) / 2; + let cell = &self.cells[mid]; + + if target < cell.cell_id.range_min() { + hi = mid; + } else if target > cell.cell_id.range_max() { + lo = mid + 1; + } else { + return Some(cell); + } + } + + None + } + + /// Returns the index cells whose ranges overlap the target cell's range: the cell containing + /// it, or the cells it contains. + pub fn cells_overlapping(&self, target: S2CellId) -> &[IndexCell] { + let start = self + .cells + .partition_point(|c| c.cell_id.range_max() < target.range_min()); + let end = self + .cells + .partition_point(|c| c.cell_id.range_min() <= target.range_max()); + &self.cells[start..end.max(start)] + } + + /// Writes the cell index to a CountingWriter with a dictionary at the end. + /// Offsets are recorded relative to the start of this write, not the global + /// stream position, because CompositeFile shares one CountingWriter across + /// all fields but the reader sees a per-field slice starting at zero. + pub fn write(&self, write: &mut CountingWriter) -> io::Result<()> { + let base = write.written_bytes(); + let mut offsets: Vec<(u64, u64)> = Vec::with_capacity(self.cells.len()); + + for cell in &self.cells { + offsets.push((cell.cell_id.0, write.written_bytes() - base)); + write_cell(write, cell, |shape| shape.geometry_id.1)?; + } + + let dir_offset = write.written_bytes() - base; + write_directory(write, &offsets, dir_offset) + } +} + +/// Writes one cell: its id, its shape count, then each shape's geometry id (as given by +/// `geometry_id`), contains_center flag, edge count and edge ids. Counts and edge ids are +/// variable length, so that neither a crowded cell nor a geometry with many vertices overflows +/// them. +pub(crate) fn write_cell( + write: &mut W, + cell: &IndexCell, + geometry_id: impl Fn(&ClippedShape) -> u32, +) -> io::Result<()> { + write.write_all(&cell.cell_id.0.to_le_bytes())?; + write_u32_vint(cell.shapes.len() as u32, write)?; + for shape in &cell.shapes { + write.write_all(&geometry_id(shape).to_le_bytes())?; + write.write_all(&[shape.contains_center as u8])?; + write_u32_vint(shape.edge_indices.len() as u32, write)?; + for &edge_id in &shape.edge_indices { + write_u32_vint(edge_id, write)?; + } + } + Ok(()) +} + +/// Writes the directory of a cell index, one `(cell_id, offset)` entry per cell in cell order, +/// followed by the footer that locates it. `dir_offset` is where the directory starts, relative to +/// the start of the cell index. +pub(crate) fn write_directory( + write: &mut W, + offsets: &[(u64, u64)], + dir_offset: u64, +) -> io::Result<()> { + for &(cell_id, offset) in offsets { + write.write_all(&cell_id.to_le_bytes())?; + write.write_all(&offset.to_le_bytes())?; + } + write.write_all(&(offsets.len() as u32).to_le_bytes())?; + write.write_all(&dir_offset.to_le_bytes())?; + super::write_format_marker(write, super::CELL_INDEX_MAGIC) +} + +/// Where a target cell lies relative to the cells of an index. Mirrors the cases of S2's +/// ShapeIndexIterator::Locate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CellRelation { + /// Target does not overlap any index cell. + Disjoint, + /// Target contains one or more index cells. The usize is the first overlapping cell's + /// position in the index. + Subdivided(usize), + /// Target is within (or equals) an index cell. The usize is that cell's position in the + /// index. + Indexed(usize), +} + +/// Builder for constructing a CellIndex from a collection of geometries. +pub struct IndexBuilder { + options: BuildOptions, + /// Tracks which cells' centers are inside the polygon as we traverse in S2CellId order. + tracker: InteriorTracker, + /// The index as an array of cells. + cells: Vec, +} + +impl IndexBuilder { + /// Creates a new IndexBuilder with the given options. + pub fn new(options: BuildOptions) -> Self { + Self { + options, + tracker: InteriorTracker::new(), + cells: Vec::new(), + } + } + + /// Build a CellIndex from smashed GeometrySets. Geometry IDs are arrival order — the first + /// member of the first set is 0, and the count increments from there. + pub fn build_from_sets(mut self, sets: &[GeometrySet]) -> CellIndex { + if sets.is_empty() { + return CellIndex { cells: Vec::new() }; + } + + // Seed the InteriorTracker from precomputed contains_hilbert_start flags. + let mut gid: u32 = 0; + for set in sets { + for member in &set.members { + if member.closed { + self.tracker.is_active = true; + if member.contains_hilbert_start { + self.tracker.toggle_shape(gid); + } + } + gid += 1; + } + } + + let mut face_edges: [Vec; 6] = Default::default(); + + gid = 0; + for set in sets { + for member in &set.members { + let closed = member.closed; + let vertices = &member.vertices; + let offsets = &member.ring_offsets; + let mut edge_id: u32 = 0; + + for ring_idx in 0..offsets.len() - 1 { + let ring_start = offsets[ring_idx]; + let ring_end = offsets[ring_idx + 1]; + let ring_len = ring_end - ring_start; + + if ring_len < 2 { + if ring_len == 1 { + let v = vertices[ring_start]; + let max_level = self.get_edge_max_level(&v, &v); + self.add_face_edge( + gid, + edge_id, + max_level, + false, + &v, + &v, + &mut face_edges, + ); + edge_id += 1; + } + continue; + } + + let edge_count = ring_len - 1; + for i in 0..edge_count { + let v0 = vertices[ring_start + i]; + let v1 = vertices[ring_start + i + 1]; + let max_level = self.get_edge_max_level(&v0, &v1); + self.add_face_edge( + gid, + edge_id, + max_level, + closed, + &v0, + &v1, + &mut face_edges, + ); + edge_id += 1; + } + + if closed { + edge_id += 1; + } + } + + gid += 1; + } + } + + for face in 0..NUM_FACES { + self.update_face_edges(face, &face_edges[face as usize]); + } + + self.cells.sort_by_key(|c| c.cell_id); + + CellIndex { cells: self.cells } + } + + /// Delegates to the free function `get_edge_max_level`. + fn get_edge_max_level(&self, v0: &[f64; 3], v1: &[f64; 3]) -> i32 { + get_edge_max_level(v0, v1) + } + + /// Clips an edge to all cube faces. + #[allow(clippy::too_many_arguments)] + fn add_face_edge( + &self, + geometry_id: u32, + edge_id: u32, + max_level: i32, + has_interior: bool, + v0: &[f64; 3], + v1: &[f64; 3], + face_edges: &mut [Vec; 6], + ) { + let a_face = get_face(v0); + let b_face = get_face(v1); + + if a_face == b_face { + let (a_u, a_v) = valid_face_xyz_to_uv(a_face, v0); + let (b_u, b_v) = valid_face_xyz_to_uv(a_face, v1); + let max_uv = 1.0 - CELL_PADDING; + if a_u.abs() <= max_uv + && a_v.abs() <= max_uv + && b_u.abs() <= max_uv + && b_v.abs() <= max_uv + { + face_edges[a_face as usize].push(FaceEdge { + geometry_id, + edge_id, + max_level, + has_interior, + a: [a_u, a_v], + b: [b_u, b_v], + v0: *v0, + v1: *v1, + }); + return; + } + } + + for face in 0..NUM_FACES { + if let Some((a_uv, b_uv)) = clip_to_padded_face(v0, v1, face, CELL_PADDING) { + face_edges[face as usize].push(FaceEdge { + geometry_id, + edge_id, + max_level, + has_interior, + a: a_uv, + b: b_uv, + v0: *v0, + v1: *v1, + }); + } + } + } + + fn update_face_edges(&mut self, face: i32, face_edges: &[FaceEdge]) { + // A face with no edges still needs a cell if shapes contain it. + if face_edges.is_empty() && self.tracker.shape_ids().is_empty() { + return; + } + + let mut clipped: Vec = Vec::with_capacity(face_edges.len()); + let mut bound = R2Rect::empty(); + + for (i, fe) in face_edges.iter().enumerate() { + let edge_bound = R2Rect::from_point_pair(fe.a, fe.b); + clipped.push(ClippedEdge { + face_edge_index: i, + bound: edge_bound, + }); + bound = bound.union(&edge_bound); + } + + let face_id = S2CellId::from_face(face); + let pcell = S2PaddedCell::new(face_id, CELL_PADDING); + + // With no edges, there is nothing to shrink to. + let shrunk_id = if face_edges.is_empty() { + face_id + } else { + pcell.shrink_to_fit(&bound) + }; + if shrunk_id != face_id { + self.skip_cell_range(face_id.range_min(), shrunk_id.range_min()); + + let pcell = S2PaddedCell::new(shrunk_id, CELL_PADDING); + self.update_edges(&pcell, &mut clipped, face_edges); + + self.skip_cell_range(shrunk_id.range_max().next(), face_id.range_max().next()); + } else { + self.update_edges(&pcell, &mut clipped, face_edges); + } + } + + fn skip_cell_range(&mut self, begin: S2CellId, end: S2CellId) { + if self.tracker.shape_ids().is_empty() { + return; + } + + for cell_id in cell_union_from_range(begin, end) { + let pcell = S2PaddedCell::new(cell_id, CELL_PADDING); + let mut empty_edges: Vec = Vec::new(); + self.update_edges(&pcell, &mut empty_edges, &[]); + } + } + + fn update_edges( + &mut self, + pcell: &S2PaddedCell, + edges: &mut [ClippedEdge], + face_edges: &[FaceEdge], + ) { + if edges.is_empty() && self.tracker.shape_ids().is_empty() { + return; + } + + if self.should_subdivide(pcell, edges, face_edges, &self.tracker) { + let middle = pcell.middle(); + + let mut child_edges: [[Vec; 2]; 2] = Default::default(); + for i in 0..2 { + for j in 0..2 { + child_edges[i][j].reserve(edges.len()); + } + } + + for edge in edges.iter() { + let fe = &face_edges[edge.face_edge_index]; + clip_to_children(&edge.bound, fe.a, fe.b, &middle, |i, j, bound| { + child_edges[i][j].push(ClippedEdge { bound, ..*edge }); + }); + } + + for pos in 0..4 { + let (i, j) = pcell.get_child_ij(pos); + + let dominated_by_geometry_with_no_edges = + !self.tracker.shape_ids().is_empty() && child_edges[i][j].is_empty(); + + if !child_edges[i][j].is_empty() || dominated_by_geometry_with_no_edges { + let child_pcell = S2PaddedCell::from_parent(pcell, i, j); + self.update_edges(&child_pcell, &mut child_edges[i][j], face_edges); + } + } + } else { + self.make_index_cell(pcell, edges, face_edges); + } + } + + /// Returns true if the cell should be subdivided. + fn should_subdivide( + &self, + pcell: &S2PaddedCell, + edges: &[ClippedEdge], + face_edges: &[FaceEdge], + tracker: &InteriorTracker, + ) -> bool { + if edges.len() <= self.options.max_edges_per_cell { + return false; + } + + // Compute the threshold for short edges. This ensures linear index size. The formula is: + // max(max_edges_per_cell, min_short_edge_fraction * (edges + shapes)) + let max_short_edges = self.options.max_edges_per_cell.max( + (self.options.min_short_edge_fraction + * (edges.len() + tracker.shape_ids().len()) as f64) as usize, + ); + + // Count short edges (those that could benefit from subdivision). + let mut short_count = 0; + for edge in edges { + let fe = &face_edges[edge.face_edge_index]; + if pcell.level() < fe.max_level { + short_count += 1; + // Early exit: if we exceed max_short_edges, keep subdividing + if short_count > max_short_edges { + return true; + } + } + } + + // If short_count <= max_short_edges, stop subdividing (make a cell) + false + } + + fn make_index_cell( + &mut self, + pcell: &S2PaddedCell, + edges: &[ClippedEdge], + face_edges: &[FaceEdge], + ) { + let cell_id = pcell.id(); + + if self.tracker.is_active && !edges.is_empty() { + if !self.tracker.at_cellid(cell_id) { + self.tracker.move_to(&pcell.get_entry_vertex()); + } + self.tracker.draw_to(&pcell.get_center()); + + for edge in edges { + let fe = &face_edges[edge.face_edge_index]; + if fe.has_interior { + self.tracker.test_edge(fe.geometry_id, &fe.v0, &fe.v1); + } + } + } + + let containing_ids: Vec = self.tracker.shape_ids().to_vec(); + + if edges.is_empty() && containing_ids.is_empty() { + return; + } + + let mut cell = IndexCell::new(cell_id); + + let mut edge_idx = 0; + let mut containing_iter = containing_ids.iter().peekable(); + + while edge_idx < edges.len() || containing_iter.peek().is_some() { + let edge_geo_id = if edge_idx < edges.len() { + face_edges[edges[edge_idx].face_edge_index].geometry_id + } else { + u32::MAX + }; + + let containing_geo_id = containing_iter.peek().copied().copied().unwrap_or(u32::MAX); + + if containing_geo_id < edge_geo_id { + cell.add_shape(ClippedShape::new((0, containing_geo_id), true)); + containing_iter.next(); + } else { + let geo_id = edge_geo_id; + let mut shape = ClippedShape::new((0, geo_id), false); + + while edge_idx < edges.len() + && face_edges[edges[edge_idx].face_edge_index].geometry_id == geo_id + { + shape.add_edge(face_edges[edges[edge_idx].face_edge_index].edge_id); + edge_idx += 1; + } + + if containing_geo_id == geo_id { + shape.contains_center = true; + containing_iter.next(); + } + + cell.add_shape(shape); + } + } + + self.cells.push(cell); + + if self.tracker.is_active && !edges.is_empty() { + self.tracker.draw_to(&pcell.get_exit_vertex()); + for edge in edges { + let fe = &face_edges[edge.face_edge_index]; + if fe.has_interior { + self.tracker.test_edge(fe.geometry_id, &fe.v0, &fe.v1); + } + } + self.tracker.set_next_cellid(cell_id.next()); + } + } +} + +/// Distributes an edge among the children of a padded cell. The edge runs from `a` to `b` in +/// (u,v) coordinates, `bound` bounds the part of it in the cell, and `middle` is the cell's middle, +/// the region that belongs to all four children. Calls `add` with each child's (i, j) the edge +/// lies in, and the bound of the part of the edge in it. Port of +/// MutableS2ShapeIndex::ClipVAxis and the subdivision of UpdateEdges. +pub(crate) fn clip_to_children( + bound: &R2Rect, + a: [f64; 2], + b: [f64; 2], + middle: &R2Rect, + mut add: impl FnMut(usize, usize, R2Rect), +) { + let (u_mid_lo, u_mid_hi) = (middle[0].lo(), middle[0].hi()); + let (v_mid_lo, v_mid_hi) = (middle[1].lo(), middle[1].hi()); + let mut clip_v_axis = |i: usize, bound: R2Rect| { + if bound[1].hi() <= v_mid_lo { + // The edge is entirely in the lower child. + add(i, 0, bound); + } else if bound[1].lo() >= v_mid_hi { + // The edge is entirely in the upper child. + add(i, 1, bound); + } else { + // The edge spans both children. + add(i, 0, clip_v_bound(&bound, a, b, v_mid_hi, true)); + add(i, 1, clip_v_bound(&bound, a, b, v_mid_lo, false)); + } + }; + if bound[0].hi() <= u_mid_lo { + // The edge is entirely in the left children. + clip_v_axis(0, *bound); + } else if bound[0].lo() >= u_mid_hi { + // The edge is entirely in the right children. + clip_v_axis(1, *bound); + } else { + // The edge spans the left and the right children. + clip_v_axis(0, clip_u_bound(bound, a, b, u_mid_hi, true)); + clip_v_axis(1, clip_u_bound(bound, a, b, u_mid_lo, false)); + } +} + +/// Clips the u-axis of the bound of an edge from `a` to `b` to `u`: its high end if `clip_hi`, +/// otherwise its low end. Shrinks the v-axis to match. Port of MutableS2ShapeIndex::ClipUBound. +fn clip_u_bound(bound: &R2Rect, a: [f64; 2], b: [f64; 2], u: f64, clip_hi: bool) -> R2Rect { + // Check if clipping is actually needed. + if clip_hi { + if bound[0].hi() <= u { + return *bound; + } + } else if bound[0].lo() >= u { + return *bound; + } + + // Interpolate to find the v-coordinate where the edge crosses u, and update the end of the + // v-bound the slope of the edge determines. + let v = bound[1].clamp(interpolate_double(u, a[0], b[0], a[1], b[1])); + let positive_slope = (a[0] > b[0]) == (a[1] > b[1]); + let u_bound = if clip_hi { + R1Interval::new(bound[0].lo(), u) + } else { + R1Interval::new(u, bound[0].hi()) + }; + let v_bound = if clip_hi == positive_slope { + R1Interval::new(bound[1].lo(), v) + } else { + R1Interval::new(v, bound[1].hi()) + }; + R2Rect::new(u_bound, v_bound) +} + +/// Clips the v-axis of the bound of an edge from `a` to `b` to `v`: its high end if `clip_hi`, +/// otherwise its low end. Shrinks the u-axis to match. Port of MutableS2ShapeIndex::ClipVBound. +fn clip_v_bound(bound: &R2Rect, a: [f64; 2], b: [f64; 2], v: f64, clip_hi: bool) -> R2Rect { + // Check if clipping is actually needed. + if clip_hi { + if bound[1].hi() <= v { + return *bound; + } + } else if bound[1].lo() >= v { + return *bound; + } + + // Interpolate to find the u-coordinate where the edge crosses v, and update the end of the + // u-bound the slope of the edge determines. + let u = bound[0].clamp(interpolate_double(v, a[1], b[1], a[0], b[0])); + let positive_slope = (a[0] > b[0]) == (a[1] > b[1]); + let u_bound = if clip_hi == positive_slope { + R1Interval::new(bound[0].lo(), u) + } else { + R1Interval::new(u, bound[0].hi()) + }; + let v_bound = if clip_hi { + R1Interval::new(bound[1].lo(), v) + } else { + R1Interval::new(v, bound[1].hi()) + }; + R2Rect::new(u_bound, v_bound) +} + +/// Interpolates a value along a line segment. Given points (x0, y0) and (x1, y1), returns the +/// y-value at the given x. +/// +/// This function makes the following guarantees: +/// - If x == x0, then result == y0 (exactly). +/// - If x == x1, then result == y1 (exactly). +/// - If x0 <= x <= x1 and y0 <= y1, then y0 <= result <= y1. +/// - More generally, if x is between x0 and x1, then result is between y0 and y1. +/// +/// Port of S2::InterpolateDouble from s2edge_clipping.h. +#[inline] +fn interpolate_double(x: f64, x0: f64, x1: f64, y0: f64, y1: f64) -> f64 { + // If x0 == x1 == x all we can return is the single point. + if x0 == x1 { + // C++ asserts: x == x0 && y0 == y1 + debug_assert!(x == x0 && y0 == y1); + return y0; + } + + // To get results that are accurate near both endpoints, we interpolate + // starting from the closer of the two points. + if (x0 - x).abs() <= (x1 - x).abs() { + y0 + (y1 - y0) * ((x - x0) / (x1 - x0)) + } else { + y1 + (y0 - y1) * ((x - x1) / (x0 - x1)) + } +} + +fn cell_union_from_range(begin: S2CellId, end: S2CellId) -> Vec { + let mut result = Vec::new(); + let mut id = begin; + + while id < end { + // Find largest cell that fits. + let mut level = id.level(); + + while level > 0 { + let parent = id.parent(level - 1); + if parent.range_min() != id || parent.range_max().next() > end { + break; + } + level -= 1; + } + + if level < id.level() { + id = id.parent(level); + } + + result.push(id); + id = id.range_max().next(); + } + + result +} diff --git a/src/spatial/cell_index_reader.rs b/src/spatial/cell_index_reader.rs new file mode 100644 index 0000000000..b38ceaf7ab --- /dev/null +++ b/src/spatial/cell_index_reader.rs @@ -0,0 +1,415 @@ +//! Cell index deserialization. +//! +//! Reads a serialized cell index from a byte slice. The footer at the end locates the dictionary, +//! and the dictionary provides random access by cell_id. The iterator walks cells in sorted order +//! for merge. + +use common::read_u32_vint_no_advance; + +use crate::error::DataCorruption; +use crate::spatial::cell_index::{ + CellRelation, ClippedShape, IndexCell, DIRECTORY_ENTRY_LEN, FOOTER_LEN, +}; +use crate::spatial::s2cell_id::S2CellId; + +/// Reads a serialized cell index from a byte slice. +pub struct CellIndexReader<'a> { + data: &'a [u8], + cell_count: u32, + dir_offset: u64, +} + +impl<'a> CellIndexReader<'a> { + /// Opens a cell index reader from the raw bytes of a serialized cell index. Fails when the + /// footer does not locate a directory that ends where the footer begins. + pub fn open(data: &'a [u8]) -> crate::Result { + let data = super::check_format_marker(data, super::CELL_INDEX_MAGIC, "cell index")?; + let corrupt = + || DataCorruption::comment_only(format!("Corrupt {}-byte cell index", data.len())); + let (body, footer) = data.split_last_chunk::().ok_or_else(corrupt)?; + let (cell_count, dir_offset) = footer.split_at(std::mem::size_of::()); + let cell_count = u32::from_le_bytes(cell_count.try_into().unwrap()); + let dir_offset = u64::from_le_bytes(dir_offset.try_into().unwrap()); + let dir_len = cell_count as u64 * DIRECTORY_ENTRY_LEN as u64; + if dir_offset.checked_add(dir_len) != Some(body.len() as u64) { + return Err(corrupt().into()); + } + Ok(CellIndexReader { + data, + cell_count, + dir_offset, + }) + } + + /// A reader over no cells, for a segment that has no geometry for the field. + pub fn empty() -> Self { + CellIndexReader { + data: &[], + cell_count: 0, + dir_offset: 0, + } + } + + /// Number of cells in the index. + pub fn cell_count(&self) -> u32 { + self.cell_count + } + + /// Finds the cell containing the given cell_id using the S2 range containment test. + pub fn find(&self, target: S2CellId) -> Option { + let mut lo = 0u32; + let mut hi = self.cell_count; + + while lo < hi { + let mid = lo + (hi - lo) / 2; + let (cell_id, _) = self.read_dir_entry(mid); + let cell_id = S2CellId(cell_id); + + if target < cell_id.range_min() { + hi = mid; + } else if target > cell_id.range_max() { + lo = mid + 1; + } else { + let (_, offset) = self.read_dir_entry(mid); + return Some(self.read_cell(offset as usize)); + } + } + + None + } + + /// Returns an iterator over all cells in sorted order. + pub fn iter(&'a self) -> CellIndexIter<'a> { + CellIndexIter { + reader: self, + pos: 0, + end_at: None, + } + } + + /// Returns an iterator over all index cells whose ranges overlap the given target cell's + /// range. This is the candidate collection step: for a covering cell, find all index cells + /// that might contain relevant geometry. + /// + /// Binary search the dictionary for the first cell whose range_max is at or past the target's + /// range_min, then yield cells forward until past the target's range_max. + pub fn scan_range(&'a self, target: S2CellId) -> CellIndexIter<'a> { + CellIndexIter { + reader: self, + pos: self.first_ending_at_or_after(target.range_min()), + end_at: Some(target.range_max()), + } + } + + /// Calls `visit` with the id of each index cell `scan_range` would yield for the target and + /// the position of each geometry in it, without decoding the cells' edge indices. + pub fn visit_range(&self, target: S2CellId, mut visit: impl FnMut(S2CellId, u32)) { + let end = target.range_max(); + let mut pos = self.first_ending_at_or_after(target.range_min()); + while pos < self.cell_count { + let (cell_id, offset) = self.read_dir_entry(pos); + let cell_id = S2CellId(cell_id); + if cell_id.range_min() > end { + break; + } + self.visit_cell(offset as usize, |position| visit(cell_id, position)); + pos += 1; + } + } + + /// Calls `visit` with the position of each geometry in the cell at the given offset. + fn visit_cell(&self, offset: usize, mut visit: impl FnMut(u32)) { + let mut pos = offset + std::mem::size_of::(); + let read_vint = |pos: &mut usize| { + let (value, len) = read_u32_vint_no_advance(&self.data[*pos..]); + *pos += len; + value + }; + for _ in 0..read_vint(&mut pos) { + visit(u32::from_le_bytes( + self.data[pos..pos + 4].try_into().unwrap(), + )); + // The geometry position and the flag for the cell's center precede the edge indices. + pos += 5; + for _ in 0..read_vint(&mut pos) { + read_vint(&mut pos); + } + } + } + + /// Locates a target cell relative to the index cells from the dictionary alone, without + /// decoding any cell. + pub fn locate(&self, target: S2CellId) -> CellRelation { + let pos = self.first_ending_at_or_after(target.range_min()); + if pos >= self.cell_count { + return CellRelation::Disjoint; + } + let (cell_id, _) = self.read_dir_entry(pos); + let cell_id = S2CellId(cell_id); + if cell_id.contains(target) { + CellRelation::Indexed(pos as usize) + } else if target.contains(cell_id) { + CellRelation::Subdivided(pos as usize) + } else { + CellRelation::Disjoint + } + } + + /// Decodes the cell at the given position in cell order. + pub fn cell(&self, position: usize) -> IndexCell { + let (_, offset) = self.read_dir_entry(position as u32); + self.read_cell(offset as usize) + } + + /// Binary search the dictionary for the position of the first cell whose range_max is at or + /// past the given cell id. + fn first_ending_at_or_after(&self, target_min: S2CellId) -> u32 { + let mut lo = 0u32; + let mut hi = self.cell_count; + while lo < hi { + let mid = lo + (hi - lo) / 2; + let (cell_id, _) = self.read_dir_entry(mid); + if S2CellId(cell_id).range_max() < target_min { + lo = mid + 1; + } else { + hi = mid; + } + } + lo + } + + fn read_dir_entry(&self, index: u32) -> (u64, u64) { + let base = self.dir_offset as usize + index as usize * DIRECTORY_ENTRY_LEN; + let cell_id = u64::from_le_bytes(self.data[base..base + 8].try_into().unwrap()); + let offset = u64::from_le_bytes(self.data[base + 8..base + 16].try_into().unwrap()); + (cell_id, offset) + } + + fn read_cell(&self, offset: usize) -> IndexCell { + let mut pos = offset; + + let cell_id = u64::from_le_bytes(self.data[pos..pos + 8].try_into().unwrap()); + pos += 8; + + let read_vint = |pos: &mut usize| { + let (value, len) = read_u32_vint_no_advance(&self.data[*pos..]); + *pos += len; + value as usize + }; + + let shape_count = read_vint(&mut pos); + + // A count read from a corrupt index must not size an allocation beyond what the bytes + // left could hold: a shape takes at least 6 bytes, an edge id at least 1. + let mut shapes = Vec::with_capacity(shape_count.min((self.data.len() - pos) / 6)); + for _ in 0..shape_count { + let geometry_id = u32::from_le_bytes(self.data[pos..pos + 4].try_into().unwrap()); + pos += 4; + + let contains_center = self.data[pos] != 0; + pos += 1; + + let edge_count = read_vint(&mut pos); + let mut edge_indices = Vec::with_capacity(edge_count.min(self.data.len() - pos)); + for _ in 0..edge_count { + edge_indices.push(read_vint(&mut pos) as u32); + } + + shapes.push(ClippedShape { + geometry_id: (0, geometry_id), + contains_center, + edge_indices, + }); + } + + IndexCell { + cell_id: S2CellId(cell_id), + shapes, + } + } +} + +/// Iterates cells in sorted order, optionally bounded by a range_max. +pub struct CellIndexIter<'a> { + reader: &'a CellIndexReader<'a>, + pos: u32, + /// When set, stop yielding cells whose range_min exceeds this value. + end_at: Option, +} + +impl<'a> Iterator for CellIndexIter<'a> { + type Item = IndexCell; + + fn next(&mut self) -> Option { + if self.pos >= self.reader.cell_count { + return None; + } + let (cell_id, offset) = self.reader.read_dir_entry(self.pos); + if let Some(max) = self.end_at { + if S2CellId(cell_id).range_min() > max { + return None; + } + } + self.pos += 1; + Some(self.reader.read_cell(offset as usize)) + } +} + +#[cfg(test)] +mod tests { + use common::CountingWriter; + + use super::*; + use crate::spatial::cell_index::CellIndex; + + fn written(cell_index: &CellIndex) -> Vec { + let mut write = CountingWriter::wrap(Vec::new()); + cell_index.write(&mut write).unwrap(); + write.finish() + } + + fn cell(face: i32, geometry_id: u32) -> IndexCell { + let mut cell = IndexCell::new(S2CellId::from_face(face)); + cell.add_shape(ClippedShape::new((0, geometry_id), true)); + cell + } + + #[test] + fn test_empty_reads_no_cells() { + let reader = CellIndexReader::empty(); + assert_eq!(reader.cell_count(), 0); + assert!(reader.iter().next().is_none()); + assert!(reader.find(S2CellId::from_face(0)).is_none()); + assert_eq!( + reader.locate(S2CellId::from_face(0)), + CellRelation::Disjoint + ); + } + + #[test] + fn test_visit_range_finds_what_scan_range_does() -> crate::Result<()> { + let mut cells = Vec::new(); + for (child, face_cell) in S2CellId::from_face(1).children().into_iter().enumerate() { + let mut cell = IndexCell::new(face_cell); + for shape in 0..child as u32 + 1 { + let mut clipped = ClippedShape::new((0, 10 * child as u32 + shape), shape == 0); + clipped.edge_indices = (0..shape * 200).collect(); + cell.add_shape(clipped); + } + cells.push(cell); + } + cells.push(cell(4, 99)); + let data = written(&CellIndex { cells }); + let reader = CellIndexReader::open(&data)?; + let face = S2CellId::from_face(1); + for target in [ + face, + face.children()[2], + S2CellId::from_face(4), + S2CellId::from_face(5), + ] { + let scanned: Vec<(S2CellId, u32)> = reader + .scan_range(target) + .flat_map(|cell| { + let id = cell.cell_id; + cell.shapes + .into_iter() + .map(move |shape| (id, shape.geometry_id.1)) + }) + .collect(); + let mut visited = Vec::new(); + reader.visit_range(target, |cell_id, position| { + visited.push((cell_id, position)) + }); + assert_eq!(visited, scanned); + } + Ok(()) + } + + #[test] + fn test_open_reads_a_written_index() -> crate::Result<()> { + let data = written(&CellIndex { + cells: vec![cell(1, 7), cell(4, 9)], + }); + let reader = CellIndexReader::open(&data)?; + assert_eq!(reader.cell_count(), 2); + let cells: Vec = reader.iter().collect(); + assert_eq!(cells, vec![cell(1, 7), cell(4, 9)]); + assert_eq!(reader.cell(1), cell(4, 9)); + + let footer_only = written(&CellIndex::new()); + assert_eq!( + footer_only.len(), + FOOTER_LEN + crate::spatial::FORMAT_MARKER_LEN + ); + assert_eq!(CellIndexReader::open(&footer_only)?.cell_count(), 0); + Ok(()) + } + + #[test] + fn test_counts_and_edge_ids_beyond_a_u16_round_trip() -> crate::Result<()> { + let mut crowded = IndexCell::new(S2CellId::from_face(2)); + for geometry_id in 0..70_000 { + crowded.add_shape(ClippedShape::new((0, geometry_id), geometry_id % 2 == 0)); + } + let mut long = ClippedShape::new((0, 3), false); + long.edge_indices = (0..70_000).chain([u32::MAX]).collect(); + crowded.add_shape(long); + let data = written(&CellIndex { + cells: vec![cell(1, 7), crowded.clone(), cell(4, 9)], + }); + let reader = CellIndexReader::open(&data)?; + let cells: Vec = reader.iter().collect(); + assert_eq!(cells, vec![cell(1, 7), crowded, cell(4, 9)]); + Ok(()) + } + + #[test] + fn test_open_rejects_a_corrupt_index() { + let data = written(&CellIndex { + cells: vec![cell(1, 7)], + }); + for len in 0..FOOTER_LEN + crate::spatial::FORMAT_MARKER_LEN { + assert!(CellIndexReader::open(&data[..len]).is_err(), "{len} bytes"); + } + assert!(CellIndexReader::open(&data[1..]).is_err()); + let mut longer = data.clone(); + longer.insert(0, 0); + assert!(CellIndexReader::open(&longer).is_err()); + } + + #[test] + fn test_open_names_an_index_of_another_format() { + let data = written(&CellIndex { + cells: vec![cell(1, 7)], + }); + // Without the marker, as older versions wrote it, and with another version. + let unmarked = &data[..data.len() - crate::spatial::FORMAT_MARKER_LEN]; + let mut other_version = data.clone(); + let version_at = data.len() - crate::spatial::FORMAT_MARKER_LEN; + other_version[version_at..version_at + 4].copy_from_slice(&99u32.to_le_bytes()); + for data in [unmarked, &other_version[..]] { + let error = CellIndexReader::open(data).err().unwrap().to_string(); + assert!(error.contains("re-index"), "{error}"); + } + } + + #[test] + fn test_locate_finds_containing_and_contained_cells() -> crate::Result<()> { + let face_1 = S2CellId::from_face(1); + let child = face_1.child(2); + let data = written(&CellIndex { + cells: vec![IndexCell::new(child)], + }); + let reader = CellIndexReader::open(&data)?; + assert_eq!(reader.locate(face_1), CellRelation::Subdivided(0)); + assert_eq!(reader.locate(child), CellRelation::Indexed(0)); + assert_eq!(reader.locate(child.child(3)), CellRelation::Indexed(0)); + assert_eq!(reader.locate(face_1.child(1)), CellRelation::Disjoint); + assert_eq!( + reader.locate(S2CellId::from_face(0)), + CellRelation::Disjoint + ); + Ok(()) + } +} diff --git a/src/spatial/cell_union.rs b/src/spatial/cell_union.rs new file mode 100644 index 0000000000..00c8e190f6 --- /dev/null +++ b/src/spatial/cell_union.rs @@ -0,0 +1,508 @@ +//! CellUnion: A normalized, sorted collection of non-overlapping S2CellIds. +//! +//! A CellUnion represents a region on the sphere as a collection of S2CellIds at various levels. +//! The collection is: +//! - Sorted in increasing order by cell ID +//! - Non-overlapping (no cell contains another) +//! - Normalized (four sibling cells are replaced by their parent) +//! +//! From s2cell_union.h and s2cell_union.cc + +use super::latlng_rect::S2LatLngRect; +use super::math::normalize; +use super::region::Region; +use super::s2cap::S2Cap; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; + +/// A normalized collection of S2CellIds representing a region. +/// +/// The cell IDs are stored sorted and non-overlapping. In a normalized CellUnion, groups of 4 +/// sibling cells are replaced by their parent cell. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CellUnion { + cell_ids: Vec, +} + +impl CellUnion { + /// Creates an empty CellUnion. + pub fn new() -> Self { + Self { + cell_ids: Vec::new(), + } + } + + /// Creates a CellUnion from a vector of S2CellIds, then normalizes. + /// + /// The input cells can be in any order and may overlap. The constructor sorts them, removes + /// duplicates and contained cells, and merges sibling groups. + pub fn from_cell_ids(mut cell_ids: Vec) -> Self { + Self::normalize_vec(&mut cell_ids); + Self { cell_ids } + } + + /// Creates a CellUnion from pre-normalized cell IDs (no normalization). + /// + /// REQUIRES: cell_ids satisfies is_normalized() requirements. + pub fn from_normalized(cell_ids: Vec) -> Self { + debug_assert!(Self::check_normalized(&cell_ids)); + Self { cell_ids } + } + + /// Creates a CellUnion from sorted, non-overlapping cell IDs. + /// + /// Unlike from_normalized, this does not require sibling merging. + /// The cells must be valid, sorted, and non-overlapping. + pub fn from_verbatim(cell_ids: Vec) -> Self { + debug_assert!(Self::check_valid(&cell_ids)); + Self { cell_ids } + } + + /// Creates a CellUnion covering the entire sphere (all 6 face cells). + pub fn whole_sphere() -> Self { + Self { + cell_ids: vec![ + S2CellId::from_face(0), + S2CellId::from_face(1), + S2CellId::from_face(2), + S2CellId::from_face(3), + S2CellId::from_face(4), + S2CellId::from_face(5), + ], + } + } + + /// Creates a CellUnion covering leaf cells from min_id to max_id inclusive. + /// + /// Uses the maximum_tile algorithm to find the largest cells that fit + /// + /// REQUIRES: min_id.is_leaf(), max_id.is_leaf(), min_id <= max_id + pub fn from_min_max(min_id: S2CellId, max_id: S2CellId) -> Self { + debug_assert!(min_id.is_leaf()); + debug_assert!(max_id.is_leaf()); + debug_assert!(min_id <= max_id); + Self::from_begin_end(min_id, max_id.next()) + } + + /// Creates a CellUnion covering leaf cells from begin (inclusive) to end (exclusive). + /// + /// Like Python ranges or STL iterator ranges. Returns empty if begin == end. + /// + /// REQUIRES: begin.is_leaf(), end.is_leaf(), begin <= end + pub fn from_begin_end(begin: S2CellId, end: S2CellId) -> Self { + debug_assert!(begin.is_leaf() || begin == S2CellId::end(S2CellId::MAX_LEVEL)); + debug_assert!(end.is_leaf() || end == S2CellId::end(S2CellId::MAX_LEVEL)); + debug_assert!(begin <= end); + + // We repeatedly add the largest cell we can. + let mut cell_ids = Vec::new(); + let mut id = begin.maximum_tile(end); + while id != end { + cell_ids.push(id); + id = id.next().maximum_tile(end); + } + // The output is already normalized. + Self { cell_ids } + } + + /// Returns the number of cells. + #[inline] + pub fn len(&self) -> usize { + self.cell_ids.len() + } + + /// Returns true if empty. + #[inline] + pub fn is_empty(&self) -> bool { + self.cell_ids.is_empty() + } + + /// Returns the cell at the given index. + #[inline] + pub fn cell_id(&self, i: usize) -> S2CellId { + self.cell_ids[i] + } + + /// Returns the underlying cell ID vector. + #[inline] + pub fn cell_ids(&self) -> &[S2CellId] { + &self.cell_ids + } + + /// Consumes self and returns the cell ID vector. + pub fn into_cell_ids(self) -> Vec { + self.cell_ids + } + + /// Clears the CellUnion. + pub fn clear(&mut self) { + self.cell_ids.clear(); + self.cell_ids.shrink_to_fit(); + } + + /// Returns true if the cell union is valid. + /// + /// Valid means: cell IDs are valid, non-overlapping, and sorted. + pub fn is_valid(&self) -> bool { + Self::check_valid(&self.cell_ids) + } + + /// Returns true if the cell union is normalized. + /// + /// Normalized means: valid and no four cells have a common parent. + pub fn is_normalized(&self) -> bool { + Self::check_normalized(&self.cell_ids) + } + + /// Checks validity of a cell ID vector. + fn check_valid(cell_ids: &[S2CellId]) -> bool { + if !cell_ids.is_empty() && !cell_ids[0].is_valid() { + return false; + } + for i in 1..cell_ids.len() { + if !cell_ids[i].is_valid() { + return false; + } + // Check non-overlapping: previous range_max < current range_min + if cell_ids[i - 1].range_max() >= cell_ids[i].range_min() { + return false; + } + } + true + } + + /// Checks normalization of a cell ID vector. + fn check_normalized(cell_ids: &[S2CellId]) -> bool { + if !cell_ids.is_empty() && !cell_ids[0].is_valid() { + return false; + } + for i in 1..cell_ids.len() { + if !cell_ids[i].is_valid() { + return false; + } + if cell_ids[i - 1].range_max() >= cell_ids[i].range_min() { + return false; + } + // Check no 4 siblings exist + if i >= 3 + && are_siblings( + cell_ids[i - 3], + cell_ids[i - 2], + cell_ids[i - 1], + cell_ids[i], + ) + { + return false; + } + } + true + } + + /// Normalizes the cell union in place. + /// + /// Sorts the cells, removes contained cells, and merges sibling groups. + pub fn normalize(&mut self) { + Self::normalize_vec(&mut self.cell_ids); + } + + /// Normalizes a vector of cell IDs in place. + pub fn normalize_vec(ids: &mut Vec) { + // Sort cells in increasing order + ids.sort(); + + let mut out = 0usize; + for i in 0..ids.len() { + let mut id = ids[i]; + debug_assert!(id.is_valid()); + + // Check whether this cell is contained by the previous output cell. + if out > 0 && ids[out - 1].contains(id) { + continue; + } + + // Discard any previous cells contained by this cell. + while out > 0 && id.contains(ids[out - 1]) { + out -= 1; + } + + // Check whether the last 3 elements plus "id" can be collapse into + // a single parent cell. + while out >= 3 && are_siblings(ids[out - 3], ids[out - 2], ids[out - 1], id) { + // Replace four children by their parent cell. + id = id.immediate_parent(); + out -= 3; + } + ids[out] = id; + out += 1; + } + + // Resize to actual output size + ids.truncate(out); + } + + /// Denormalizes to satisfy min_level and level_mod constraints. + /// + /// Replaces cells whose level is less than min_level or where + /// (level - min_level) is not a multiple of level_mod with their children. + /// + /// This allows coverings generated with min_level/level_mod constraints + /// to be stored normalized and then restored to the original form. + pub fn denormalize(&self, min_level: i32, level_mod: i32) -> Vec { + debug_assert!(min_level <= S2CellId::MAX_LEVEL); + debug_assert!((1..=3).contains(&level_mod)); + + let mut output = Vec::with_capacity(self.cell_ids.len()); + + for &id in &self.cell_ids { + let level = id.level(); + let mut new_level = level.max(min_level); + + if level_mod > 1 { + // Round up so that (new_level - min_level) is a multiple of level_mod. + // (Note that MAX_LEVEL is a multiple of 1, 2, and 3.) + let remainder = (S2CellId::MAX_LEVEL - (new_level - min_level)) % level_mod; + new_level += remainder; + new_level = new_level.min(S2CellId::MAX_LEVEL); + } + + if new_level == level { + output.push(id); + } else { + // Expand to children at new_level + let end = id.child_end_at_level(new_level); + let mut child = id.child_begin_at_level(new_level); + while child != end { + output.push(child); + child = child.next(); + } + } + } + output + } + + /// Returns true if the cell union contains the given cell ID. + /// + /// Uses binary search with the EntirelyPrecedes comparator. + pub fn contains(&self, id: S2CellId) -> bool { + debug_assert!(id.is_valid()); + + // Binary search for the first cell that might contain id + // (i.e., the first cell that does not entirely precede id). + let i = self.lower_bound(id); + i < self.cell_ids.len() && self.cell_ids[i].contains(id) + } + + /// Returns true if the cell union intersects the given cell ID. + /// + /// Uses binary search with the EntirelyPrecedes comparator. + pub fn intersects(&self, id: S2CellId) -> bool { + debug_assert!(id.is_valid()); + + // Binary search for the first cell that might intersect id. + let i = self.lower_bound(id); + i < self.cell_ids.len() && self.cell_ids[i].intersects(id) + } + + /// Returns true if the cell union contains the given point. + /// + /// The point does not need to be normalized. + pub fn contains_point(&self, p: &[f64; 3]) -> bool { + self.contains(S2CellId::from_point(p)) + } + + /// Binary search for the first cell that does not entirely precede target. + /// + /// Uses the EntirelyPrecedes comparator: a cell entirely precedes another + /// if its range_max < the other's range_min. + fn lower_bound(&self, target: S2CellId) -> usize { + self.cell_ids + .partition_point(|&id| entirely_precedes(id, target)) + } + + /// Shrinks the capacity if there is excess allocation. + /// + /// Useful when holding many CellUnions in memory. + pub fn pack(&mut self, excess: usize) { + if self.cell_ids.capacity() - self.cell_ids.len() > excess { + self.cell_ids.shrink_to_fit(); + } + } +} + +/// Returns true if cell `a` entirely precedes cell `b` on the Hilbert curve. +/// +/// This means a.range_max() < b.range_min(), i.e., no overlap is possible +/// +/// Note: This is not a total ordering. For sorted, disjoint cell IDs, it can +/// be used with partition_point to find the first cell that might intersect. +#[inline] +fn entirely_precedes(a: S2CellId, b: S2CellId) -> bool { + a.range_max() < b.range_min() +} + +/// Returns true if the four cells are siblings (have a common parent). +/// +/// Uses the XOR trick: a ^ b ^ c == d is a necessary (but not sufficient) +/// condition. The mask test provides the exact check. +/// +/// REQUIRES: The four cells are distinct. +fn are_siblings(a: S2CellId, b: S2CellId, c: S2CellId, d: S2CellId) -> bool { + // A necessary (but not sufficient) condition is that the XOR of the + // four cells must be zero. This is also very fast to test. + if (a.id() ^ b.id() ^ c.id()) != d.id() { + return false; + } + + // Now we do a slightly more expensive but exact test. First, compute + // mask that blocks out the two bits that encode the child position of + // "d" with respect to its parent, then check that the other three + // children all agree with "mask". + let lsb = d.lsb(); + let mask = lsb << 1; + let mask = !(mask + (mask << 1)); + let id_masked = d.id() & mask; + + (a.id() & mask) == id_masked + && (b.id() & mask) == id_masked + && (c.id() & mask) == id_masked + && !d.is_face() // Face cells have no parent +} + +impl std::ops::Index for CellUnion { + type Output = S2CellId; + + fn index(&self, index: usize) -> &Self::Output { + &self.cell_ids[index] + } +} + +impl IntoIterator for CellUnion { + type Item = S2CellId; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.cell_ids.into_iter() + } +} + +impl<'a> IntoIterator for &'a CellUnion { + type Item = &'a S2CellId; + type IntoIter = std::slice::Iter<'a, S2CellId>; + + fn into_iter(self) -> Self::IntoIter { + self.cell_ids.iter() + } +} + +impl FromIterator for CellUnion { + fn from_iter>(iter: I) -> Self { + Self::from_cell_ids(iter.into_iter().collect()) + } +} + +impl std::fmt::Display for CellUnion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + const MAX_DISPLAY: usize = 20; + write!(f, "CellUnion(size={}, cells=[", self.len())?; + for (i, id) in self.cell_ids.iter().take(MAX_DISPLAY).enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{}", id.to_token())?; + } + if self.len() > MAX_DISPLAY { + write!(f, ", ...]")?; + } else { + write!(f, "]")?; + } + write!(f, ")") + } +} + +impl Region for CellUnion { + /// Returns a bounding spherical cap for this cell union. + /// + /// Computes an area-weighted centroid of the cells, then expands a cap + /// from that center to include all cells' bounding caps. + /// + /// Note: This won't produce the bounding cap of minimal area, but + /// it should be close enough for covering purposes. + fn get_cap_bound(&self) -> S2Cap { + if self.cell_ids.is_empty() { + return S2Cap::empty(); + } + + // Compute the approximate centroid of the region, weighted by cell area. + let mut centroid = [0.0, 0.0, 0.0]; + for &id in &self.cell_ids { + let area = S2Cell::average_area(id.level()); + let p = id.to_point(); + centroid[0] += area * p[0]; + centroid[1] += area * p[1]; + centroid[2] += area * p[2]; + } + + // Handle the case where centroid is zero (can happen if cells are + // symmetrically distributed around the origin). + let centroid = if centroid[0] == 0.0 && centroid[1] == 0.0 && centroid[2] == 0.0 { + [1.0, 0.0, 0.0] + } else { + normalize(¢roid) + }; + + // Use the centroid as the cap axis, and expand the cap angle so that it + // contains the bounding caps of all the individual cells. Note that it is + // *not* sufficient to just bound all the cell vertices because the bounding + // cap may be concave (i.e. cover more than one hemisphere). + let mut cap = S2Cap::from_point(centroid); + for &id in &self.cell_ids { + let cell = S2Cell::new(id); + cap.add_cap(&cell.get_cap_bound()); + } + cap + } + + /// Returns a bounding latitude-longitude rectangle for this cell union. + /// + /// Computes the union of all cells' rect bounds. + fn get_rect_bound(&self) -> S2LatLngRect { + let mut bound = S2LatLngRect::empty(); + for &id in &self.cell_ids { + let cell = S2Cell::new(id); + bound = bound.union(&cell.get_rect_bound()); + } + bound + } + + /// Returns a small cell covering for this cell union. + /// + /// Delegates to get_cap_bound().get_cell_union_bound() rather than + /// returning the cells directly, because get_cell_union_bound should + /// return a small covering. + fn get_cell_union_bound(&self, cell_ids: &mut Vec) { + self.get_cap_bound().get_cell_union_bound(cell_ids); + } + + /// Returns true if the cell union completely contains the given cell. + fn contains_cell(&self, cell: &S2Cell) -> bool { + self.contains(cell.id()) + } + + /// Returns true if the cell union may intersect the given cell. + /// + /// Because Intersects is exact, false means definitely no intersection. + fn may_intersect(&self, cell: &S2Cell) -> bool { + self.intersects(cell.id()) + } + + /// Returns true if the cell union contains the given point. + /// + /// The point does not need to be normalized. + fn contains_point(&self, p: &[f64; 3]) -> bool { + self.contains_point(p) + } +} + +#[cfg(test)] +#[path = "tests/cell_union_tests.rs"] +mod tests; diff --git a/src/spatial/closest_edge_query.rs b/src/spatial/closest_edge_query.rs new file mode 100644 index 0000000000..6d4a8ac601 --- /dev/null +++ b/src/spatial/closest_edge_query.rs @@ -0,0 +1,515 @@ +//! Branch-and-bound closest edge query over the cell index. +//! +//! Finds geometries closest to a query geometry by traversing the S2 cell hierarchy in distance +//! order. The query geometry has its own CellIndex for bilateral pruning -- a stored cell is only +//! checked against the query cells whose bounding caps could be closer than the best distance so +//! far, not against every query edge. +//! +//! Interiors count, as with S2ClosestEdgeQuery's include_interiors: a polygon that contains part +//! of the other geometry is at distance zero, however far away its edges are. +//! +//! Three modes from one algorithm: +//! kNN -- dynamic threshold, tightens to Kth-best +//! within-D -- fixed threshold, collect all within D +//! boolean -- fixed threshold, exit on first witness + +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet}; + +use common::BitSet; + +use super::cell_index::{ + BuildOptions, CellIndex, CellRelation, GeometryId, IndexBuilder, IndexCell, +}; +use super::cell_index_reader::CellIndexReader; +use super::contains_query::QueryEdgeProvider; +use super::edge_cache::EdgeCache; +use super::geometry_set::GeometrySet; +use super::s1chord_angle::S1ChordAngle; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; +use super::s2edge_distances::{update_edge_pair_min_distance, update_min_distance}; +use super::shape_index_region::{any_shape_contains_point, clipped_shape_contains_point}; +use super::sphere::Sphere; + +/// A result from the closest edge query: document ID and distance. +#[derive(Clone, Debug)] +pub struct ClosestEdgeResult { + /// The document containing the closest geometry. + pub doc_id: u32, + /// The minimum distance from the query geometry to this document's geometry. + pub distance: S1ChordAngle, +} + +/// Prepared closest-edge query. The query geometry is smashed into a GeometrySet and spatially +/// indexed. The same CellIndex that contains and intersects build for the query polygon. +pub struct ClosestEdgeQuery { + query_index: CellIndex, + query_edges: QueryEdgeProvider, + /// Bounding cap of each query index cell, parallel to `query_index.cells`. + query_bounds: Vec, + max_results: usize, + max_distance: S1ChordAngle, + min_distance: S1ChordAngle, + first_only: bool, +} + +impl ClosestEdgeQuery { + /// Build a kNN query from a smashed query geometry. + pub fn knn(set: GeometrySet, k: usize) -> Self { + Self::build(set, k, S1ChordAngle::infinity(), false) + } + + /// Build a within-D query from a smashed query geometry. + pub fn within(set: GeometrySet, max_distance: S1ChordAngle) -> Self { + Self::build(set, usize::MAX, max_distance, false) + } + + /// Build a boolean distance query. Returns at most one result. + pub fn any_within(set: GeometrySet, max_distance: S1ChordAngle) -> Self { + Self::build(set, 1, max_distance, true) + } + + /// Build a range distance query: results between inner and outer distance. + /// + /// There is no boolean form that stops at the first witness. A geometry's distance is final + /// only once every closer cell has been visited, and the first geometry found may turn out + /// to lie inside the inner distance while another lies in range. + pub fn between( + set: GeometrySet, + min_distance: S1ChordAngle, + max_distance: S1ChordAngle, + ) -> Self { + let mut q = Self::build(set, usize::MAX, max_distance, false); + q.min_distance = min_distance; + q + } + + fn build( + set: GeometrySet, + max_results: usize, + max_distance: S1ChordAngle, + first_only: bool, + ) -> Self { + let builder = IndexBuilder::new(BuildOptions::default()); + let query_index = builder.build_from_sets(std::slice::from_ref(&set)); + let query_bounds = query_index + .cells + .iter() + .map(|cell| CellBound::new(&S2Cell::new(cell.cell_id))) + .collect(); + let query_edges = QueryEdgeProvider { set }; + Self { + query_index, + query_edges, + query_bounds, + max_results, + max_distance, + min_distance: S1ChordAngle::zero(), + first_only, + } + } + + /// Search one segment. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + self.search_impl(cell_reader, edge_cache, None) + } + + /// Search one segment, skipping shapes whose doc_id is not in the filter bitset. + pub fn search_segment_filtered( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: &BitSet, + ) -> Vec { + self.search_impl(cell_reader, edge_cache, Some(terms_filter)) + } + + fn search_impl( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: Option<&BitSet>, + ) -> Vec { + if cell_reader.cell_count() == 0 + || self.max_results == 0 + || self.max_distance <= S1ChordAngle::zero() + || self.min_distance >= self.max_distance + { + return Vec::new(); + } + + let mut best = Best::new(self.max_results, self.max_distance); + + self.add_containing_polygons(cell_reader, edge_cache, terms_filter, &mut best); + if best.limit == S1ChordAngle::zero() || (self.first_only && !best.is_empty()) { + return self.collect_results(&best); + } + + // Stored geometries already tested for lying inside the query geometry. + let mut tested_interiors: HashSet = HashSet::new(); + let mut queue: BinaryHeap = BinaryHeap::new(); + + // Seed with the six face-level cells. + for face in 0..6 { + self.enqueue( + &S2Cell::from_face(face), + cell_reader, + best.limit, + &mut queue, + ); + } + + while let Some(entry) = queue.pop() { + if entry.distance >= best.limit { + break; + } + + match entry.relation { + // The queue cell is the index cell: any coarser index cell containing it would + // have been processed in its place. + CellRelation::Indexed(position) => { + self.process_edges( + &cell_reader.cell(position), + edge_cache, + terms_filter, + &mut tested_interiors, + &mut best, + ); + if self.first_only && !best.is_empty() { + return self.collect_results(&best); + } + } + CellRelation::Subdivided(_) => { + let cell = S2Cell::new(entry.cell_id); + if let Some(children) = cell.subdivide() { + for child in &children { + self.enqueue(child, cell_reader, best.limit, &mut queue); + } + } + } + CellRelation::Disjoint => {} + } + } + + self.collect_results(&best) + } + + /// Queues a cell that overlaps the index and may be closer to the query than the limit. + fn enqueue( + &self, + cell: &S2Cell, + cell_reader: &CellIndexReader, + distance_limit: S1ChordAngle, + queue: &mut BinaryHeap, + ) { + let relation = cell_reader.locate(cell.id()); + if relation == CellRelation::Disjoint { + return; + } + let distance = self.cell_to_query_distance(cell); + if distance < distance_limit { + queue.push(QueueEntry { + distance, + cell_id: cell.id(), + relation, + }); + } + } + + /// Stored polygons that contain part of the query geometry are at distance zero. Testing + /// one vertex of each query member suffices: a polygon that contains part of a member but + /// not that vertex has an edge that crosses the member, which the edge distances find. + fn add_containing_polygons( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: Option<&BitSet>, + best: &mut Best, + ) { + for member in &self.query_edges.set.members { + let Some(point) = member.vertices.first() else { + continue; + }; + let Some(index_cell) = cell_reader.find(S2CellId::from_point(point)) else { + continue; + }; + for clipped in &index_cell.shapes { + if let Some(filter) = terms_filter { + if !filter.contains(edge_cache.doc_id_for(clipped.geometry_id)) { + continue; + } + } + let entry = edge_cache.get(clipped.geometry_id); + if !entry.edge_set().closed { + continue; + } + // Use the already decoded shape; looking it up again would decode every shape + // in the cell once per polygon, making overlapping polygons quadratic. + if clipped_shape_contains_point(index_cell.cell_id, clipped, point, |edge_idx| { + entry.edge(edge_idx) + }) { + best.update(entry.doc_id(), S1ChordAngle::zero()); + // K distinct zero-distance hits cannot be improved by any later geometry. + if best.limit == S1ChordAngle::zero() { + return; + } + } + } + } + } + + /// Whether a closed member of the query geometry contains the point. + fn query_contains(&self, point: &[f64; 3]) -> bool { + any_shape_contains_point(&self.query_index, &self.query_edges, point) + } + + /// Minimum distance from a stored cell to the query geometry, its interior included: a + /// lower bound on the distance of anything stored in the cell. + fn cell_to_query_distance(&self, cell: &S2Cell) -> S1ChordAngle { + if self.query_contains(&cell.get_center()) { + return S1ChordAngle::zero(); + } + let bound = CellBound::new(cell); + let mut min_dist = S1ChordAngle::infinity(); + for (query_cell, query_bound) in self.query_index.cells.iter().zip(&self.query_bounds) { + if bound.distance(query_bound) >= min_dist { + continue; + } + for clipped in &query_cell.shapes { + let vertices = &self.query_edges.get_edge_set(clipped.geometry_id).vertices; + for &edge_idx in &clipped.edge_indices { + let distance = match edge(vertices, edge_idx) { + (point, other) if point == other => cell.get_distance_to_point(point), + (v0, v1) => cell.get_distance_to_edge(v0, v1), + }; + min_dist = min_dist.min(distance); + } + } + } + min_dist + } + + /// Lowers `min_dist` to the distance between the stored edge (a0, a1), clipped to the cell + /// bounded by `bound`, and the query edges in query cells that could be closer. + fn update_min_distance_to_query( + &self, + bound: &CellBound, + a0: &[f64; 3], + a1: &[f64; 3], + min_dist: &mut S1ChordAngle, + ) { + for (query_cell, query_bound) in self.query_index.cells.iter().zip(&self.query_bounds) { + if bound.distance(query_bound) >= *min_dist { + continue; + } + for clipped in &query_cell.shapes { + let vertices = &self.query_edges.get_edge_set(clipped.geometry_id).vertices; + for &edge_idx in &clipped.edge_indices { + let (b0, b1) = edge(vertices, edge_idx); + if a0 == a1 { + update_min_distance(a0, b0, b1, min_dist); + } else if b0 == b1 { + update_min_distance(b0, a0, a1, min_dist); + } else { + update_edge_pair_min_distance(a0, a1, b0, b1, min_dist); + } + } + } + } + } + + fn process_edges( + &self, + index_cell: &IndexCell, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: Option<&BitSet>, + tested_interiors: &mut HashSet, + best: &mut Best, + ) { + let bound = CellBound::new(&S2Cell::new(index_cell.cell_id)); + + for clipped in &index_cell.shapes { + let geometry_id = clipped.geometry_id; + + if let Some(filter) = terms_filter { + let doc_id = edge_cache.doc_id_for(geometry_id); + if !filter.contains(doc_id) { + continue; + } + } + + let cache_entry = edge_cache.get(geometry_id); + let doc_id = cache_entry.doc_id(); + let candidate_vertices = cache_entry.vertices(); + if candidate_vertices.is_empty() { + continue; + } + + let mut min_dist = best.limit; + if let Some(prev) = best.get(doc_id) { + if prev < min_dist { + min_dist = prev; + } + } + + // A stored geometry inside the query polygon is at distance zero wherever its edges + // are. If only part of it is inside, its edges cross the query's, which the edge + // distances find. + if tested_interiors.insert(geometry_id) && self.query_contains(&candidate_vertices[0]) { + min_dist = S1ChordAngle::zero(); + } else { + // A shape with no edges in this cell covers it. Its distance comes from the cells + // its edges cross, or from add_containing_polygons when it contains the query. + for &edge_idx in &clipped.edge_indices { + let (cv0, cv1) = cache_entry.edge(edge_idx); + self.update_min_distance_to_query(&bound, &cv0, &cv1, &mut min_dist); + } + } + + if min_dist < best.limit { + best.update(doc_id, min_dist); + } + } + } + + fn collect_results(&self, best: &Best) -> Vec { + let mut out: Vec = best + .distances + .iter() + .filter(|(_, &distance)| distance >= self.min_distance) + .map(|(&doc_id, &distance)| ClosestEdgeResult { doc_id, distance }) + .collect(); + // Ties go to the lower document, so that a query always returns the same documents. + out.sort_by_key(|result| (result.distance, result.doc_id)); + if self.max_results < out.len() { + out.truncate(self.max_results); + } + out + } +} + +/// The best distance found so far for each document, and the distance a geometry must be nearer +/// than to count. When the results are bounded to K, only the K nearest documents are kept, and +/// once there are K the limit is the farthest of them. +struct Best { + distances: HashMap, + /// The kept documents by distance, farthest on top, when the results are bounded. Entries for + /// a document whose distance has since improved are stale, and dropped when they surface. + nearest: BinaryHeap<(S1ChordAngle, u32)>, + max_results: usize, + limit: S1ChordAngle, +} + +impl Best { + fn new(max_results: usize, max_distance: S1ChordAngle) -> Self { + Best { + distances: HashMap::new(), + nearest: BinaryHeap::new(), + max_results, + limit: max_distance, + } + } + + fn get(&self, doc_id: u32) -> Option { + self.distances.get(&doc_id).copied() + } + + fn is_empty(&self) -> bool { + self.distances.is_empty() + } + + /// Records a document's distance, unless it already has one at least as near. + fn update(&mut self, doc_id: u32, distance: S1ChordAngle) { + if self.get(doc_id).is_some_and(|prev| distance >= prev) { + return; + } + self.distances.insert(doc_id, distance); + if self.max_results == usize::MAX { + return; + } + self.nearest.push((distance, doc_id)); + // Drop stale entries, and the farthest document while there are more than K. + while let Some(&(top_distance, top_doc)) = self.nearest.peek() { + let current = self.get(top_doc) == Some(top_distance); + if current && self.distances.len() <= self.max_results { + break; + } + self.nearest.pop(); + if current { + self.distances.remove(&top_doc); + } + } + if self.distances.len() == self.max_results { + if let Some(&(farthest, _)) = self.nearest.peek() { + self.limit = farthest; + } + } + } +} + +/// The endpoints of an edge in a smashed vertex array. A point is an edge from itself to itself. +fn edge(vertices: &[[f64; 3]], edge_idx: u32) -> (&[f64; 3], &[f64; 3]) { + let i = edge_idx as usize; + (&vertices[i], vertices.get(i + 1).unwrap_or(&vertices[i])) +} + +/// A cell's bounding cap, for a cheap lower bound on the distance between two cells. +struct CellBound { + center: [f64; 3], + radius: f64, +} + +impl CellBound { + fn new(cell: &S2Cell) -> Self { + let cap = cell.get_cap_bound(); + CellBound { + center: cap.center(), + radius: cap.get_radius_radians(), + } + } + + /// Lower bound on the distance between any point of one cap and any point of the other. + /// Rounded down by more than the error of the computation, so that it stays a lower bound. + fn distance(&self, other: &CellBound) -> S1ChordAngle { + let centers = S1ChordAngle::from_points(&self.center, &other.center).to_radians(); + let gap = centers - self.radius - other.radius - 1e-14; + if gap > 0.0 { + S1ChordAngle::from_radians(gap) + } else { + S1ChordAngle::zero() + } + } +} + +struct QueueEntry { + distance: S1ChordAngle, + cell_id: S2CellId, + /// Where the cell lies relative to the index cells. + relation: CellRelation, +} + +impl PartialEq for QueueEntry { + fn eq(&self, other: &Self) -> bool { + self.distance == other.distance + } +} + +impl Eq for QueueEntry {} + +impl PartialOrd for QueueEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +// Reversed ordering for min-heap (BinaryHeap is a max-heap). +impl Ord for QueueEntry { + fn cmp(&self, other: &Self) -> Ordering { + other.distance.cmp(&self.distance) + } +} diff --git a/src/spatial/containment.rs b/src/spatial/containment.rs new file mode 100644 index 0000000000..9b6dcdd7d3 --- /dev/null +++ b/src/spatial/containment.rs @@ -0,0 +1,312 @@ +//! Loop containment and intersection predicates. +use super::crossings::{ordered_ccw, ref_dir, robust_cross, sign, S2EdgeCrosser}; +use super::exact::expensive_sign; +use super::math::{cross, dot, norm2}; +use super::sphere::ORIGIN; + +/// A version of contains that does not use the shape index. +pub fn brute_force_contains(p: &[f64; 3], vertices: &[[f64; 3]], origin_inside: bool) -> bool { + // Handle degenerate cases + let n = vertices.len(); + if n < 3 { + return origin_inside; + } + + let origin = ORIGIN; + let mut crosser = S2EdgeCrosser::new(&origin, p); + crosser.restart_at(&vertices[0]); + + let mut inside = origin_inside; + + // Test each edge of the loop + for i in 1..=n { + let vertex = &vertices[i % n]; + inside ^= crosser.edge_or_vertex_crossing(vertex); + } + + inside +} + +/// Computes whether the reference origin is inside a loop. Compares angle containment at vertex 1 +/// with brute force result to determine the answer. +pub fn compute_origin_inside(vertices: &[[f64; 3]]) -> bool { + let n = vertices.len(); + if n < 3 { + return false; + } + + // Check if vertex(1) is inside the angle formed by edges 0->1 and 1->2. This is true by CCW + // convention for non-degenerate loops. A repeated vertex makes the angle degenerate, and S2 + // takes it to contain nothing. + let v1_inside = vertices[0] != vertices[1] + && vertices[2] != vertices[1] + && angle_contains_vertex(&vertices[0], &vertices[1], &vertices[2]); + + // Compute containment of v1 assuming origin_inside is false + let contains_v1 = brute_force_contains(&vertices[1], vertices, false); + + // If the results disagree, Origin must be inside + v1_inside != contains_v1 +} + +/// Removes the degeneracies of a ring: repeated vertices, and spikes, where an edge doubles back +/// along the one before it. Returns the loop without a closing vertex, or an empty loop when +/// nothing but degeneracies remain, and the spikes it removed, as (base, tip) edges. Port of +/// PruneDegeneracies from s2loop_measures.cc, which drops the spikes. +pub fn prune_degeneracies(ring: &[[f64; 3]]) -> (Vec<[f64; 3]>, Vec<[[f64; 3]; 2]>) { + let mut vertices: Vec<[f64; 3]> = Vec::with_capacity(ring.len()); + let mut spikes = Vec::new(); + for v in ring { + if vertices.last() == Some(v) { + continue; + } + if vertices.len() >= 2 && vertices[vertices.len() - 2] == *v { + // Remove the "B" of a spike "ABA". + let tip = vertices.pop().unwrap(); + spikes.push([*v, tip]); + } else { + vertices.push(*v); + } + } + if vertices.len() < 3 { + return (Vec::new(), spikes); + } + if vertices[0] == vertices[vertices.len() - 1] { + vertices.pop(); + } + // Remove the spikes across the start of the loop: it begins with "BA" and ends with "A". + let n = vertices.len(); + let mut k = 0; + while 2 * k + 3 <= n && vertices[k + 1] == vertices[n - 1 - k] { + spikes.push([vertices[k + 1], vertices[k]]); + k += 1; + } + if n - 2 * k < 3 { + return (Vec::new(), spikes); + } + vertices.truncate(n - k); + vertices.drain(..k); + (vertices, spikes) +} + +/// Whether every vertex of a loop lies exactly on one great circle, so that the loop bounds no +/// area however many distinct vertices it has. Exact, as S2's predicates are. +pub fn on_one_great_circle(vertices: &[[f64; 3]]) -> bool { + let (a, b) = (&vertices[0], &vertices[1]); + vertices[2..] + .iter() + .all(|v| v == a || v == b || expensive_sign(a, b, v, false) == 0) +} + +/// The geodesic curvature of a loop without degeneracies or a closing vertex: the sum of its +/// turn angles, which is 2π minus its area. It is positive when the loop turns left around a +/// region smaller than a hemisphere. Port of S2::GetCurvature. +pub fn curvature(vertices: &[[f64; 3]]) -> f64 { + // Kahan summation, as S2's, so that the error stays within curvature_max_error. + let n = vertices.len(); + let (mut sum, mut compensation) = (0.0f64, 0.0f64); + for i in 0..n { + let angle = turn_angle( + &vertices[(i + n - 1) % n], + &vertices[i], + &vertices[(i + 1) % n], + ) + compensation; + let old_sum = sum; + sum += angle; + compensation = (old_sum - sum) + angle; + } + sum + compensation +} + +/// The maximum error in the curvature of a loop of `n` vertices. Port of +/// S2::GetCurvatureMaxError. +pub fn curvature_max_error(n: usize) -> f64 { + 11.25 * f64::EPSILON * n as f64 +} + +/// The exterior angle at vertex B of the path ABC: positive for a left turn, negative for a right +/// one. Port of S2::TurnAngle. +fn turn_angle(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3]) -> f64 { + let x = robust_cross(a, b); + let y = robust_cross(b, c); + let angle = norm2(&cross(&x, &y)).sqrt().atan2(dot(&x, &y)); + if sign(a, b, c) > 0 { + angle + } else { + -angle + } +} + +fn angle_contains_vertex(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3]) -> bool { + // B is inside angle ABC if C is to the left of ray AB + // This is equivalent to Sign(A, B, C) > 0 + !ordered_ccw(&ref_dir(b), c, a, b) +} + +/// Returns true if wedge A contains wedge B. Both wedges share the common vertex ab1. +pub fn wedge_contains( + a0: &[f64; 3], + ab1: &[f64; 3], + a2: &[f64; 3], + b0: &[f64; 3], + b2: &[f64; 3], +) -> bool { + // For A to contain B, the CCW edge order around ab1 must be a2 b2 b0 a0 + ordered_ccw(a2, b2, b0, ab1) && ordered_ccw(b0, a0, a2, ab1) +} + +/// Returns true if wedge A intersects wedge B. Both wedges share the common vertex ab1. +pub fn wedge_intersects( + a0: &[f64; 3], + ab1: &[f64; 3], + a2: &[f64; 3], + b0: &[f64; 3], + b2: &[f64; 3], +) -> bool { + // A and B are disjoint if CCW order is a0 b2 b0 a2 + !(ordered_ccw(a0, b2, b0, ab1) && ordered_ccw(b0, a2, a0, ab1)) +} + +/// Returns true if loop A contains loop B. +pub fn loop_contains( + a_vertices: &[[f64; 3]], + a_origin_inside: bool, + b_vertices: &[[f64; 3]], + b_origin_inside: bool, +) -> bool { + let n_a = a_vertices.len(); + let n_b = b_vertices.len(); + + // Handle degenerate cases + if n_a < 3 || n_b < 3 { + // Empty loop contains nothing, full loop contains everything + return a_origin_inside || !b_origin_inside; + } + + let mut found_shared_vertex = false; + + // Check all edge pairs + for i in 0..n_a { + let a0 = &a_vertices[i]; + let a1 = &a_vertices[(i + 1) % n_a]; + + let mut crosser = S2EdgeCrosser::new(a0, a1); + + for j in 0..n_b { + let b0 = &b_vertices[j]; + let b1 = &b_vertices[(j + 1) % n_b]; + + let crossing = crosser.crossing_sign_two(b0, b1); + + if crossing > 0 { + // Interior crossing: A does not contain B + return false; + } + + if crossing == 0 { + // Shared vertex: check wedge containment + // Find which vertex is shared + if a1 == b1 { + found_shared_vertex = true; + let a_prev = a0; + let a_next = &a_vertices[(i + 2) % n_a]; + let b_prev = b0; + let b_next = &b_vertices[(j + 2) % n_b]; + + // A must contain B at this vertex + if !wedge_contains(a_prev, a1, a_next, b_prev, b_next) { + return false; + } + } + } + } + } + + // If we found shared vertices and all wedges passed, A contains B + if found_shared_vertex { + return true; + } + + // No shared vertices: check point containment + // A must contain a vertex of B + if !brute_force_contains(&b_vertices[0], a_vertices, a_origin_inside) { + return false; + } + + // B must not contain a vertex of A (handles sphere-covering case) + if brute_force_contains(&a_vertices[0], b_vertices, b_origin_inside) { + return false; + } + + true +} + +/// Returns true if loop A intersects loop B. +pub fn loop_intersects( + a_vertices: &[[f64; 3]], + a_origin_inside: bool, + b_vertices: &[[f64; 3]], + b_origin_inside: bool, +) -> bool { + let n_a = a_vertices.len(); + let n_b = b_vertices.len(); + + // Handle degenerate cases + if n_a < 3 || n_b < 3 { + // Empty loop intersects nothing, full loop intersects everything non-empty + return a_origin_inside && b_origin_inside; + } + + let mut found_shared_vertex = false; + + // Check all edge pairs + for i in 0..n_a { + let a0 = &a_vertices[i]; + let a1 = &a_vertices[(i + 1) % n_a]; + + let mut crosser = S2EdgeCrosser::new(a0, a1); + + for j in 0..n_b { + let b0 = &b_vertices[j]; + let b1 = &b_vertices[(j + 1) % n_b]; + + let crossing = crosser.crossing_sign_two(b0, b1); + + if crossing > 0 { + // Interior crossing: loops intersect + return true; + } + + if crossing == 0 { + // Shared vertex: check wedge intersection + if a1 == b1 { + found_shared_vertex = true; + let a_prev = a0; + let a_next = &a_vertices[(i + 2) % n_a]; + let b_prev = b0; + let b_next = &b_vertices[(j + 2) % n_b]; + + if wedge_intersects(a_prev, a1, a_next, b_prev, b_next) { + return true; + } + } + } + } + } + + // If we found shared vertices but no wedge intersection, loops don't intersect + if found_shared_vertex { + return false; + } + + // No crossings or shared vertices: check containment + // If A contains B or B contains A, they intersect + if brute_force_contains(&b_vertices[0], a_vertices, a_origin_inside) { + return true; + } + if brute_force_contains(&a_vertices[0], b_vertices, b_origin_inside) { + return true; + } + + false +} diff --git a/src/spatial/contains_query.rs b/src/spatial/contains_query.rs new file mode 100644 index 0000000000..942d13f13f --- /dev/null +++ b/src/spatial/contains_query.rs @@ -0,0 +1,745 @@ +//! Polygon-contains-polygon query using the verification cascade. +//! +//! Given a query polygon, finds all indexed geometries contained by it. Builds a query-local +//! CellIndex using the same IndexBuilder as ingest, produces a covering via RegionCoverer, then +//! verifies candidates per segment using bilateral edge narrowing. +//! +//! This module does not port any single S2 class. It composes S2 primitives with the serialized +//! segment index and query-local CellIndex. + +use std::collections::{HashMap, HashSet}; +use std::io; +use std::ops::Range; +use std::sync::OnceLock; + +use common::{BitSet, CountingWriter}; + +use super::cell_index::{BuildOptions, CellIndex, GeometryId, IndexBuilder}; +use super::cell_index_reader::CellIndexReader; +use super::crossings::{triage_sign, S2EdgeCrosser}; +use super::edge_cache::EdgeCache; +use super::edge_reader::EdgeReader; +use super::edge_writer::EdgeWriter; +use super::exact::expensive_sign; +use super::geometry_set::{EdgeSet, GeometrySet}; +use super::math::{cross, dot, normalize}; +use super::polygon_union::{member_edges, split_edge, union_boundary, PolygonEdge}; +use super::region::Region; +use super::region_coverer::{CovererOptions, RegionCoverer}; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; +use super::serializer::EDGE_SKIP_INTERVAL; +use super::shape_index_region::{ + any_shape_contains_point, contains_point, index_contains_point, CellIndexRegion, EdgeProvider, + SegmentIndex, +}; +use super::sphere::Sphere; + +/// In-memory edge provider wrapping a smashed GeometrySet. Resolves edge indices into the +/// smashed vertex arrays. No modulo wrapping — smashed vertices include the closure vertex. +pub struct QueryEdgeProvider { + pub(crate) set: GeometrySet, +} + +impl QueryEdgeProvider { + /// Returns the EdgeSet for the given geometry_id. + pub fn get_edge_set(&self, geometry_id: GeometryId) -> &EdgeSet { + &self.set.members[geometry_id.1 as usize] + } +} + +impl EdgeProvider for QueryEdgeProvider { + fn get_edge(&self, geometry_id: GeometryId, edge_idx: u32) -> ([f64; 3], [f64; 3]) { + let vertices = &self.get_edge_set(geometry_id).vertices; + let i = edge_idx as usize; + (vertices[i], *vertices.get(i + 1).unwrap_or(&vertices[i])) + } + + fn is_closed(&self, geometry_id: GeometryId) -> bool { + self.get_edge_set(geometry_id).closed + } +} + +/// A query geometry prepared for searching segments: its cell index, its covering, and for each +/// covering cell whether the geometry contains it and which query edges lie in it. Shared by the +/// contains and intersects queries. +pub(crate) struct QueryRegion { + query_index: CellIndex, + query_edges: QueryEdgeProvider, + covering: Vec, + interior: Vec, + /// The query edges in each covering cell, parallel to `covering`. + cell_edges: Vec>, + /// The exposed boundary of the union of the query's members, computed when first needed. + boundary: OnceLock>, +} + +/// An edge of the query geometry, and the member and edge index that identify it. A point is an +/// edge from itself to itself. +struct QueryEdge { + member: u32, + edge_idx: u32, + v0: [f64; 3], + v1: [f64; 3], +} + +/// How a candidate's boundary meets the query geometry's boundary. +pub(crate) enum Contact { + /// An edge of each crosses the other at a point interior to both. + Cross, + /// The boundaries meet at most at shared vertices and along shared edges. + Touch(Touch), +} + +/// Where the boundaries of a candidate and the query geometry meet without crossing. A ring of +/// either that meets nothing lies wholly inside or wholly outside the other geometry, as does a +/// touching edge, but for its endpoints. Vertices where they meet may lie on either side, so they +/// decide nothing. +#[derive(Default)] +pub(crate) struct Touch { + /// Candidate edges that share a vertex with a query edge, other than those in `along`. + pub(crate) touching: HashSet, + /// Candidate edges that run along query edges. Values identify the shared query edges + /// as (member, edge, reversed direction). + pub(crate) along: HashMap>, + /// Query edges, as (member, edge index), that share a vertex with a candidate edge, other + /// than those in `query_along`. + pub(crate) query_touching: HashSet<(u32, u32)>, + /// Query edges that run along a candidate edge. + pub(crate) query_along: HashSet<(u32, u32)>, + /// Whether a vertex of either lies inside an edge of the other: a T-junction, or edges + /// that overlap in part. The pieces of the edges, split where they meet, then decide. + pub(crate) irregular: bool, +} + +impl Touch { + /// Whether an edge of the candidate ring meets the query geometry's boundary. + pub(crate) fn candidate_ring_meets(&self, ring: Range) -> bool { + self.touching + .iter() + .chain(self.along.keys()) + .any(|&edge_idx| ring.contains(&(edge_idx as usize))) + } + + /// The edges of the query member's ring that touch the candidate's boundary. + fn query_ring_touching( + &self, + member: u32, + ring: Range, + ) -> impl Iterator + '_ { + self.query_touching + .iter() + .filter(move |&&(m, edge_idx)| m == member && ring.contains(&(edge_idx as usize))) + .map(|&(_, edge_idx)| edge_idx) + } + + /// Whether an edge of the query member's ring meets the candidate's boundary. + pub(crate) fn query_ring_meets(&self, member: u32, ring: Range) -> bool { + self.query_touching + .iter() + .chain(&self.query_along) + .any(|&(m, edge_idx)| m == member && ring.contains(&(edge_idx as usize))) + } +} + +/// The midpoint of an edge. +pub(crate) fn midpoint((v0, v1): ([f64; 3], [f64; 3])) -> [f64; 3] { + normalize(&[v0[0] + v1[0], v0[1] + v1[1], v0[2] + v1[2]]) +} + +/// Whether point P lies on the edge from A to B, endpoints included. Exact, as S2's predicates +/// are: P must lie on the great circle through A and B, not merely near it. +fn point_on_edge(p: &[f64; 3], a: &[f64; 3], b: &[f64; 3]) -> bool { + if p == a || p == b { + return true; + } + if a == b || expensive_sign(a, b, p, false) != 0 { + return false; + } + // On the great circle, P lies between A and B if it turns the same way from each. + let normal = cross(a, b); + dot(&cross(a, p), &normal) > 0.0 && dot(&cross(p, b), &normal) > 0.0 +} + +/// Whether point P lies on the edge from A to B, strictly between its endpoints. Exact, but cheap +/// for the many points that are clearly off the edge's great circle. +fn point_inside_edge(p: &[f64; 3], a: &[f64; 3], b: &[f64; 3], a_cross_b: &[f64; 3]) -> bool { + p != a && p != b && triage_sign(a, b, p, a_cross_b) == 0 && point_on_edge(p, a, b) +} + +/// Whether a vertex of one edge lies strictly inside the other edge. +fn edges_meet_inside(a0: &[f64; 3], a1: &[f64; 3], b0: &[f64; 3], b1: &[f64; 3]) -> bool { + let (a_cross, b_cross) = (cross(a0, a1), cross(b0, b1)); + point_inside_edge(b0, a0, a1, &a_cross) + || point_inside_edge(b1, a0, a1, &a_cross) + || point_inside_edge(a0, b0, b1, &b_cross) + || point_inside_edge(a1, b0, b1, &b_cross) +} + +/// Smashed polygon rings are CCW. The polygon lies to the left of an exterior edge, and to +/// the right of a hole edge. This classification needs no offset from a shared boundary. +pub(crate) fn interior_is_left(candidate: &EdgeSet, edge_idx: u32) -> bool { + (edge_idx as usize) < candidate.ring_offsets[1] +} + +impl QueryRegion { + /// Index and cover a smashed query geometry. + pub(crate) fn new(set: GeometrySet, options: CovererOptions) -> Self { + let builder = IndexBuilder::new(BuildOptions::default()); + let query_index = builder.build_from_sets(std::slice::from_ref(&set)); + let query_edges = QueryEdgeProvider { set }; + + let region = CellIndexRegion::new(&query_index, &query_edges); + let coverer = RegionCoverer::new(options); + let covering = coverer.get_covering(®ion); + + let interior: Vec = covering + .cell_ids() + .iter() + .map(|&id| region.contains_cell(&S2Cell::new(id))) + .collect(); + + let covering = covering.into_cell_ids(); + + // A covering cell is usually coarser than the query index cells, so it holds the edges of + // every index cell it overlaps. + let cell_edges = covering + .iter() + .map(|&cell_id| { + let mut edges = Vec::new(); + for cell in query_index.cells_overlapping(cell_id) { + for clipped in &cell.shapes { + for &edge_idx in &clipped.edge_indices { + let (v0, v1) = query_edges.get_edge(clipped.geometry_id, edge_idx); + edges.push(QueryEdge { + member: clipped.geometry_id.1, + edge_idx, + v0, + v1, + }); + } + } + } + edges + }) + .collect(); + + Self { + query_index, + query_edges, + covering, + interior, + cell_edges, + boundary: OnceLock::new(), + } + } + + /// The exposed boundary of the union of the query's members, as pieces that know the side + /// of the query's interior. + pub(crate) fn boundary(&self) -> &[PolygonEdge] { + self.boundary + .get_or_init(|| union_boundary(&self.query_index, &self.query_edges)) + } + + /// The members of the query geometry. + pub(crate) fn members(&self) -> &[EdgeSet] { + &self.query_edges.set.members + } + + /// The endpoints of an edge of a query member. + pub(crate) fn query_edge(&self, member: u32, edge_idx: u32) -> ([f64; 3], [f64; 3]) { + self.query_edges.get_edge((0, member), edge_idx) + } + + /// Whether a member of the query geometry that has an interior contains the point. + pub(crate) fn contains(&self, point: &[f64; 3]) -> bool { + any_shape_contains_point(&self.query_index, &self.query_edges, point) + } + + /// Whether the query covers the candidate's interior beside a shared edge. A member + /// sharing the edge contributes its interior side; a member containing its midpoint + /// without sharing it covers both sides. + pub(crate) fn covers_shared_interior( + &self, + candidate: &EdgeSet, + edge_idx: u32, + touch: &Touch, + ) -> bool { + let left = interior_is_left(candidate, edge_idx); + let i = edge_idx as usize; + let point = midpoint((candidate.vertices[i], candidate.vertices[i + 1])); + let shared_edges = &touch.along[&edge_idx]; + self.members().iter().enumerate().any(|(member, edges)| { + if !edges.closed { + return false; + } + let mut shared = shared_edges + .iter() + .filter(|&&(m, _, _)| m == member as u32) + .peekable(); + if shared.peek().is_some() { + shared.any(|&(_, qi, reversed)| (interior_is_left(edges, qi) ^ reversed) == left) + } else { + index_contains_point( + &self.query_index, + &self.query_edges, + (0, member as u32), + &point, + ) + } + }) + } + + /// Collect every stored geometry in the covering, recording which covering cells it lies in. + /// Skips geometries whose doc_id is not in the filter bitset. + pub(crate) fn collect_candidates( + &self, + reader: &CellIndexReader, + terms_filter: Option<&BitSet>, + edge_cache: &EdgeCache<'_, Sphere>, + ) -> HashMap { + let mut candidates: HashMap = HashMap::new(); + + for (covering_index, &covering_cell_id) in self.covering.iter().enumerate() { + let is_interior = self.interior[covering_index]; + + for index_cell in reader.scan_range(covering_cell_id) { + let cell_is_interior = is_interior && covering_cell_id.contains(index_cell.cell_id); + + for clipped in index_cell.shapes { + if let Some(filter) = terms_filter { + let doc_id = edge_cache.doc_id_for(clipped.geometry_id); + if !filter.contains(doc_id) { + continue; + } + } + + let info = candidates + .entry(clipped.geometry_id) + .or_insert_with(CandidateInfo::new); + + if !cell_is_interior && !clipped.edge_indices.is_empty() { + info.boundary_edges.push(BoundaryEdges { + covering_index, + edge_indices: clipped.edge_indices, + }); + } + } + } + } + + candidates + } + + /// Whether a point of the candidate lies on an edge or point of the query geometry, or a point + /// of the query geometry on an edge or point of the candidate, where either has no interior. + /// A point is an edge from itself to itself, which the boundary contact skips. Where both + /// have interiors, the containment tests decide. + pub(crate) fn point_contact(&self, info: &CandidateInfo, candidate: &EdgeSet) -> bool { + let vertices = &candidate.vertices; + for boundary in &info.boundary_edges { + for query_edge in &self.cell_edges[boundary.covering_index] { + if candidate.closed && self.members()[query_edge.member as usize].closed { + continue; + } + let (q0, q1) = (&query_edge.v0, &query_edge.v1); + for &edge_idx in &boundary.edge_indices { + let i = edge_idx as usize; + let (c0, c1) = (&vertices[i], vertices.get(i + 1).unwrap_or(&vertices[i])); + let touches = if c0 == c1 { + point_on_edge(c0, q0, q1) + } else { + q0 == q1 && point_on_edge(q0, c0, c1) + }; + if touches { + return true; + } + } + } + } + false + } + + /// Tests the candidate's edges against the query edges, narrowed to edges sharing a boundary + /// covering cell. + pub(crate) fn boundary_contact(&self, info: &CandidateInfo, candidate: &EdgeSet) -> Contact { + let mut touch = Touch::default(); + let vertices = &candidate.vertices; + if vertices.len() < 2 { + return Contact::Touch(touch); + } + + for boundary in &info.boundary_edges { + let query_edges = &self.cell_edges[boundary.covering_index]; + if query_edges.is_empty() { + continue; + } + + for &candidate_edge_idx in &boundary.edge_indices { + let ci = candidate_edge_idx as usize; + let (cv0, cv1) = (&vertices[ci], &vertices[ci + 1]); + if cv0 == cv1 { + continue; + } + let mut crosser = S2EdgeCrosser::new(cv0, cv1); + + for query_edge in query_edges { + let (qv0, qv1) = (&query_edge.v0, &query_edge.v1); + if qv0 == qv1 { + continue; + } + let sign = crosser.crossing_sign_two(qv0, qv1); + if sign != 0 && edges_meet_inside(cv0, cv1, qv0, qv1) { + touch.irregular = true; + touch.touching.insert(candidate_edge_idx); + touch + .query_touching + .insert((query_edge.member, query_edge.edge_idx)); + continue; + } + match sign { + 1 => return Contact::Cross, + 0 => { + let query_edge = (query_edge.member, query_edge.edge_idx); + if (cv0, cv1) == (qv0, qv1) || (cv0, cv1) == (qv1, qv0) { + touch.query_along.insert(query_edge); + touch.along.entry(candidate_edge_idx).or_default().insert(( + query_edge.0, + query_edge.1, + cv0 != qv0, + )); + } else { + touch.touching.insert(candidate_edge_idx); + touch.query_touching.insert(query_edge); + } + } + _ => {} + } + } + } + } + + let Touch { + touching, + along, + query_touching, + query_along, + .. + } = &mut touch; + touching.retain(|edge_idx| !along.contains_key(edge_idx)); + query_touching.retain(|edge| !query_along.contains(edge)); + Contact::Touch(touch) + } +} + +/// Prepared contains query, built once from a query polygon and applied per-segment. Holds the +/// query-local CellIndex, the covering, and the interior/boundary classification. +pub struct ContainsQuery { + region: QueryRegion, +} + +impl ContainsQuery { + /// Build the query from a smashed GeometrySet. + pub fn new(set: GeometrySet, options: CovererOptions) -> Self { + Self { + region: QueryRegion::new(set, options), + } + } + + /// Search one segment for geometries contained by the query polygon. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + let candidates = self + .region + .collect_candidates(cell_reader, None, edge_cache); + self.verify_candidates(candidates, cell_reader, edge_cache) + } + + /// Search one segment, skipping shapes whose doc_id is not in the filter bitset. + pub fn search_segment_filtered( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: &BitSet, + ) -> Vec { + let candidates = + self.region + .collect_candidates(cell_reader, Some(terms_filter), edge_cache); + self.verify_candidates(candidates, cell_reader, edge_cache) + } + + /// Verify each candidate using the cascade. A geometry set is contained only when every one + /// of its members is: a member that is not a candidate lies outside the covering. + fn verify_candidates( + &self, + candidates: HashMap, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + // Keyed by the set's head: its document and the members not yet verified. + let mut unverified: HashMap = HashMap::new(); + let mut rejected: HashSet = HashSet::new(); + + for (geometry_id, info) in candidates { + let entry = edge_cache.get(geometry_id); + let head = (geometry_id.0, entry.head()); + if rejected.contains(&head) { + continue; + } + if !self.verify_one(geometry_id, &info, cell_reader, edge_cache) { + rejected.insert(head); + continue; + } + unverified + .entry(head) + .or_insert((entry.doc_id(), entry.member_count())) + .1 -= 1; + } + + unverified + .into_iter() + .filter(|(head, (_, members))| *members == 0 && !rejected.contains(head)) + .map(|(_, (doc_id, _))| doc_id) + .collect() + } + + /// Verify a single candidate geometry. If the candidate has edges in boundary covering cells, + /// test those edges against the query polygon's edges in the same cells. Any interior crossing + /// means the candidate is not contained. Otherwise every part of the candidate's boundary + /// must lie inside the query polygon, and no hole of the query polygon inside the candidate. + fn verify_one( + &self, + geometry_id: GeometryId, + info: &CandidateInfo, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> bool { + let entry = edge_cache.get(geometry_id); + let candidate = entry.edge_set(); + if candidate.vertices.is_empty() { + return false; + } + // Multiple members can cover each other's boundaries or jointly enclose a hole, so + // candidates are verified against their exposed union boundary, as they are where the + // boundaries meet other than at shared vertices and edges. + if self.region.members().len() > 1 { + let boundary = self.region.boundary(); + return self.verify_union(geometry_id, candidate, boundary, cell_reader, edge_cache); + } + let touch = match self.region.boundary_contact(info, candidate) { + Contact::Cross => return false, + Contact::Touch(touch) => touch, + }; + if touch.irregular { + let boundary = self.region.boundary(); + return self.verify_union(geometry_id, candidate, boundary, cell_reader, edge_cache); + } + + let touches_outside = touch + .touching + .iter() + .any(|&edge_idx| !self.region.contains(&midpoint(entry.edge(edge_idx)))); + if touches_outside { + return false; + } + for ring in candidate.ring_offsets.windows(2) { + if !touch.candidate_ring_meets(ring[0]..ring[1]) + && !self.region.contains(&candidate.vertices[ring[0]]) + { + return false; + } + } + if !candidate.closed { + return true; + } + + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + let mut candidate_contains = + |point: &[f64; 3]| contains_point(&mut segment, geometry_id, point); + + // Where the boundaries run together, the candidate's interior must lie on the query + // polygon's side, however narrow either polygon is. + for &edge_idx in touch.along.keys() { + if !self + .region + .covers_shared_interior(candidate, edge_idx, &touch) + { + return false; + } + } + + // A ring of the query polygon inside the candidate leaves part of it uncovered: a hole, + // or the exterior of a polygon larger than a hemisphere, which the candidate surrounds. + for (member, edge_set) in self.region.members().iter().enumerate() { + let member = member as u32; + if !edge_set.closed { + continue; + } + for hole in edge_set.ring_offsets.windows(2) { + let hole_inside = if touch.query_ring_meets(member, hole[0]..hole[1]) { + touch + .query_ring_touching(member, hole[0]..hole[1]) + .any(|edge_idx| { + candidate_contains(&midpoint(self.region.query_edge(member, edge_idx))) + }) + } else { + candidate_contains(&edge_set.vertices[hole[0]]) + }; + if hole_inside { + return false; + } + } + } + true + } + + /// Every piece of the candidate's boundary must be covered. For a polygon, no exposed + /// query boundary may lie inside it: that would leave an uncovered part of its interior. + fn verify_union( + &self, + geometry_id: GeometryId, + candidate: &EdgeSet, + boundary: &[PolygonEdge], + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> bool { + let edges: Vec<_> = member_edges(candidate).map(|(_, edge)| edge).collect(); + if edges.is_empty() { + return candidate + .vertices + .iter() + .all(|point| self.region.contains(point)); + } + for edge in &edges { + for piece in split_edge(*edge, boundary) { + let shared = piece.shared_sides.iter().any(|&side| side); + let covered = if shared { + !candidate.closed || piece.shared_sides[usize::from(piece.edge.interior_left)] + } else { + self.region.contains(&piece.edge.midpoint()) + }; + if !covered { + return false; + } + } + } + if !candidate.closed { + return true; + } + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + for edge in boundary { + for piece in split_edge(*edge, &edges) { + if !piece.shared_sides.iter().any(|&side| side) + && contains_point(&mut segment, geometry_id, &piece.edge.midpoint()) + { + return false; + } + } + } + true + } +} + +/// Prepared query for the geometries that contain a query geometry: a contains query run the +/// other way around. A geometry contains the query geometry only if it meets it, so the query's +/// covering finds the candidates. A contains query built from each candidate's geometry set then +/// searches an index of the query geometry alone, as a segment stores it, so that the members of +/// a set may jointly contain it. +pub struct ContainingQuery { + region: QueryRegion, + options: CovererOptions, + /// Whether a member of the query geometry has an interior, which no point or line contains. + has_interior: bool, + /// The query geometry's cell index. + cells: Vec, + /// The query geometry's edge index. + edges: Vec, +} + +impl ContainingQuery { + /// Build the query from a smashed GeometrySet. + pub fn new(set: GeometrySet, options: CovererOptions) -> io::Result { + let cell_index = + IndexBuilder::new(BuildOptions::default()).build_from_sets(std::slice::from_ref(&set)); + let mut cells = CountingWriter::wrap(Vec::new()); + cell_index.write(&mut cells)?; + let mut edges = CountingWriter::wrap(Vec::new()); + let mut edge_writer = EdgeWriter::::new(&mut edges, EDGE_SKIP_INTERVAL); + edge_writer.insert(&set)?; + edge_writer.finish()?; + Ok(Self { + has_interior: set.members.iter().any(|member| member.closed), + region: QueryRegion::new(set, options.clone()), + options, + cells: cells.finish(), + edges: edges.finish(), + }) + } + + /// Search one segment for geometries that contain the query geometry. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + const WRITTEN: &str = "the query geometry's index was written by the query"; + let query_cells = CellIndexReader::open(&self.cells).expect(WRITTEN); + let query_edges = EdgeReader::::open(&self.edges).expect(WRITTEN); + let mut query_cache = EdgeCache::new(vec![query_edges], usize::MAX); + + let candidates = self + .region + .collect_candidates(cell_reader, None, edge_cache); + let mut heads: HashSet = HashSet::new(); + let mut doc_ids = Vec::new(); + for geometry_id in candidates.into_keys() { + let entry = edge_cache.get(geometry_id); + // The contains query tests the whole set, once. + if (self.has_interior && !entry.edge_set().closed) + || !heads.insert((geometry_id.0, entry.head())) + { + continue; + } + let contains = ContainsQuery::new(entry.geometry_set().clone(), self.options.clone()); + if !contains + .search_segment(&query_cells, &mut query_cache) + .is_empty() + { + doc_ids.push(entry.doc_id()); + } + } + doc_ids + } +} + +/// Tracks what we know about a candidate geometry from the covering scan. +pub(crate) struct CandidateInfo { + /// Candidate edges grouped by the boundary covering cell they appeared in. + pub(crate) boundary_edges: Vec, +} + +impl CandidateInfo { + pub(crate) fn new() -> Self { + Self { + boundary_edges: Vec::new(), + } + } +} + +/// Candidate edge indices from a single boundary covering cell. +pub(crate) struct BoundaryEdges { + /// The position in the covering of the cell these edges appeared in, used to look up the + /// query polygon's edges in the same cell. + pub(crate) covering_index: usize, + /// Edge indices into the candidate geometry's vertex array. + pub(crate) edge_indices: Vec, +} diff --git a/src/spatial/crossings.rs b/src/spatial/crossings.rs new file mode 100644 index 0000000000..b917959e5b --- /dev/null +++ b/src/spatial/crossings.rs @@ -0,0 +1,850 @@ +//! Edge crossing predicates and the S2EdgeCrosser for efficient edge intersection testing. +//! +//! The predicates in this module are robust, meaning they produce correct, consistent results even +//! in pathological cases. +//! +//! From s2predicates.h, s2predicates.cc, s2edge_crossings.h and s2edge_crossings.cc. +use num_traits::Zero; + +use super::exact::{exact_cross, expensive_sign, to_exact, ExactPoint}; +use super::math::{cross, dot, ldexp, norm2}; + +/// Error bound for tangent plane rejection test in edge crossing. +/// +/// The maximum error in computing tangent planes for fast rejection is: +/// - Error in RobustCrossProd() is insignificant +/// - Error in CrossProd() call: (0.5 + 1/sqrt(3)) * DBL_EPSILON +/// - Error in each DotProd() call: DBL_EPSILON +/// - Total: (1.5 + 1/sqrt(3)) * DBL_EPSILON +/// +/// CrossingSignInternal2 in s2edge_crosser.cc +const TANGENT_ERROR: f64 = (1.5 + 1.0 / 1.732050807568877_f64) * f64::EPSILON; + +/// sqrt(3) constant for error bound computations. +/// +/// Google S2 value 1.7320508075688772935274463415058 is truncated to fit f64. +/// +/// kSqrt3 in s2predicates_internal.cc +const SQRT3: f64 = 1.732_050_807_568_877_2_f64; + +/// Rounding error for f64 arithmetic: 0.5 * f64::EPSILON = 2^(-53) using the simplifed method +/// mentioned in s2predicates_internal.h. +/// +/// DBL_ERR in s2predicates_internal.h +const DBL_ERR: f64 = 0.5 * f64::EPSILON; + +/// Upper bound on the angle (in radians) between the vector returned by robust_cross(a, b) and the +/// true cross product of a and b. +/// +/// This value can be set somewhat arbitrarily because the algorithm uses more precision as needed +/// to achieve the specified error. The only strict requirement is kRobustCrossProdError >= DBL_ERR +/// (since this is the minimum error even when using exact arithmetic). We set the error to +/// 6 * DBL_ERR (= 3 * DBL_EPSILON) so that virtually all cases can be handled using ordinary +/// double-precision arithmetic. +/// +/// kRobustCrossProdError in s2edge_crossings.h +/// static_assert on kRobustCrossProdError in s2edge_crossings.cc +pub const ROBUST_CROSS_PROD_ERROR: f64 = 6.0 * DBL_ERR; + +/// Minimum norm threshold for the stable cross product computation. +/// +/// The directional error in computing (a-b) x (a+b) for unit vectors a, b is: +/// +/// ```text +/// (1 + 2*sqrt(3) + 32*sqrt(3)*DBL_ERR/||N||) * T_ERR +/// ```` +/// +/// where ||N|| is the norm of the result and T_ERR is the rounding error for the precision +/// type T (for f64, T_ERR = DBL_ERR = 0.5 * EPSILON). +/// +/// To keep this error to at most ROBUST_CROSS_PROD_ERROR, we require: +/// +/// ```text +/// ||N|| >= 32*sqrt(3)*DBL_ERR / (ROBUST_CROSS_PROD_ERROR/T_ERR - (1 + 2*sqrt(3))) +/// ```` +/// +/// For double precision where T_ERR = DBL_ERR: +/// +/// ```text +/// kMinNorm = (32 * sqrt(3) * DBL_ERR) / (6.0 - (1 + 2*sqrt(3))) +/// = (32 * sqrt(3) * DBL_ERR) / (5 - 2*sqrt(3)) +/// ~= 4.00644e-15 (~= 18.04 * f64::EPSILON) +/// ```` +/// +/// GetStableCrossProd in s2edge_crossings.cc +pub const ROBUST_CROSS_PROD_MIN_NORM: f64 = + (32.0 * SQRT3 * DBL_ERR) / (ROBUST_CROSS_PROD_ERROR / DBL_ERR - (1.0 + 2.0 * SQRT3)); + +const MAX_DET_ERROR: f64 = 3.6548 * f64::EPSILON; + +/// This version of Sign returns +1 if the points are definitely CCW, -1 if they are definitely CW, +/// and 0 if two points are identical or the result is uncertain. Uncertain cases can be resolved, +/// if desired, by calling ExpensiveSign. +/// +/// The purpose of this method is to allow additional cheap tests to be done, where possible, in +/// order to avoid calling ExpensiveSign unnecessarily. +/// +/// REQUIRES: a_cross_b == cross(a, b) +pub fn triage_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], a_cross_b: &[f64; 3]) -> i32 { + debug_assert!(norm2(a) <= 2.0 + f64::EPSILON); + debug_assert!(norm2(b) <= 2.0 + f64::EPSILON); + debug_assert!(norm2(c) <= 2.0 + f64::EPSILON); + + let det = dot(a_cross_b, c); + + if det > MAX_DET_ERROR { + 1 + } else if det < -MAX_DET_ERROR { + -1 + } else { + 0 + } +} + +/// Returns +1 if the points A, B, C are counterclockwise, -1 if the points are clockwise, and 0 if +/// any two points are the same. This function is essentially like taking the sign of the +/// determinant of ABC, except that it has additional logic to make sure that the above properties +/// hold even when the three points are coplanar, and to deal with the limitations of +/// floating-point arithmetic. +/// +/// Sign satisfies the following conditions: +/// +/// (1) Sign(a,b,c) == 0 if and only if a == b, b == c, or c == a +/// (2) Sign(b,c,a) == Sign(a,b,c) for all a,b,c +/// (3) Sign(c,b,a) == -Sign(a,b,c) for all a,b,c +/// +/// In other words: +/// +/// (1) The result is zero if and only if two points are the same. +/// (2) Rotating the order of the arguments does not affect the result. +/// (3) Exchanging any two arguments inverts the result. +/// +/// On the other hand, note that it is not true in general that Sign(-a,b,c) == -Sign(a,b,c), or +/// any similar identities involving antipodal points. +pub fn sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3]) -> i32 { + // We don't need RobustCrossProd() here because Sign() does its own error estimation and calls + // ExpensiveSign() if there is any uncertainty about the result. + sign_with_cross(a, b, c, &cross(a, b)) +} + +/// Like sign(), but allows the precomputed cross-product of A and B to be specified. +pub fn sign_with_cross(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], a_cross_b: &[f64; 3]) -> i32 { + let result = triage_sign(a, b, c, a_cross_b); + if result != 0 { + return result; + } + expensive_sign(a, b, c, true) +} + +/// Given 4 points on the unit sphere, return true if the edges OA, OB, and OC are encountered in +/// that order while sweeping CCW around the point O. You can think of this as testing whether A <= +/// B <= C with respect to the CCW ordering around O that starts at A, or equivalently, whether B +/// is contained in the range of angles (inclusive) that starts at A and extends +/// CCW to C. +/// +/// Properties: +/// +/// (1) If OrderedCCW(a,b,c,o) && OrderedCCW(b,a,c,o), then a == b +/// (2) If OrderedCCW(a,b,c,o) && OrderedCCW(a,c,b,o), then b == c +/// (3) If OrderedCCW(a,b,c,o) && OrderedCCW(c,b,a,o), then a == b == c +/// (4) If a == b or b == c, then OrderedCCW(a,b,c,o) is true +/// (5) Otherwise if a == c, then OrderedCCW(a,b,c,o) is false +/// +/// REQUIRES: a != o && b != o && c != o +pub fn ordered_ccw(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], o: &[f64; 3]) -> bool { + debug_assert!(a != o); + debug_assert!(b != o); + debug_assert!(c != o); + + // Count how many of the three "is_left_of" tests pass. + // The last test uses > instead of >= so that: + // - We return true if A == B or B == C + // - We return false if A == C (and A != B) + let mut sum = 0; + if sign(b, o, a) >= 0 { + sum += 1; + } + if sign(c, o, b) >= 0 { + sum += 1; + } + if sign(a, o, c) > 0 { + sum += 1; + } + sum >= 2 +} + +/// Returns a vector whose direction is guaranteed to be very close to the exact mathematical cross +/// product of the given unit-length vectors "a" and "b", but whose magnitude is arbitrary. Unlike +/// a.CrossProd(b), this statement is true even when "a" and "b" are very nearly +/// parallel (i.e., a ~= b or a ~= -b). Specifically, the direction of the result vector differs +/// from the exact cross product by at most kRobustCrossProdError radians (see below). +/// +/// When a == -b exactly, the result is consistent with the symbolic perturbation model used by +/// S2::Sign (see s2predicates.h). In other words, even antipodal point pairs have a consistent +/// and well-defined edge between them. (In fact this is true for any pair of distinct points +/// whose vectors are parallel.) +/// +/// When a == b exactly, an arbitrary vector orthogonal to "a" is returned. [From a strict +/// mathematical viewpoint it would be better to return (0, 0, 0), but this behavior helps to avoid +/// special cases in client code.] +/// +/// This function has the following properties (RCP == RobustCrossProd): +/// +/// (1) RCP(a, b) != 0 for all a, b +/// (2) RCP(b, a) == -RCP(a, b) unless a == b +/// (3) RCP(-a, b) == -RCP(a, b) unless a and b are exactly proportional +/// (4) RCP(a, -b) == -RCP(a, b) unless a and b are exactly proportional +/// +/// Note that if you want the result to be unit-length, you must call Normalize() explicitly. (The +/// result is always scaled such that Normalize() can be called without precision loss due to +/// floating-point underflow.) +pub fn robust_cross(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + debug_assert!((norm2(a) - 1.0).abs() < 5.0 * f64::EPSILON); + debug_assert!((norm2(b) - 1.0).abs() < 5.0 * f64::EPSILON); + + // (a - b) x (a + b) = 2 * (a x b), more stable when nearly parallel + let diff = [a[0] - b[0], a[1] - b[1], a[2] - b[2]]; + let sum = [a[0] + b[0], a[1] + b[1], a[2] + b[2]]; + + let result = cross(&diff, &sum); + + let norm2 = norm2(&result); + if norm2 >= ROBUST_CROSS_PROD_MIN_NORM * ROBUST_CROSS_PROD_MIN_NORM { + return result; + } + + if a == b { + return ortho(a); + } + + exact_cross_prod(a, b) +} + +// Computes the cross product of two unit-length vectors using exact arithmetic. If the exact +// result is non-zero, converts it to a normalizable f64 vector. Otherwise falls back to symbolic +// perturbation. +fn exact_cross_prod(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + debug_assert!(a != b); + let xa = to_exact(a); + let xb = to_exact(b); + let result = exact_cross(&xa, &xb); + if !result[0].is_zero() || !result[1].is_zero() || !result[2].is_zero() { + normalizable_from_exact(&result) + } else { + symbolic_cross_prod(a, b) + } +} + +// Converts a BigRational vector to f64, scaling as necessary to ensure the result can be +// normalized without loss of precision due to underflow. When the exact components are too +// small for direct f64 conversion, we scale in the exact domain first -- dividing by the +// largest absolute component so the largest becomes +/-1.0, then converting. The division +// is exact in BigRational; only the final to_f64() rounds. +fn normalizable_from_exact(xf: &ExactPoint) -> [f64; 3] { + use num_traits::{Signed, ToPrimitive}; + + // Try direct conversion. + let x = [ + xf[0].to_f64().unwrap_or(0.0), + xf[1].to_f64().unwrap_or(0.0), + xf[2].to_f64().unwrap_or(0.0), + ]; + if is_normalizable(&x) { + return x; + } + + // Direct conversion underflowed. Scale in the exact domain by dividing + // by the largest absolute component, then convert. + let mut abs_max = xf[0].abs(); + for i in 1..3 { + let a = xf[i].abs(); + if a > abs_max { + abs_max = a; + } + } + + if abs_max.is_zero() { + return [0.0, 0.0, 0.0]; + } + + [ + (&xf[0] / &abs_max) + .to_f64() + .expect("scaled component in [-1, 1]"), + (&xf[1] / &abs_max) + .to_f64() + .expect("scaled component in [-1, 1]"), + (&xf[2] / &abs_max) + .to_f64() + .expect("scaled component in [-1, 1]"), + ] +} + +// In S2, ortho is a general-purpose geometric utility that returns any orthogonal vector, while +// ref_dir is a semantic construct for the semi-open boundary model used in vertex containment +// queries. They share an implementation, but code calling ref_dir is participating in the +// containment protocol, while code calling ortho is just doing geometry. S2 separated them for +// clarity of intent, and the one-line wrapper is intentional. + +/// Returns a unit-length vector that is orthogonal to "a". Satisfies Ortho(-a) = -Ortho(a) for +/// all a. +/// +/// Note that Vector3_d also defines an "Ortho" method, but this one is preferred for use in S2 +/// code because it explicitly tries to avoid result coordinates that are zero. (This is a +/// performance optimization that reduces the amount of time spent in functions that handle +/// degeneracies.) +/// +/// S2::Ortho in s2pointutil.h and s2pointutil.cc +fn ortho(a: &[f64; 3]) -> [f64; 3] { + let ax = a[0].abs(); + let ay = a[1].abs(); + let az = a[2].abs(); + + // Strict > mirrors LargestAbsComponent() in vector.h. + let largest = if ax > ay { + if ax > az { + 0 + } else { + 2 + } + } else if ay > az { + 1 + } else { + 2 + }; + + // s2pointutil.cc: k = LargestAbsComponent() - 1, wrapping to 2. + let k = if largest == 0 { 2 } else { largest - 1 }; + let mut temp = [0.012, 0.0053, 0.00457]; + temp[k] = 1.0; + + let cross = cross(a, &temp); + + // Normalize + let norm = norm2(&cross).sqrt(); + [cross[0] / norm, cross[1] / norm, cross[2] / norm] +} + +/// Returns a unit-length vector used as the reference direction for deciding whether a polygon +/// with semi-open boundaries contains the given vertex "a" (see S2ContainsVertexQuery). The +/// result is unit length and is guaranteed to be different from the given point "a". +/// +/// S2::RefDir in s2pointutil.h +#[inline] +pub fn ref_dir(a: &[f64; 3]) -> [f64; 3] { + ortho(a) +} + +/// The exponent of x, as C's ilogb gives it for a finite, nonzero x. +#[inline] +pub(crate) fn ilogb(x: f64) -> i32 { + let bits = x.to_bits(); + let exp = ((bits >> 52) & 0x7FF) as i32; + if exp == 0 { + // Subnormal: true exponent depends on leading zeros in mantissa. + let mantissa = bits & 0x000F_FFFF_FFFF_FFFF; + -1011 - mantissa.leading_zeros() as i32 + } else { + exp - 1023 + } +} + +/// Returns true if the given vector's magnitude is large enough such that the angle to another +/// vector of the same magnitude can be measured using Angle() without loss of precision due to +/// floating-point underflow. (This requirement is also sufficient to ensure that Normalize() can +/// be called without risk of precision loss.) +/// +/// Let ab = RobustCrossProd(a, b) and cd = RobustCrossProd(cd). In order for ab.Angle(cd) to not +/// lose precision, the squared magnitudes of ab and cd must each be at least 2**-484. This +/// ensures that the sum of the squared magnitudes of ab.CrossProd(cd) and ab.DotProd(cd) is at +/// least 2**-968, which ensures that any denormalized terms in these two calculations do not +/// affect the accuracy of the result (since all denormalized numbers are smaller than 2**-1022, +/// which is less than DBL_ERR * 2**-968). +/// +/// The fastest way to ensure this is to test whether the largest component of the result has a +/// magnitude of at least 2**-242. +/// +/// IsNormalizable in s2edge_crossings.cc +#[inline] +fn is_normalizable(p: &[f64; 3]) -> bool { + const MIN_COMPONENT: f64 = 1.0_f64 + / (1_u64 << 62) as f64 + / (1_u64 << 62) as f64 + / (1_u64 << 62) as f64 + / (1_u64 << 56) as f64; // 2^(-242) + p[0].abs().max(p[1].abs().max(p[2].abs())) >= MIN_COMPONENT +} + +/// Scales a 3-vector as necessary to ensure that the result can be normalized +/// without loss of precision due to floating-point underflow. +/// +/// REQUIRES: p != (0, 0, 0) +/// +/// We can't just scale by a fixed factor because the smallest representable double is 2**-1074, so +/// if we multiplied by 2**(1074 - 242) then the result might be so large that we couldn't square +/// it without overflow. +/// +/// Note that we must scale by a power of two to avoid rounding errors, and that the calculation of +/// "pmax" is free because IsNormalizable() is inline. The code below scales "p" such that the +/// largest component is +/// in the range [1, 2). +/// +/// EnsureNormalizable in s2edge_crossings.cc +#[inline] +fn ensure_normalizable(p: [f64; 3]) -> [f64; 3] { + debug_assert!(p != [0.0, 0.0, 0.0]); + if !is_normalizable(&p) { + let p_max = p[0].abs().max(p[1].abs().max(p[2].abs())); + // The expression below avoids signed overflow for any value of ilogb(). + // ldexp(2, -1 - ilogb(p_max)) scales p so largest component is in [1, 2). + let scale = ldexp(2.0, -1 - ilogb(p_max)); + [p[0] * scale, p[1] * scale, p[2] * scale] + } else { + p + } +} + +/// Computes the cross product of a and b using symbolic perturbations. +/// +/// REQUIRES: a < b (lexicographic ordering) +/// REQUIRES: The exact cross product a x b is zero +/// +/// The following code uses the same symbolic perturbation model as S2::Sign. The particular +/// sequence of tests below was obtained using Mathematica (although it would be easy to do it by +/// hand for this simple case). +/// +/// Just like the function SymbolicallyPerturbedSign() in s2predicates.cc, every input coordinate +/// x[i] is assigned a symbolic perturbation dx[i]. We then compute the cross product +/// +/// ```text +/// (a + da).CrossProd(b + db) . +/// ``` +/// +/// The result is a polynomial in the perturbation symbols. For example if we did this in one +/// dimension, the result would be +/// +/// ```text +/// a * b + b * da + a * db + da * db +/// ``` +/// +/// where "a" and "b" have numerical values and "da" and "db" are symbols. In 3 dimensions the +/// result is similar except that the coefficients are 3-vectors rather than scalars. +/// +/// Every possible S2Point has its own symbolic perturbation in each coordinate (i.e., there are +/// about 3 * 2**192 symbols). The magnitudes of the perturbations are chosen such that if x < y +/// lexicographically, the perturbations for "y" are much smaller than the perturbations for "x". +/// Similarly, the perturbations for the coordinates of a given point x are chosen such that dx[0] +/// is much smaller than dx[1] which is much smaller than dx[2]. Putting this together with fact +/// the inputs to this function have been sorted so that a < b lexicographically, this tells us +/// that +/// +/// ```text +/// da[2] > da[1] > da[0] > db[2] > db[1] > db[0] +/// ``` +/// +/// where each perturbation is so much smaller than the previous one that we don't even need to +/// consider it unless the coefficients of all previous perturbations are zero. In fact, each +/// succeeding perturbation is so small that we don't need to consider it unless the coefficient of +/// all products of the previous perturbations are zero. For example, we don't need to consider +/// the coefficient of db[1] unless the coefficient of db[2]*da[0] is zero. +/// +/// The follow code simply enumerates the coefficients of the perturbations (and products of +/// perturbations) that appear in the cross product above, in order of decreasing perturbation +/// magnitude. The first non-zero coefficient determines the result. The easiest way to enumerate +/// the coefficients in the correct order is to pretend that each perturbation is some tiny value +/// "eps" raised to a power of two: +/// +/// ```text +/// eps** 1 2 4 8 16 32 +/// da[2] da[1] da[0] db[2] db[1] db[0] +/// ``` +/// +/// Essentially we can then just count in binary and test the corresponding subset of perturbations +/// at each step. So for example, we must test the coefficient of db[2]*da[0] before db[1] because +/// eps**12 > eps**16. +/// +/// SymbolicCrossProdSorted in s2edge_crossings.cc +fn symbolic_cross_prod_sorted(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + debug_assert!(a < b); + debug_assert!({ + let xa = to_exact(a); + let xb = to_exact(b); + let c = exact_cross(&xa, &xb); + c[0].is_zero() && c[1].is_zero() && c[2].is_zero() + }); + + if b[0] != 0.0 || b[1] != 0.0 { + // da[2] + return [-b[1], b[0], 0.0]; + } + if b[2] != 0.0 { + // da[1] + return [b[2], 0.0, 0.0]; // Note that b[0] == 0. + } + + // None of the remaining cases can occur in practice, because we can only get + // to this point if b = (0, 0, 0). Nevertheless, even (0, 0, 0) has a + // well-defined direction under the symbolic perturbation model. + debug_assert!(b[1] == 0.0 && b[2] == 0.0); // da[0] coefficients (always zero) + + if a[0] != 0.0 || a[1] != 0.0 { + // db[2] + return [a[1], -a[0], 0.0]; + } + + // The following coefficient is always non-zero, so we can stop here. + // + // It may seem strange that we are returning (1, 0, 0) as the cross product + // without even looking at the sign of a[2]. (Wouldn't you expect + // (0, 0, -1) x (0, 0, 0) and (0, 0, 1) x (0, 0, 0) to point in opposite + // directions?) It's worth pointing out that in this function there is *no + // relationship whatsoever* between the vectors "a" and "-a", because the + // perturbations applied to these vectors may be entirely different. This is + // why the identity "RobustCrossProd(-a, b) == -RobustCrossProd(a, b)" does + // not hold whenever "a" and "b" are linearly dependent (i.e., proportional). + // [As it happens the two cross products above actually do point in opposite + // directions, but for example (1, 1, 1) x (2, 2, 2) = (-2, 2, 0) and + // (-1, -1, -1) x (2, 2, 2) = (-2, 2, 0) do not.] + [1.0, 0.0, 0.0] // db[2] * da[1] +} + +/// Computes the cross product of a and b using symbolic perturbations. +/// +/// REQUIRES: a != b +/// +/// SymbolicCrossProdSorted() requires that a < b. +/// +/// SymbolicCrossProd in s2edge_crossings.cc +fn symbolic_cross_prod(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + debug_assert!(a != b); + // SymbolicCrossProdSorted() requires that a < b. + if a < b { + ensure_normalizable(symbolic_cross_prod_sorted(a, b)) + } else { + let result = ensure_normalizable(symbolic_cross_prod_sorted(b, a)); + [-result[0], -result[1], -result[2]] + } +} + +/// Compute the determinant in a numerically stable way. Unlike triage_sign(), this method can +/// usually compute the correct determinant sign even when all three points are as collinear as +/// possible. For example if three points are spaced 1km apart along a random line on the Earth's +/// surface using the nearest representable points, there is only a 0.4% chance that this method +/// will not be able to find the determinant sign. The probability of failure decreases as the +/// points get closer together; if the collinear points are +/// 1 meter apart, the failure rate drops to 0.0004%. +/// +/// This method could be extended to also handle nearly-antipodal points (and in fact an earlier +/// version of this code did exactly that), but antipodal points are rare in practice so it seems +/// better to simply fall back to exact arithmetic in that case. +/// +/// StableSign in s2predicates.cc +pub fn stable_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3]) -> i32 { + let ab = [b[0] - a[0], b[1] - a[1], b[2] - a[2]]; + let bc = [c[0] - b[0], c[1] - b[1], c[2] - b[2]]; + let ca = [a[0] - c[0], a[1] - c[1], a[2] - c[2]]; + let ab2 = norm2(&ab); + let bc2 = norm2(&bc); + let ca2 = norm2(&ca); + + // Now compute the determinant ((A-C)x(B-C)).C, where the vertices have been + // cyclically permuted if necessary so that AB is the longest edge. (This + // minimizes the magnitude of cross product.) At the same time we also + // compute the maximum error in the determinant. Using a similar technique + // to the one used for kMaxDetError, the error is at most + // + // |d| <= (3 + 6/sqrt(3)) * |A-C| * |B-C| * e + // + // where e = 0.5 * DBL_EPSILON. If the determinant magnitude is larger than + // this value then we know its sign with certainty. + const DET_ERROR_MULTIPLIER: f64 = 3.2321 * f64::EPSILON; // see above + let (det, max_error) = if ab2 >= bc2 && ab2 >= ca2 { + // AB is the longest edge, so compute (A-C)x(B-C).C. + let det = -dot(&cross(&ca, &bc), c); + let max_error = DET_ERROR_MULTIPLIER * (ca2 * bc2).sqrt(); + (det, max_error) + } else if bc2 >= ca2 { + // BC is the longest edge, so compute (B-A)x(C-A).A. + let det = -dot(&cross(&ab, &ca), a); + let max_error = DET_ERROR_MULTIPLIER * (ab2 * ca2).sqrt(); + (det, max_error) + } else { + // CA is the longest edge, so compute (C-B)x(A-B).B. + let det = -dot(&cross(&bc, &ab), b); + let max_error = DET_ERROR_MULTIPLIER * (bc2 * ab2).sqrt(); + (det, max_error) + }; + + // Precomputed: 3.2321 * f64::EPSILON * f64::MIN_POSITIVE.sqrt() + // = 7.176703675781937e-16 * 1.4916681462400413e-154 + const MIN_NO_UNDERFLOW_ERROR: f64 = 1.0705260268167732e-169; + if max_error < MIN_NO_UNDERFLOW_ERROR { + return 0; + } + + if det.abs() <= max_error { + 0 + } else if det > 0.0 { + 1 + } else { + -1 + } +} + +/// Given two edges AB and CD where at least two vertices are identical +/// (i.e. CrossingSign(a,b,c,d) == 0), this function defines whether the two edges "cross" in such +/// a way that point-in-polygon containment tests can be implemented by counting the number of edge +/// crossings. The basic rule is that a "crossing" occurs if AB is encountered after CD during a +/// CCW sweep around the shared vertex starting from a fixed reference point. +/// +/// Note that according to this rule, if AB crosses CD then in general CD does not cross AB. +/// However, this leads to the correct result when counting polygon edge crossings. For example, +/// suppose that A,B,C are three consecutive vertices of a CCW polygon. If we now consider the +/// edge crossings of a segment BP as P sweeps around B, the crossing number +/// changes parity exactly when BP crosses BA or BC. +/// +/// Useful properties of VertexCrossing (VC): +/// +/// (1) VC(a,a,c,d) == VC(a,b,c,c) == false +/// (2) VC(a,b,a,b) == VC(a,b,b,a) == true +/// (3) VC(a,b,c,d) == VC(a,b,d,c) == VC(b,a,c,d) == VC(b,a,d,c) +/// (3) If exactly one of a,b equals one of c,d, then exactly one of +/// VC(a,b,c,d) and VC(c,d,a,b) is true +/// +/// It is an error to call this method with 4 distinct vertices. +pub fn vertex_crossing(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], d: &[f64; 3]) -> bool { + if a == b || c == d { + return false; + } + + if a == c { + return (b == d) || ordered_ccw(&ref_dir(a), d, b, a); + } + if b == d { + return ordered_ccw(&ref_dir(b), c, a, b); + } + if a == d { + return (b == c) || ordered_ccw(&ref_dir(a), c, b, a); + } + if b == c { + return ordered_ccw(&ref_dir(b), d, a, b); + } + + debug_assert!(false, "vertex_crossing called with 4 distinct vertices"); + false +} + +/// Efficiently tests edges against a fixed edge AB. +/// +/// This is especially efficient when testing against an edge chain connecting vertices +/// v0, v1, v2, ... Use restart_at() to jump to a new location in the chain. +#[derive(Clone)] +pub struct S2EdgeCrosser { + // Fixed edge AB + a: [f64; 3], + b: [f64; 3], + a_cross_b: [f64; 3], + + // Tangent planes for fast rejection (computed lazily) + have_tangents: bool, + a_tangent: [f64; 3], + b_tangent: [f64; 3], + + // Edge chain state + c: [f64; 3], + acb: i32, // Orientation of triangle ACB + + // Temporary for CrossingSignInternal + bda: i32, // Orientation of triangle BDA +} + +impl S2EdgeCrosser { + /// Creates a new edge crosser for testing edges against the fixed edge AB. + pub fn new(a: &[f64; 3], b: &[f64; 3]) -> Self { + debug_assert!((norm2(a) - 1.0).abs() < 5.0 * f64::EPSILON); + debug_assert!((norm2(b) - 1.0).abs() < 5.0 * f64::EPSILON); + + Self { + a: *a, + b: *b, + a_cross_b: cross(a, b), + have_tangents: false, + a_tangent: [0.0, 0.0, 0.0], + b_tangent: [0.0, 0.0, 0.0], + c: [0.0, 0.0, 0.0], + acb: 0, + bda: 0, + } + } + + /// Returns the first endpoint of the fixed edge. + pub fn a(&self) -> &[f64; 3] { + &self.a + } + + /// Returns the second endpoint of the fixed edge. + pub fn b(&self) -> &[f64; 3] { + &self.b + } + + /// Returns the current chain vertex C. + pub fn c(&self) -> &[f64; 3] { + &self.c + } + + /// Tests whether the fixed edge AB crosses edge CD. + /// + /// This is the two-argument version for testing disconnected edges. + pub fn crossing_sign_two(&mut self, c: &[f64; 3], d: &[f64; 3]) -> i32 { + if *c != self.c { + self.restart_at(c); + } + self.crossing_sign(d) + } + + /// Tests whether the fixed edge AB crosses edge CD using the vertex crossing rule. + pub fn edge_or_vertex_crossing_two(&mut self, c: &[f64; 3], d: &[f64; 3]) -> bool { + if *c != self.c { + self.restart_at(c); + } + self.edge_or_vertex_crossing(d) + } + + /// Sets the current chain vertex to C. + /// + /// Call this to "jump" to a new location in an edge chain. + pub fn restart_at(&mut self, c: &[f64; 3]) { + debug_assert!((norm2(c) - 1.0).abs() < 5.0 * f64::EPSILON); + self.c = *c; + self.acb = -triage_sign(&self.a, &self.b, c, &self.a_cross_b); + } + + /// Tests whether the fixed edge AB crosses edge CD, where C is the previous vertex from the + /// last call. + /// + /// This is the single-argument chain version. Use `restart_at` to set C. + pub fn crossing_sign(&mut self, d: &[f64; 3]) -> i32 { + debug_assert!((norm2(d) - 1.0).abs() < 5.0 * f64::EPSILON); + + // Fast path: test if BDA has opposite orientation from ACB + // TriageSign is invariant under rotation, so ABD has same orientation as BDA + let bda = triage_sign(&self.a, &self.b, d, &self.a_cross_b); + if self.acb == -bda && bda != 0 { + // Most common case: C and D are on the same side of AB + // Save D as the new C for the next call + self.c = *d; + self.acb = -bda; + return -1; + } + self.bda = bda; + self.crossing_sign_internal(d) + } + + /// Returns the crossing sign from the last interior crossing. + /// + /// When `crossing_sign` returns +1, this returns the sign of the crossing: + /// +1 if AB crosses CD from left to right, -1 if right to left. + pub fn last_interior_crossing_sign(&self) -> i32 { + // The crossing sign equals Sign(ABC) which equals acb with sign flipped + self.acb + } + + /// Internal implementation for when fast path doesn't work. + fn crossing_sign_internal(&mut self, d: &[f64; 3]) -> i32 { + // At this point it is still very likely that CD does not cross AB. Two common situations + // are (1) CD crosses the great circle through AB but does not cross AB itself, or (2) + // A,B,C,D are four points on a line such that AB does not overlap CD. For example, the + // latter happens when a line or curve is sampled finely, or when geometry is constructed + // by computing the union of S2CellIds. + // + // Most of the time, we can determine that AB and CD do not intersect by computing the two + // outward-facing tangents at A and B (parallel to AB) and testing whether AB and CD are on + // opposite sides of the plane perpendicular to one of these tangents. This is somewhat + // expensive but still much cheaper than s2pred::ExpensiveSign. + if !self.have_tangents { + let norm = robust_cross(&self.a, &self.b); + self.a_tangent = cross(&self.a, &norm); + self.b_tangent = cross(&norm, &self.b); + self.have_tangents = true; + } + + // The error in RobustCrossProd() is insignificant. The maximum error in the call to + // CrossProd() (i.e., the maximum norm of the error vector) is + // (0.5 + 1/sqrt(3)) * DBL_EPSILON. The maximum error in each call to DotProd() below is + // DBL_EPSILON. (There is also a small relative error term that is insignificant because + // we are comparing the result against a constant that is very close to zero.) + if (dot(&self.c, &self.a_tangent) > TANGENT_ERROR + && dot(d, &self.a_tangent) > TANGENT_ERROR) + || (dot(&self.c, &self.b_tangent) > TANGENT_ERROR + && dot(d, &self.b_tangent) > TANGENT_ERROR) + { + self.c = *d; + self.acb = -self.bda; + return -1; + } + + // Otherwise, eliminate the cases where two vertices from different edges are equal. + // (These cases could be handled in the code below, but we would rather avoid calling + // ExpensiveSign whenever possible.) + if self.a == self.c || self.a == *d || self.b == self.c || self.b == *d { + self.c = *d; + self.acb = -self.bda; + return 0; + } + + // Eliminate cases where an input edge is degenerate. (Note that in most cases, if CD is + // degenerate then this method is not even called because acb_ and bda have different + // signs.) + if self.a == self.b || self.c == *d { + self.c = *d; + self.acb = -self.bda; + return -1; + } + + // Otherwise it's time to break out the big guns. + if self.acb == 0 { + self.acb = -expensive_sign(&self.a, &self.b, &self.c, true); + } + debug_assert!(self.acb != 0); + + if self.bda == 0 { + self.bda = expensive_sign(&self.a, &self.b, d, true); + } + debug_assert!(self.bda != 0); + + if self.bda != self.acb { + self.c = *d; + self.acb = -self.bda; + return -1; + } + + let c_cross_d = cross(&self.c, d); + let cbd = -sign_with_cross(&self.c, d, &self.b, &c_cross_d); + if cbd != self.acb { + self.c = *d; + self.acb = -self.bda; + return -1; + } + + let dac = sign_with_cross(&self.c, d, &self.a, &c_cross_d); + let result = if dac != self.acb { -1 } else { 1 }; + + self.c = *d; + self.acb = -self.bda; + result + } + + /// Tests edge crossing with vertex crossing rule included. + pub fn edge_or_vertex_crossing(&mut self, d: &[f64; 3]) -> bool { + // Save c since crossing_sign will overwrite it + let c = self.c; + let crossing = self.crossing_sign(d); + if crossing < 0 { + return false; + } + if crossing > 0 { + return true; + } + vertex_crossing(&self.a, &self.b, &c, d) + } +} diff --git a/src/spatial/edge_cache.rs b/src/spatial/edge_cache.rs new file mode 100644 index 0000000000..4f7de01d7f --- /dev/null +++ b/src/spatial/edge_cache.rs @@ -0,0 +1,228 @@ +//! Cached geometry set access over one or more edge readers. +//! +//! Wraps EdgeReaders with an LRU cache keyed by (segment, head_position). The query path uses +//! segment 0 with a single reader. The merge path uses real segment indices with one reader per +//! source segment and a shared eviction budget. +//! +//! The cache uses interior mutability so callers take &EdgeCache. Returns Rc so the +//! data survives eviction. The budget, in bytes of heap the cached sets hold, is advisory: it +//! grows to four times the largest set the cache holds, so that a few sets larger than the +//! budget are not decoded again for every cell they share with each other and other geometries. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; + +use lru::LruCache; + +use super::cell_index::GeometryId; +use super::edge_reader::{EdgeReader, MemberSummary}; +use super::geometry_set::{EdgeSet, GeometrySet}; +use super::surface::Surface; + +struct CachedSet { + head_position: u32, + set: Rc, +} + +struct CacheInner<'a, S: Surface> { + readers: Vec>, + index: HashMap<(usize, u32), (usize, u32)>, + sets: LruCache<(usize, u32), CachedSet>, + cached_bytes: usize, + max_bytes: usize, +} + +/// A geometry set entry returned from the cache. Holds an Rc to the set and knows its member +/// index within the set. The Rc keeps the data alive even if the cache evicts the entry. +pub struct GeometryEntry { + set: Rc, + head: u32, + member_idx: usize, +} + +impl GeometryEntry { + /// The vertices for this member. + pub fn vertices(&self) -> &Vec<[f64; 3]> { + &self.set.members[self.member_idx].vertices + } + + /// The full edge set for this member. + pub fn edge_set(&self) -> &EdgeSet { + &self.set.members[self.member_idx] + } + + /// The document id. + pub fn doc_id(&self) -> u32 { + self.set.doc_id + } + + /// The head position of the set this member belongs to. + pub fn head(&self) -> u32 { + self.head + } + + /// The number of members in the geometry set. + pub fn member_count(&self) -> u32 { + self.set.members.len() as u32 + } + + /// The full geometry set. + pub fn geometry_set(&self) -> &GeometrySet { + &self.set + } + + /// The two endpoints of the given edge. For point geometries (single vertex), returns the + /// vertex twice. + pub fn edge(&self, index: u32) -> ([f64; 3], [f64; 3]) { + let vertices = self.vertices(); + let i = index as usize; + let v0 = vertices[i]; + let v1 = if i + 1 < vertices.len() { + vertices[i + 1] + } else { + debug_assert_eq!( + i, 0, + "edge index past vertex count is only valid for points" + ); + v0 + }; + (v0, v1) + } +} + +/// Cached geometry set access. Keyed by (segment, position) so the query path and merge path use +/// the same type. The query path passes segment 0. The merge path passes the real segment index. +/// Uses interior mutability so callers take &EdgeCache. +pub struct EdgeCache<'a, S: Surface> { + inner: RefCell>, +} + +impl<'a, S: Surface> EdgeCache<'a, S> { + /// Create a cache over the given readers. The max_bytes budget is shared across all + /// segments. Eviction removes the least recently used set regardless of which segment it came + /// from. + pub fn new(readers: Vec>, max_bytes: usize) -> Self { + EdgeCache { + inner: RefCell::new(CacheInner { + readers, + index: HashMap::new(), + sets: LruCache::unbounded(), + cached_bytes: 0, + max_bytes, + }), + } + } + + /// Total number of geometry entries in the given segment. + pub fn geometry_count(&self, segment: usize) -> u32 { + self.inner.borrow().readers[segment].geometry_count() + } + + /// Resolve a geometry to its doc_id without decoding vertices or touching the cache. + pub fn doc_id_for(&self, id: GeometryId) -> u32 { + self.inner.borrow().readers[id.0 as usize].doc_id_for(id.1) + } + + /// Summarize a geometry from its entry and its set's head, without decoding its set or + /// touching the cache. + pub fn summary(&self, id: GeometryId) -> MemberSummary { + self.inner.borrow().readers[id.0 as usize].summary(id.1) + } + + /// Retrieve the geometry entry for the given geometry id. The returned entry holds an Rc to + /// the geometry set and knows its member index. The Rc keeps the data alive regardless of + /// cache eviction. + pub fn get(&self, id: GeometryId) -> GeometryEntry { + let mut inner = self.inner.borrow_mut(); + let segment = id.0 as usize; + let position = id.1; + let key = (segment, position); + + let head_key = if let Some(&hk) = inner.index.get(&key) { + hk + } else { + let (head_position, set) = inner.readers[segment].read_geometry_set(position); + let head_key = (segment, head_position); + + let member_count = set.members.len() as u32; + let bytes = set.mem_usage(); + inner.cached_bytes += bytes; + // Leave room for the geometries around a large set, or the set and they would evict + // each other in turn, and the set would be decoded once per cell. + inner.max_bytes = inner.max_bytes.max(4 * bytes); + + for i in 0..member_count { + inner.index.insert((segment, head_position + i), head_key); + } + + inner.sets.put( + head_key, + CachedSet { + head_position, + set: Rc::new(set), + }, + ); + + while inner.cached_bytes > inner.max_bytes && inner.sets.len() > 1 { + if let Some((evicted_key, evicted_set)) = inner.sets.pop_lru() { + inner.cached_bytes -= evicted_set.set.mem_usage(); + for i in 0..evicted_set.set.members.len() as u32 { + inner.index.remove(&(evicted_key.0, evicted_key.1 + i)); + } + } + } + + head_key + }; + + let cached = inner.sets.get(&head_key).unwrap(); + let head = cached.head_position; + let member_idx = (position - head) as usize; + GeometryEntry { + set: Rc::clone(&cached.set), + head, + member_idx, + } + } +} + +#[cfg(test)] +mod tests { + use common::CountingWriter; + + use super::*; + use crate::spatial::edge_writer::EdgeWriter; + use crate::spatial::geometry::Geometry; + use crate::spatial::geometry_set::to_geometry_set; + use crate::spatial::plane::Plane; + use crate::spatial::sphere::Sphere; + + #[test] + fn test_a_set_larger_than_the_budget_stays_cached_among_others() -> crate::Result<()> { + let line = |doc_id: u32, len: usize| { + let points = (0..len) + .map(|i| [i as f64 / 1000.0, doc_id as f64]) + .collect(); + to_geometry_set( + &Geometry::::LineString(points).project::(), + doc_id, + ) + }; + let mut write = CountingWriter::wrap(Vec::new()); + let mut edge_writer = EdgeWriter::::new(&mut write, 16); + for set in [line(0, 1000), line(1, 10), line(2, 10), line(3, 10)] { + edge_writer.insert(&set)?; + } + edge_writer.finish()?; + let data = write.finish(); + let cache = EdgeCache::new(vec![EdgeReader::::open(&data)?], 100); + + let large = cache.get((0, 0)).set; + for small in 1..4 { + cache.get((0, small)); + assert!(Rc::ptr_eq(&large, &cache.get((0, 0)).set), "decoded again"); + } + Ok(()) + } +} diff --git a/src/spatial/edge_reader.rs b/src/spatial/edge_reader.rs new file mode 100644 index 0000000000..d546b18614 --- /dev/null +++ b/src/spatial/edge_reader.rs @@ -0,0 +1,432 @@ +//! Edge index deserialization. +//! +//! Reads geometry entries from a byte slice, typically mmap'd segment data. The 16-byte footer +//! locates the skip list directory. The directory provides random access by geometry position, +//! and records for each entry the head of the set it belongs to, so that finding a head never +//! walks more than two skip intervals. Between skip list entries, a forward walk bridges the gap. +//! +//! Geometry entry format per 004_edge_index.md: +//! flags (u8): bit 0 = closed, bit 1 = contains_hilbert_start, bit 2 = has_holes, bit 3 = is_head +//! len (u32): byte length of data section +//! Head: [flags] [len] [doc_id] [ring boundaries?] [vertex data] +//! Member: [flags] [len] [ring boundaries?] [vertex data] + +use std::marker::PhantomData; + +use super::geometry_set::{EdgeSet, GeometrySet}; +use super::surface::Surface; +use crate::error::DataCorruption; + +/// Bytes in one skip list directory entry: the offset (u64) of a geometry entry and the position +/// (u32) of the head of its set. +const DIRECTORY_ENTRY_LEN: usize = std::mem::size_of::() + std::mem::size_of::(); + +/// Bytes in the footer: the geometry count (u32), the skip interval (u32) and the directory +/// offset (u64). +const FOOTER_LEN: usize = 2 * std::mem::size_of::() + std::mem::size_of::(); + +struct EntryHeader { + data_len: u32, + closed: bool, + contains_hilbert_start: bool, + has_holes: bool, + is_head: bool, +} + +impl EntryHeader { + /// Bytes before the data section: the flags, the data length and, for a head, the doc id. + fn header_len(&self) -> u64 { + if self.is_head { + 9 + } else { + 5 + } + } + + /// Bytes in the whole entry. + fn entry_len(&self) -> u64 { + self.header_len() + self.data_len as u64 + } +} + +/// A member of a geometry set, read from its entry and its set's head without decoding the set. +pub struct MemberSummary { + /// The document of the member's set. + pub doc_id: u32, + /// The position of the head of the member's set. + pub head: u32, + /// Whether the member is the only one of its set. + pub sole: bool, + /// The member's vertex, when it is a point. + pub point: Option<[f64; 3]>, +} + +/// Reads geometry entries from a serialized edge index segment. Pure deserialization with no +/// mutable state. The Surface type parameter determines how many f64 coordinates are read per +/// vertex. +pub struct EdgeReader<'a, S: Surface> { + data: &'a [u8], + geometry_count: u32, + skip_interval: u32, + dir_offset: u64, + _surface: PhantomData, +} + +impl<'a, S: Surface> EdgeReader<'a, S> { + /// Construct a reader from the raw bytes of an edge index segment. Fails when the footer does + /// not locate a skip list directory that ends where the footer begins. + pub fn open(data: &'a [u8]) -> crate::Result { + let data = super::check_format_marker(data, super::EDGE_INDEX_MAGIC, "edge index")?; + let corrupt = + || DataCorruption::comment_only(format!("Corrupt {}-byte edge index", data.len())); + let (body, footer) = data.split_last_chunk::().ok_or_else(corrupt)?; + let geometry_count = u32::from_le_bytes(footer[0..4].try_into().unwrap()); + let skip_interval = u32::from_le_bytes(footer[4..8].try_into().unwrap()); + let dir_offset = u64::from_le_bytes(footer[8..16].try_into().unwrap()); + if skip_interval == 0 { + return Err(corrupt().into()); + } + let dir_len = geometry_count.div_ceil(skip_interval) as u64 * DIRECTORY_ENTRY_LEN as u64; + if dir_offset.checked_add(dir_len) != Some(body.len() as u64) { + return Err(corrupt().into()); + } + + Ok(EdgeReader { + data, + geometry_count, + skip_interval, + dir_offset, + _surface: PhantomData, + }) + } + + /// A reader over no geometries, for a segment that has no geometry for the field. + pub fn empty() -> Self { + EdgeReader { + data: &[], + geometry_count: 0, + skip_interval: 1, + dir_offset: 0, + _surface: PhantomData, + } + } + + /// Total number of geometry entries in this segment. + pub fn geometry_count(&self) -> u32 { + self.geometry_count + } + + /// Resolve a geometry position to its doc_id without decoding vertices. + pub fn doc_id_for(&self, position: u32) -> u32 { + let (_, head_offset) = self.find_head(position); + self.read_doc_id(head_offset) + } + + /// What a query needs of the member at a position to decide it without decoding its set: + /// its document, the head of its set, whether it is the set's only member, and its vertex + /// when it is a point. + pub fn summary(&self, position: u32) -> MemberSummary { + debug_assert!(position < self.geometry_count); + let skip_index = position / self.skip_interval; + let (mut offset, block_head) = self.read_dir_entry(skip_index); + let mut head = None; + let mut header = self.read_header(offset); + for pos in skip_index * self.skip_interval..position { + if header.is_head { + head = Some((pos, offset)); + } + offset += header.entry_len(); + header = self.read_header(offset); + } + let (head, head_offset) = if header.is_head { + (position, offset) + } else { + head.unwrap_or_else(|| (block_head, self.offset_of(block_head))) + }; + let next = position + 1; + let sole = header.is_head + && (next == self.geometry_count + || self.read_header(offset + header.entry_len()).is_head); + let vertex_len = S::DIMENSIONS * 8; + let point = (!header.has_holes && header.data_len as usize == vertex_len).then(|| { + let start = (offset + header.header_len()) as usize; + Self::decode_vertex(&self.data[start..start + vertex_len]) + }); + MemberSummary { + doc_id: self.read_doc_id(head_offset), + head, + sole, + point, + } + } + + /// Read the full geometry set containing the given position. Returns (head_position, set). No + /// caching. Every call decodes from the byte stream. + pub fn read_geometry_set(&self, position: u32) -> (u32, GeometrySet) { + debug_assert!(position < self.geometry_count); + + let (head_position, head_offset) = self.find_head(position); + let doc_id = self.read_doc_id(head_offset); + + let mut members: Vec = Vec::new(); + let mut cur = head_offset; + let mut pos = head_position; + loop { + let h = self.read_header(cur); + if h.is_head && !members.is_empty() { + break; + } + let data_start = cur + h.header_len(); + + let (hole_starts, vertex_start) = if h.has_holes { + self.decode_ring_boundaries(data_start as usize, h.data_len as usize) + } else { + (Vec::new(), data_start as usize) + }; + + let vertex_end = data_start as usize + h.data_len as usize; + let vertices = self.decode_vertices(&self.data[vertex_start..vertex_end]); + + let mut offsets = vec![0usize]; + for &hole_start in &hole_starts { + offsets.push(hole_start); + } + offsets.push(vertices.len()); + + members.push(EdgeSet { + vertices, + closed: h.closed, + contains_hilbert_start: h.contains_hilbert_start, + ring_offsets: offsets, + }); + + cur = data_start + h.data_len as u64; + pos += 1; + + if pos >= self.geometry_count { + break; + } + } + + (head_position, GeometrySet { members, doc_id }) + } + + /// The position and offset of the head of the set containing the given position: the last + /// head at or before it in its skip block, or else the head the directory records for the + /// block, which starts inside that head's set. + fn find_head(&self, position: u32) -> (u32, u64) { + let skip_index = position / self.skip_interval; + let (mut offset, block_head) = self.read_dir_entry(skip_index); + let mut head = None; + for pos in skip_index * self.skip_interval..=position { + let h = self.read_header(offset); + if h.is_head { + head = Some((pos, offset)); + } + offset += h.entry_len(); + } + head.unwrap_or_else(|| (block_head, self.offset_of(block_head))) + } + + /// The offset of the entry at the given position. + fn offset_of(&self, position: u32) -> u64 { + let skip_index = position / self.skip_interval; + let (mut offset, _) = self.read_dir_entry(skip_index); + for _ in skip_index * self.skip_interval..position { + offset += self.read_header(offset).entry_len(); + } + offset + } + + /// The offset of the given skip block's first entry and the position of its set's head. + fn read_dir_entry(&self, skip_index: u32) -> (u64, u32) { + let at = self.dir_offset as usize + skip_index as usize * DIRECTORY_ENTRY_LEN; + let offset = u64::from_le_bytes(self.data[at..at + 8].try_into().unwrap()); + let head_position = u32::from_le_bytes(self.data[at + 8..at + 12].try_into().unwrap()); + (offset, head_position) + } + + /// The doc id of the head entry at the given offset. + fn read_doc_id(&self, head_offset: u64) -> u32 { + let at = head_offset as usize + 5; + u32::from_le_bytes(self.data[at..at + 4].try_into().unwrap()) + } + + fn read_header(&self, offset: u64) -> EntryHeader { + let o = offset as usize; + let flags = self.data[o]; + let data_len = u32::from_le_bytes(self.data[o + 1..o + 5].try_into().unwrap()); + EntryHeader { + data_len, + closed: flags & 0x01 != 0, + contains_hilbert_start: flags & 0x02 != 0, + has_holes: flags & 0x04 != 0, + is_head: flags & 0x08 != 0, + } + } + + fn decode_ring_boundaries(&self, data_start: usize, data_len: usize) -> (Vec, usize) { + let mut hole_starts = Vec::new(); + let mut pos = data_start; + let data_end = data_start + data_len; + loop { + debug_assert!(pos + 4 <= data_end, "ring boundary overruns data section"); + let entry = u32::from_le_bytes(self.data[pos..pos + 4].try_into().unwrap()); + let vertex_index = (entry & 0x7FFFFFFF) as usize; + let continuation = entry & 0x80000000 != 0; + hole_starts.push(vertex_index); + pos += 4; + if !continuation { + break; + } + } + (hole_starts, pos) + } + + fn decode_vertices(&self, data: &[u8]) -> Vec<[f64; 3]> { + data.chunks_exact(S::DIMENSIONS * 8) + .map(Self::decode_vertex) + .collect() + } + + fn decode_vertex(chunk: &[u8]) -> [f64; 3] { + let mut point = [0.0f64; 3]; + for i in 0..S::DIMENSIONS { + point[i] = f64::from_le_bytes(chunk[i * 8..(i + 1) * 8].try_into().unwrap()); + } + point + } +} + +#[cfg(test)] +mod tests { + use std::path::Path; + + use common::{CountingWriter, TerminatingWrite}; + + use super::*; + use crate::directory::{Directory, RamDirectory}; + use crate::spatial::edge_writer::EdgeWriter; + use crate::spatial::geometry::Geometry; + use crate::spatial::geometry_set::to_geometry_set; + use crate::spatial::plane::Plane; + use crate::spatial::sphere::Sphere; + + fn written(sets: &[GeometrySet]) -> crate::Result> { + let directory = RamDirectory::create(); + let path = Path::new("edges"); + let mut write = CountingWriter::wrap(directory.open_write(path)?); + let mut edge_writer = EdgeWriter::::new(&mut write, 2); + for set in sets { + edge_writer.insert(set)?; + } + edge_writer.finish()?; + write.finish().terminate()?; + Ok(directory.open_read(path)?.read_bytes()?.to_vec()) + } + + fn point(doc_id: u32) -> GeometrySet { + let point = Geometry::::Point([doc_id as f64, 0.0]).project::(); + to_geometry_set(&point, doc_id) + } + + #[test] + fn test_open_reads_a_written_index() -> crate::Result<()> { + let data = written(&[point(3), point(5), point(8)])?; + let reader = EdgeReader::::open(&data)?; + assert_eq!(reader.geometry_count(), 3); + assert_eq!(reader.doc_id_for(2), 8); + let footer_only = written(&[])?; + assert_eq!( + footer_only.len(), + FOOTER_LEN + crate::spatial::FORMAT_MARKER_LEN + ); + assert_eq!( + EdgeReader::::open(&footer_only)?.geometry_count(), + 0 + ); + Ok(()) + } + + fn points(doc_id: u32, count: usize) -> GeometrySet { + let points = (0..count).map(|i| [i as f64, doc_id as f64]).collect(); + to_geometry_set( + &Geometry::::MultiPoint(points).project::(), + doc_id, + ) + } + + #[test] + fn test_find_the_set_of_every_position() -> crate::Result<()> { + // With a skip interval of 2, the long set spans several skip blocks. + let sets = [points(3, 1), points(5, 9), points(8, 2), points(9, 1)]; + let data = written(&sets)?; + let reader = EdgeReader::::open(&data)?; + let mut position = 0; + for set in &sets { + let head = position; + for member in &set.members { + assert_eq!( + reader.doc_id_for(position), + set.doc_id, + "position {position}" + ); + let (head_position, read) = reader.read_geometry_set(position); + assert_eq!(head_position, head); + assert_eq!(read.doc_id, set.doc_id); + assert_eq!(read.members.len(), set.members.len()); + assert_eq!( + read.members[(position - head) as usize].vertices, + member.vertices + ); + position += 1; + } + } + assert_eq!(reader.geometry_count(), position); + Ok(()) + } + + #[test] + fn test_summarize_every_position() -> crate::Result<()> { + let line = to_geometry_set( + &Geometry::::LineString(vec![[0.0, 0.0], [1.0, 1.0]]).project::(), + 4, + ); + // With a skip interval of 2, the long set spans several skip blocks. + let sets = [points(3, 1), points(5, 9), line, points(8, 2), points(9, 1)]; + let data = written(&sets)?; + let reader = EdgeReader::::open(&data)?; + let mut position = 0; + for set in &sets { + let head = position; + for member in &set.members { + let summary = reader.summary(position); + assert_eq!(summary.doc_id, set.doc_id, "position {position}"); + assert_eq!(summary.head, head, "position {position}"); + assert_eq!(summary.sole, set.members.len() == 1, "position {position}"); + let point = (member.vertices.len() == 1).then(|| member.vertices[0]); + assert_eq!(summary.point, point, "position {position}"); + position += 1; + } + } + Ok(()) + } + + #[test] + fn test_open_rejects_a_corrupt_index() -> crate::Result<()> { + let data = written(&[point(3), point(5), point(8)])?; + for len in 0..FOOTER_LEN + crate::spatial::FORMAT_MARKER_LEN { + assert!( + EdgeReader::::open(&data[..len]).is_err(), + "{len} bytes" + ); + } + assert!(EdgeReader::::open(&data[1..]).is_err()); + let mut no_skip_interval = data.clone(); + let interval_at = data.len() - crate::spatial::FORMAT_MARKER_LEN - FOOTER_LEN + 4; + no_skip_interval[interval_at..interval_at + 4].copy_from_slice(&0u32.to_le_bytes()); + assert!(EdgeReader::::open(&no_skip_interval).is_err()); + assert_eq!(EdgeReader::::empty().geometry_count(), 0); + Ok(()) + } +} diff --git a/src/spatial/edge_writer.rs b/src/spatial/edge_writer.rs new file mode 100644 index 0000000000..803aef158f --- /dev/null +++ b/src/spatial/edge_writer.rs @@ -0,0 +1,129 @@ +//! Edge index serialization. +//! +//! Stores flattened polygon vertices for query-time crossing tests. Vertices stream through +//! immediately on insert because merge cannot accumulate an entire segment's vertex data in +//! memory. The skip list directory records every Nth geometry offset, with the position of the +//! head of the set that geometry belongs to. + +use std::io::{self, Write}; +use std::marker::PhantomData; + +use common::CountingWriter; + +use crate::directory::WritePtr; +use crate::spatial::geometry_set::GeometrySet; +use crate::spatial::surface::Surface; + +/// Streams geometry entries to disk, or to any other output. The Surface type parameter +/// determines how many f64 coordinates are written per vertex. +pub struct EdgeWriter<'a, S: Surface, W: Write = WritePtr> { + write: &'a mut CountingWriter, + /// Every Nth geometry's offset, and the position of its set's head. + directory: Vec<(u64, u32)>, + geometry_count: u32, + skip_interval: u32, + base: u64, + _surface: PhantomData, +} + +impl<'a, S: Surface, W: Write> EdgeWriter<'a, S, W> { + /// Creates a new writer backed by the given output. + pub fn new(write: &'a mut CountingWriter, skip_interval: u32) -> Self { + let base = write.written_bytes(); + EdgeWriter { + write, + directory: Vec::new(), + geometry_count: 0, + skip_interval, + base, + _surface: PhantomData, + } + } + + /// Write all geometries for one document from a smashed GeometrySet. + pub fn insert(&mut self, set: &GeometrySet) -> io::Result<()> { + self.insert_for_doc(set, set.doc_id) + } + + /// Write all geometries of a set under the given doc id rather than the set's own, as a merge + /// does when it renumbers documents. + pub fn insert_for_doc(&mut self, set: &GeometrySet, doc_id: u32) -> io::Result<()> { + let head_position = self.geometry_count; + for (member_idx, member) in set.members.iter().enumerate() { + let pos = self.write.written_bytes() - self.base; + + if self.geometry_count.is_multiple_of(self.skip_interval) { + self.directory.push((pos, head_position)); + } + self.geometry_count += 1; + + let is_head = member_idx == 0; + let has_holes = member.ring_offsets.len() > 2; + + let flags: u8 = (if member.closed { 0x01 } else { 0 }) + | (if member.contains_hilbert_start { + 0x02 + } else { + 0 + }) + | (if has_holes { 0x04 } else { 0 }) + | (if is_head { 0x08 } else { 0 }); + + let ring_boundary_bytes = if has_holes { + (member.ring_offsets.len() - 2) * 4 + } else { + 0 + }; + let vertex_bytes = member.vertices.len() * S::DIMENSIONS * 8; + let data_bytes = u32::try_from(ring_boundary_bytes + vertex_bytes).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "A geometry of {} vertices is too large", + member.vertices.len() + ), + ) + })?; + + self.write.write_all(&[flags])?; + self.write.write_all(&data_bytes.to_le_bytes())?; + if is_head { + self.write.write_all(&doc_id.to_le_bytes())?; + } + + if has_holes { + let hole_count = member.ring_offsets.len() - 2; + for (h, &offset) in member.ring_offsets[1..member.ring_offsets.len() - 1] + .iter() + .enumerate() + { + let continuation = if h < hole_count - 1 { 0x80000000u32 } else { 0 }; + let entry = continuation | (offset as u32); + self.write.write_all(&entry.to_le_bytes())?; + } + } + + for v in &member.vertices { + for &coord in &v[..S::DIMENSIONS] { + self.write.write_all(&coord.to_le_bytes())?; + } + } + } + Ok(()) + } + + /// Write the skip list directory and footer. + pub fn finish(&mut self) -> io::Result<()> { + let dir_offset = self.write.written_bytes() - self.base; + + for &(offset, head_position) in &self.directory { + self.write.write_all(&offset.to_le_bytes())?; + self.write.write_all(&head_position.to_le_bytes())?; + } + + self.write.write_all(&self.geometry_count.to_le_bytes())?; + self.write.write_all(&self.skip_interval.to_le_bytes())?; + self.write.write_all(&dir_offset.to_le_bytes())?; + crate::spatial::write_format_marker(self.write, crate::spatial::EDGE_INDEX_MAGIC) + } +} diff --git a/src/spatial/exact.rs b/src/spatial/exact.rs new file mode 100644 index 0000000000..a3c9de6b5b --- /dev/null +++ b/src/spatial/exact.rs @@ -0,0 +1,288 @@ +//! Exact arithmetic predicates for spherical geometry. +//! +//! This module provides robust geometric predicates that handle degenerate cases using exact +//! arithmetic and symbolic perturbation. +//! +//! s2predicates.cc + +use num_rational::BigRational; +use num_traits::{FromPrimitive, Signed}; + +use super::math::cross; + +// Exact arithmetic types and operations for the ExactSign path. +// +// Every f64 is exactly representable as a rational number (mantissa times a power of two), so +// BigRational::from_f64 loses nothing. Multiply two rationals, get an exact rational. Add them, +// exact. Check sign, exact. +// +// These replace C++ ExactFloat / Vector3_xf from exactfloat.h. + +pub(super) type ExactPoint = [BigRational; 3]; + +pub(super) fn to_exact(p: &[f64; 3]) -> ExactPoint { + [ + BigRational::from_f64(p[0]).unwrap(), + BigRational::from_f64(p[1]).unwrap(), + BigRational::from_f64(p[2]).unwrap(), + ] +} + +pub(super) fn exact_cross(a: &ExactPoint, b: &ExactPoint) -> ExactPoint { + [ + &a[1] * &b[2] - &a[2] * &b[1], + &a[2] * &b[0] - &a[0] * &b[2], + &a[0] * &b[1] - &a[1] * &b[0], + ] +} + +fn exact_dot(a: &ExactPoint, b: &ExactPoint) -> BigRational { + &a[0] * &b[0] + &a[1] * &b[1] + &a[2] * &b[2] +} + +fn exact_sgn(x: &BigRational) -> i32 { + if x.is_positive() { + 1 + } else if x.is_negative() { + -1 + } else { + 0 + } +} + +/// Uses arbitrary-precision arithmetic and the "simulation of simplicity" technique in order to be +/// completely robust (i.e., to return consistent results for all possible inputs). +pub fn expensive_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], perturb: bool) -> i32 { + // Return zero if and only if two points are the same. This ensures (1). + if a == b || b == c || c == a { + return 0; + } + + // Next we try recomputing the determinant still using floating-point arithmetic but in a more + // precise way. This is more expensive than the simple calculation done by TriageSign(), but + // it is still *much* cheaper than using arbitrary-precision arithmetic. This optimization is + // able to compute the correct determinant sign in virtually all cases except when the three + // points are truly collinear (e.g., three points on the equator). + let det_sign = stable_sign(a, b, c); + if det_sign != 0 { + return det_sign; + } + + // Otherwise fall back to exact arithmetic and symbolic permutations. + exact_sign(a, b, c, perturb) +} + +// Compute the determinant using exact arithmetic and/or symbolic permutations. Requires that the +// three points are distinct. +fn exact_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], perturb: bool) -> i32 { + debug_assert!(a != b && b != c && c != a); + + // Sort the three points in lexicographic order, keeping track of the sign of the permutation. + // (Each exchange inverts the sign of the determinant.) + let mut points = [*a, *b, *c]; + let perm_sign = sort_with_parity(&mut points); + + // Construct multiple-precision versions of the sorted points and compute their exact 3x3 + // determinant. + let xa = to_exact(&points[0]); + let xb = to_exact(&points[1]); + let xc = to_exact(&points[2]); + let xb_cross_xc = exact_cross(&xb, &xc); + let det = exact_dot(&xa, &xb_cross_xc); + + // C++ checks !isnan(det) and det.prec() < det.max_prec() here. + // BigRational cannot produce NaN or overflow, so these are unnecessary. + + // If the exact determinant is non-zero, we're done. + let det_sign = exact_sgn(&det); + if det_sign == 0 && perturb { + // Otherwise, we need to resort to symbolic perturbations to resolve the + // sign of the determinant. + let det_sign = symbolically_perturbed_sign(&xa, &xb, &xc, &xb_cross_xc); + debug_assert_ne!(0, det_sign); + perm_sign * det_sign + } else { + perm_sign * det_sign + } +} + +// Compute the determinant in a numerically stable way. Unlike TriageSign(), this method can +// usually compute the correct determinant sign even when all three points are as collinear as +// possible. For example if three points are spaced 1km apart along a random line on the Earth's +// surface using the nearest representable points, there is only a 0.4% chance that this method +// will not be able to find the determinant sign. The probability of failure decreases as the +// points get closer together; if the collinear points are +// 1 meter apart, the failure rate drops to 0.0004%. +// +// This method could be extended to also handle nearly-antipodal points (and in fact an earlier +// version of this code did exactly that), but antipodal points are rare in practice so it seems +// better to simply fall back to exact arithmetic in that case. +fn stable_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3]) -> i32 { + let ab = [b[0] - a[0], b[1] - a[1], b[2] - a[2]]; + let bc = [c[0] - b[0], c[1] - b[1], c[2] - b[2]]; + let ca = [a[0] - c[0], a[1] - c[1], a[2] - c[2]]; + let ab2 = ab[0] * ab[0] + ab[1] * ab[1] + ab[2] * ab[2]; + let bc2 = bc[0] * bc[0] + bc[1] * bc[1] + bc[2] * bc[2]; + let ca2 = ca[0] * ca[0] + ca[1] * ca[1] + ca[2] * ca[2]; + + // Now compute the determinant ((A-C)x(B-C)).C, where the vertices have been cyclically + // permuted if necessary so that AB is the longest edge. (This minimizes the magnitude of + // cross product.) At the same time we also compute the maximum error in the determinant. + // Using a similar technique to the one used for kMaxDetError, the error is at most + // + // |d| <= (3 + 6/sqrt(3)) * |A-C| * |B-C| * e + // + // where e = 0.5 * DBL_EPSILON. If the determinant magnitude is larger than this value then we + // know its sign with certainty. + const DET_ERROR_MULTIPLIER: f64 = 3.2321 * f64::EPSILON; // see above + let (det, max_error) = if ab2 >= bc2 && ab2 >= ca2 { + // AB is the longest edge, so compute (A-C)x(B-C).C. + let cross = cross(&ca, &bc); + let det = -(cross[0] * c[0] + cross[1] * c[1] + cross[2] * c[2]); + let max_error = DET_ERROR_MULTIPLIER * (ca2 * bc2).sqrt(); + (det, max_error) + } else if bc2 >= ca2 { + // BC is the longest edge, so compute (B-A)x(C-A).A. + let cross = cross(&ab, &ca); + let det = -(cross[0] * a[0] + cross[1] * a[1] + cross[2] * a[2]); + let max_error = DET_ERROR_MULTIPLIER * (ab2 * ca2).sqrt(); + (det, max_error) + } else { + // CA is the longest edge, so compute (C-B)x(A-B).B. + let cross = cross(&bc, &ab); + let det = -(cross[0] * b[0] + cross[1] * b[1] + cross[2] * b[2]); + let max_error = DET_ERROR_MULTIPLIER * (bc2 * ab2).sqrt(); + (det, max_error) + }; + + // Errors smaller than this value may not be accurate due to underflow. + const MIN_NO_UNDERFLOW_ERROR: f64 = 3.2321 * f64::EPSILON * 1.4916681462400413e-154; // sqrt(DBL_MIN) + if max_error < MIN_NO_UNDERFLOW_ERROR { + return 0; + } + + if det.abs() <= max_error { + 0 + } else if det > 0.0 { + 1 + } else { + -1 + } +} + +/// The following function returns the sign of the determinant of three points A, B, C under a +/// model where every possible S2Point is slightly perturbed by a unique infinitesmal amount such +/// that no three perturbed points are collinear and no four points are coplanar. The +/// perturbations are so small that they do not change the sign of any determinant that was +/// non-zero before the perturbations, and therefore can be safely ignored unless the determinant +/// of three points is exactly zero (using multiple-precision arithmetic). +/// +/// Since the symbolic perturbation of a given point is fixed (i.e., the perturbation is the same +/// for all calls to this method and does not depend on the other two arguments), the results of +/// this method are always self-consistent. It will never return results that would correspond to +/// an "impossible" configuration of non-degenerate points. +/// +/// Requirements: +/// - The 3x3 determinant of A, B, C must be exactly zero. +/// - The points must be distinct, with A < B < C in lexicographic order. +/// +/// Returns: +/// - +1 or -1 according to the sign of the determinant after the symbolic perturbations are taken +/// into account. +/// +/// Reference: +/// - "Simulation of Simplicity" (Edelsbrunner and Muecke, ACM Transactions on Graphics, 1990). +fn symbolically_perturbed_sign( + a: &ExactPoint, + b: &ExactPoint, + c: &ExactPoint, + b_cross_c: &ExactPoint, +) -> i32 { + // This method requires that the points are sorted in lexicographically increasing order. This + // is because every possible S2Point has its own symbolic perturbation such that if A < B then + // the symbolic perturbation for A is much larger than the perturbation for B. + // + // Alternatively, we could sort the points in this method and keep track of the sign of the + // permutation, but it is more efficient to do this before converting the inputs to the + // multi-precision representation, and this also lets us re-use the result of the + // cross product B x C. + debug_assert!(a < b && b < c); + + let det_sign = exact_sgn(&b_cross_c[2]); // da[2] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&b_cross_c[1]); // da[1] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&b_cross_c[0]); // da[0] + if det_sign != 0 { + return det_sign; + } + + let det_sign = exact_sgn(&(&c[0] * &a[1] - &c[1] * &a[0])); // db[2] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&c[0]); // db[2] * da[1] + if det_sign != 0 { + return det_sign; + } + let det_sign = -exact_sgn(&c[1]); // db[2] * da[0] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&(&c[2] * &a[0] - &c[0] * &a[2])); // db[1] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&c[2]); // db[1] * da[0] + if det_sign != 0 { + return det_sign; + } + + // The following test is listed in the paper, but it is redundant because + // the previous tests guarantee that C == (0, 0, 0). + debug_assert_eq!(0, exact_sgn(&(&c[1] * &a[2] - &c[2] * &a[1]))); // db[0] + + let det_sign = exact_sgn(&(&a[0] * &b[1] - &a[1] * &b[0])); // dc[2] + if det_sign != 0 { + return det_sign; + } + let det_sign = -exact_sgn(&b[0]); // dc[2] * da[1] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&b[1]); // dc[2] * da[0] + if det_sign != 0 { + return det_sign; + } + let det_sign = exact_sgn(&a[0]); // dc[2] * db[1] + if det_sign != 0 { + return det_sign; + } + + 1 // dc[2] * db[1] * da[0] +} + +// Sort the three points in lexicographic order, keeping track of the sign of the permutation. +// (Each exchange inverts the sign of the determinant.) +fn sort_with_parity(points: &mut [[f64; 3]; 3]) -> i32 { + let mut parity = 1; + + if points[1] < points[0] { + points.swap(0, 1); + parity = -parity; + } + if points[2] < points[1] { + points.swap(1, 2); + parity = -parity; + } + if points[1] < points[0] { + points.swap(0, 1); + parity = -parity; + } + + parity +} diff --git a/src/spatial/executor.rs b/src/spatial/executor.rs new file mode 100644 index 0000000000..bd233dac4a --- /dev/null +++ b/src/spatial/executor.rs @@ -0,0 +1,875 @@ +//! Spatial query executor. +//! +//! Evaluates multi-stage spatial queries during Weight construction, producing per-segment +//! bitsets or scored result sets that per-segment scorers iterate. + +use std::collections::HashMap; +use std::fmt; + +use common::BitSet; + +use super::cell_index_reader::CellIndexReader; +use super::closest_edge_query::ClosestEdgeQuery; +use super::contains_query::ContainsQuery; +use super::edge_cache::EdgeCache; +use super::edge_reader::EdgeReader; +use super::geometry::Geometry; +use super::geometry_set::{to_geometry_set, GeometrySet}; +use super::intersects_query::IntersectsQuery; +use super::plane::Plane; +use super::region_coverer::CovererOptions; +use super::s1chord_angle::S1ChordAngle; +use super::sphere::Sphere; +use crate::core::Searcher; +use crate::docset::{DocSet, TERMINATED}; +use crate::index::SegmentId; +use crate::query::{ + does_not_match, AllQuery, BitSetDocSet, BooleanQuery, EnableScoring, Explanation, Query, + Scorer, Weight, +}; +use crate::schema::Field; +use crate::{DocId, Score, SegmentReader, TantivyError}; + +/// Spatial relationship for a join predicate. +#[derive(Clone, Debug)] +pub enum SpatialRelation { + /// Within a distance. + Near(f64), + /// Between two distances. + Between(f64, f64), + /// Geometries intersect. + Intersects, + /// The outer geometry contains the inner one. + Contains, +} + +/// A node in the execution plan tree. +pub enum PlanNode { + /// A tantivy query evaluated to per-segment bitsets. + Query(Box), + + /// Find K nearest documents to a query geometry. + Knn { + /// The spatial field. + field: Field, + /// The query geometry. + geometry: GeometrySet, + /// Number of nearest neighbors. + k: usize, + /// If present, only documents passing this node are considered. + filter: Option>, + }, + + /// Intersects query filtered by a terms bitset during traversal. + Intersects { + /// The spatial field. + field: Field, + /// The query geometry (smashed). + geometry: GeometrySet, + /// Documents passing this node are eligible candidates. + filter: Box, + }, + + /// Distance query filtered by a terms bitset during traversal. + Within { + /// The spatial field. + field: Field, + /// The query geometry (smashed). + geometry: GeometrySet, + /// Maximum distance (chord angle length2). + radius: f64, + /// Documents passing this node are eligible candidates. + filter: Box, + }, + + /// Spatial join between an outer and inner result set. + Join { + /// The spatial field. + field: Field, + /// Documents to test. + outer: Box, + /// Documents to test against. + inner: Box, + /// The spatial relationship. + relation: SpatialRelation, + }, +} + +/// The output of evaluating a plan node. Per-segment results keyed by segment ID. +pub struct StageOutput { + /// Per-segment results. + pub results: HashMap, +} + +/// The precomputed results for a single segment. +pub enum SegmentResult { + /// A bitset of matching doc_ids with a constant score (join, distance, text queries). + Match(BitSet), + /// A scored list of (doc_id, score) pairs sorted by doc_id (kNN). + Scored(Vec<(DocId, f32)>), +} + +impl StageOutput { + /// Extract the bitset for a segment, building one from scored results if necessary. + fn bitset_for(&self, segment_id: &SegmentId, max_doc: u32) -> BitSet { + match self.results.get(segment_id) { + Some(SegmentResult::Match(bitset)) => bitset.clone(), + Some(SegmentResult::Scored(results)) => { + let mut bitset = BitSet::with_max_value(max_doc); + for &(doc_id, _) in results { + bitset.insert(doc_id); + } + bitset + } + None => BitSet::with_max_value(max_doc), + } + } +} + +/// Bytes of geometry each edge cache of a query may hold before it evicts geometries. +pub(crate) const QUERY_EDGE_CACHE_BYTES: usize = 8 << 20; + +/// Evaluate a plan node recursively, producing per-segment results. Every stage's documents are +/// alive: a Query stage leaves out deleted documents, and every other stage draws its documents +/// from stages below it, or applies the deletions itself. +fn evaluate( + node: &PlanNode, + searcher: &Searcher, + segments: &[SegmentReader], +) -> crate::Result { + match node { + PlanNode::Query(query) => { + // Run the tantivy query to completion per segment, collect bitsets. Deleted documents + // must not become the inner or outer documents of a join, or eligible candidates. + let scoring = EnableScoring::disabled_from_searcher(searcher); + let weight = query.weight(scoring)?; + let mut results = HashMap::new(); + for reader in segments { + let mut bitset = BitSet::with_max_value(reader.max_doc()); + let alive_bitset = reader.alive_bitset(); + weight.for_each_no_score(reader, &mut |docs| { + for &doc in docs { + if alive_bitset.is_none_or(|alive| alive.is_alive(doc)) { + bitset.insert(doc); + } + } + })?; + results.insert(reader.segment_id(), SegmentResult::Match(bitset)); + } + Ok(StageOutput { results }) + } + + PlanNode::Intersects { + field, + geometry, + filter, + } => { + // Evaluate the filter node -> per-segment bitsets. + let filter_output = evaluate(filter, searcher, segments)?; + + // Build the IntersectsQuery once from the smashed geometry. + let intersects = IntersectsQuery::new(geometry.clone(), CovererOptions::default()); + + // For each segment, run the filtered spatial traversal. + let mut results = HashMap::new(); + for reader in segments { + let spatial = reader.spatial_fields().get_field(*field)?; + if let Some(spatial_reader) = spatial { + let cell_reader = CellIndexReader::open(spatial_reader.cells_bytes())?; + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + let mut edge_cache = EdgeCache::new(vec![er], QUERY_EDGE_CACHE_BYTES); + + let filter_bitset = + filter_output.bitset_for(&reader.segment_id(), reader.max_doc()); + + let doc_ids = intersects.search_segment_filtered( + &cell_reader, + &mut edge_cache, + &filter_bitset, + ); + + let mut bitset = BitSet::with_max_value(reader.max_doc()); + for doc_id in doc_ids { + bitset.insert(doc_id); + } + results.insert(reader.segment_id(), SegmentResult::Match(bitset)); + } + } + Ok(StageOutput { results }) + } + + PlanNode::Within { + field, + geometry, + radius, + filter, + } => { + let filter_output = evaluate(filter, searcher, segments)?; + + let query = + ClosestEdgeQuery::within(geometry.clone(), S1ChordAngle::from_length2(*radius)); + + let mut results = HashMap::new(); + for reader in segments { + let spatial = reader.spatial_fields().get_field(*field)?; + if let Some(spatial_reader) = spatial { + let cell_reader = CellIndexReader::open(spatial_reader.cells_bytes())?; + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + let mut edge_cache = EdgeCache::new(vec![er], QUERY_EDGE_CACHE_BYTES); + + let filter_bitset = + filter_output.bitset_for(&reader.segment_id(), reader.max_doc()); + + let hits = query.search_segment_filtered( + &cell_reader, + &mut edge_cache, + &filter_bitset, + ); + + let mut bitset = BitSet::with_max_value(reader.max_doc()); + for result in hits { + bitset.insert(result.doc_id); + } + results.insert(reader.segment_id(), SegmentResult::Match(bitset)); + } + } + Ok(StageOutput { results }) + } + + PlanNode::Knn { + field, + geometry, + k, + filter, + } => { + let filter_output = match filter { + Some(node) => Some(evaluate(node, searcher, segments)?), + None => None, + }; + + let query = ClosestEdgeQuery::knn(geometry.clone(), *k); + + // Run branch-and-bound per segment, collecting up to K results with + // distances. Each result is tagged with its segment ID for redistribution + // after the cross-segment merge. + let mut all_results: Vec<(SegmentId, u32, f64, usize)> = Vec::new(); + for (segment_ord, reader) in segments.iter().enumerate() { + let filter_bitset = filter_output + .as_ref() + .map(|fo| fo.bitset_for(&reader.segment_id(), reader.max_doc())); + for (doc_id, distance) in knn_segment_hits(&query, *field, reader, filter_bitset)? { + all_results.push((reader.segment_id(), doc_id, distance, segment_ord)); + } + } + + // Merge across segments: sort by distance, keep global top K. Ties go to the lower + // document address, so that a query always returns the same documents. + all_results.sort_by(|a, b| a.2.total_cmp(&b.2).then((a.3, a.1).cmp(&(b.3, b.1)))); + all_results.truncate(*k); + + // Redistribute to per-segment scored results. + let mut results: HashMap> = HashMap::new(); + let distances: Vec = all_results.iter().map(|hit| hit.2).collect(); + for ((seg_id, doc_id, _, _), score) in + all_results.iter().zip(rank_scores(&distances, *k)) + { + results.entry(*seg_id).or_default().push((*doc_id, score)); + } + + // Sort each segment's results by doc_id for DocSet ordering. + let mut final_results = HashMap::new(); + for (seg_id, mut docs) in results { + docs.sort_by_key(|&(doc_id, _)| doc_id); + final_results.insert(seg_id, SegmentResult::Scored(docs)); + } + + Ok(StageOutput { + results: final_results, + }) + } + + PlanNode::Join { + field, + outer, + inner, + relation, + } => { + // Phase 1: collect inner bitsets across all segments. + let inner_output = evaluate(inner, searcher, segments)?; + + // Phase 2: collect outer bitsets across all segments. + let outer_output = evaluate(outer, searcher, segments)?; + + let spatial_readers = segments + .iter() + .map(|reader| reader.spatial_fields().get_field(*field)) + .collect::>>()?; + + // Phase 3: open each segment that holds inner documents once, for every probe. The + // probes' caches live for the whole join, so they share the query's budget. + let probe_cache_bytes = (QUERY_EDGE_CACHE_BYTES / segments.len().max(1)).max(1 << 20); + let mut probes = Vec::new(); + for (segment_ord, (reader, spatial_reader)) in + segments.iter().zip(&spatial_readers).enumerate() + { + let Some(spatial_reader) = spatial_reader else { + continue; + }; + let inner = inner_output.bitset_for(&reader.segment_id(), reader.max_doc()); + if inner.len() == 0 { + continue; + } + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + probes.push(ProbeSegment { + segment_ord, + cell_reader: CellIndexReader::open(spatial_reader.cells_bytes())?, + edge_cache: EdgeCache::new(vec![er], probe_cache_bytes), + inner, + }); + } + + // Phase 4: for each segment, walk its cell index for outer documents, and probe the + // segments for inner matches with each one's geometry. + let mut results = HashMap::new(); + for (segment_ord, (reader, spatial_reader)) in + segments.iter().zip(&spatial_readers).enumerate() + { + let mut result_bitset = BitSet::with_max_value(reader.max_doc()); + // With no inner documents anywhere, nothing matches. + if let Some(spatial_reader) = spatial_reader.as_ref().filter(|_| !probes.is_empty()) + { + let outer_bitset = + outer_output.bitset_for(&reader.segment_id(), reader.max_doc()); + let cell_reader = CellIndexReader::open(spatial_reader.cells_bytes())?; + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + let edge_cache = EdgeCache::new(vec![er], QUERY_EDGE_CACHE_BYTES); + // Each outer document is probed once, with all of its geometry. + let mut visited = BitSet::with_max_value(reader.max_doc()); + + for cell in cell_reader.iter() { + for clipped in &cell.shapes { + let doc_id = edge_cache.doc_id_for(clipped.geometry_id); + if visited.contains(doc_id) || !outer_bitset.contains(doc_id) { + continue; + } + visited.insert(doc_id); + + let outer_geometry = + edge_cache.get(clipped.geometry_id).geometry_set().clone(); + let probe = JoinProbe::new(relation, outer_geometry); + let found = probes.iter_mut().any(|segment| { + // A document is not joined to itself. + let is_self = segment.segment_ord == segment_ord + && segment.inner.contains(doc_id); + if is_self { + segment.inner.remove(doc_id); + } + let found = probe.matches(segment); + if is_self { + segment.inner.insert(doc_id); + } + found + }); + if found { + result_bitset.insert(doc_id); + } + } + } + } + results.insert(reader.segment_id(), SegmentResult::Match(result_bitset)); + } + + Ok(StageOutput { results }) + } + } +} + +/// The nearest documents of one segment and their distances, as squared chord lengths, at most +/// the query's K of them. Only documents in the filter are considered, and deleted documents +/// never are: they must not take the K slots before ranking. A filter's documents are alive +/// already. +fn knn_segment_hits( + query: &ClosestEdgeQuery, + field: Field, + reader: &SegmentReader, + filter: Option, +) -> crate::Result> { + let Some(spatial_reader) = reader.spatial_fields().get_field(field)? else { + return Ok(Vec::new()); + }; + let cell_reader = CellIndexReader::open(spatial_reader.cells_bytes())?; + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + let mut edge_cache = EdgeCache::new(vec![er], QUERY_EDGE_CACHE_BYTES); + let filter = filter.or_else(|| { + reader.alive_bitset().map(|alive| { + let mut bitset = BitSet::with_max_value_and_full(reader.max_doc()); + bitset.intersect_update(alive.bitset()); + bitset + }) + }); + let hits = match &filter { + Some(filter) => query.search_segment_filtered(&cell_reader, &mut edge_cache, filter), + None => query.search_segment(&cell_reader, &mut edge_cache), + }; + Ok(hits + .into_iter() + .map(|hit| (hit.doc_id, hit.distance.length2())) + .collect()) +} + +/// The scores of kNN hits, given their distances in ascending order. A hit scores by its rank: +/// 1 for the nearest, down to 1/K for the Kth, with hits at the same distance sharing a rank. The +/// scores are positive, so that a kNN clause lifts the nearest documents of a boolean query, and +/// they order the hits exactly, however close their distances. +fn rank_scores(distances: &[f64], k: usize) -> Vec { + let mut rank = 0; + distances + .iter() + .enumerate() + .map(|(i, distance)| { + if i > 0 && distances[i - 1] != *distance { + rank = i; + } + (k - rank) as f32 / k as f32 + }) + .collect() +} + +/// A segment a join probes for inner matches: its readers, and its inner documents. +struct ProbeSegment<'a> { + segment_ord: usize, + cell_reader: CellIndexReader<'a>, + edge_cache: EdgeCache<'a, Sphere>, + inner: BitSet, +} + +/// The query a join runs on the probed segments for one outer geometry, built once per outer +/// document. +enum JoinProbe { + Distance(ClosestEdgeQuery), + Intersects(IntersectsQuery), + Contains(ContainsQuery), +} + +impl JoinProbe { + fn new(relation: &SpatialRelation, outer_geometry: GeometrySet) -> Self { + match relation { + SpatialRelation::Near(r) => JoinProbe::Distance(ClosestEdgeQuery::any_within( + outer_geometry, + S1ChordAngle::from_radians(*r), + )), + SpatialRelation::Between(inner_r, outer_r) => { + JoinProbe::Distance(ClosestEdgeQuery::between( + outer_geometry, + S1ChordAngle::from_radians(*inner_r), + S1ChordAngle::from_radians(*outer_r), + )) + } + SpatialRelation::Intersects => JoinProbe::Intersects(IntersectsQuery::new( + outer_geometry, + CovererOptions::default(), + )), + SpatialRelation::Contains => JoinProbe::Contains(ContainsQuery::new( + outer_geometry, + CovererOptions::default(), + )), + } + } + + /// Whether any inner document of the segment matches. + fn matches(&self, segment: &mut ProbeSegment<'_>) -> bool { + let (cell_reader, edge_cache, inner) = ( + &segment.cell_reader, + &mut segment.edge_cache, + &segment.inner, + ); + match self { + JoinProbe::Distance(query) => !query + .search_segment_filtered(cell_reader, edge_cache, inner) + .is_empty(), + JoinProbe::Intersects(query) => !query + .search_segment_filtered(cell_reader, edge_cache, inner) + .is_empty(), + JoinProbe::Contains(query) => !query + .search_segment_filtered(cell_reader, edge_cache, inner) + .is_empty(), + } + } +} + +/// Query that evaluates a spatial execution plan during Weight construction. +pub struct SpatialExecutor { + /// The root of the plan tree. + root: PlanNode, +} + +impl SpatialExecutor { + /// Create a new executor from a plan tree. + pub fn new(root: PlanNode) -> Self { + Self { root } + } + + /// Create the plan for the K nearest documents to a lon/lat point. + pub fn knn(field: Field, point: [f64; 2], k: usize) -> Self { + let point = Geometry::::Point(point).project::(); + Self::new(PlanNode::Knn { + field, + geometry: to_geometry_set(&point, 0), + k, + filter: None, + }) + } + + /// Narrow the documents the root of the plan considers to those the query also matches: the + /// outer documents of a join, or the candidates of a kNN query, which then finds the nearest + /// of those. Restrictions add up. Outer documents or a filter given as another plan node than + /// a query are left as they are. + pub fn restrict(&mut self, query: Box) { + fn intersect(node: &mut PlanNode, query: Box) { + if let PlanNode::Query(slot) = node { + let current = std::mem::replace(slot, Box::new(AllQuery)); + *slot = if current.is::() { + query + } else { + Box::new(BooleanQuery::intersection(vec![current, query])) + }; + } + } + match &mut self.root { + PlanNode::Join { outer, .. } => intersect(outer, query), + PlanNode::Knn { filter, .. } => match filter { + Some(filter) => intersect(filter, query), + None => *filter = Some(Box::new(PlanNode::Query(query))), + }, + _ => {} + } + } + + /// Replace the outer documents of the root join with those the query matches. + #[deprecated(note = "use `restrict`, which also narrows a kNN query")] + pub fn set_outer(&mut self, query: Box) { + if let PlanNode::Join { outer, .. } = &mut self.root { + **outer = PlanNode::Query(query); + } + } +} + +impl fmt::Debug for SpatialExecutor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SpatialExecutor") + } +} + +impl Clone for PlanNode { + fn clone(&self) -> Self { + match self { + PlanNode::Query(q) => PlanNode::Query(q.box_clone()), + PlanNode::Knn { + field, + geometry, + k, + filter, + } => PlanNode::Knn { + field: *field, + geometry: geometry.clone(), + k: *k, + filter: filter.clone(), + }, + PlanNode::Intersects { + field, + geometry, + filter, + } => PlanNode::Intersects { + field: *field, + geometry: geometry.clone(), + filter: filter.clone(), + }, + PlanNode::Within { + field, + geometry, + radius, + filter, + } => PlanNode::Within { + field: *field, + geometry: geometry.clone(), + radius: *radius, + filter: filter.clone(), + }, + PlanNode::Join { + field, + outer, + inner, + relation, + } => PlanNode::Join { + field: *field, + outer: outer.clone(), + inner: inner.clone(), + relation: relation.clone(), + }, + } + } +} + +impl Clone for SpatialExecutor { + fn clone(&self) -> Self { + Self { + root: self.root.clone(), + } + } +} + +impl Query for SpatialExecutor { + fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result> { + let Some(searcher) = enable_scoring.searcher() else { + // Without a searcher, as when deleting by query or in a filter aggregation, the + // segments are not known up front. A kNN query then finds the K nearest of each + // segment. A join needs the inner documents of every segment. + if let PlanNode::Knn { + field, + geometry, + k, + filter, + } = &self.root + { + let filter = match filter.as_deref() { + None => None, + Some(PlanNode::Query(query)) => Some(query.weight(enable_scoring)?), + Some(_) => return Err(needs_a_searcher()), + }; + return Ok(Box::new(SegmentKnnWeight { + field: *field, + query: ClosestEdgeQuery::knn(geometry.clone(), *k), + k: *k, + filter, + })); + } + return Err(needs_a_searcher()); + }; + let segments = searcher.segment_readers(); + + let output = evaluate(&self.root, searcher, segments)?; + + Ok(Box::new(ExecutorWeight { + results: output.results, + })) + } +} + +fn needs_a_searcher() -> TantivyError { + TantivyError::InvalidArgument( + "A spatial join needs a searcher: it evaluates every segment at once".to_string(), + ) +} + +/// A kNN query without a searcher: each segment keeps the K nearest of its own documents. +struct SegmentKnnWeight { + field: Field, + query: ClosestEdgeQuery, + k: usize, + /// Only documents this weight matches are candidates. + filter: Option>, +} + +impl Weight for SegmentKnnWeight { + fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + let filter = match &self.filter { + Some(weight) => { + let mut bitset = BitSet::with_max_value(reader.max_doc()); + let alive_bitset = reader.alive_bitset(); + weight.for_each_no_score(reader, &mut |docs| { + for &doc in docs { + if alive_bitset.is_none_or(|alive| alive.is_alive(doc)) { + bitset.insert(doc); + } + } + })?; + Some(bitset) + } + None => None, + }; + let mut hits = knn_segment_hits(&self.query, self.field, reader, filter)?; + hits.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); + hits.truncate(self.k); + let distances: Vec = hits.iter().map(|hit| hit.1).collect(); + let mut scored: Vec<(DocId, f32)> = hits + .iter() + .zip(rank_scores(&distances, self.k)) + .map(|((doc_id, _), score)| (*doc_id, score)) + .collect(); + scored.sort_by_key(|&(doc_id, _)| doc_id); + Ok(Box::new(ReplayScorer::from_scored(scored, boost))) + } + + fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { + let mut scorer = self.scorer(reader, 1.0)?; + // Seeking backward is not allowed: a scorer may start past the document. + if scorer.doc() > doc || scorer.seek(doc) != doc { + return Err(does_not_match(doc)); + } + Ok(Explanation::new("SpatialExecutor", scorer.score())) + } +} + +/// Holds precomputed per-segment results. Scorers iterate them. +struct ExecutorWeight { + /// Precomputed results keyed by segment ID. + results: HashMap, +} + +impl Weight for ExecutorWeight { + fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + match self.results.get(&reader.segment_id()) { + Some(SegmentResult::Match(bitset)) => { + Ok(Box::new(ReplayScorer::from_bitset(bitset.clone(), boost))) + } + Some(SegmentResult::Scored(results)) => { + Ok(Box::new(ReplayScorer::from_scored(results.clone(), boost))) + } + None => Ok(Box::new(ReplayScorer::empty())), + } + } + + fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { + let mut scorer = self.scorer(reader, 1.0)?; + // Seeking backward is not allowed: a scorer may start past the document. + if scorer.doc() > doc || scorer.seek(doc) != doc { + return Err(does_not_match(doc)); + } + Ok(Explanation::new("SpatialExecutor", scorer.score())) + } +} + +/// Iterates precomputed results as a scorer. +struct ReplayScorer { + /// The iteration state. + inner: ReplayInner, +} + +enum ReplayInner { + /// Iterate a bitset with a constant score. + BitSet { + docs: BitSetDocSet, + doc_id: DocId, + score: Score, + }, + /// Iterate a sorted vec of (doc_id, score). + Scored { + results: Vec<(DocId, f32)>, + index: usize, + boost: Score, + }, + /// No results for this segment. + Empty, +} + +impl ReplayScorer { + fn from_bitset(bitset: BitSet, score: Score) -> Self { + let docs = BitSetDocSet::from(bitset); + let doc_id = docs.doc(); + Self { + inner: ReplayInner::BitSet { + docs, + doc_id, + score, + }, + } + } + + fn from_scored(results: Vec<(DocId, f32)>, boost: Score) -> Self { + Self { + inner: ReplayInner::Scored { + results, + index: 0, + boost, + }, + } + } + + fn empty() -> Self { + Self { + inner: ReplayInner::Empty, + } + } +} + +impl Scorer for ReplayScorer { + fn score(&mut self) -> Score { + match &self.inner { + ReplayInner::BitSet { score, .. } => *score, + ReplayInner::Scored { + results, + index, + boost, + } => { + if *index < results.len() { + results[*index].1 * *boost + } else { + 0.0 + } + } + ReplayInner::Empty => 0.0, + } + } +} + +impl DocSet for ReplayScorer { + fn advance(&mut self) -> DocId { + match &mut self.inner { + ReplayInner::BitSet { docs, doc_id, .. } => { + if *doc_id == TERMINATED { + return TERMINATED; + } + *doc_id = docs.advance(); + *doc_id + } + ReplayInner::Scored { results, index, .. } => { + *index += 1; + if *index < results.len() { + results[*index].0 + } else { + TERMINATED + } + } + ReplayInner::Empty => TERMINATED, + } + } + + fn seek(&mut self, target: DocId) -> DocId { + match &mut self.inner { + ReplayInner::BitSet { docs, doc_id, .. } => { + if *doc_id < target { + *doc_id = docs.seek(target); + } + *doc_id + } + ReplayInner::Scored { results, index, .. } => { + let start = (*index).min(results.len()); + *index = start + results[start..].partition_point(|&(doc_id, _)| doc_id < target); + self.doc() + } + ReplayInner::Empty => TERMINATED, + } + } + + fn doc(&self) -> DocId { + match &self.inner { + ReplayInner::BitSet { doc_id, .. } => *doc_id, + ReplayInner::Scored { results, index, .. } => { + if *index < results.len() { + results[*index].0 + } else { + TERMINATED + } + } + ReplayInner::Empty => TERMINATED, + } + } + + fn size_hint(&self) -> u32 { + match &self.inner { + ReplayInner::BitSet { docs, .. } => docs.size_hint(), + ReplayInner::Scored { results, .. } => results.len() as u32, + ReplayInner::Empty => 0, + } + } +} diff --git a/src/spatial/geometry.rs b/src/spatial/geometry.rs new file mode 100644 index 0000000000..4a9a0b8260 --- /dev/null +++ b/src/spatial/geometry.rs @@ -0,0 +1,566 @@ +//! GeoJSON geometry types, parsing, serialization, and projection. + +use std::io::{self, Read, Write}; + +use common::{BinarySerializable, VInt}; +use serde_json::{json, Map, Value}; + +use crate::spatial::plane::Plane; +use crate::spatial::surface::Surface; +use crate::spatial::xor::{compress_f64, decompress_f64}; + +/// Errors from parsing GeoJSON geometry. +#[derive(Debug)] +pub enum GeometryError { + /// The "type" field is missing. + MissingType, + /// A required field is missing. + MissingField(String), + /// The geometry type is not supported. + UnsupportedType(String), + /// A coordinate value is out of range or not a number. + InvalidCoordinate(String), + /// The coordinate structure is wrong (bad nesting, too few points, unclosed ring). + InvalidStructure(String), +} + +/// A GeoJSON geometry, parameterized by surface. On the Plane, coordinates are lon/lat in degrees. +/// On the Sphere, coordinates are unit vectors. +#[derive(Debug, Clone, PartialEq)] +pub enum Geometry { + /// A single position. + Point(S::Point), + /// Multiple positions. + MultiPoint(Vec), + /// An ordered sequence of positions forming a line. + LineString(Vec), + /// Multiple line strings. + MultiLineString(Vec>), + /// A closed ring (outer) with optional holes (inner rings). + Polygon(Vec>), + /// Multiple polygons. + MultiPolygon(Vec>>), + /// A heterogeneous collection of geometries. + GeometryCollection(Vec), +} + +impl Geometry { + /// Project this geometry onto a surface. For Sphere, converts lon/lat to unit sphere + /// coordinates. For Plane, returns the geometry unchanged. + pub fn project(&self) -> Geometry { + match self { + Geometry::Point(p) => Geometry::Point(S::project(p[0], p[1])), + Geometry::MultiPoint(points) => { + Geometry::MultiPoint(points.iter().map(|p| S::project(p[0], p[1])).collect()) + } + Geometry::LineString(line) => Geometry::LineString(Self::project_ring::(line)), + Geometry::MultiLineString(lines) => Geometry::MultiLineString( + lines.iter().map(|l| Self::project_ring::(l)).collect(), + ), + Geometry::Polygon(rings) => { + Geometry::Polygon(rings.iter().map(|r| Self::project_ring::(r)).collect()) + } + Geometry::MultiPolygon(polygons) => Geometry::MultiPolygon( + polygons + .iter() + .map(|p| p.iter().map(|r| Self::project_ring::(r)).collect()) + .collect(), + ), + Geometry::GeometryCollection(geometries) => { + Geometry::GeometryCollection(geometries.iter().map(|g| g.project::()).collect()) + } + } + } + + fn project_ring(ring: &[[f64; 2]]) -> Vec { + ring.iter().map(|p| S::project(p[0], p[1])).collect() + } + + /// Parse a GeoJSON geometry object. + pub fn from_geojson(object: &Map) -> Result { + let geometry_type = object + .get("type") + .and_then(|v| v.as_str()) + .ok_or(GeometryError::MissingType)?; + match geometry_type { + "Point" => { + let coordinates = get_coordinates(object)?; + let point = to_point(coordinates)?; + Ok(Geometry::Point(point)) + } + "MultiPoint" => { + let coordinates = get_coordinates(object)?; + let multi_point = to_line_string(coordinates)?; + Ok(Geometry::MultiPoint(multi_point)) + } + "LineString" => { + let coordinates = get_coordinates(object)?; + let line_string = to_line_string(coordinates)?; + if line_string.len() < 2 { + return Err(GeometryError::InvalidStructure( + "a line string contains at least 2 points".to_string(), + )); + } + Ok(Geometry::LineString(line_string)) + } + "MultiLineString" => { + let coordinates = get_coordinates(object)?; + let multi_line_string = to_multi_line_string(coordinates)?; + for line_string in &multi_line_string { + if line_string.len() < 2 { + return Err(GeometryError::InvalidStructure( + "a line string contains at least 2 points".to_string(), + )); + } + } + Ok(Geometry::MultiLineString(multi_line_string)) + } + "Polygon" => { + let coordinates = get_coordinates(object)?; + let polygon = to_multi_line_string(coordinates)?; + for ring in &polygon { + // RFC 7946: a linear ring has four or more positions, + // first and last identical. + if ring.len() < 4 { + return Err(GeometryError::InvalidStructure( + "a linear ring has four or more positions".to_string(), + )); + } + if ring.first() != ring.last() { + return Err(GeometryError::InvalidStructure( + "a linear ring must be closed (first == last)".to_string(), + )); + } + } + Ok(Geometry::Polygon(polygon)) + } + "MultiPolygon" => { + let mut result = Vec::new(); + let multi_polygons = get_coordinates(object)?; + let multi_polygons = + multi_polygons + .as_array() + .ok_or(GeometryError::InvalidStructure( + "expected an array of polygons".to_string(), + ))?; + for polygon in multi_polygons { + let polygon = to_multi_line_string(polygon)?; + for ring in &polygon { + if ring.len() < 4 { + return Err(GeometryError::InvalidStructure( + "a linear ring has four or more positions".to_string(), + )); + } + if ring.first() != ring.last() { + return Err(GeometryError::InvalidStructure( + "a linear ring must be closed (first == last)".to_string(), + )); + } + } + result.push(polygon); + } + Ok(Geometry::MultiPolygon(result)) + } + // Older versions read only the misspelling "GeometriesCollection". + "GeometryCollection" | "GeometriesCollection" => { + let geometries = object + .get("geometries") + .ok_or(GeometryError::MissingField("geometries".to_string()))?; + let geometries = geometries + .as_array() + .ok_or(GeometryError::InvalidStructure( + "geometries is not an array".to_string(), + ))?; + let mut result = Vec::new(); + for geometry in geometries { + let object = geometry.as_object().ok_or(GeometryError::InvalidStructure( + "geometry is not an object".to_string(), + ))?; + result.push(Geometry::from_geojson(object)?); + } + Ok(Geometry::GeometryCollection(result)) + } + _ => Err(GeometryError::UnsupportedType(geometry_type.to_string())), + } + } + + /// Serialize the geometry to GeoJSON format. + /// https://fr.wikipedia.org/wiki/GeoJSON + pub fn to_geojson(&self) -> Map { + let mut map = Map::new(); + match self { + Geometry::Point(point) => { + map.insert("type".to_string(), Value::String("Point".to_string())); + let coords = json!([point[0], point[1]]); + map.insert("coordinates".to_string(), coords); + } + Geometry::MultiPoint(points) => { + map.insert("type".to_string(), Value::String("MultiPoint".to_string())); + let coords: Vec = points.iter().map(|p| json!([p[0], p[1]])).collect(); + map.insert("coordinates".to_string(), Value::Array(coords)); + } + Geometry::LineString(line) => { + map.insert("type".to_string(), Value::String("LineString".to_string())); + let coords: Vec = line.iter().map(|p| json!([p[0], p[1]])).collect(); + map.insert("coordinates".to_string(), Value::Array(coords)); + } + Geometry::MultiLineString(lines) => { + map.insert( + "type".to_string(), + Value::String("MultiLineString".to_string()), + ); + let coords: Vec = lines + .iter() + .map(|line| Value::Array(line.iter().map(|p| json!([p[0], p[1]])).collect())) + .collect(); + map.insert("coordinates".to_string(), Value::Array(coords)); + } + Geometry::Polygon(rings) => { + map.insert("type".to_string(), Value::String("Polygon".to_string())); + let coords: Vec = rings + .iter() + .map(|ring| Value::Array(ring.iter().map(|p| json!([p[0], p[1]])).collect())) + .collect(); + map.insert("coordinates".to_string(), Value::Array(coords)); + } + Geometry::MultiPolygon(polygons) => { + map.insert( + "type".to_string(), + Value::String("MultiPolygon".to_string()), + ); + let coords: Vec = polygons + .iter() + .map(|polygon| { + Value::Array( + polygon + .iter() + .map(|ring| { + Value::Array(ring.iter().map(|p| json!([p[0], p[1]])).collect()) + }) + .collect(), + ) + }) + .collect(); + map.insert("coordinates".to_string(), Value::Array(coords)); + } + Geometry::GeometryCollection(geometries) => { + map.insert( + "type".to_string(), + Value::String("GeometryCollection".to_string()), + ); + let geoms: Vec = geometries + .iter() + .map(|g| Value::Object(g.to_geojson())) + .collect(); + map.insert("geometries".to_string(), Value::Array(geoms)); + } + } + map + } +} + +fn get_coordinates(object: &Map) -> Result<&Value, GeometryError> { + let coordinates = object + .get("coordinates") + .ok_or(GeometryError::MissingField("coordinates".to_string()))?; + Ok(coordinates) +} + +/// Parse a position. RFC 7946 allows elements after lon and lat, such as an altitude. They are +/// not indexed, so they are dropped. +fn to_point(value: &Value) -> Result<[f64; 2], GeometryError> { + let lonlat = value.as_array().ok_or(GeometryError::InvalidStructure( + "expected an array of lon, lat and optional altitude".to_string(), + ))?; + if lonlat.len() < 2 { + return Err(GeometryError::InvalidStructure( + "expected an array of lon, lat and optional altitude".to_string(), + )); + } + let lon = lonlat[0].as_f64().ok_or(GeometryError::InvalidCoordinate( + "longitude must be f64".to_string(), + ))?; + if !lon.is_finite() || !(-180.0..=180.0).contains(&lon) { + return Err(GeometryError::InvalidCoordinate(format!( + "invalid longitude: {}", + lon + ))); + } + let lat = lonlat[1].as_f64().ok_or(GeometryError::InvalidCoordinate( + "latitude must be f64".to_string(), + ))?; + if !lat.is_finite() || !(-90.0..=90.0).contains(&lat) { + return Err(GeometryError::InvalidCoordinate(format!( + "invalid latitude: {}", + lat + ))); + } + Ok([lon, lat]) +} + +fn to_line_string(value: &Value) -> Result, GeometryError> { + let mut result = Vec::new(); + let coordinates = value.as_array().ok_or(GeometryError::InvalidStructure( + "expected an array of lon/lat arrays".to_string(), + ))?; + for coordinate in coordinates { + result.push(to_point(coordinate)?); + } + Ok(result) +} + +fn to_multi_line_string(value: &Value) -> Result>, GeometryError> { + let mut result = Vec::new(); + let coordinates = value.as_array().ok_or(GeometryError::InvalidStructure( + "expected an array of an array of lon/lat arrays".to_string(), + ))?; + for coordinate in coordinates { + result.push(to_line_string(coordinate)?); + } + Ok(result) +} + +impl BinarySerializable for Geometry { + fn serialize(&self, writer: &mut W) -> io::Result<()> { + match self { + Geometry::Point(point) => { + 0u8.serialize(writer)?; + point[0].serialize(writer)?; + point[1].serialize(writer)?; + Ok(()) + } + Geometry::MultiPoint(points) => { + 1u8.serialize(writer)?; + serialize_line_string(points, writer) + } + Geometry::LineString(line_string) => { + 2u8.serialize(writer)?; + serialize_line_string(line_string, writer) + } + Geometry::MultiLineString(multi_line_string) => { + 3u8.serialize(writer)?; + serialize_polygon(&multi_line_string[..], writer) + } + Geometry::Polygon(polygon) => { + 4u8.serialize(writer)?; + serialize_polygon(polygon, writer) + } + Geometry::MultiPolygon(multi_polygon) => { + 5u8.serialize(writer)?; + BinarySerializable::serialize(&VInt(multi_polygon.len() as u64), writer)?; + for polygon in multi_polygon { + BinarySerializable::serialize(&VInt(polygon.len() as u64), writer)?; + for ring in polygon { + BinarySerializable::serialize(&VInt(ring.len() as u64), writer)?; + } + } + let mut lon = Vec::new(); + let mut lat = Vec::new(); + for polygon in multi_polygon { + for ring in polygon { + for point in ring { + lon.push(point[0]); + lat.push(point[1]); + } + } + } + let lon = compress_f64(&lon); + let lat = compress_f64(&lat); + VInt(lon.len() as u64).serialize(writer)?; + writer.write_all(&lon)?; + VInt(lat.len() as u64).serialize(writer)?; + writer.write_all(&lat)?; + Ok(()) + } + Geometry::GeometryCollection(geometries) => { + 6u8.serialize(writer)?; + BinarySerializable::serialize(&VInt(geometries.len() as u64), writer)?; + for geometry in geometries { + geometry.serialize(writer)?; + } + Ok(()) + } + } + } + + fn deserialize(reader: &mut R) -> io::Result { + let discriminant: u8 = BinarySerializable::deserialize(reader)?; + match discriminant { + 0 => { + let lon = BinarySerializable::deserialize(reader)?; + let lat = BinarySerializable::deserialize(reader)?; + Ok(Geometry::Point([lon, lat])) + } + 1 => Ok(Geometry::MultiPoint(deserialize_line_string(reader)?)), + 2 => Ok(Geometry::LineString(deserialize_line_string(reader)?)), + 3 => Ok(Geometry::MultiLineString(deserialize_polygon(reader)?)), + 4 => Ok(Geometry::Polygon(deserialize_polygon(reader)?)), + 5 => { + let polygon_count = VInt::deserialize(reader)?.0 as usize; + let mut polygons = Vec::new(); + let mut count = 0; + for _ in 0..polygon_count { + let ring_count = VInt::deserialize(reader)?.0 as usize; + let mut rings = Vec::new(); + for _ in 0..ring_count { + let point_count = VInt::deserialize(reader)?.0 as usize; + rings.push(point_count); + count += point_count; + } + polygons.push(rings); + } + let lon_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lat_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lon = decompress_f64(&lon_bytes, count); + let lat = decompress_f64(&lat_bytes, count); + let mut multi_polygon = Vec::new(); + let mut offset = 0; + for rings in polygons { + let mut polygon = Vec::new(); + for point_count in rings { + let mut ring = Vec::new(); + for _ in 0..point_count { + ring.push([lon[offset], lat[offset]]); + offset += 1; + } + polygon.push(ring); + } + multi_polygon.push(polygon); + } + Ok(Geometry::MultiPolygon(multi_polygon)) + } + 6 => { + let geometry_count = VInt::deserialize(reader)?.0 as usize; + let mut geometries = Vec::new(); + for _ in 0..geometry_count { + geometries.push(Geometry::deserialize(reader)?); + } + Ok(Geometry::GeometryCollection(geometries)) + } + _ => Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid geometry type", + )), + } + } +} + +fn serialize_line_string(line: &[[f64; 2]], writer: &mut W) -> io::Result<()> { + BinarySerializable::serialize(&VInt(line.len() as u64), writer)?; + let mut lon = Vec::new(); + let mut lat = Vec::new(); + for point in line { + lon.push(point[0]); + lat.push(point[1]); + } + let lon = compress_f64(&lon); + let lat = compress_f64(&lat); + VInt(lon.len() as u64).serialize(writer)?; + writer.write_all(&lon)?; + VInt(lat.len() as u64).serialize(writer)?; + writer.write_all(&lat)?; + Ok(()) +} + +fn serialize_polygon( + line_string: &[Vec<[f64; 2]>], + writer: &mut W, +) -> io::Result<()> { + BinarySerializable::serialize(&VInt(line_string.len() as u64), writer)?; + for ring in line_string { + BinarySerializable::serialize(&VInt(ring.len() as u64), writer)?; + } + let mut lon: Vec = Vec::new(); + let mut lat: Vec = Vec::new(); + for ring in line_string { + for point in ring { + lon.push(point[0]); + lat.push(point[1]); + } + } + let lon: Vec = compress_f64(&lon); + let lat: Vec = compress_f64(&lat); + VInt(lon.len() as u64).serialize(writer)?; + writer.write_all(&lon)?; + VInt(lat.len() as u64).serialize(writer)?; + writer.write_all(&lat)?; + Ok(()) +} + +fn deserialize_line_string(reader: &mut R) -> io::Result> { + let point_count = VInt::deserialize(reader)?.0 as usize; + let lon_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lat_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lon: Vec = decompress_f64(&lon_bytes, point_count); + let lat: Vec = decompress_f64(&lat_bytes, point_count); + let mut line_string: Vec<[f64; 2]> = Vec::new(); + for offset in 0..point_count { + line_string.push([lon[offset], lat[offset]]); + } + Ok(line_string) +} + +fn deserialize_polygon(reader: &mut R) -> io::Result>> { + let ring_count = VInt::deserialize(reader)?.0 as usize; + let mut rings = Vec::new(); + let mut count = 0; + for _ in 0..ring_count { + let point_count = VInt::deserialize(reader)?.0 as usize; + rings.push(point_count); + count += point_count; + } + let lon_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lat_bytes: Vec = BinarySerializable::deserialize(reader)?; + let lon: Vec = decompress_f64(&lon_bytes, count); + let lat: Vec = decompress_f64(&lat_bytes, count); + let mut polygon: Vec> = Vec::new(); + let mut offset = 0; + for point_count in rings { + let mut ring = Vec::new(); + for _ in 0..point_count { + ring.push([lon[offset], lat[offset]]); + offset += 1; + } + polygon.push(ring); + } + Ok(polygon) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parse(json: Value) -> Result, GeometryError> { + Geometry::from_geojson(json.as_object().unwrap()) + } + + #[test] + fn test_positions_with_altitude() { + let point = parse(json!({"type": "Point", "coordinates": [1.0, 2.0, 30.0]})).unwrap(); + assert_eq!(point, Geometry::Point([1.0, 2.0])); + let ring = json!([ + [0.0, 0.0, 5.0], + [1.0, 0.0, 5.0], + [1.0, 1.0, 5.0], + [0.0, 0.0, 5.0] + ]); + let polygon = parse(json!({"type": "Polygon", "coordinates": [ring]})).unwrap(); + assert_eq!( + polygon, + Geometry::Polygon(vec![vec![[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 0.0]]]) + ); + assert!(parse(json!({"type": "Point", "coordinates": [1.0]})).is_err()); + } + + #[test] + fn test_collections_parse_by_either_name() { + let point = json!({"type": "Point", "coordinates": [1.0, 2.0]}); + for name in ["GeometryCollection", "GeometriesCollection"] { + let collection = parse(json!({"type": name, "geometries": [point]})).unwrap(); + assert_eq!( + collection, + Geometry::GeometryCollection(vec![Geometry::Point([1.0, 2.0])]) + ); + } + } +} diff --git a/src/spatial/geometry_set.rs b/src/spatial/geometry_set.rs new file mode 100644 index 0000000000..3b33db3557 --- /dev/null +++ b/src/spatial/geometry_set.rs @@ -0,0 +1,403 @@ +//! The universal currency for stored-format geometry. +//! +//! A GeometrySet holds all members of a document's geometry in stored format: polygon rings closed +//! per RFC 7946 (first == last), flattened into a single vertex array per member, edge_id gaps at +//! ring boundaries. After smashing, the ring structure is gone. Edge access is always +//! vertices[edge_idx] and vertices[edge_idx + 1]. One Geometry in, one GeometrySet out. + +use std::sync::LazyLock; + +use super::containment::{ + brute_force_contains, compute_origin_inside, curvature, curvature_max_error, + on_one_great_circle, prune_degeneracies, +}; +use super::geometry::Geometry; +use super::math::normalize; +use super::s2coords::face_uv_to_xyz; +use super::sphere::Sphere; + +/// The Hilbert curve start point: face 0, UV (-1, -1), normalized to the unit sphere. +static HILBERT_START: LazyLock<[f64; 3]> = + LazyLock::new(|| normalize(&face_uv_to_xyz(0, -1.0, -1.0))); + +/// One smashed geometry's edges. After smashing, the type distinctions (point, line string, +/// polygon) are gone. Everything is consecutive vertex pairs and a closed flag. +#[derive(Clone)] +pub struct EdgeSet { + /// Flattened vertex array. Edges are consecutive pairs: vertices[i] to vertices[i+1]. + pub vertices: Vec<[f64; 3]>, + /// Whether this geometry has interior (polygon). + pub closed: bool, + /// Whether the Hilbert curve start point is inside this geometry's interior. + pub contains_hilbert_start: bool, + /// Ring boundary offsets. Always starts with 0 and ends with the vertex count. + pub ring_offsets: Vec, +} + +/// Smashed geometry set for one field of one document. The geometry_id is assigned by the +/// builder or writer that receives the set, not carried here. +#[derive(Clone)] +pub struct GeometrySet { + /// The smashed members of this document's geometry field. + pub members: Vec, + /// The document this set belongs to. + pub doc_id: u32, +} + +impl GeometrySet { + /// Heap held by the set: its members and their vertex and ring offset arrays, spare capacity + /// included. + pub fn mem_usage(&self) -> usize { + heap_size(&self.members) + + self + .members + .iter() + .map(|member| heap_size(&member.vertices) + heap_size(&member.ring_offsets)) + .sum::() + } +} + +/// Heap held by a Vec's buffer, spare capacity included. +fn heap_size(vec: &Vec) -> usize { + vec.capacity() * std::mem::size_of::() +} + +/// Which of the two regions a polygon ring separates on the sphere is the polygon's. A ring and +/// its reverse have the same vertices, so only the direction it runs can tell a small polygon from +/// the rest of the sphere around it. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RingOrientation { + /// The smaller region, whichever way the ring runs. RFC 7946 has parsers accept rings of + /// either direction, and much data has them, but no polygon is larger than a hemisphere. + #[default] + Smaller, + /// The right-hand rule of RFC 7946: an exterior ring runs counter-clockwise around the + /// polygon and a hole clockwise, so the polygon lies to the left of every ring. A polygon may + /// be larger than a hemisphere, but a ring running the wrong way bounds the rest of the + /// sphere. + RightHand, +} + +/// Smash a projected geometry into a GeometrySet in stored format, with rings bounding the +/// smaller region. See `to_geometry_set_with_orientation`. +pub fn to_geometry_set(geometry: &Geometry, doc_id: u32) -> GeometrySet { + to_geometry_set_with_orientation(geometry, doc_id, RingOrientation::Smaller) +} + +/// Smash a projected geometry into a GeometrySet in stored format. Orients polygon rings CCW +/// around the regions the orientation has them bound, computes contains_hilbert_start, and +/// flattens rings into a single vertex array per member. One Geometry in, one GeometrySet out. +/// Every array is allocated at its exact length, but for degenerate polygon rings, since the +/// writer holds the sets until the segment is flushed. +pub fn to_geometry_set_with_orientation( + geometry: &Geometry, + doc_id: u32, + orientation: RingOrientation, +) -> GeometrySet { + let mut members = Vec::with_capacity(member_count(geometry)); + smash(&mut members, geometry, orientation); + GeometrySet { members, doc_id } +} + +/// Number of members a geometry smashes into. +fn member_count(geometry: &Geometry) -> usize { + match geometry { + Geometry::Point(_) | Geometry::LineString(_) | Geometry::Polygon(_) => 1, + Geometry::MultiPoint(points) => points.len(), + Geometry::MultiLineString(lines) => lines.len(), + Geometry::MultiPolygon(polygons) => polygons.len(), + Geometry::GeometryCollection(collection) => collection.iter().map(member_count).sum(), + } +} + +fn smash(members: &mut Vec, geometry: &Geometry, orientation: RingOrientation) { + match geometry { + Geometry::Point(v) => { + members.push(EdgeSet { + vertices: vec![*v], + closed: false, + contains_hilbert_start: false, + ring_offsets: vec![0, 1], + }); + } + Geometry::MultiPoint(points) => { + for v in points { + members.push(EdgeSet { + vertices: vec![*v], + closed: false, + contains_hilbert_start: false, + ring_offsets: vec![0, 1], + }); + } + } + // A line with no vertices is empty. + Geometry::LineString(line) => { + if !line.is_empty() { + members.push(open_member(line.clone())); + } + } + Geometry::MultiLineString(lines) => { + for line in lines.iter().filter(|line| !line.is_empty()) { + members.push(open_member(line.clone())); + } + } + Geometry::Polygon(polygon) => { + if !polygon.is_empty() { + members.extend(smash_polygon(polygon, orientation)); + } + } + Geometry::MultiPolygon(multi_polygon) => { + for polygon in multi_polygon.iter().filter(|polygon| !polygon.is_empty()) { + members.extend(smash_polygon(polygon, orientation)); + } + } + Geometry::GeometryCollection(collection) => { + for geo in collection { + smash(members, geo, orientation); + } + } + } +} + +/// Smash one polygon (outer ring + holes) into stored-format members. Takes already-projected +/// sphere coordinates with rings closed per RFC 7946 (first == last), or unclosed. Orients every +/// ring CCW around the region it bounds, computes whether the Hilbert curve start point is inside, +/// and flattens into a single vertex array of closed rings. The orientation decides which region +/// a ring bounds. +/// +/// Degeneracies are made valid as ST_MakeValid does. Repeated vertices are removed. A spike, +/// where an edge doubles back along the one before it, bounds no area: it leaves the ring, and +/// follows the polygon as a line member, where queries still meet it. An exterior that bounds no +/// area at all is the point or line it is, without its holes, and a hole that bounds none is left +/// out. A polygon with no rings is empty. +/// +/// Ring offsets always start with 0 and end with the total vertex count. +fn smash_polygon(rings: &[Vec<[f64; 3]>], orientation: RingOrientation) -> Vec { + let mut members = Vec::new(); + // Room for every ring closed, as GeoJSON rings are. + let closed_len = |ring: &Vec<[f64; 3]>| ring.len() + usize::from(ring.first() != ring.last()); + let mut flat = Vec::with_capacity(rings.iter().map(closed_len).sum()); + let mut hilbert_inside = false; + let mut offsets = Vec::with_capacity(rings.len() + 1); + offsets.push(0); + + for (i, ring) in rings.iter().enumerate() { + let (mut ring, spikes) = prune_degeneracies(ring); + if ring.is_empty() || on_one_great_circle(&ring) { + if i == 0 { + return vec![open_member(remove_repeated_vertices(&rings[0]))]; + } + continue; + } + members.extend(spikes.iter().map(|spike| open_member(spike.to_vec()))); + let curvature = curvature(&ring); + let reverse = match orientation { + RingOrientation::Smaller => curvature < -curvature_max_error(ring.len()), + // A hole runs clockwise around the region it bounds. + RingOrientation::RightHand => i > 0, + }; + if reverse { + // Keep the first vertex first. + ring[1..].reverse(); + } + ring.push(ring[0]); + let mut origin_inside = compute_origin_inside(&ring); + if orientation == RingOrientation::Smaller && curvature.abs() < std::f64::consts::PI { + // A ring this close to a hemisphere may cross itself, which defeats the angle test + // at a vertex. Take the inside to be the smaller side. + let inside = |origin_inside| { + SAMPLE_POINTS + .iter() + .filter(|point| brute_force_contains(point, &ring, origin_inside)) + .count() + }; + if inside(origin_inside) > SAMPLE_POINTS.len() * 3 / 4 { + origin_inside = !origin_inside; + } + } + if brute_force_contains(&HILBERT_START, &ring, origin_inside) { + hilbert_inside = !hilbert_inside; + } + flat.extend_from_slice(&ring); + offsets.push(flat.len()); + } + + let polygon = EdgeSet { + vertices: flat, + closed: true, + contains_hilbert_start: hilbert_inside, + ring_offsets: offsets, + }; + members.insert(0, polygon); + members +} + +/// Directions spread over the sphere: the centers, edge midpoints and corners of the cube. +static SAMPLE_POINTS: LazyLock> = LazyLock::new(|| { + let mut points = Vec::new(); + for x in -1..=1 { + for y in -1..=1 { + for z in -1..=1 { + if (x, y, z) != (0, 0, 0) { + points.push(normalize(&[x as f64, y as f64, z as f64])); + } + } + } + } + points +}); + +/// A path without consecutive repeated vertices. A ring of a single position becomes that +/// position alone. +fn remove_repeated_vertices(ring: &[[f64; 3]]) -> Vec<[f64; 3]> { + let mut vertices: Vec<[f64; 3]> = Vec::with_capacity(ring.len()); + for v in ring { + if vertices.last() != Some(v) { + vertices.push(*v); + } + } + vertices +} + +/// A member with no interior: a point for a single vertex, otherwise a line through the vertices. +fn open_member(vertices: Vec<[f64; 3]>) -> EdgeSet { + let n = vertices.len(); + EdgeSet { + vertices, + closed: false, + contains_hilbert_start: false, + ring_offsets: vec![0, n], + } +} + +#[cfg(test)] +mod tests { + use std::mem::size_of; + + use super::*; + use crate::spatial::plane::Plane; + + fn smashed(geometry: Geometry) -> GeometrySet { + to_geometry_set(&geometry.project::(), 0) + } + + #[test] + fn test_mem_usage_counts_exactly_sized_arrays() { + let point = smashed(Geometry::Point([19.94, 50.06])); + assert_eq!( + point.mem_usage(), + size_of::() + size_of::<[f64; 3]>() + 2 * size_of::() + ); + + let square = vec![[0.0, 0.0], [4.0, 0.0], [4.0, 4.0], [0.0, 4.0], [0.0, 0.0]]; + let hole = vec![[1.0, 1.0], [1.0, 2.0], [2.0, 2.0], [2.0, 1.0], [1.0, 1.0]]; + let polygons = smashed(Geometry::MultiPolygon(vec![ + vec![square.clone(), hole], + vec![square], + ])); + assert_eq!( + polygons.mem_usage(), + 2 * size_of::() + 15 * size_of::<[f64; 3]>() + 5 * size_of::() + ); + } + + #[test] + fn test_smash_polygon_reverses_holes() { + let square = vec![[0.0, 0.0], [4.0, 0.0], [4.0, 4.0], [0.0, 4.0], [0.0, 0.0]]; + let hole = vec![[1.0, 1.0], [1.0, 2.0], [2.0, 2.0], [2.0, 1.0], [1.0, 1.0]]; + let projected = + Geometry::::Polygon(vec![square.clone(), hole.clone()]).project::(); + let Geometry::Polygon(rings) = &projected else { + unreachable!() + }; + let set = to_geometry_set(&projected, 0); + let member = &set.members[0]; + assert_eq!(member.ring_offsets, [0, 5, 10]); + assert_eq!(member.vertices[..5], rings[0][..]); + let reversed: Vec<[f64; 3]> = rings[1].iter().rev().copied().collect(); + assert_eq!(member.vertices[5..], reversed[..]); + } + + #[test] + fn test_smash_polygon_orients_rings_ccw() { + // A clockwise exterior is reversed. A counter-clockwise hole, which RFC 7946 orders + // the other way, already runs CCW around the hole, as the stored format has it. + let square = vec![[0.0, 0.0], [0.0, 4.0], [4.0, 4.0], [4.0, 0.0], [0.0, 0.0]]; + let hole = vec![[1.0, 1.0], [2.0, 1.0], [2.0, 2.0], [1.0, 2.0], [1.0, 1.0]]; + let projected = + Geometry::::Polygon(vec![square.clone(), hole.clone()]).project::(); + let set = to_geometry_set(&projected, 0); + let member = &set.members[0]; + let expected = Geometry::::Polygon(vec![ + vec![[0.0, 0.0], [4.0, 0.0], [4.0, 4.0], [0.0, 4.0], [0.0, 0.0]], + hole, + ]) + .project::(); + let Geometry::Polygon(rings) = &expected else { + unreachable!() + }; + assert_eq!(member.vertices[..5], rings[0][..]); + assert_eq!(member.vertices[5..], rings[1][..]); + assert!(!member.contains_hilbert_start); + } + + #[test] + fn test_smash_polygon_keeps_what_degeneracies_leave() { + // A repeated leading vertex is removed. A spike leaves the ring, and stays as a line. + let spiked = vec![ + [0.0, 0.0], + [0.0, 0.0], + [10.0, 0.0], + [12.0, 0.0], + [10.0, 0.0], + [10.0, 10.0], + [0.0, 10.0], + [0.0, 0.0], + ]; + let set = smashed(Geometry::Polygon(vec![spiked])); + let square = vec![ + [0.0, 0.0], + [10.0, 0.0], + [10.0, 10.0], + [0.0, 10.0], + [0.0, 0.0], + ]; + let unspiked = smashed(Geometry::Polygon(vec![square.clone()])); + assert_eq!(set.members[0].vertices, unspiked.members[0].vertices); + assert_eq!( + set.members[0].contains_hilbert_start, + unspiked.members[0].contains_hilbert_start + ); + let spike = smashed(Geometry::LineString(vec![[10.0, 0.0], [12.0, 0.0]])); + assert_eq!(set.members.len(), 2); + assert!(!set.members[1].closed); + assert_eq!(set.members[1].vertices, spike.members[0].vertices); + + // An unclosed ring is closed. + let unclosed = smashed(Geometry::Polygon(vec![square[..4].to_vec()])); + assert_eq!(unclosed.members[0].vertices, unspiked.members[0].vertices); + + // A polygon with no area is the line or point it is. A hole with no area is left out. + let flat = vec![[0.0, 0.0], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0]]; + let set = smashed(Geometry::Polygon(vec![flat.clone()])); + let line = smashed(Geometry::LineString(flat[..3].to_vec())); + assert!(!set.members[0].closed); + assert_eq!(set.members[0].vertices, line.members[0].vertices); + let collinear = vec![[0.0, 0.0], [2.0, 0.0], [1.0, 0.0], [0.0, 0.0]]; + assert!(!smashed(Geometry::Polygon(vec![collinear])).members[0].closed); + let point = smashed(Geometry::Polygon(vec![vec![[5.0, 5.0]; 4]])); + assert_eq!(point.members[0].vertices.len(), 1); + let square = vec![[0.0, 0.0], [4.0, 0.0], [4.0, 4.0], [0.0, 4.0], [0.0, 0.0]]; + let set = smashed(Geometry::Polygon(vec![square, flat])); + assert_eq!(set.members[0].ring_offsets, [0, 5]); + + // Empty polygons and lines are empty. + assert!(smashed(Geometry::Polygon(vec![])).members.is_empty()); + let multi = smashed(Geometry::MultiPolygon(vec![ + vec![], + vec![vec![[5.0, 5.0]; 4]], + ])); + assert_eq!(multi.members.len(), 1); + assert!(smashed(Geometry::LineString(vec![])).members.is_empty()); + } +} diff --git a/src/spatial/interleave.rs b/src/spatial/interleave.rs new file mode 100644 index 0000000000..398fb62d99 --- /dev/null +++ b/src/spatial/interleave.rs @@ -0,0 +1,609 @@ +//! Heap-based merge of N source cell indexes. Produces cells in Hilbert order. Subdivides coarse +//! cells on demand when they overlap finer cells or exceed the edge limit. + +use std::cmp::Ordering; +use std::collections::BinaryHeap; + +use common::ReadOnlyBitSet; +use rustc_hash::{FxHashMap, FxHashSet}; + +use super::cell_index::{ + clip_to_children, get_edge_max_level, BuildOptions, ClippedShape, GeometryId, IndexCell, + CELL_PADDING, +}; +use super::cell_index_reader::CellIndexIter; +use super::edge_cache::EdgeCache; +use super::r2rect::R2Rect; +use super::s2cell_id::S2CellId; +use super::s2edge_clipping::{clip_edge_bound, clip_to_padded_face}; +use super::s2padded_cell::S2PaddedCell; +use super::sphere::Sphere; + +struct CoarseEdge { + edge_index: u32, + max_level: i32, + v0: [f64; 3], + v1: [f64; 3], +} + +struct CoarseGeometry { + edges: Vec, + /// Whether the geometry has an interior. Only then can it contain a cell's center. + closed: bool, +} + +struct CoarseIndexCell { + cell: IndexCell, + geometries: FxHashMap, + geometry_ids: Vec, +} + +struct HeapEntry { + cell_id: S2CellId, + source: Option, + cells: Vec, + geometry_ids: FxHashSet, +} + +impl CoarseIndexCell { + pub fn new(cell: IndexCell, edge_cache: &EdgeCache<'_, Sphere>) -> Self { + let mut geometries = FxHashMap::default(); + let mut geometry_ids = Vec::with_capacity(cell.shapes.len()); + for shape in &cell.shapes { + geometry_ids.push(shape.geometry_id); + + if shape.edge_indices.is_empty() { + continue; + } + + let entry = edge_cache.get(shape.geometry_id); + + let level_geometry = + geometries + .entry(shape.geometry_id) + .or_insert_with(|| CoarseGeometry { + edges: Vec::new(), + closed: entry.edge_set().closed, + }); + + for &edge_index in &shape.edge_indices { + let (v0, v1) = entry.edge(edge_index); + let max_level = get_edge_max_level(&v0, &v1); + level_geometry.edges.push(CoarseEdge { + edge_index, + max_level, + v0, + v1, + }); + } + } + CoarseIndexCell { + cell, + geometries, + geometry_ids, + } + } +} + +impl HeapEntry { + fn range_min(&self) -> S2CellId { + self.cell_id.range_min() + } + + fn level(&self) -> i32 { + self.cell_id.level() + } +} + +impl PartialEq for HeapEntry { + fn eq(&self, other: &Self) -> bool { + self.range_min() == other.range_min() && self.level() == other.level() + } +} + +impl Eq for HeapEntry {} + +impl PartialOrd for HeapEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for HeapEntry { + fn cmp(&self, other: &Self) -> Ordering { + other + .range_min() + .cmp(&self.range_min()) + .then(other.level().cmp(&self.level())) + } +} + +/// Heap-based merge of N source cell indexes. Produces cells in Hilbert order. +pub struct HeapInterleave<'a> { + sources: Vec>, + edge_cache: &'a EdgeCache<'a, Sphere>, + edge_counter: FxHashSet<(GeometryId, u32)>, + alive_bitsets: &'a [Option], + heap: BinaryHeap, + max_edges: usize, + min_short_edge_fraction: f64, + cells_from_source: u64, + absorptions: u64, + splits: u64, + cells_emitted: u64, + short_edge_checks: u64, + max_sponge_cells: u64, + crossing_checks: u64, +} + +impl<'a> HeapInterleave<'a> { + /// Creates a heap merge from source iterators and a shared edge cache. + pub fn new( + iters: Vec>, + edge_cache: &'a EdgeCache<'a, Sphere>, + alive_bitsets: &'a [Option], + ) -> Self { + let options = BuildOptions::default(); + let mut merge = HeapInterleave { + sources: iters, + edge_cache, + edge_counter: FxHashSet::default(), + alive_bitsets, + heap: BinaryHeap::new(), + max_edges: options.max_edges_per_cell, + min_short_edge_fraction: options.min_short_edge_fraction, + cells_from_source: 0, + absorptions: 0, + splits: 0, + cells_emitted: 0, + short_edge_checks: 0, + max_sponge_cells: 0, + crossing_checks: 0, + }; + let n = merge.sources.len(); + for i in 0..n { + merge.push_next_from_source(i); + } + merge + } + + /// Log interleave metrics at debug level. + pub fn report(&self) { + debug!( + "interleave: {} from source, {} absorbed, {} splits, {} emitted, {} edge checks, {} \ + crossing checks, {} max sponge", + self.cells_from_source, + self.absorptions, + self.splits, + self.cells_emitted, + self.short_edge_checks, + self.crossing_checks, + self.max_sponge_cells, + ); + } + + fn push_next_from_source(&mut self, source_idx: usize) { + let segment = source_idx as u16; + while let Some(mut cell) = self.sources[source_idx].next() { + if let Some(bitset) = &self.alive_bitsets[source_idx] { + cell.shapes.retain(|shape| { + let id = (segment, shape.geometry_id.1); + let doc_id = self.edge_cache.doc_id_for(id); + bitset.contains(doc_id) + }); + } + for shape in &mut cell.shapes { + shape.geometry_id.0 = segment; + } + if !cell.shapes.is_empty() { + self.cells_from_source += 1; + let mut geometries = FxHashSet::default(); + geometries.extend(cell.shapes.iter().map(|s| s.geometry_id)); + self.heap.push(HeapEntry { + cell_id: cell.cell_id, + source: Some(source_idx), + cells: vec![CoarseIndexCell::new(cell, self.edge_cache)], + geometry_ids: geometries, + }); + return; + } + } + } + + fn coarse_split(&mut self, sponge: HeapEntry) { + self.splits += 1; + let mut cells = sponge.cells; + let head = cells.remove(0); + let children = subdivide_cell(&head, &mut self.crossing_checks); + + // Build four HeapEntries from the subdivided children. + let mut entries: Vec = children + .into_iter() + .map(|child_cell| { + let child_cell_id = child_cell.cell_id; + let mut geometry_ids = FxHashSet::default(); + geometry_ids.extend(child_cell.shapes.iter().map(|s| s.geometry_id)); + HeapEntry { + cell_id: child_cell_id, + source: None, + cells: vec![CoarseIndexCell::new(child_cell, self.edge_cache)], + geometry_ids, + } + }) + .collect(); + + // Distribute remaining coarse cells to their parent child. If a coarse cell is at the + // same level as the children it can't be assigned as a descendant. Push it onto the heap + // as its own entry. + for coarse_cell in cells { + let coarse_cell_id = coarse_cell.cell.cell_id; + if let Some(parent_entry) = entries + .iter_mut() + .find(|e| e.cell_id.contains(coarse_cell_id)) + { + parent_entry.geometry_ids.extend(&coarse_cell.geometry_ids); + parent_entry.cells.push(coarse_cell); + } else { + let mut geometries = FxHashSet::default(); + geometries.extend(&coarse_cell.geometry_ids); + self.heap.push(HeapEntry { + cell_id: coarse_cell_id, + source: None, + cells: vec![coarse_cell], + geometry_ids: geometries, + }); + } + } + + // Push non-empty entries onto the heap. + for entry in entries { + let has_geometries = entry.cells.iter().any(|c| !c.geometry_ids.is_empty()); + if has_geometries { + self.heap.push(entry); + } + } + } +} + +impl<'a> Iterator for HeapInterleave<'a> { + type Item = IndexCell; + + fn next(&mut self) -> Option { + loop { + let mut sponge = self.heap.pop()?; + if let Some(source) = sponge.source { + sponge.source = None; + self.push_next_from_source(source); + } + if let Some(mut entry) = self.heap.pop() { + if let Some(source) = entry.source { + entry.source = None; + self.push_next_from_source(source); + } + if sponge.cell_id.contains(entry.cell_id) || sponge.cell_id == entry.cell_id { + self.absorptions += 1; + sponge.geometry_ids.extend(&entry.geometry_ids); + sponge.cells.extend(entry.cells); + let sponge_size = sponge.cells.len() as u64; + if sponge_size > self.max_sponge_cells { + self.max_sponge_cells = sponge_size; + } + let mut edges = std::mem::take(&mut self.edge_counter); + edges.clear(); + self.short_edge_checks += + count_short_edges(&sponge.cells, sponge.cell_id.level(), &mut edges); + let threshold = self.max_edges.max( + (self.min_short_edge_fraction + * (edges.len() + sponge.geometry_ids.len()) as f64) + as usize, + ); + if edges.len() > threshold { + self.coarse_split(sponge); + } else { + self.heap.push(sponge); + } + self.edge_counter = edges; + } else { + self.cells_emitted += 1; + self.heap.push(entry); + return flatten(sponge); + } + } else { + self.cells_emitted += 1; + return flatten(sponge); + } + } + } +} + +/// A geometry's edges across the cells of a sponge, with the contains_center flag and the center +/// of the first cell it appears in: the anchor of the ray-cast to the sponge center. +struct EdgeUnion { + /// The edges as (edge index, v0, v1). + edges: Vec<(u32, [f64; 3], [f64; 3])>, + /// Whether the geometry has an interior. A geometry with no edges in any cell is there only + /// because it contains a cell's center, so it does. + closed: bool, + anchor_flag: bool, + anchor_center: [f64; 3], +} + +fn flatten(sponge: HeapEntry) -> Option { + if sponge.cells.len() == 1 { + return Some(sponge.cells.into_iter().next().unwrap().cell); + } + + let sponge_pcell = S2PaddedCell::new(sponge.cell_id, CELL_PADDING); + let sponge_center = sponge_pcell.get_center(); + + // Build per-geometry union of edges with anchor info for the ray-cast. + let mut edge_unions: FxHashMap = FxHashMap::default(); + + for coarse_cell in &sponge.cells { + let cell_center = S2PaddedCell::new(coarse_cell.cell.cell_id, CELL_PADDING).get_center(); + + for shape in &coarse_cell.cell.shapes { + let union_entry = edge_unions + .entry(shape.geometry_id) + .or_insert_with(|| EdgeUnion { + edges: Vec::new(), + closed: true, + anchor_flag: shape.contains_center, + anchor_center: cell_center, + }); + if let Some(geo) = coarse_cell.geometries.get(&shape.geometry_id) { + union_entry.closed = geo.closed; + for &edge_idx in &shape.edge_indices { + if let Some(ce) = geo.edges.iter().find(|e| e.edge_index == edge_idx) { + union_entry.edges.push((edge_idx, ce.v0, ce.v1)); + } + } + } + } + } + + for union in edge_unions.values_mut() { + union.edges.sort_unstable_by_key(|(idx, _, _)| *idx); + union.edges.dedup_by_key(|(idx, _, _)| *idx); + } + + let mut result = IndexCell::new(sponge.cell_id); + + for (geometry_id, union) in &edge_unions { + if union.edges.is_empty() { + result.add_shape(ClippedShape::new(*geometry_id, union.anchor_flag)); + continue; + } + + // Only a geometry with an interior can contain the center. For a line, the parity of + // its crossings means nothing. + let mut contains = false; + if union.closed { + let mut crosser = + super::crossings::S2EdgeCrosser::new(&union.anchor_center, &sponge_center); + let mut crossings = 0u32; + for (_, v0, v1) in &union.edges { + if crosser.edge_or_vertex_crossing_two(v0, v1) { + crossings += 1; + } + } + contains = union.anchor_flag ^ !crossings.is_multiple_of(2); + } + + let mut shape = ClippedShape::new(*geometry_id, contains); + shape.edge_indices = union.edges.iter().map(|(idx, _, _)| *idx).collect(); + result.add_shape(shape); + } + + Some(result) +} + +fn count_short_edges( + cells: &[CoarseIndexCell], + level: i32, + counter: &mut FxHashSet<(GeometryId, u32)>, +) -> u64 { + let mut checked = 0u64; + for cell in cells { + for (&gid, geo) in &cell.geometries { + for edge in &geo.edges { + checked += 1; + if level < edge.max_level { + counter.insert((gid, edge.edge_index)); + } + } + } + } + checked +} + +struct MergeEdge { + geometry_id: GeometryId, + edge_index: u32, + v0: [f64; 3], + v1: [f64; 3], + a_uv: [f64; 2], + b_uv: [f64; 2], +} + +#[derive(Clone)] +struct MergeClippedEdge { + edge_index: usize, + bound: R2Rect, +} + +fn compute_contains_center( + from: &[f64; 3], + to: &[f64; 3], + from_flag: bool, + all_edges: &[MergeEdge], + clipped: &[MergeClippedEdge], + geometry_id: GeometryId, + crossing_checks: &mut u64, +) -> bool { + use super::crossings::S2EdgeCrosser; + let mut crosser = S2EdgeCrosser::new(from, to); + let mut crossings = 0u32; + for ce in clipped { + let me = &all_edges[ce.edge_index]; + if me.geometry_id == geometry_id { + *crossing_checks += 1; + if crosser.edge_or_vertex_crossing_two(&me.v0, &me.v1) { + crossings += 1; + } + } + } + from_flag ^ !crossings.is_multiple_of(2) +} + +fn subdivide_cell(coarse_cell: &CoarseIndexCell, crossing_checks: &mut u64) -> Vec { + let cell = &coarse_cell.cell; + let face = cell.cell_id.face(); + let pcell = S2PaddedCell::new(cell.cell_id, CELL_PADDING); + let parent_center = pcell.get_center(); + + let mut merge_edges: Vec = Vec::new(); + let mut clipped: Vec = Vec::new(); + + // The cell holds the part of each edge on its face, clipped as the builder clips it. An edge + // that misses the padded cell crosses none of its children, nor any path inside it. + for (&gid, geo) in &coarse_cell.geometries { + for edge in &geo.edges { + let Some((a_uv, b_uv)) = clip_to_padded_face(&edge.v0, &edge.v1, face, CELL_PADDING) + else { + continue; + }; + // The bound of the part of the edge in the padded cell, as tight as the builder's. + let mut bound = R2Rect::from_point_pair(a_uv, b_uv); + if !clip_edge_bound(a_uv, b_uv, &pcell.bound(), &mut bound) { + continue; + } + clipped.push(MergeClippedEdge { + edge_index: merge_edges.len(), + bound, + }); + merge_edges.push(MergeEdge { + geometry_id: gid, + edge_index: edge.edge_index, + v0: edge.v0, + v1: edge.v1, + a_uv, + b_uv, + }); + } + } + + // Only a geometry with an interior can contain a child's center. One with no edges in the + // cell is there only because it contains the cell's center, so it has one. + let parent_contains: Vec<(GeometryId, bool)> = cell + .shapes + .iter() + .filter(|s| { + coarse_cell + .geometries + .get(&s.geometry_id) + .is_none_or(|geo| geo.closed) + }) + .map(|s| (s.geometry_id, s.contains_center)) + .collect(); + + let middle = pcell.middle(); + let mut child_edges: [[Vec; 2]; 2] = Default::default(); + for ce in &clipped { + let me = &merge_edges[ce.edge_index]; + clip_to_children(&ce.bound, me.a_uv, me.b_uv, &middle, |i, j, bound| { + child_edges[i][j].push(MergeClippedEdge { + edge_index: ce.edge_index, + bound, + }); + }); + } + + let mut children = Vec::new(); + + for pos in 0..4i32 { + let (i, j) = pcell.get_child_ij(pos); + let child_pcell = S2PaddedCell::from_parent(&pcell, i, j); + let child_center = child_pcell.get_center(); + let mut index_cell = IndexCell::new(child_pcell.id()); + + if child_edges[i][j].is_empty() { + for &(gid, parent_flag) in &parent_contains { + let contains = compute_contains_center( + &parent_center, + &child_center, + parent_flag, + &merge_edges, + &clipped, + gid, + crossing_checks, + ); + if contains { + index_cell.add_shape(ClippedShape::new(gid, true)); + } + } + } else { + let mut edge_map: Vec<(GeometryId, Vec)> = Vec::new(); + for ce in &child_edges[i][j] { + let me = &merge_edges[ce.edge_index]; + if let Some((_, edges)) = + edge_map.iter_mut().find(|(gid, _)| *gid == me.geometry_id) + { + edges.push(me.edge_index); + } else { + edge_map.push((me.geometry_id, vec![me.edge_index])); + } + } + + let mut emitted_gids: Vec = Vec::new(); + for (geometry_id, edge_indices) in &edge_map { + let parent_flag = parent_contains + .iter() + .find(|(gid, _)| gid == geometry_id) + .map(|(_, c)| *c); + + // A geometry missing from parent_contains has no interior. + let contains = parent_flag.is_some_and(|parent_flag| { + compute_contains_center( + &parent_center, + &child_center, + parent_flag, + &merge_edges, + &clipped, + *geometry_id, + crossing_checks, + ) + }); + + let mut shape = ClippedShape::new(*geometry_id, contains); + shape.edge_indices = edge_indices.clone(); + index_cell.add_shape(shape); + emitted_gids.push(*geometry_id); + } + + for &(gid, parent_flag) in &parent_contains { + if emitted_gids.contains(&gid) { + continue; + } + let contains = compute_contains_center( + &parent_center, + &child_center, + parent_flag, + &merge_edges, + &clipped, + gid, + crossing_checks, + ); + if contains { + index_cell.add_shape(ClippedShape::new(gid, true)); + } + } + } + + if !index_cell.is_empty() { + children.push(index_cell); + } + } + + children +} diff --git a/src/spatial/intersects_query.rs b/src/spatial/intersects_query.rs new file mode 100644 index 0000000000..8af67395fa --- /dev/null +++ b/src/spatial/intersects_query.rs @@ -0,0 +1,197 @@ +//! Polygon-intersects-polygon query. +//! +//! Shares infrastructure with contains_query. Differs in verification: any overlap between the +//! query polygon and a candidate is sufficient. + +use std::collections::{HashMap, HashSet}; + +use common::BitSet; + +use super::cell_index::GeometryId; +use super::cell_index_reader::CellIndexReader; +use super::contains_query::{midpoint, CandidateInfo, Contact, QueryRegion}; +use super::edge_cache::EdgeCache; +use super::geometry_set::{EdgeSet, GeometrySet}; +use super::polygon_union::{member_edges, split_edge, PolygonEdge}; +use super::region_coverer::CovererOptions; +use super::shape_index_region::{contains_point, SegmentIndex}; +use super::sphere::Sphere; + +/// Prepared intersects query, built once from a query polygon and applied per-segment. +pub struct IntersectsQuery { + region: QueryRegion, +} + +impl IntersectsQuery { + /// Build the query from a smashed GeometrySet. + pub fn new(set: GeometrySet, options: CovererOptions) -> Self { + Self { + region: QueryRegion::new(set, options), + } + } + + /// Search one segment for geometries that intersect the query polygon. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + let candidates = self + .region + .collect_candidates(cell_reader, None, edge_cache); + self.verify_candidates(candidates, cell_reader, edge_cache) + } + + /// Search one segment, skipping shapes whose doc_id is not in the filter bitset. + pub fn search_segment_filtered( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + terms_filter: &BitSet, + ) -> Vec { + let candidates = + self.region + .collect_candidates(cell_reader, Some(terms_filter), edge_cache); + self.verify_candidates(candidates, cell_reader, edge_cache) + } + + fn verify_candidates( + &self, + candidates: HashMap, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> Vec { + let mut doc_ids = HashSet::new(); + + for (geometry_id, info) in candidates { + // A document intersects once, however many of its members do. + let doc_id = edge_cache.get(geometry_id).doc_id(); + if !doc_ids.contains(&doc_id) + && self.verify_one(geometry_id, &info, cell_reader, edge_cache) + { + doc_ids.insert(doc_id); + } + } + + doc_ids.into_iter().collect() + } + + /// The candidate intersects the query geometry if their edges cross, if part of the + /// candidate's boundary lies inside the query geometry, if part of the query geometry's + /// boundary lies inside the candidate, or if they share an edge with both interiors on the same + /// side of it. + fn verify_one( + &self, + geometry_id: GeometryId, + info: &CandidateInfo, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> bool { + let entry = edge_cache.get(geometry_id); + let candidate = entry.edge_set(); + if candidate.vertices.is_empty() { + return false; + } + let touch = match self.region.boundary_contact(info, candidate) { + Contact::Cross => return true, + Contact::Touch(touch) => touch, + }; + // Points and lines have no interior, so they intersect what they meet: at a shared + // vertex, along a shared edge, or where a point lies on an edge. Within a polygon's + // boundary is within the polygon, as its containment tests have it. + let members = self.region.members(); + let open_contact = touch + .query_touching + .iter() + .chain(&touch.query_along) + .any(|&(member, _)| !candidate.closed || !members[member as usize].closed); + if open_contact || self.region.point_contact(info, candidate) { + return true; + } + if touch.irregular { + return self.verify_pieces(geometry_id, candidate, cell_reader, edge_cache); + } + + let touches_inside = touch + .touching + .iter() + .any(|&edge_idx| self.region.contains(&midpoint(entry.edge(edge_idx)))); + if touches_inside { + return true; + } + let ring_inside = candidate.ring_offsets.windows(2).any(|ring| { + !touch.candidate_ring_meets(ring[0]..ring[1]) + && self.region.contains(&candidate.vertices[ring[0]]) + }); + if ring_inside { + return true; + } + if !candidate.closed { + return false; + } + + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + let mut candidate_contains = + |point: &[f64; 3]| contains_point(&mut segment, geometry_id, point); + + let query_touches_inside = touch.query_touching.iter().any(|&(member, edge_idx)| { + candidate_contains(&midpoint(self.region.query_edge(member, edge_idx))) + }); + if query_touches_inside { + return true; + } + for (member, edge_set) in self.region.members().iter().enumerate() { + for ring in edge_set.ring_offsets.windows(2) { + if !touch.query_ring_meets(member as u32, ring[0]..ring[1]) + && candidate_contains(&edge_set.vertices[ring[0]]) + { + return true; + } + } + } + touch.along.keys().any(|&edge_idx| { + self.region + .covers_shared_interior(candidate, edge_idx, &touch) + }) + } + + /// Where the boundaries meet other than at shared vertices and edges, split the edges of + /// each where the other's meet them. The candidate intersects the query geometry if a piece + /// of its boundary lies inside the query's, if one runs along the query's boundary with + /// both interiors on the same side, or if a piece of the query's boundary lies inside it. + fn verify_pieces( + &self, + geometry_id: GeometryId, + candidate: &EdgeSet, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> bool { + let boundary = self.region.boundary(); + let edges: Vec = member_edges(candidate).map(|(_, edge)| edge).collect(); + for edge in &edges { + for piece in split_edge(*edge, boundary) { + let inside = if piece.shared_sides.iter().any(|&side| side) { + piece.shared_sides[usize::from(piece.edge.interior_left)] + } else { + self.region.contains(&piece.edge.midpoint()) + }; + if inside { + return true; + } + } + } + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + boundary.iter().any(|edge| { + split_edge(*edge, &edges).iter().any(|piece| { + !piece.shared_sides.iter().any(|&side| side) + && contains_point(&mut segment, geometry_id, &piece.edge.midpoint()) + }) + }) + } +} diff --git a/src/spatial/latlng_rect.rs b/src/spatial/latlng_rect.rs new file mode 100644 index 0000000000..9cd4588438 --- /dev/null +++ b/src/spatial/latlng_rect.rs @@ -0,0 +1,477 @@ +//! An S2LatLngRect represents a closed latitude-longitude rectangle. It is capable of +//! representing the empty and full rectangles as well as single points. Note that the +//! latitude-longitude space is considered to have a *cylindrical* topology rather than a spherical +//! one, i.e. the poles have multiple lat/lng representations. An S2LatLngRect may be defined so +//! that includes some representations of a pole but not others. Use the PolarClosure() method if +//! you want to expand a rectangle so that it contains all possible representations of any +//! contained poles. +//! +//! Ported from s2/s2latlng_rect.h (Eric Veach, Google). + +use std::f64::consts::{FRAC_PI_2, PI}; + +use super::crossings::{robust_cross, S2EdgeCrosser}; +use super::math::{cross, dot, normalize}; +use super::r1interval::R1Interval; +use super::region::Region; +use super::s1chord_angle::S1ChordAngle; +use super::s1interval::S1Interval; +use super::s2cap::S2Cap; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; + +/// A closed latitude-longitude rectangle. +#[derive(Clone, Copy, Debug)] +pub struct S2LatLngRect { + lat: R1Interval, + lng: S1Interval, +} + +impl S2LatLngRect { + /// Construct a rectangle from latitude and longitude intervals. The two intervals must + /// either be both empty or both non-empty, and the latitude interval must not extend outside + /// [-90, +90] degrees. Note that both intervals (and hence the rectangle) are closed. + pub fn new(lat: R1Interval, lng: S1Interval) -> Self { + let result = Self { lat, lng }; + debug_assert!(result.is_valid()); + result + } + + /// Construct a rectangle from minimum and maximum latitudes and longitudes. If lo.lng() > + /// hi.lng(), the rectangle spans the 180 degree longitude line. Both points must be + /// normalized, with lo.lat() <= hi.lat(). + pub fn from_lo_hi(lo: &[f64; 2], hi: &[f64; 2]) -> Self { + Self::new(R1Interval::new(lo[1], hi[1]), S1Interval::new(lo[0], hi[0])) + } + + /// Returns the empty rectangle (contains no points). + pub fn empty() -> Self { + Self { + lat: R1Interval::empty(), + lng: S1Interval::empty(), + } + } + + /// Returns the full rectangle (contains all points). + pub fn full() -> Self { + Self { + lat: Self::full_lat(), + lng: S1Interval::full(), + } + } + + /// The full allowable range of latitudes. + pub fn full_lat() -> R1Interval { + R1Interval::new(-FRAC_PI_2, FRAC_PI_2) + } + + /// Construct a rectangle containing a single (normalized) point. + pub fn from_point(p: &[f64; 2]) -> Self { + Self { + lat: R1Interval::from_point(p[1]), + lng: S1Interval::from_point(p[0]), + } + } + + /// Construct the minimal bounding rectangle containing the two given normalized points. + /// This is equivalent to starting with an empty rectangle and calling AddPoint() twice. + /// Note that it is different than the S2LatLngRect(lo, hi) constructor, where the first + /// point is always used as the lower-left corner of the resulting rectangle. + pub fn from_point_pair(p1: &[f64; 2], p2: &[f64; 2]) -> Self { + Self { + lat: R1Interval::from_point_pair(p1[1], p2[1]), + lng: S1Interval::from_point_pair(p1[0], p2[0]), + } + } + + /// Construct a rectangle of the given size centered around the given point. The latitude + /// interval of the result is clamped to [-90,90] degrees, and the longitude interval of the + /// result is Full() if and only if the longitude size is 360 degrees or more. + pub fn from_center_size(center: &[f64; 2], size: &[f64; 2]) -> Self { + Self::from_point(center).expanded(size[1] * 0.5, size[0] * 0.5) + } + + /// The latitude interval. + pub fn lat(&self) -> &R1Interval { + &self.lat + } + + /// The longitude interval. + pub fn lng(&self) -> &S1Interval { + &self.lng + } + + /// The low latitude. + pub fn lat_lo(&self) -> f64 { + self.lat.lo() + } + + /// The high latitude. + pub fn lat_hi(&self) -> f64 { + self.lat.hi() + } + + /// The low longitude. + pub fn lng_lo(&self) -> f64 { + self.lng.lo() + } + + /// The high longitude. + pub fn lng_hi(&self) -> f64 { + self.lng.hi() + } + + /// The low corner as [lng, lat]. + pub fn lo(&self) -> [f64; 2] { + [self.lng.lo(), self.lat.lo()] + } + + /// The high corner as [lng, lat]. + pub fn hi(&self) -> [f64; 2] { + [self.lng.hi(), self.lat.hi()] + } + + /// Returns true if the rectangle is valid, which essentially just means that the latitude + /// bounds do not exceed Pi/2 in absolute value and the longitude bounds do not exceed Pi in + /// absolute value. Also, if either the latitude or longitude bound is empty then both must + /// be. + pub fn is_valid(&self) -> bool { + self.lat.lo().abs() <= FRAC_PI_2 + && self.lat.hi().abs() <= FRAC_PI_2 + && self.lng.is_valid() + && self.lat.is_empty() == self.lng.is_empty() + } + + /// Returns true if the rectangle is empty, i.e. it contains no points at all. + pub fn is_empty(&self) -> bool { + self.lat.is_empty() + } + + /// Returns true if the rectangle is full, i.e. it contains all points. + pub fn is_full(&self) -> bool { + self.lat == Self::full_lat() && self.lng.is_full() + } + + /// Returns true if the rectangle is a point, i.e. lo() == hi() + pub fn is_point(&self) -> bool { + self.lat.lo() == self.lat.hi() && self.lng.lo() == self.lng.hi() + } + + /// Returns true if lng_.lo() > lng_.hi(), i.e. the rectangle crosses the 180 degree + /// longitude line. + pub fn is_inverted(&self) -> bool { + self.lng.is_inverted() + } + + /// Returns the k-th vertex of the rectangle (k = 0,1,2,3) in CCW order (lower left, lower + /// right, upper right, upper left). For convenience, the argument is reduced modulo 4 to + /// the range [0..3]. + pub fn get_vertex(&self, k: usize) -> [f64; 2] { + let i = (k >> 1) & 1; + let lat = if i == 0 { self.lat.lo() } else { self.lat.hi() }; + let j = i ^ (k & 1); + let lng = if j == 0 { self.lng.lo() } else { self.lng.hi() }; + [lng, lat] + } + + /// Returns the center of the rectangle in latitude-longitude space (in general this is not + /// the center of the region on the sphere). + pub fn get_center(&self) -> [f64; 2] { + [self.lng.center(), self.lat.center()] + } + + /// Returns the width and height of this rectangle in latitude-longitude space. Empty + /// rectangles have a negative width and height. + pub fn get_size(&self) -> [f64; 2] { + [self.lng.length(), self.lat.length()] + } + + /// Returns the surface area of this rectangle on the unit sphere. + pub fn area(&self) -> f64 { + if self.is_empty() { + 0.0 + } else { + self.lng.length() * (self.lat_hi().sin() - self.lat_lo().sin()) + } + } + + /// More efficient version of Contains() that accepts a lat/lng rather than an S2Point. + pub fn contains_latlng(&self, ll: &[f64; 2]) -> bool { + self.lat.contains_point(ll[1]) && self.lng.contains_point(ll[0]) + } + + /// Returns true if and only if the given point is contained by the region. + pub fn contains_point(&self, p: &[f64; 3]) -> bool { + let ll = [ + p[1].atan2(p[0]), + p[2].atan2((p[0] * p[0] + p[1] * p[1]).sqrt()), + ]; + self.contains_latlng(&ll) + } + + /// Returns true if and only if the given point is contained in the interior of the region + /// (i.e. the region excluding its boundary). + pub fn interior_contains_latlng(&self, ll: &[f64; 2]) -> bool { + self.lat.interior_contains_point(ll[1]) && self.lng.interior_contains_point(ll[0]) + } + + /// Returns true if and only if the rectangle contains the given other rectangle. + pub fn contains(&self, other: &S2LatLngRect) -> bool { + self.lat.contains(&other.lat) && self.lng.contains(&other.lng) + } + + /// Returns true if and only if the interior of this rectangle contains all points of the + /// given other rectangle (including its boundary). + pub fn interior_contains(&self, other: &S2LatLngRect) -> bool { + self.lat.interior_contains(&other.lat) && self.lng.interior_contains(&other.lng) + } + + /// Returns true if this rectangle and the given other rectangle have any points in common. + pub fn intersects(&self, other: &S2LatLngRect) -> bool { + self.lat.intersects(&other.lat) && self.lng.intersects(&other.lng) + } + + /// Returns true if and only if the interior of this rectangle intersects any point + /// (including the boundary) of the given other rectangle. + pub fn interior_intersects(&self, other: &S2LatLngRect) -> bool { + self.lat.interior_intersects(&other.lat) && self.lng.interior_intersects(&other.lng) + } + + /// Increase the size of the bounding rectangle to include the given point. The rectangle + /// is expanded by the minimum amount possible. + pub fn add_point_3d(&mut self, p: &[f64; 3]) { + let ll = [ + p[1].atan2(p[0]), + p[2].atan2((p[0] * p[0] + p[1] * p[1]).sqrt()), + ]; + self.add_point(&ll); + } + + /// Increase the size of the bounding rectangle to include the given point. + pub fn add_point(&mut self, ll: &[f64; 2]) { + self.lat.add_point(ll[1]); + self.lng.add_point(ll[0]); + } + + /// Returns a rectangle that has been expanded by the given margin on each side in the + /// latitude direction, and by the given margin on each side in the longitude direction. + /// If either margin is negative, then shrinks the rectangle on the corresponding sides + /// instead. The resulting rectangle may be empty. + /// + /// If either the latitude or longitude interval becomes empty after expansion by a negative + /// margin, the result is empty. + pub fn expanded(&self, lat_margin: f64, lng_margin: f64) -> S2LatLngRect { + let lat = self.lat.expanded(lat_margin); + let lng = self.lng.expanded(lng_margin); + if lat.is_empty() || lng.is_empty() { + return S2LatLngRect::empty(); + } + S2LatLngRect::new(lat.intersection(&Self::full_lat()), lng) + } + + /// If the rectangle does not include either pole, returns it unmodified. Otherwise expands + /// the longitude range to Full() so that the rectangle contains all possible representations + /// of the contained pole(s). + pub fn polar_closure(&self) -> S2LatLngRect { + if self.lat.lo() == -FRAC_PI_2 || self.lat.hi() == FRAC_PI_2 { + S2LatLngRect::new(self.lat, S1Interval::full()) + } else { + *self + } + } + + /// Returns the smallest rectangle containing the union of this rectangle and the given + /// rectangle. + pub fn union(&self, other: &S2LatLngRect) -> S2LatLngRect { + S2LatLngRect::new(self.lat.union(&other.lat), self.lng.union(&other.lng)) + } + + /// Returns the smallest rectangle containing the intersection of this rectangle and the + /// given rectangle. Note that the region of intersection may consist of two disjoint + /// rectangles, in which case a single rectangle spanning both of them is returned. + pub fn intersection(&self, other: &S2LatLngRect) -> S2LatLngRect { + let lat = self.lat.intersection(&other.lat); + let lng = self.lng.intersection(&other.lng); + if lat.is_empty() || lng.is_empty() { + S2LatLngRect::empty() + } else { + S2LatLngRect::new(lat, lng) + } + } + + /// Returns true if the latitude and longitude intervals of the two rectangles are the same + /// up to the given tolerance. + pub fn approx_equals(&self, other: &S2LatLngRect, max_error: f64) -> bool { + self.lat.approx_equals(&other.lat, max_error) + && self.lng.approx_equals(&other.lng, max_error) + } + + /// Returns true if the boundary of this rectangle intersects the given geodesic edge + /// (v0, v1). + pub fn boundary_intersects(&self, v0: &[f64; 3], v1: &[f64; 3]) -> bool { + if self.is_empty() { + return false; + } + if !self.lng.is_full() + && (Self::intersects_lng_edge(v0, v1, &self.lat, self.lng.lo()) + || Self::intersects_lng_edge(v0, v1, &self.lat, self.lng.hi())) + { + return true; + } + (self.lat.lo() != -FRAC_PI_2 && Self::intersects_lat_edge(v0, v1, self.lat.lo(), &self.lng)) + || (self.lat.hi() != FRAC_PI_2 + && Self::intersects_lat_edge(v0, v1, self.lat.hi(), &self.lng)) + } + + /// Returns true if the edge AB intersects the given edge of constant longitude. The nice + /// thing about edges of constant longitude is that they are straight lines on the sphere + /// (geodesics). + pub fn intersects_lng_edge(a: &[f64; 3], b: &[f64; 3], lat: &R1Interval, lng: f64) -> bool { + let point = |lat: f64| [lat.cos() * lng.cos(), lat.cos() * lng.sin(), lat.sin()]; + S2EdgeCrosser::new(a, b).crossing_sign_two(&point(lat.lo()), &point(lat.hi())) > 0 + } + + /// Returns true if the edge AB intersects the given edge of constant latitude. + /// Unfortunately, lines of constant latitude are curves on the sphere. They can intersect a + /// straight edge in 0, 1, or 2 points. + pub fn intersects_lat_edge(a: &[f64; 3], b: &[f64; 3], lat: f64, lng: &S1Interval) -> bool { + // First, compute the normal to the plane AB that points vaguely north. + let mut z = normalize(&robust_cross(a, b)); + if z[2] < 0.0 { + z = [-z[0], -z[1], -z[2]]; + } + + // Extend this to an orthonormal frame (x,y,z) where x is the direction where the great + // circle through AB achieves its maximium latitude. + let y = normalize(&robust_cross(&z, &[0.0, 0.0, 1.0])); + let x = cross(&y, &z); + + // Compute the angle "theta" from the x-axis (in the x-y plane defined above) where the + // great circle intersects the given line of latitude. + let sin_lat = lat.sin(); + if sin_lat.abs() >= x[2] { + return false; // The great circle does not reach the given latitude. + } + let cos_theta = sin_lat / x[2]; + let sin_theta = (1.0 - cos_theta * cos_theta).sqrt(); + let theta = sin_theta.atan2(cos_theta); + + // The candidate intersection points are located +/- theta in the x-y plane. For an + // intersection to be valid, we need to check that the intersection point is contained in + // the interior of the edge AB and also that it is contained within the given longitude + // interval "lng". + + // Compute the range of theta values spanned by the edge AB. + let ab_theta = + S1Interval::from_point_pair(dot(a, &y).atan2(dot(a, &x)), dot(b, &y).atan2(dot(b, &x))); + + let within = |sin_theta: f64| { + let isect = [ + x[0] * cos_theta + y[0] * sin_theta, + x[1] * cos_theta + y[1] * sin_theta, + ]; + lng.contains_point(isect[1].atan2(isect[0])) + }; + (ab_theta.contains_point(theta) && within(sin_theta)) + || (ab_theta.contains_point(-theta) && within(-sin_theta)) + } + + /// Returns a bounding spherical cap for this rectangle. We consider two possible bounding + /// caps, one whose axis passes through the center of the rectangle and one whose axis is the + /// north or south pole, and return the smaller of the two. + pub fn get_cap_bound_impl(&self) -> S2Cap { + if self.is_empty() { + return S2Cap::empty(); + } + + let (pole_z, pole_angle) = if self.lat.lo() + self.lat.hi() < 0.0 { + (-1.0, FRAC_PI_2 + self.lat.hi()) + } else { + (1.0, FRAC_PI_2 - self.lat.lo()) + }; + + let pole_cap = S2Cap::new( + [0.0, 0.0, pole_z], + S1ChordAngle::from_radians((1.0 + 2.0 * f64::EPSILON) * pole_angle), + ); + + if self.lng.length() < 2.0 * PI { + let center = self.get_center(); + let cos_lat = center[1].cos(); + let center_point = [ + cos_lat * center[0].cos(), + cos_lat * center[0].sin(), + center[1].sin(), + ]; + + let mut mid_cap = S2Cap::from_point(center_point); + for k in 0..4 { + let vertex = self.get_vertex(k); + let cos_lat_v = vertex[1].cos(); + let vertex_point = [ + cos_lat_v * vertex[0].cos(), + cos_lat_v * vertex[0].sin(), + vertex[1].sin(), + ]; + mid_cap.add_point(&vertex_point); + } + + // C++ compares height(), which is 0.5 * length2(). The factor + // cancels in the comparison so length2() is equivalent. + if mid_cap.radius().length2() < pole_cap.radius().length2() { + return mid_cap; + } + } + + pole_cap + } +} + +impl Region for S2LatLngRect { + fn get_cap_bound(&self) -> S2Cap { + self.get_cap_bound_impl() + } + + fn get_rect_bound(&self) -> S2LatLngRect { + *self + } + + fn get_cell_union_bound(&self, cell_ids: &mut Vec) { + self.get_cap_bound().get_cell_union_bound(cell_ids); + } + + /// Returns true if this rectangle contains the given cell. This is an exact test assuming + /// that S2Cell::GetRectBound() returns tight bounds. + fn contains_cell(&self, cell: &S2Cell) -> bool { + self.contains(&cell.get_rect_bound()) + } + + /// This test is cheap but is NOT exact. Use Intersects() if you want a more accurate and + /// more expensive test. Note that when this method is used by an S2RegionCoverer, the + /// accuracy isn't all that important since if a cell may intersect the region then it is + /// subdivided, and the accuracy of this method goes up as the cells get smaller. + fn may_intersect(&self, cell: &S2Cell) -> bool { + self.intersects(&cell.get_rect_bound()) + } + + /// The point 'p' does not need to be normalized. + fn contains_point(&self, p: &[f64; 3]) -> bool { + self.contains_point(p) + } +} + +impl PartialEq for S2LatLngRect { + fn eq(&self, other: &Self) -> bool { + self.lat == other.lat && self.lng == other.lng + } +} + +impl std::fmt::Display for S2LatLngRect { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[Lo{:?}, Hi{:?}]", self.lo(), self.hi()) + } +} + +#[cfg(test)] +#[path = "tests/latlng_rect_tests.rs"] +mod tests; diff --git a/src/spatial/math.rs b/src/spatial/math.rs new file mode 100644 index 0000000000..8e0a8ea781 --- /dev/null +++ b/src/spatial/math.rs @@ -0,0 +1,54 @@ +//! Vector math utilities for 3D point operations. + +/// Returns the dot product of two 3D vectors. +#[inline] +pub fn dot(a: &[f64; 3], b: &[f64; 3]) -> f64 { + a[0] * b[0] + a[1] * b[1] + a[2] * b[2] +} + +/// Returns the unit vector in the same direction as p. +#[inline] +pub fn normalize(p: &[f64; 3]) -> [f64; 3] { + let norm = (p[0] * p[0] + p[1] * p[1] + p[2] * p[2]).sqrt(); + [p[0] / norm, p[1] / norm, p[2] / norm] +} + +/// Returns the negation of a 3D vector. +#[inline] +pub fn neg(p: [f64; 3]) -> [f64; 3] { + [-p[0], -p[1], -p[2]] +} + +/// Returns the squared Euclidean norm of a 3D vector. +#[inline] +pub fn norm2(p: &[f64; 3]) -> f64 { + p[0] * p[0] + p[1] * p[1] + p[2] * p[2] +} + +/// Returns the cross product of two 3D vectors. +#[inline] +pub fn cross(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + [ + a[1] * b[2] - a[2] * b[1], + a[2] * b[0] - a[0] * b[2], + a[0] * b[1] - a[1] * b[0], + ] +} + +/// Returns the difference of two 3D vectors. +#[inline] +pub fn sub(a: &[f64; 3], b: &[f64; 3]) -> [f64; 3] { + [a[0] - b[0], a[1] - b[1], a[2] - b[2]] +} + +/// Returns x * 2^exp. Equivalent to C's ldexp. +#[inline] +pub fn ldexp(x: f64, exp: i32) -> f64 { + x * 2.0_f64.powi(exp) +} + +/// IEEE 754 remainder. Equivalent to C's remainder(). +#[inline] +pub fn remainder(x: f64, y: f64) -> f64 { + x - (x / y).round_ties_even() * y +} diff --git a/src/spatial/mod.rs b/src/spatial/mod.rs new file mode 100644 index 0000000000..a1954d0ce1 --- /dev/null +++ b/src/spatial/mod.rs @@ -0,0 +1,117 @@ +//! Spatial module. +//! +//! Geospatial indexing using S2 cell partitioning. Polygon edges are stored in a streaming +//! binary format backed by an S2 cell index, enabling polygon containment queries at search +//! engine scale. +//! +//! References: +//! +//! https://s2geometry.io/ +//! https://s2geometry.io/about/overview +//! https://s2geometry.io/devguide/s2cell_hierarchy.html +//! https://proj.org/en/stable/operations/projections/s2.html +//! https://s2geometry.io/resources/earthcube.html + +mod area_query; +pub mod cell_index; +pub mod cell_index_reader; +pub mod cell_union; +pub mod closest_edge_query; +pub mod containment; +pub mod contains_query; +pub mod crossings; +pub mod edge_cache; +pub mod edge_reader; +pub mod edge_writer; +pub mod exact; +pub mod executor; +pub mod geometry; +pub mod geometry_set; +pub mod interleave; +pub mod intersects_query; +pub mod latlng_rect; +pub mod math; +pub mod plane; +pub mod plugin; +mod polygon_union; +pub mod r1interval; +pub mod r2rect; +pub mod reader; +pub mod rect_query; +pub mod region; +pub mod region_coverer; +pub mod s1chord_angle; +pub mod s1interval; +pub mod s2cap; +pub mod s2cell; +pub mod s2cell_id; +pub mod s2coords; +pub mod s2edge_clipping; +pub mod s2edge_distances; +pub mod s2metrics; +pub mod s2padded_cell; +pub mod serializer; +pub mod shape_index_region; +pub mod spatial_index_manager; +pub mod sphere; +pub mod surface; +pub mod within_query; +pub mod writer; +pub mod xor; + +pub use geometry::Geometry; +pub use geometry_set::RingOrientation; + +/// Version of the spatial cell and edge index formats. It follows each file's footer, with a +/// magic number that names the file, so that an index written in another format fails to open +/// with an error that says so, rather than as corrupt. +pub(crate) const FORMAT_VERSION: u32 = 1; +/// Magic number of a cell index. +pub(crate) const CELL_INDEX_MAGIC: u32 = u32::from_le_bytes(*b"S2CI"); +/// Magic number of an edge index. +pub(crate) const EDGE_INDEX_MAGIC: u32 = u32::from_le_bytes(*b"S2EI"); +/// Bytes of the format version and magic number at the end of a file. +pub(crate) const FORMAT_MARKER_LEN: usize = 2 * std::mem::size_of::(); + +/// Write the format version and magic number that end a file. +pub(crate) fn write_format_marker( + write: &mut W, + magic: u32, +) -> std::io::Result<()> { + write.write_all(&FORMAT_VERSION.to_le_bytes())?; + write.write_all(&magic.to_le_bytes()) +} + +/// The file without its format marker, or an error naming what is wrong with the marker. +pub(crate) fn check_format_marker<'a>( + data: &'a [u8], + magic: u32, + file: &str, +) -> crate::Result<&'a [u8]> { + let wrong = |reason: String| { + crate::error::DataCorruption::comment_only(format!( + "The {}-byte spatial {file} {reason}; re-index the spatial field", + data.len() + )) + }; + let Some((body, marker)) = data.split_last_chunk::() else { + return Err(wrong("is too short to hold a format marker".to_string()).into()); + }; + let version = u32::from_le_bytes(marker[..4].try_into().unwrap()); + let found_magic = u32::from_le_bytes(marker[4..].try_into().unwrap()); + if found_magic != magic { + return Err(wrong( + "has no format marker: an older version of tantivy wrote it, or it is corrupt" + .to_string(), + ) + .into()); + } + if version != FORMAT_VERSION { + return Err(wrong(format!( + "has format version {version}, and this version of tantivy reads version \ + {FORMAT_VERSION}" + )) + .into()); + } + Ok(body) +} diff --git a/src/spatial/plane/math.rs b/src/spatial/plane/math.rs new file mode 100644 index 0000000000..c43bdb495e --- /dev/null +++ b/src/spatial/plane/math.rs @@ -0,0 +1,90 @@ +//! Planar geometry predicates. +//! +//! 2D equivalents of the spherical predicates in the S2 port. The orientation test uses Shewchuk's +//! adaptive precision orient2d from the robust crate. Everything else is Euclidean arithmetic. + +/// Orientation of point C relative to the directed line from A to B. Returns positive if C is to +/// the left (counter-clockwise), negative if to the right (clockwise), zero if collinear. Uses +/// Shewchuk's adaptive precision predicate. +pub fn orient2d(a: &[f64; 2], b: &[f64; 2], c: &[f64; 2]) -> f64 { + robust::orient2d( + robust::Coord { x: a[0], y: a[1] }, + robust::Coord { x: b[0], y: b[1] }, + robust::Coord { x: c[0], y: c[1] }, + ) +} + +/// Returns true if segments AB and CD cross. Does not count shared endpoints as crossings. +pub fn segments_cross(a: &[f64; 2], b: &[f64; 2], c: &[f64; 2], d: &[f64; 2]) -> bool { + let d1 = orient2d(a, b, c); + let d2 = orient2d(a, b, d); + let d3 = orient2d(c, d, a); + let d4 = orient2d(c, d, b); + + if ((d1 > 0.0 && d2 < 0.0) || (d1 < 0.0 && d2 > 0.0)) + && ((d3 > 0.0 && d4 < 0.0) || (d3 < 0.0 && d4 > 0.0)) + { + return true; + } + + false +} + +/// Returns true if point P is inside the polygon defined by the given vertices. The polygon is +/// assumed closed (first vertex equals last vertex). Uses the ray casting algorithm with orient2d +/// for robust crossing decisions. +pub fn point_in_polygon(p: &[f64; 2], vertices: &[[f64; 2]]) -> bool { + if vertices.len() < 4 { + return false; + } + let mut inside = false; + let n = vertices.len() - 1; // last vertex is the closure duplicate + let mut j = n - 1; + for i in 0..n { + let vi = &vertices[i]; + let vj = &vertices[j]; + if ((vi[1] > p[1]) != (vj[1] > p[1])) && (orient2d(vi, vj, p) > 0.0) == (vi[1] > vj[1]) { + inside = !inside; + } + j = i; + } + inside +} + +/// Minimum distance from point P to segment AB. Euclidean. +pub fn point_to_segment_distance(p: &[f64; 2], a: &[f64; 2], b: &[f64; 2]) -> f64 { + let dx = b[0] - a[0]; + let dy = b[1] - a[1]; + let len_sq = dx * dx + dy * dy; + if len_sq == 0.0 { + let ex = p[0] - a[0]; + let ey = p[1] - a[1]; + return (ex * ex + ey * ey).sqrt(); + } + let t = ((p[0] - a[0]) * dx + (p[1] - a[1]) * dy) / len_sq; + let t = t.clamp(0.0, 1.0); + let proj_x = a[0] + t * dx; + let proj_y = a[1] + t * dy; + let ex = p[0] - proj_x; + let ey = p[1] - proj_y; + (ex * ex + ey * ey).sqrt() +} + +/// Minimum distance between segments AB and CD. Euclidean. +pub fn segment_to_segment_distance(a: &[f64; 2], b: &[f64; 2], c: &[f64; 2], d: &[f64; 2]) -> f64 { + if segments_cross(a, b, c, d) { + return 0.0; + } + let d1 = point_to_segment_distance(a, c, d); + let d2 = point_to_segment_distance(b, c, d); + let d3 = point_to_segment_distance(c, a, b); + let d4 = point_to_segment_distance(d, a, b); + d1.min(d2).min(d3).min(d4) +} + +/// Euclidean distance between two points. +pub fn edge_length(a: &[f64; 2], b: &[f64; 2]) -> f64 { + let dx = b[0] - a[0]; + let dy = b[1] - a[1]; + (dx * dx + dy * dy).sqrt() +} diff --git a/src/spatial/plane/mod.rs b/src/spatial/plane/mod.rs new file mode 100644 index 0000000000..fb5673c48d --- /dev/null +++ b/src/spatial/plane/mod.rs @@ -0,0 +1,22 @@ +//! Planar surface implementation and predicates for 2D spatial indexing. + +pub mod math; + +use crate::spatial::surface::Surface; + +/// Implementation of a planar surface. +#[derive(Debug, Clone, PartialEq)] +pub struct Plane; + +impl Surface for Plane { + const DIMENSIONS: usize = 2; + type Point = [f64; 2]; + + fn project(lon: f64, lat: f64) -> [f64; 2] { + [lon, lat] + } + + fn origin_inside(_ring: &[Self::Point]) -> bool { + false + } +} diff --git a/src/spatial/plugin.rs b/src/spatial/plugin.rs new file mode 100644 index 0000000000..ba7da2284e --- /dev/null +++ b/src/spatial/plugin.rs @@ -0,0 +1,303 @@ +//! Spatial fields as a [`SegmentPlugin`]. +//! +//! A segment's spatial fields live in two composite files, the cell indexes and the edge +//! indexes, one section per field. The plugin writes them when a segment is flushed, merges them, +//! and reports their space usage. A segment of a schema with no spatial field has neither file. + +use std::any::Any; +use std::collections::BTreeMap; +use std::io::Write; + +use super::cell_index::{write_cell, write_directory}; +use super::cell_index_reader::CellIndexReader; +use super::edge_cache::EdgeCache; +use super::edge_reader::EdgeReader; +use super::edge_writer::EdgeWriter; +use super::interleave::HeapInterleave; +use super::serializer::{SpatialSerializer, EDGE_SKIP_INTERVAL}; +use super::sphere::Sphere; +use super::writer::SpatialWriter; +use crate::directory::CompositeWrite; +use crate::index::{SegmentComponent, SegmentReader}; +use crate::indexer::doc_id_mapping::DocIdMapping; +use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; +use crate::schema::document::{Document, Value}; +use crate::schema::{FieldType, Schema}; +use crate::space_usage::{ComponentSpaceUsage, SPATIAL_CELLS, SPATIAL_EDGES}; +use crate::{DocId, Segment}; + +/// Extension of a segment's cell index file. +pub(crate) const CELLS_EXTENSION: &str = "cells"; +/// Extension of a segment's edge index file. +pub(crate) const EDGES_EXTENSION: &str = "edges"; + +/// The component of a segment's cell indexes. +pub(crate) fn cells_component() -> SegmentComponent { + SegmentComponent::Custom(CELLS_EXTENSION.to_string()) +} + +/// The component of a segment's edge indexes. +pub(crate) fn edges_component() -> SegmentComponent { + SegmentComponent::Custom(EDGES_EXTENSION.to_string()) +} + +/// Bytes of geometry the edge cache of a spatial merge may hold before it evicts geometries: +/// enough for the geometries around the merge's position on the Hilbert curve, without holding a +/// segment's worth of vertex data in memory. +const MERGE_EDGE_CACHE_BYTES: usize = 64 << 20; + +/// Built-in segment plugin that indexes, stores and merges spatial fields. +pub struct SpatialPlugin; + +impl SegmentPlugin for SpatialPlugin { + fn extensions(&self) -> &[&str] { + &[CELLS_EXTENSION, EDGES_EXTENSION] + } + + fn create_writer(&self, ctx: &PluginWriterContext) -> crate::Result> { + Ok(Box::new(SpatialPluginWriter::new(ctx))) + } + + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()> { + merge_spatial_fields(&ctx) + } + + fn space_usage( + &self, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let spatial = segment_reader.spatial_fields(); + let schema = segment_reader.schema(); + Ok(BTreeMap::from([ + ( + SPATIAL_CELLS.to_string(), + ComponentSpaceUsage::PerField(spatial.cells_space_usage(schema)), + ), + ( + SPATIAL_EDGES.to_string(), + ComponentSpaceUsage::PerField(spatial.edges_space_usage(schema)), + ), + ])) + } +} + +/// Accumulates the geometry of a segment's spatial fields, and writes their indexes when the +/// segment is flushed. +pub struct SpatialPluginWriter { + writer: SpatialWriter, + has_spatial_field: bool, +} + +impl SpatialPluginWriter { + pub(crate) fn new(ctx: &PluginWriterContext) -> Self { + SpatialPluginWriter { + writer: SpatialWriter::default(), + has_spatial_field: ctx.segment.schema().contains_spatial_field(), + } + } + + /// Generic, zero-copy document ingestion, called directly by the `SegmentWriter`. + pub(crate) fn index_document( + &mut self, + doc_id: DocId, + doc: &D, + schema: &Schema, + ) -> crate::Result<()> { + if !self.has_spatial_field { + return Ok(()); + } + for (field, value) in doc.iter_fields_and_values() { + let FieldType::Spatial(options) = schema.get_field_entry(field).field_type() else { + continue; + }; + if let Some(geometry) = value.as_value().into_geometry() { + self.writer + .add_geometry(doc_id, field, *geometry, options.ring_orientation()); + } + } + Ok(()) + } +} + +impl PluginWriter for SpatialPluginWriter { + fn serialize( + mut self: Box, + segment: &Segment, + doc_id_map: Option<&DocIdMapping>, + ) -> crate::Result<()> { + if !self.has_spatial_field { + return Ok(()); + } + if let Some(mapping) = doc_id_map { + self.writer + .remap_doc_ids(|doc_id| mapping.get_new_doc_id(doc_id)); + } + let cells = segment.open_write(cells_component())?; + let edges = segment.open_write(edges_component())?; + self.writer + .serialize(SpatialSerializer::new(cells, edges)?)?; + Ok(()) + } + + fn mem_usage(&self) -> usize { + self.writer.mem_usage() + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +/// Merge the spatial fields of the source segments into the target segment. Each field's cells +/// are interleaved in Hilbert order, and its geometry sets are renumbered in the order the cells +/// first reach them, under their documents' new ids. +fn merge_spatial_fields(ctx: &PluginMergeContext) -> crate::Result<()> { + if !ctx.schema.contains_spatial_field() { + return Ok(()); + } + let target_segment = ctx.target_segment; + let mut cells_composite = CompositeWrite::wrap(target_segment.open_write(cells_component())?); + let mut edges_composite = CompositeWrite::wrap(target_segment.open_write(edges_component())?); + + // The spatial fields that have geometry in some segment, with every segment's reader. + // Collected first so they live long enough for the readers they produce. + let mut fields = Vec::new(); + for (field, field_entry) in ctx.schema.fields() { + if !matches!(field_entry.field_type(), FieldType::Spatial(_)) { + continue; + } + let spatial_readers = ctx + .readers + .iter() + .map(|reader| reader.spatial_fields().get_field(field)) + .collect::>>()?; + if spatial_readers.iter().any(Option::is_some) { + fields.push((field, spatial_readers)); + } + } + + // Build inverted doc_id mapping: (segment_ord, old_doc_id) → new_doc_id. Only segments + // with geometry get a row, since only their doc ids are looked up. + let mut doc_id_inverse: Vec>> = Vec::new(); + if !fields.is_empty() { + for (segment_ord, reader) in ctx.readers.iter().enumerate() { + let has_geometry = fields + .iter() + .any(|(_, spatial_readers)| spatial_readers[segment_ord].is_some()); + doc_id_inverse.push(if has_geometry { + vec![None; reader.max_doc() as usize] + } else { + Vec::new() + }); + } + for (new_doc_id, old_addr) in ctx.doc_id_mapping.iter_old_doc_addrs().enumerate() { + let row = &mut doc_id_inverse[old_addr.segment_ord as usize]; + if let Some(slot) = row.get_mut(old_addr.doc_id as usize) { + *slot = Some(new_doc_id as DocId); + } + } + } + for (field, spatial_readers) in &fields { + // Open source cell index readers and edge readers. A segment without geometry for the + // field gets empty readers, so that geometry ids keep their segment ordinals. + let mut cell_readers: Vec = Vec::new(); + let mut edge_readers: Vec> = Vec::new(); + let mut geometry_counts: Vec = Vec::new(); + + for sr in spatial_readers { + if let Some(spatial_reader) = sr { + let cr = CellIndexReader::open(spatial_reader.cells_bytes())?; + let er = EdgeReader::::open(spatial_reader.edges_bytes())?; + geometry_counts.push(er.geometry_count()); + cell_readers.push(cr); + edge_readers.push(er); + } else { + geometry_counts.push(0); + cell_readers.push(CellIndexReader::empty()); + edge_readers.push(EdgeReader::::empty()); + } + } + + // Like the segment writer, leave a field without geometry out of the composites. + let total_source_geometries: u32 = geometry_counts.iter().sum(); + if total_source_geometries == 0 { + continue; + } + + let edge_cache = EdgeCache::new(edge_readers, MERGE_EDGE_CACHE_BYTES); + let iters = cell_readers.iter().map(|cr| cr.iter()).collect(); + let mut interleave = + HeapInterleave::new(iters, &edge_cache, &ctx.doc_id_mapping.alive_bitsets); + + const NOT_ASSIGNED: u32 = u32::MAX; + let mut old_to_new: Vec> = geometry_counts + .iter() + .map(|&count| vec![NOT_ASSIGNED; count as usize]) + .collect(); + let mut next_new_id: u32 = 0; + + let merge_start = std::time::Instant::now(); + + let edges_out = edges_composite.for_field(*field); + let mut edge_writer = EdgeWriter::::new(edges_out, EDGE_SKIP_INTERVAL); + let cells_out = cells_composite.for_field(*field); + let cells_base = cells_out.written_bytes(); + let mut offsets: Vec<(u64, u64)> = Vec::new(); + let mut cell_count: u64 = 0; + + for cell in interleave.by_ref() { + cell_count += 1; + for shape in &cell.shapes { + let segment = shape.geometry_id.0 as usize; + let position = shape.geometry_id.1 as usize; + if old_to_new[segment][position] != NOT_ASSIGNED { + continue; + } + let entry = edge_cache.get(shape.geometry_id); + let set_size = entry.member_count(); + for i in 0..set_size { + let pos = (entry.head() + i) as usize; + if old_to_new[segment][pos] == NOT_ASSIGNED { + old_to_new[segment][pos] = next_new_id; + next_new_id += 1; + } + } + let old_doc_id = entry.doc_id(); + let new_doc_id = doc_id_inverse[segment][old_doc_id as usize].unwrap(); + edge_writer.insert_for_doc(entry.geometry_set(), new_doc_id)?; + } + + offsets.push((cell.cell_id.0, cells_out.written_bytes() - cells_base)); + write_cell(cells_out, &cell, |shape| { + let new_id = old_to_new[shape.geometry_id.0 as usize][shape.geometry_id.1 as usize]; + debug_assert_ne!(new_id, NOT_ASSIGNED); + new_id + })?; + } + + edge_writer.finish()?; + + debug!( + "spatial merge: {} source geometries, {} output geometries, {} cells, {:.2}s", + total_source_geometries, + next_new_id, + cell_count, + merge_start.elapsed().as_secs_f64() + ); + interleave.report(); + + // Cell index dictionary and footer. + let dir_offset = cells_out.written_bytes() - cells_base; + write_directory(cells_out, &offsets, dir_offset)?; + cells_out.flush()?; + edges_out.flush()?; + } + + cells_composite.close()?; + edges_composite.close()?; + Ok(()) +} diff --git a/src/spatial/point.rs b/src/spatial/point.rs new file mode 100644 index 0000000000..04b1664b01 --- /dev/null +++ b/src/spatial/point.rs @@ -0,0 +1,10 @@ +//! A point in the geographical coordinate system. + +/// A point in the geographical coordinate system. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct GeoPoint { + /// Longitude + pub lon: f64, + /// Latitude + pub lat: f64, +} diff --git a/src/spatial/polygon_union.rs b/src/spatial/polygon_union.rs new file mode 100644 index 0000000000..0625a66935 --- /dev/null +++ b/src/spatial/polygon_union.rs @@ -0,0 +1,354 @@ +//! Exposed boundary of a union of polygon members, including boundaries of holes created by +//! their arrangement. Edges are split at crossings and coincident endpoints before classifying +//! each piece, so another member can cover part or all of a member's boundary. + +use std::collections::{HashMap, HashSet}; + +use num_traits::{Signed, ToPrimitive}; + +use super::cell_index::CellIndex; +use super::contains_query::{interior_is_left, midpoint, QueryEdgeProvider}; +use super::crossings::{robust_cross, S2EdgeCrosser}; +use super::exact::{exact_cross, expensive_sign, to_exact}; +use super::geometry_set::EdgeSet; +use super::math::{cross, dot, neg, normalize}; +use super::s1chord_angle::S1ChordAngle; +use super::s2cap::S2Cap; +use super::s2cell_id::S2CellId; +use super::shape_index_region::{clipped_shape_contains_point, EdgeProvider}; + +#[derive(Clone, Copy)] +pub(crate) struct PolygonEdge { + pub(crate) v0: [f64; 3], + pub(crate) v1: [f64; 3], + pub(crate) interior_left: bool, + /// Original endpoints define the great circle even after its cuts have been rounded. + source: ([f64; 3], [f64; 3]), + bound: S2Cap, +} + +impl PolygonEdge { + fn new(v0: [f64; 3], v1: [f64; 3], interior_left: bool) -> Self { + let center = midpoint((v0, v1)); + let radius = S1ChordAngle::from_points(¢er, &v0).to_radians(); + Self { + v0, + v1, + interior_left, + source: (v0, v1), + // Account for rounding in the midpoint and cap radius when rejecting edge pairs. + bound: S2Cap::from_center_angle(center, radius + 1e-14), + } + } + + pub(crate) fn midpoint(&self) -> [f64; 3] { + midpoint((self.v0, self.v1)) + } + + /// Position along the minor arc, with endpoints at zero and one. + fn position(&self, point: &[f64; 3]) -> f64 { + if *point == self.v0 { + return 0.0; + } + if *point == self.v1 { + return 1.0; + } + let (tangent, length) = self.frame(); + dot(point, &tangent).atan2(dot(point, &self.v0)) / length + } + + fn frame(&self) -> ([f64; 3], f64) { + let tangent = normalize(&cross( + &robust_cross(&self.source.0, &self.source.1), + &self.v0, + )); + let length = dot(&self.v1, &tangent).atan2(dot(&self.v1, &self.v0)); + (tangent, length) + } + + fn piece(&self, v0: [f64; 3], v1: [f64; 3]) -> Self { + let mut piece = Self::new(v0, v1, self.interior_left); + piece.source = self.source; + piece + } +} + +/// Edges of a member, excluding closure-to-next-ring links and zero-length edges. For an open +/// member the interior side is unused. +pub(crate) fn member_edges(member: &EdgeSet) -> impl Iterator + '_ { + member.ring_offsets.windows(2).flat_map(move |ring| { + (ring[0]..ring[1].saturating_sub(1)) + .filter(|&i| member.vertices[i] != member.vertices[i + 1]) + .map(move |i| { + ( + i as u32, + PolygonEdge::new( + member.vertices[i], + member.vertices[i + 1], + interior_is_left(member, i as u32), + ), + ) + }) + }) +} + +struct Overlap { + start: f64, + end: f64, + member: u32, + interior_left: bool, +} + +struct SplitEdge { + edge: PolygonEdge, + cuts: Vec<(f64, [f64; 3])>, + overlaps: Vec, +} + +impl SplitEdge { + fn new(edge: PolygonEdge) -> Self { + Self { + edge, + cuts: vec![(0.0, edge.v0), (1.0, edge.v1)], + overlaps: Vec::new(), + } + } + + fn add_cut(&mut self, point: [f64; 3]) { + let position = self.edge.position(&point); + if position > 0.0 && position < 1.0 { + self.cuts.push((position, point)); + } + } + + fn split_at(&mut self, other: PolygonEdge, member: u32) { + if !self.edge.bound.intersects_cap(&other.bound) { + return; + } + let edge = self.edge; + // Treat coincident arcs geometrically, before the crossing predicate's symbolic + // perturbation distinguishes otherwise collinear edges. + let identical = + edge.source == other.source || edge.source == (other.source.1, other.source.0); + let collinear = identical + || (expensive_sign(&edge.source.0, &edge.source.1, &other.source.0, false) == 0 + && expensive_sign(&edge.source.0, &edge.source.1, &other.source.1, false) == 0); + if collinear { + let a = edge.position(&other.v0); + let mut b = edge.position(&other.v1); + let same_direction = dot( + &robust_cross(&edge.source.0, &edge.source.1), + &robust_cross(&other.source.0, &other.source.1), + ) > 0.0; + let period = std::f64::consts::TAU / edge.frame().1; + // atan2 wraps at pi. Unwrap the other minor arc using its direction before + // intersecting it with this arc, including its equivalent positions a turn away. + if same_direction && b < a { + b += period; + } else if !same_direction && b > a { + b -= period; + } + for shift in [-period, 0.0, period] { + let start = (a.min(b) + shift).max(0.0); + let end = (a.max(b) + shift).min(1.0); + if start < end { + self.add_cut(other.v0); + self.add_cut(other.v1); + self.overlaps.push(Overlap { + start, + end, + member, + interior_left: other.interior_left ^ !same_direction, + }); + } + } + } else if S2EdgeCrosser::new(&edge.v0, &edge.v1).crossing_sign_two(&other.v0, &other.v1) > 0 + { + self.add_cut(intersection(edge, other)); + } + } + + fn pieces(&mut self) -> Vec<(PolygonEdge, f64)> { + self.cuts.sort_unstable_by(|a, b| a.0.total_cmp(&b.0)); + self.cuts.dedup_by(|a, b| a.0 == b.0); + self.cuts + .windows(2) + .filter(|pair| pair[0].1 != pair[1].1) + .map(|pair| { + ( + self.edge.piece(pair[0].1, pair[1].1), + 0.5 * (pair[0].0 + pair[1].0), + ) + }) + .collect() + } + + fn overlaps_at(&self, position: f64) -> impl Iterator { + self.overlaps + .iter() + .filter(move |overlap| overlap.start < position && position < overlap.end) + } +} + +/// Intersection of two crossing minor arcs. Exact products avoid moving the cut across a +/// narrow gap when the edges are nearly parallel; the resulting direction is rounded only once. +fn intersection(a: PolygonEdge, b: PolygonEdge) -> [f64; 3] { + // Reuse an existing vertex at a T-junction. Renormalizing it can otherwise introduce a + // tiny extra piece between two representations of the same intersection. + for (point, other) in [ + (a.source.0, b), + (a.source.1, b), + (b.source.0, a), + (b.source.1, a), + ] { + if expensive_sign(&other.source.0, &other.source.1, &point, false) == 0 + && dot(&point, &a.midpoint()) > 0.0 + && dot(&point, &b.midpoint()) > 0.0 + { + return point; + } + } + let ab = exact_cross(&to_exact(&a.source.0), &to_exact(&a.source.1)); + let cd = exact_cross(&to_exact(&b.source.0), &to_exact(&b.source.1)); + let point = exact_cross(&ab, &cd); + let scale = point.iter().map(Signed::abs).max().unwrap(); + let point = normalize(&point.map(|coordinate| (coordinate / &scale).to_f64().unwrap())); + let middle = a.midpoint(); + if dot(&point, &middle) < 0.0 { + neg(point) + } else { + point + } +} + +/// Prepare the union's exposed boundary once per query. The query index narrows the edge +/// pairs that can meet, and each pair is split only once even if it occurs in several cells. +pub(crate) fn union_boundary(index: &CellIndex, provider: &QueryEdgeProvider) -> Vec { + let mut positions = HashMap::new(); + let mut edges = Vec::new(); + for (member, set) in provider.set.members.iter().enumerate() { + if !set.closed { + continue; + } + for (edge_idx, edge) in member_edges(set) { + positions.insert((member as u32, edge_idx), edges.len()); + edges.push((member as u32, SplitEdge::new(edge))); + } + } + if edges.is_empty() { + return Vec::new(); + } + + let mut visited = HashSet::new(); + for cell in &index.cells { + let shapes: Vec<_> = cell + .shapes + .iter() + .filter(|shape| !shape.edge_indices.is_empty() && provider.is_closed(shape.geometry_id)) + .collect(); + for (i, a) in shapes.iter().enumerate() { + for b in &shapes[i + 1..] { + for ai in &a.edge_indices { + let Some(&ai) = positions.get(&(a.geometry_id.1, *ai)) else { + continue; + }; + for bi in &b.edge_indices { + let Some(&bi) = positions.get(&(b.geometry_id.1, *bi)) else { + continue; + }; + let (ai, bi) = (ai.min(bi), ai.max(bi)); + if !visited.insert((ai, bi)) { + continue; + } + let (before, after) = edges.split_at_mut(bi); + let (am, a) = &mut before[ai]; + let (bm, b) = &mut after[0]; + a.split_at(b.edge, *bm); + b.split_at(a.edge, *am); + } + } + } + } + } + + let mut boundary = Vec::new(); + for (member, mut edge) in edges { + for (piece, position) in edge.pieces() { + let mut covered = [false; 2]; + covered[usize::from(piece.interior_left)] = true; + for overlap in edge.overlaps_at(position) { + covered[usize::from(overlap.interior_left)] = true; + } + let point = piece.midpoint(); + if let Some(cell) = index.find_cell(S2CellId::from_point(&point)) { + for shape in &cell.shapes { + let other = shape.geometry_id.1; + if other == member + || !provider.is_closed(shape.geometry_id) + || edge + .overlaps_at(position) + .any(|overlap| overlap.member == other) + { + continue; + } + if clipped_shape_contains_point(cell.cell_id, shape, &point, |edge_idx| { + provider.get_edge(shape.geometry_id, edge_idx) + }) { + covered = [true; 2]; + break; + } + } + } + if covered[0] != covered[1] { + boundary.push(piece); + } + } + } + boundary +} + +pub(crate) struct EdgePiece { + pub(crate) edge: PolygonEdge, + /// The interior sides contributed by edges running along this piece, indexed by leftness. + pub(crate) shared_sides: [bool; 2], +} + +/// Split an edge against another boundary, retaining which sides of coincident pieces it covers. +pub(crate) fn split_edge(edge: PolygonEdge, boundary: &[PolygonEdge]) -> Vec { + let mut split = SplitEdge::new(edge); + for other in boundary { + split.split_at(*other, 0); + } + split + .pieces() + .into_iter() + .map(|(edge, position)| { + let mut shared_sides = [false; 2]; + for overlap in split.overlaps_at(position) { + shared_sides[usize::from(overlap.interior_left)] = true; + } + EdgePiece { edge, shared_sides } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spatial::sphere::Sphere; + use crate::spatial::surface::Surface; + + #[test] + fn test_coincident_minor_arcs_across_the_angle_wrap() { + let point = |longitude| Sphere::project(longitude, 0.0); + let edge = PolygonEdge::new(point(10.0), point(100.0), true); + for (start, end, left) in [(80.0, 200.0, true), (200.0, 80.0, false)] { + let other = PolygonEdge::new(point(start), point(end), true); + let pieces = split_edge(edge, &[other]); + assert_eq!(pieces.len(), 2); + assert_eq!(pieces[0].shared_sides, [false; 2]); + assert!(pieces[1].shared_sides[usize::from(left)]); + assert!(!pieces[1].shared_sides[usize::from(!left)]); + } + } +} diff --git a/src/spatial/r1interval.rs b/src/spatial/r1interval.rs new file mode 100644 index 0000000000..a0c918725f --- /dev/null +++ b/src/spatial/r1interval.rs @@ -0,0 +1,213 @@ +//! An R1Interval represents a closed, bounded interval on the real line. It is capable of +//! representing the empty interval (containing no points) and zero-length intervals (containing a +//! single point). + +/// A closed interval on the real line [lo, hi]. +/// +/// The interval is empty if lo > hi. This is used for latitude bounds and other 1D ranges where +/// wraparound is not needed. +#[derive(Clone, Copy, Debug)] +pub struct R1Interval { + lo: f64, + hi: f64, +} + +impl PartialEq for R1Interval { + fn eq(&self, other: &Self) -> bool { + (self.lo == other.lo && self.hi == other.hi) || (self.is_empty() && other.is_empty()) + } +} + +impl R1Interval { + /// Creates a new interval [lo, hi]. If lo > hi, the interval is empty. + pub fn new(lo: f64, hi: f64) -> Self { + Self { lo, hi } + } + + /// Returns an empty interval. An empty interval contains no points. + pub fn empty() -> Self { + Self { lo: 1.0, hi: 0.0 } + } + + /// Creates an interval containing a single point. + pub fn from_point(p: f64) -> Self { + Self { lo: p, hi: p } + } + + /// Creates the minimal interval containing two points. + pub fn from_point_pair(p1: f64, p2: f64) -> Self { + if p1 <= p2 { + Self { lo: p1, hi: p2 } + } else { + Self { lo: p2, hi: p1 } + } + } + + /// The low bound of the interval. + pub fn lo(&self) -> f64 { + self.lo + } + + /// The high bound of the interval. + pub fn hi(&self) -> f64 { + self.hi + } + + /// Returns true if the interval is empty (contains no points). + pub fn is_empty(&self) -> bool { + self.lo > self.hi + } + + /// Returns the center of the interval. For empty intervals, the result is arbitrary. + pub fn center(&self) -> f64 { + 0.5 * (self.lo + self.hi) + } + + /// Returns the length of the interval. The length of an empty interval is negative. + pub fn length(&self) -> f64 { + self.hi - self.lo + } + + /// Returns true if the given point is in the closed interval [lo, hi]. + pub fn contains_point(&self, p: f64) -> bool { + p >= self.lo && p <= self.hi + } + + /// Returns true if the given point is in the open interval (lo, hi). The interior excludes the + /// endpoints. + pub fn interior_contains_point(&self, p: f64) -> bool { + p > self.lo && p < self.hi + } + + /// Returns true if this interval contains the given interval. + pub fn contains(&self, other: &R1Interval) -> bool { + if other.is_empty() { + true + } else { + other.lo >= self.lo && other.hi <= self.hi + } + } + + /// Returns true if the interior of this interval contains the entire + /// interval 'other' (including its boundary). + pub fn interior_contains(&self, other: &R1Interval) -> bool { + if other.is_empty() { + true + } else { + other.lo > self.lo && other.hi < self.hi + } + } + + /// Returns true if this interval intersects the given interval. + pub fn intersects(&self, other: &R1Interval) -> bool { + if self.lo <= other.lo { + other.lo <= self.hi && other.lo <= other.hi + } else { + self.lo <= other.hi && self.lo <= self.hi + } + } + + /// Returns true if the interior of this interval intersects any point of the given interval + /// (including its boundary). + pub fn interior_intersects(&self, other: &R1Interval) -> bool { + other.lo < self.hi && self.lo < other.hi && self.lo < self.hi && other.lo <= other.hi + } + + /// Expands the interval to contain the given point. + pub fn add_point(&mut self, p: f64) { + if self.is_empty() { + self.lo = p; + self.hi = p; + } else if p < self.lo { + self.lo = p; + } else if p > self.hi { + self.hi = p; + } + } + + /// Expands the interval to contain the given interval. + pub fn add_interval(&mut self, other: &R1Interval) { + if other.is_empty() { + return; + } + if self.is_empty() { + *self = *other; + return; + } + if other.lo < self.lo { + self.lo = other.lo; + } + if other.hi > self.hi { + self.hi = other.hi; + } + } + + /// Returns the closest point in the interval to the given point. The interval must be + /// non-empty. + pub fn project(&self, p: f64) -> f64 { + debug_assert!(!self.is_empty()); + p.clamp(self.lo, self.hi) + } + + /// Returns an interval expanded by the given margin on each side. If margin is negative, + /// shrinks instead. Any expansion of an empty interval remains empty. + pub fn expanded(&self, margin: f64) -> R1Interval { + if self.is_empty() { + *self + } else { + R1Interval::new(self.lo - margin, self.hi + margin) + } + } + + /// Returns the smallest interval containing both this and 'other'. + pub fn union(&self, other: &R1Interval) -> R1Interval { + if self.is_empty() { + *other + } else if other.is_empty() { + *self + } else { + R1Interval::new(self.lo.min(other.lo), self.hi.max(other.hi)) + } + } + + /// Returns the intersection of this interval with 'other'. Empty intervals do not need + /// special-casing. + pub fn intersection(&self, other: &R1Interval) -> R1Interval { + R1Interval::new(self.lo.max(other.lo), self.hi.min(other.hi)) + } + + /// Returns true if this interval equals 'other' within the given tolerance. Empty intervals + /// are considered to match any interval with length <= 2*max_error. + pub fn approx_equals(&self, other: &R1Interval, max_error: f64) -> bool { + if self.is_empty() { + other.length() <= 2.0 * max_error + } else if other.is_empty() { + self.length() <= 2.0 * max_error + } else { + (other.lo - self.lo).abs() <= max_error && (other.hi - self.hi).abs() <= max_error + } + } + + /// Methods that allow the R1Interval to be accessed as a vector. (The recommended style is to + /// use lo() and hi() whenever possible, but these methods are useful when the endpoint to be + /// selected is not constant.) + pub fn get_bound(&self, i: usize) -> f64 { + if i == 0 { + self.lo() + } else { + self.hi() + } + } + + /// Clamps the given value to this interval. + #[inline] + pub fn clamp(&self, v: f64) -> f64 { + v.max(self.lo()).min(self.hi()) + } +} + +impl std::fmt::Display for R1Interval { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[{}, {}]", self.lo, self.hi) + } +} diff --git a/src/spatial/r2rect.rs b/src/spatial/r2rect.rs new file mode 100644 index 0000000000..5128c84f78 --- /dev/null +++ b/src/spatial/r2rect.rs @@ -0,0 +1,253 @@ +//! Axis-aligned rectangles in the (x,y) plane. +//! +//! R2Rect represents closed, bounded rectangles using intervals for each axis. +//! Used for UV-space bounds during cell operations and edge clipping. +//! +//! R2Rect in r2rect.h and r2rect.cc. + +use super::r1interval::R1Interval; + +/// A closed axis-aligned rectangle defined by intervals on the x and y axes. +/// +/// Empty rectangles have empty intervals on both axes. Use `is_empty()` to test, since empty +/// rectangles have multiple representations. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct R2Rect { + /// The x-axis interval. + pub x: R1Interval, + /// The y-axis interval. + pub y: R1Interval, +} + +impl R2Rect { + /// Construct a rectangle from the given intervals in x and y. The two intervals must either + /// be both empty or both non-empty. + #[inline] + pub fn new(x: R1Interval, y: R1Interval) -> Self { + debug_assert!(x.is_empty() == y.is_empty()); + Self { x, y } + } + + /// Construct a rectangle from the given lower-left and upper-right points. + #[inline] + pub fn from_points(lo: [f64; 2], hi: [f64; 2]) -> Self { + let x = R1Interval::new(lo[0], hi[0]); + let y = R1Interval::new(lo[1], hi[1]); + debug_assert!(x.is_empty() == y.is_empty()); + Self { x, y } + } + + /// The canonical empty rectangle. Use is_empty() to test for empty rectangles, since they + /// have more than one representation. + #[inline] + pub fn empty() -> Self { + Self { + x: R1Interval::empty(), + y: R1Interval::empty(), + } + } + + /// Convenience method to construct a rectangle containing a single point. + #[inline] + pub fn from_point(p: [f64; 2]) -> Self { + Self::from_points(p, p) + } + + /// Convenience method to construct the minimal bounding rectangle containing the two given + /// points. This is equivalent to starting with an empty rectangle and calling AddPoint() + /// twice. Note that it is different than the R2Rect(lo, hi) constructor, where the first + /// point is always used as the lower-left corner of the resulting rectangle. + #[inline] + pub fn from_point_pair(p1: [f64; 2], p2: [f64; 2]) -> Self { + Self { + x: R1Interval::from_point_pair(p1[0], p2[0]), + y: R1Interval::from_point_pair(p1[1], p2[1]), + } + } + + /// Accessor. + #[inline] + pub fn lo(&self) -> [f64; 2] { + [self.x.lo(), self.y.lo()] + } + + /// Accessor. + #[inline] + pub fn hi(&self) -> [f64; 2] { + [self.x.hi(), self.y.hi()] + } + + /// Return true if the rectangle is empty, i.e. it contains no points at all. + #[inline] + pub fn is_empty(&self) -> bool { + self.x.is_empty() + } + + /// Return the vertex in direction "i" along the x-axis (0=left, 1=right) and direction "j" + /// along the y-axis (0=down, 1=up). Equivalently, return the vertex constructed by selecting + /// endpoint "i" of the x-interval (0=lo, 1=hi) and vertex "j" of the y-interval. + #[inline] + pub fn get_vertex(&self, k: usize) -> [f64; 2] { + let j = (k >> 1) & 1; + self.get_vertex_ij(j ^ (k & 1), j) + } + + /// Return the vertex in direction "i" along the x-axis (0=left, 1=right) and direction "j" + /// along the y-axis (0=down, 1=up). Equivalently, return the vertex constructed by selecting + /// endpoint "i" of the x-interval (0=lo, 1=hi) and vertex "j" of the y-interval. + #[inline] + pub fn get_vertex_ij(&self, i: usize, j: usize) -> [f64; 2] { + [ + if i == 0 { self.x.lo() } else { self.x.hi() }, + if j == 0 { self.y.lo() } else { self.y.hi() }, + ] + } + + /// Return the center of the rectangle in (x,y)-space. + #[inline] + pub fn get_center(&self) -> [f64; 2] { + [self.x.center(), self.y.center()] + } + + /// Return the width and height of this rectangle in (x,y)-space. Empty rectangles have a + /// negative width and height. + #[inline] + pub fn get_size(&self) -> [f64; 2] { + [self.x.length(), self.y.length()] + } + + /// Return true if the rectangle contains the given point. Note that rectangles are closed + /// regions, i.e. they contain their boundary. + #[inline] + pub fn contains_point(&self, p: [f64; 2]) -> bool { + self.x.contains_point(p[0]) && self.y.contains_point(p[1]) + } + + /// Return true if and only if the given point is contained in the interior of the region (i.e. + /// the region excluding its boundary). + #[inline] + pub fn interior_contains_point(&self, p: [f64; 2]) -> bool { + self.x.interior_contains_point(p[0]) && self.y.interior_contains_point(p[1]) + } + + /// Return true if and only if the rectangle contains the given other rectangle. + #[inline] + pub fn contains(&self, other: &R2Rect) -> bool { + self.x.contains(&other.x) && self.y.contains(&other.y) + } + + /// Return true if this rectangle and the given other rectangle have any points in common. + #[inline] + pub fn intersects(&self, other: &R2Rect) -> bool { + self.x.intersects(&other.x) && self.y.intersects(&other.y) + } + + /// Return a rectangle that has been expanded on each side in the x-direction by margin.x(), + /// and on each side in the y-direction by margin.y(). If either margin is negative, then + /// shrink the interval on the corresponding sides instead. The resulting rectangle may be + /// empty. Any expansion of an empty rectangle remains empty. + #[inline] + pub fn expanded(&self, margin: f64) -> Self { + let x = self.x.expanded(margin); + let y = self.y.expanded(margin); + if x.is_empty() || y.is_empty() { + return Self::empty(); + } + Self { x, y } + } + + /// Returns the smallest rectangle containing both this and the given rectangle. + #[inline] + pub fn expanded_xy(&self, margin: [f64; 2]) -> Self { + let x = self.x.expanded(margin[0]); + let y = self.y.expanded(margin[1]); + if x.is_empty() || y.is_empty() { + return Self::empty(); + } + Self { x, y } + } + + /// Return the smallest rectangle containing the union of this rectangle and the given + /// rectangle. + #[inline] + pub fn union(&self, other: &R2Rect) -> Self { + Self { + x: self.x.union(&other.x), + y: self.y.union(&other.y), + } + } + + /// Return the smallest rectangle containing the intersection of this rectangle and the given + /// rectangle. An empty intersection is the canonical empty rectangle, which is_empty() relies + /// on. + #[inline] + pub fn intersection(&self, other: &R2Rect) -> Self { + let x = self.x.intersection(&other.x); + let y = self.y.intersection(&other.y); + if x.is_empty() || y.is_empty() { + return Self::empty(); + } + Self { x, y } + } + + /// Expand the rectangle to include the given point. The rectangle is expanded by the minimum + /// amount possible. + #[inline] + pub fn add_point(&mut self, p: [f64; 2]) { + self.x.add_point(p[0]); + self.y.add_point(p[1]); + } + + /// Expand the rectangle to include the given other rectangle. This is the same as replacing + /// the rectangle by the union of the two rectangles, but is somewhat more efficient. + #[inline] + pub fn add_rect(&mut self, other: &R2Rect) { + self.x = self.x.union(&other.x); + self.y = self.y.union(&other.y); + } + + /// Return the closest point in the rectangle to the given point "p". The rectangle must be + /// non-empty. + #[inline] + pub fn project(&self, p: [f64; 2]) -> [f64; 2] { + [self.x.project(p[0]), self.y.project(p[1])] + } +} + +impl std::ops::Index for R2Rect { + type Output = R1Interval; + #[inline] + fn index(&self, i: usize) -> &R1Interval { + match i { + 0 => &self.x, + 1 => &self.y, + _ => panic!("R2Rect index out of bounds"), + } + } +} + +impl Default for R2Rect { + /// The default constructor creates an empty R2Rect. + fn default() -> Self { + Self::empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_intersection_that_misses_on_one_axis_is_empty() { + let a = R2Rect::from_point_pair([0.0, 0.0], [2.0, 2.0]); + let below = R2Rect::from_point_pair([1.0, -3.0], [3.0, -1.0]); + assert!(a.intersection(&below).is_empty()); + let left = R2Rect::from_point_pair([-3.0, 1.0], [-1.0, 3.0]); + assert!(a.intersection(&left).is_empty()); + let overlap = R2Rect::from_point_pair([1.0, 1.0], [3.0, 3.0]); + assert_eq!( + a.intersection(&overlap), + R2Rect::from_point_pair([1.0, 1.0], [2.0, 2.0]) + ); + } +} diff --git a/src/spatial/reader.rs b/src/spatial/reader.rs new file mode 100644 index 0000000000..2fd7ccd24d --- /dev/null +++ b/src/spatial/reader.rs @@ -0,0 +1,148 @@ +//! Spatial segment readers. +//! +//! Provides per-field access to cell index and edge index data stored in a segment's +//! cell index and edge index composite files. +use std::io; +use std::sync::Arc; + +use common::file_slice::FileSlice; +use common::OwnedBytes; + +use crate::directory::CompositeFile; +use crate::error::DataCorruption; +use crate::schema::{Field, Schema}; +use crate::space_usage::PerFieldSpaceUsage; + +/// Per-field spatial index readers for a segment. +#[derive(Clone)] +pub struct SpatialReaders { + cells: Arc, + edges: Arc, +} + +impl SpatialReaders { + /// Returns an empty SpatialReaders for segments with no spatial fields. + pub fn empty() -> SpatialReaders { + SpatialReaders { + cells: Arc::new(CompositeFile::empty()), + edges: Arc::new(CompositeFile::empty()), + } + } + + /// Opens readers from cell index and edge index file slices. + pub fn open(cells_file: FileSlice, edges_file: FileSlice) -> crate::Result { + let cells = CompositeFile::open(&cells_file)?; + let edges = CompositeFile::open(&edges_file)?; + Ok(SpatialReaders { + cells: Arc::new(cells), + edges: Arc::new(edges), + }) + } + + /// Returns the per-field spatial reader, or None if the segment has no geometry for the + /// field. A field must have both a cell index and an edge index, or neither. + pub fn get_field(&self, field: Field) -> crate::Result> { + let cells_file = self.cells.open_read(field); + let edges_file = self.edges.open_read(field); + match (cells_file, edges_file) { + (Some(c), Some(e)) => { + let reader = SpatialReader::open(c, e)?; + Ok(Some(reader)) + } + (None, None) => Ok(None), + (Some(_), None) | (None, Some(_)) => Err(DataCorruption::comment_only(format!( + "Spatial field {field:?} has only one of a cell index and an edge index" + )) + .into()), + } + } + + /// Return a break down of the cell index space usage per field. + pub fn cells_space_usage(&self, schema: &Schema) -> PerFieldSpaceUsage { + self.cells.space_usage(schema) + } + + /// Return a break down of the edge index space usage per field. + pub fn edges_space_usage(&self, schema: &Schema) -> PerFieldSpaceUsage { + self.edges.space_usage(schema) + } +} + +/// Per-field spatial reader holding cell index and edge index bytes. +#[derive(Clone)] +pub struct SpatialReader { + cells_data: OwnedBytes, + edges_data: OwnedBytes, +} + +impl SpatialReader { + /// Opens the spatial reader from cell and edge file slices. + pub fn open(cells_file: FileSlice, edges_file: FileSlice) -> io::Result { + let cells_data = cells_file.read_bytes()?; + let edges_data = edges_file.read_bytes()?; + Ok(SpatialReader { + cells_data, + edges_data, + }) + } + + /// Returns the cell index bytes. + pub fn cells_bytes(&self) -> &[u8] { + self.cells_data.as_ref() + } + + /// Returns the edge index bytes. + pub fn edges_bytes(&self) -> &[u8] { + self.edges_data.as_ref() + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + use std::path::Path; + + use super::*; + use crate::directory::{CompositeWrite, Directory, RamDirectory}; + + /// A composite file with the given fields, each holding one byte. + fn composite( + directory: &RamDirectory, + name: &str, + fields: &[Field], + ) -> crate::Result { + let path = Path::new(name); + let mut write = CompositeWrite::wrap(directory.open_write(path)?); + for &field in fields { + write.for_field(field).write_all(&[0])?; + } + write.close()?; + Ok(directory.open_read(path)?) + } + + #[test] + fn test_get_field_rejects_a_field_missing_from_one_index() -> crate::Result<()> { + let directory = RamDirectory::create(); + let (both, cells_only, edges_only, neither) = ( + Field::from_field_id(0), + Field::from_field_id(1), + Field::from_field_id(2), + Field::from_field_id(3), + ); + let readers = SpatialReaders::open( + composite(&directory, "cells", &[both, cells_only])?, + composite(&directory, "edges", &[both, edges_only])?, + )?; + assert!(readers.get_field(both)?.is_some()); + assert!(readers.get_field(neither)?.is_none()); + assert!(matches!( + readers.get_field(cells_only), + Err(crate::TantivyError::DataCorruption(_)) + )); + assert!(matches!( + readers.get_field(edges_only), + Err(crate::TantivyError::DataCorruption(_)) + )); + Ok(()) + } +} diff --git a/src/spatial/rect_query.rs b/src/spatial/rect_query.rs new file mode 100644 index 0000000000..171fca33df --- /dev/null +++ b/src/spatial/rect_query.rs @@ -0,0 +1,351 @@ +//! Latitude-longitude rectangle query. +//! +//! Finds the geometries that a rectangle of latitude and longitude covers, or that it meets. +//! Two sides of such a rectangle are meridians, which are geodesics. The other two are parallels, +//! which are not: the polygon through the four corners has edges that bow toward the nearer +//! pole, by degrees where the rectangle spans a continent, and that take the short way around +//! where it is wider than a hemisphere. The rectangle is therefore tested as what it is. +//! +//! A covering of the rectangle finds the candidates, as `area_query` has it. A covering cell +//! within the rectangle decides what it holds without decoding it, for the geometries the +//! rectangle covers too. A geometry found only in such cells lies within the rectangle: it is +//! connected, so to leave the rectangle it would pass through a covering cell that the +//! rectangle's boundary crosses. +//! +//! The rectangle is closed, and a point within rounding of its boundary is inside it: the +//! vertices a segment stores are unit vectors, whose latitude and longitude are not exactly +//! those they were made from. +//! +//! This module does not port any single S2 class. It composes S2LatLngRect and the region +//! coverer with the serialized segment index. + +use std::collections::hash_map::Entry; +use std::collections::HashMap; +use std::f64::consts::{FRAC_PI_2, PI}; + +use common::BitSet; + +use super::area_query::{edges, search_meeting, Area, Covering}; +use super::cell_index_reader::CellIndexReader; +use super::crossings::robust_cross; +use super::edge_cache::EdgeCache; +use super::latlng_rect::S2LatLngRect; +use super::math::{cross, dot, normalize}; +use super::r1interval::R1Interval; +use super::region::Region; +use super::s1interval::S1Interval; +use super::shape_index_region::{contains_point, SegmentIndex}; +use super::sphere::Sphere; + +/// How far outside the rectangle a unit vector may lie and still be inside it. A stored vertex +/// is within a few units of rounding of the coordinates it was made from, and a geometry in an +/// index cell may lie outside it by the cell's padding. +const SLACK: f64 = 1e-15; + +/// How a geometry must relate to the rectangle to match. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RectRelation { + /// The rectangle covers every part of every value of the document's field. + CoveredBy, + /// The rectangle and a value of the document's field have a point in common. + Intersects, +} + +/// The longitudes of a rectangle. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Span { + /// Every longitude. + Full, + /// Half of them at most: the rectangle lies between its west and east meridians. + Narrow, + /// More than half: the rectangle is all but what lies between its east and west meridians. + Wide, +} + +/// A rectangle prepared for testing unit vectors against it, with no trigonometry for a point. +struct Bounds { + rect: S2LatLngRect, + /// The sines of the south and north latitudes: the heights the rectangle lies between. + south: f64, + north: f64, + span: Span, + /// The cosine and sine of the west, east and middle longitudes. + west: [f64; 2], + east: [f64; 2], + middle: [f64; 2], + /// A point outside the rectangle in each part of the sphere that it leaves out. + outside: Vec<[f64; 3]>, + /// A point of the rectangle. + inside: [f64; 3], +} + +impl Bounds { + /// The rectangle from its west and south, and its east and north, in degrees. It crosses the + /// antimeridian when its west is greater than its east. + fn new([west, south]: [f64; 2], [east, north]: [f64; 2]) -> Self { + let lng = if west == -180.0 && east == 180.0 { + S1Interval::full() + } else { + // From 180 to -180 is the antimeridian, though S1Interval has it as empty. + let east = if west == 180.0 && east == -180.0 { + 180.0 + } else { + east + }; + S1Interval::new( + west.to_radians().clamp(-PI, PI), + east.to_radians().clamp(-PI, PI), + ) + }; + let west = lng.lo(); + let south = south.to_radians().clamp(-FRAC_PI_2, FRAC_PI_2); + let north = north.to_radians().clamp(-FRAC_PI_2, FRAC_PI_2); + let rect = S2LatLngRect::new(R1Interval::new(south, north), lng); + let width = rect.lng().length(); + let span = if rect.lng().is_full() { + Span::Full + } else if width <= PI { + Span::Narrow + } else { + Span::Wide + }; + let direction = |lng: f64| [lng.cos(), lng.sin()]; + let east = lng.hi(); + let middle = west + width / 2.0; + let point = |lng: f64, lat: f64| [lat.cos() * lng.cos(), lat.cos() * lng.sin(), lat.sin()]; + let mut outside = Vec::new(); + if span == Span::Full { + if north.sin() < 1.0 { + outside.push([0.0, 0.0, 1.0]); + } + if south.sin() > -1.0 { + outside.push([0.0, 0.0, -1.0]); + } + } else { + outside.push(point(rect.lng().complement_center(), 0.0)); + } + Bounds { + rect, + // As a vertex's height is computed, so that a vertex at a bound is at its height. + south: south.sin(), + north: north.sin(), + span, + west: direction(west), + east: direction(east), + middle: direction(middle), + outside, + inside: point(middle, (south + north) / 2.0), + } + } + + /// Whether the rectangle contains the point. A pole has every longitude. + fn contains(&self, p: &[f64; 3]) -> bool { + if p[2] < self.south - SLACK || p[2] > self.north + SLACK { + return false; + } + // The sines of the angles from the west meridian to the point and from the point to the + // east meridian, scaled by the cosine of the point's latitude. + let east_of_west = self.west[0] * p[1] - self.west[1] * p[0] >= -SLACK; + let west_of_east = self.east[1] * p[0] - self.east[0] * p[1] >= -SLACK; + match self.span { + Span::Full => true, + // A point opposite a rectangle with no width is east of its west and west of its + // east too, but not on the rectangle's side of the sphere. + Span::Narrow => { + east_of_west + && west_of_east + && self.middle[0] * p[0] + self.middle[1] * p[1] >= -SLACK + } + Span::Wide => east_of_west || west_of_east, + } + } + + /// Whether the rectangle covers the geodesic edge between two points that it contains. The + /// edge may rise above the rectangle's north or dip below its south between them, and in a + /// rectangle wider than a hemisphere it may cut across the longitudes left out. + fn covers_edge(&self, a: &[f64; 3], b: &[f64; 3]) -> bool { + let normal = normalize(&robust_cross(a, b)); + let level = normal[0] * normal[0] + normal[1] * normal[1]; + // An edge along the equator neither rises nor dips. + if level > 0.0 { + // Where the edge's great circle is highest. It is lowest at the opposite point. + let top = [-normal[0] * normal[2], -normal[1] * normal[2], level]; + let after_a = dot(&cross(a, &top), &normal); + let before_b = dot(&cross(&top, b), &normal); + let height = level.sqrt(); + if after_a > 0.0 && before_b > 0.0 && height > self.north + SLACK { + return false; + } + if after_a < 0.0 && before_b < 0.0 && -height < self.south - SLACK { + return false; + } + } + // An edge along a meridian, or through a pole, has only the longitudes of its ends. + if self.span == Span::Wide && normal[2] != 0.0 { + let lng = |p: &[f64; 3]| p[1].atan2(p[0]); + return self + .rect + .lng() + .contains(&S1Interval::from_point_pair(lng(a), lng(b))); + } + true + } +} + +impl Area for Bounds { + fn contains(&self, point: &[f64; 3]) -> bool { + Bounds::contains(self, point) + } + + /// An edge between two points outside the rectangle meets it where it crosses its boundary. + fn meets_edge(&self, a: &[f64; 3], b: &[f64; 3]) -> bool { + self.rect.boundary_intersects(a, b) + } + + fn inside(&self) -> &[f64; 3] { + &self.inside + } +} + +/// What is known of a member that is not the only point of its document. +struct Member { + doc_id: u32, + head: u32, + sole: bool, + /// Whether the rectangle covers the member. A member found only in covering cells within + /// the rectangle is covered untested. + covered: bool, + /// Whether the member was found in a covering cell that the rectangle's boundary crosses, + /// and so tested. + tested: bool, +} + +/// Prepared rectangle query, built once and applied per segment. +pub struct RectQuery { + bounds: Bounds, + covering: Covering, + relation: RectRelation, +} + +impl RectQuery { + /// Build the query from the rectangle's west and south, and its east and north, in degrees + /// within [-180, 180] and [-90, 90]. The south must not be above the north. A west greater + /// than the east is a rectangle that crosses the antimeridian, and a west of -180 with an + /// east of 180 is every longitude. + pub fn new(west_south: [f64; 2], east_north: [f64; 2], relation: RectRelation) -> Self { + let bounds = Bounds::new(west_south, east_north); + let covering = Covering::new(&bounds.rect, |cell| bounds.rect.contains_cell(cell)); + RectQuery { + bounds, + covering, + relation, + } + } + + /// Add the documents of one segment that match to `docs`. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + match self.relation { + RectRelation::CoveredBy => self.search_covered(cell_reader, edge_cache, docs), + RectRelation::Intersects => { + search_meeting(&self.bounds, &self.covering, cell_reader, edge_cache, docs) + } + } + } + + /// A document is covered when every member of its set is. A member the covering does not + /// find lies outside the rectangle. + fn search_covered( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + let mut members: HashMap = HashMap::new(); + self.covering.visit(cell_reader, |position, within| { + let summary = edge_cache.summary((0, position)); + if let (true, Some(point)) = (summary.sole, &summary.point) { + if within || self.bounds.contains(point) { + docs.insert(summary.doc_id); + } + return; + } + let member = match members.entry(position) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(Member { + doc_id: summary.doc_id, + head: summary.head, + sole: summary.sole, + covered: true, + tested: false, + }), + }; + if !within && !member.tested { + member.tested = true; + member.covered = match &summary.point { + Some(point) => self.bounds.contains(point), + None => self.member_covered(position, cell_reader, edge_cache), + }; + } + }); + + // The members of each set with other members: those left to find, unless one is not + // covered. + let mut sets: HashMap)> = HashMap::new(); + for (position, member) in members { + if member.sole { + if member.covered { + docs.insert(member.doc_id); + } + continue; + } + let (_, left) = sets.entry(member.head).or_insert_with(|| { + let count = edge_cache.get((0, position)).member_count(); + (member.doc_id, Some(count)) + }); + *left = match *left { + Some(left) if member.covered => Some(left - 1), + _ => None, + }; + } + for (doc_id, left) in sets.into_values() { + if left == Some(0) { + docs.insert(doc_id); + } + } + } + + /// Whether the rectangle covers a line or polygon: its vertices and its edges between them, + /// and for a polygon its interior, which then holds no point outside the rectangle. + fn member_covered( + &self, + position: u32, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + ) -> bool { + let entry = edge_cache.get((0, position)); + let member = entry.edge_set(); + if member.vertices.is_empty() + || !member.vertices.iter().all(|v| self.bounds.contains(v)) + || !edges(member).all(|(a, b)| self.bounds.covers_edge(a, b)) + { + return false; + } + if !member.closed { + return true; + } + let mut segment = SegmentIndex { + cell_reader, + edge_cache, + }; + !self + .bounds + .outside + .iter() + .any(|point| contains_point(&mut segment, (0, position), point)) + } +} diff --git a/src/spatial/region.rs b/src/spatial/region.rs new file mode 100644 index 0000000000..5aebd8117a --- /dev/null +++ b/src/spatial/region.rs @@ -0,0 +1,47 @@ +//! An S2Region represents a two-dimensional region over the unit sphere. It is an abstract +//! interface with various concrete subtypes. +//! +//! The main purpose of this interface is to allow complex regions to be approximated as simpler +//! regions. So rather than having a wide variety of virtual methods that are implemented by all +//! subtypes, the interface is restricted to methods that are useful for computing approximations. +//! +//! Ported from s2/s2region.h (Eric Veach, Google). + +use super::latlng_rect::S2LatLngRect; +use super::s2cap::S2Cap; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; + +/// A two-dimensional region over the unit sphere. +pub trait Region { + /// Returns a bounding spherical cap that contains the region. The bound may not be tight. + fn get_cap_bound(&self) -> S2Cap; + + /// Returns a bounding latitude-longitude rectangle that contains the region. The bound may + /// not be tight. + fn get_rect_bound(&self) -> S2LatLngRect; + + /// Returns a small collection of S2CellIds whose union covers the region. The cells are not + /// sorted, may have redundancies (such as cells that contain other cells), and may cover much + /// more area than necessary. + /// + /// Implementations should attempt to return a small covering (ideally 4 cells or fewer) that + /// covers the region and can be computed quickly. The result is used by S2RegionCoverer as + /// a starting point for further refinement. + fn get_cell_union_bound(&self, cell_ids: &mut Vec) { + self.get_cap_bound().get_cell_union_bound(cell_ids); + } + + /// Returns true if the region completely contains the given cell, otherwise either the region + /// does not contain the cell or the containment relationship could not be determined. + fn contains_cell(&self, cell: &S2Cell) -> bool; + + /// If this method returns false, the region does not intersect the given cell. Otherwise, + /// either the region intersects the cell, or the intersection relationship could not be + /// determined. + fn may_intersect(&self, cell: &S2Cell) -> bool; + + /// Returns true if and only if the given point is contained by the region. The point 'p' is + /// generally required to be unit length, although some subtypes may relax this restriction. + fn contains_point(&self, p: &[f64; 3]) -> bool; +} diff --git a/src/spatial/region_coverer.rs b/src/spatial/region_coverer.rs new file mode 100644 index 0000000000..2c93766b33 --- /dev/null +++ b/src/spatial/region_coverer.rs @@ -0,0 +1,691 @@ +//! RegionCoverer: Approximates regions as cell unions. +//! +//! Produces coverings (cells that cover the region) using a priority-queue based refinement +//! algorithm. The output satisfies configurable constraints on cell count, levels, and level +//! spacing. + +use std::collections::BinaryHeap; + +use super::cell_union::CellUnion; +use super::region::Region; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; + +/// Options for controlling the covering algorithm. +#[derive(Clone, Debug)] +pub struct CovererOptions { + /// Maximum number of cells in the output. + /// + /// Note: Up to 6 cells may be returned if the region intersects all faces. Up to 3 cells may + /// be returned for tiny regions at cube face intersections. + /// + /// Default: 8 + max_cells: i32, + + /// Minimum cell level to use. + /// + /// Takes priority over max_cells: cells below this level are never used even if this causes + /// more cells to be returned. + /// + /// Default: 0 + min_level: i32, + + /// Maximum cell level to use. + /// + /// Default: 30 (leaf level) + max_level: i32, + + /// Level spacing constraint. + /// + /// Only cells where (level - min_level) % level_mod == 0 are used. + /// - level_mod = 1: branching factor 4 (default) + /// - level_mod = 2: branching factor 16 + /// - level_mod = 3: branching factor 64 + /// + /// Default: 1 + level_mod: i32, +} + +impl Default for CovererOptions { + fn default() -> Self { + Self { + max_cells: 8, + min_level: 0, + max_level: S2CellId::MAX_LEVEL, + level_mod: 1, + } + } +} + +impl CovererOptions { + /// Creates options with default values. + pub fn new() -> Self { + Self::default() + } + + /// Sets the maximum number of cells. + pub fn max_cells(mut self, max_cells: i32) -> Self { + self.max_cells = max_cells; + self + } + + /// Sets the minimum cell level. + pub fn min_level(mut self, min_level: i32) -> Self { + debug_assert!(min_level <= S2CellId::MAX_LEVEL); + self.min_level = min_level.min(S2CellId::MAX_LEVEL); + self + } + + /// Sets the maximum cell level. + pub fn max_level(mut self, max_level: i32) -> Self { + debug_assert!(max_level <= S2CellId::MAX_LEVEL); + self.max_level = max_level.min(S2CellId::MAX_LEVEL); + self + } + + /// Sets both min and max level to the same value. + pub fn fixed_level(mut self, level: i32) -> Self { + self.min_level = level.min(S2CellId::MAX_LEVEL); + self.max_level = level.min(S2CellId::MAX_LEVEL); + self + } + + /// Sets the level spacing constraint. + pub fn level_mod(mut self, level_mod: i32) -> Self { + debug_assert!((1..=3).contains(&level_mod)); + self.level_mod = level_mod.clamp(1, 3); + self + } + + /// Returns the maximum level that will actually be used. + /// + /// This is the largest level <= max_level where (level - min_level) % level_mod == 0. + pub fn true_max_level(&self) -> i32 { + if self.level_mod == 1 { + return self.max_level; + } + self.max_level - (self.max_level - self.min_level) % self.level_mod + } + + /// Returns the number of children per candidate (4^level_mod). + #[inline] + fn max_children(&self) -> usize { + 1 << self.max_children_shift() + } + + /// Returns log base 2 of max_children. + #[inline] + fn max_children_shift(&self) -> usize { + 2 * self.level_mod as usize + } +} + +/// A candidate cell being considered for the covering. +struct Candidate { + cell: S2Cell, + is_terminal: bool, + num_children: usize, + children: Vec, +} + +impl Candidate { + /// Creates a new candidate. + fn new(cell: S2Cell, is_terminal: bool, max_children: usize) -> Self { + Self { + cell, + is_terminal, + num_children: 0, + children: if is_terminal { + Vec::new() + } else { + Vec::with_capacity(max_children) + }, + } + } +} + +/// Priority queue entry. +#[derive(Eq, PartialEq)] +struct QueueEntry { + priority: i32, + candidate_idx: usize, +} + +/// Orders the max-heap of candidates, as S2's priority queue is: the highest priority first, and +/// of equal priorities, the candidate queued first. +impl Ord for QueueEntry { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.priority + .cmp(&other.priority) + .then_with(|| other.candidate_idx.cmp(&self.candidate_idx)) + } +} + +impl PartialOrd for QueueEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +/// Generates cell coverings for arbitrary regions. +/// +/// The covering algorithm uses a priority queue to iteratively refine cells. It starts with +/// initial candidates from the region's cell union bound, then repeatedly expands the largest +/// cells until max_cells is reached or no further refinement is possible. +pub struct RegionCoverer { + options: CovererOptions, +} + +impl RegionCoverer { + /// Creates a new RegionCoverer with the given options. + pub fn new(options: CovererOptions) -> Self { + debug_assert!(options.min_level <= options.max_level); + Self { options } + } + + /// Creates a new RegionCoverer with default options. + pub fn with_defaults() -> Self { + Self::new(CovererOptions::default()) + } + + /// Returns a reference to the options. + pub fn options(&self) -> &CovererOptions { + &self.options + } + + /// Returns a mutable reference to the options. + pub fn options_mut(&mut self) -> &mut CovererOptions { + &mut self.options + } + + /// Returns a CellUnion that covers the given region. + /// + /// The covering satisfies all constraints in options(). If min_level > 0 or level_mod > 1, the + /// result may not be normalized (siblings may exist). + pub fn get_covering(&self, region: &R) -> CellUnion { + let mut cell_ids = Vec::new(); + self.get_covering_cells(region, &mut cell_ids); + CellUnion::from_verbatim(cell_ids) + } + + /// Like get_covering but writes directly to a vector. + pub fn get_covering_cells(&self, region: &R, covering: &mut Vec) { + covering.clear(); + self.get_covering_internal(region, covering); + } + + /// Returns a fast, loose covering of the region. + /// + /// Much faster than get_covering but less tight. Uses the region's cell union bound and + /// canonicalizes it. + pub fn get_fast_covering(&self, region: &R, covering: &mut Vec) { + covering.clear(); + region.get_cell_union_bound(covering); + self.canonicalize_covering(covering); + } + + /// Internal covering algorithm. + fn get_covering_internal(&self, region: &R, result: &mut Vec) { + debug_assert!(self.options.min_level <= self.options.max_level); + + let mut pq: BinaryHeap = BinaryHeap::new(); + let mut candidates: Vec> = Vec::new(); + + // Get initial candidates from the region's cell union bound. + let mut initial_cells = Vec::new(); + self.get_initial_candidates(region, &mut initial_cells); + + for cell_id in initial_cells { + if let Some(candidate) = self.new_candidate(region, &S2Cell::new(cell_id)) { + self.add_candidate(region, candidate, result, &mut pq, &mut candidates); + } + } + + while let Some(entry) = pq.pop() { + let Some(mut candidate) = candidates[entry.candidate_idx].take() else { + continue; + }; + + // For exterior coverings, the decision logic is: + // - If below min_level: must expand + // - If only 1 child: no harm in expanding + // - If adding children wouldn't exceed max_cells: expand + // - Otherwise: mark as terminal and add + if candidate.cell.level() < self.options.min_level + || candidate.num_children == 1 + || (result.len() + pq.len() + candidate.num_children + <= self.options.max_cells as usize) + { + // Expand this candidate into its children. + for child in candidate.children.drain(..) { + self.add_candidate(region, child, result, &mut pq, &mut candidates); + } + } else { + // Mark as terminal and process. + candidate.is_terminal = true; + self.add_candidate(region, candidate, result, &mut pq, &mut candidates); + } + } + + // Rather than just returning the raw list of cell ids, we construct a cell union and then + // denormalize it. This has the effect of replacing four child cells with their parent + // whenever this does not violate the covering parameters. + CellUnion::normalize_vec(result); + if self.options.min_level > 0 || self.options.level_mod > 1 { + let normalized = std::mem::take(result); + let cu = CellUnion::from_verbatim(normalized); + *result = cu.denormalize(self.options.min_level, self.options.level_mod); + } + debug_assert!(self.is_canonical(result)); + } + + /// Computes a set of initial candidates that cover the given region. + fn get_initial_candidates(&self, region: &R, cells: &mut Vec) { + // Optimization: start with a small (usually 4 cell) covering of the + // region's bounding cap. + let tmp_options = CovererOptions { + max_cells: self.options.max_cells.min(4), + min_level: 0, + max_level: self.options.max_level, + level_mod: 1, + }; + let tmp_coverer = RegionCoverer::new(tmp_options); + tmp_coverer.get_fast_covering(region, cells); + + // Adjust cell levels for level_mod constraint. + self.adjust_cell_levels(cells); + } + + /// Processes a candidate by either adding it to the result_ vector or expanding its children + /// and inserting it into the priority queue. Passing an argument of nullptr does nothing. + fn add_candidate( + &self, + region: &R, + mut candidate: Candidate, + result: &mut Vec, + pq: &mut BinaryHeap, + candidates: &mut Vec>, + ) { + // If terminal, add to result immediately. + if candidate.is_terminal { + result.push(candidate.cell.id()); + return; + } + debug_assert_eq!(0, candidate.num_children); + + // Expand one level at a time until we hit min_level() to ensure that we don't skip over + // it. + let num_levels = if candidate.cell.level() < self.options.min_level { + 1 + } else { + self.options.level_mod as usize + }; + let num_terminals = self.expand_children(region, &mut candidate, num_levels); + + if candidate.num_children == 0 { + // No children intersect the region. + return; + } + + // Optimization: add the parent cell rather than all of its children. + if num_terminals == self.options.max_children() + && candidate.cell.level() >= self.options.min_level + { + candidate.is_terminal = true; + self.add_candidate(region, candidate, result, pq, candidates); + return; + } + + // Add candidate to priority queue. + self.add_candidate_to_queue(candidate, num_terminals, pq, candidates); + } + + /// Adds a candidate to the priority queue. + fn add_candidate_to_queue( + &self, + candidate: Candidate, + num_terminals: usize, + pq: &mut BinaryHeap, + candidates: &mut Vec>, + ) { + // We negate the priority so that smaller absolute priorities are returned first. The + // heuristic is designed to refine the largest cells first, since those are where we have + // the largest potential gain. + let priority = self.compute_priority(&candidate, num_terminals); + let candidate_idx = candidates.len(); + candidates.push(Some(candidate)); + pq.push(QueueEntry { + priority, + candidate_idx, + }); + } + + /// Creates a new candidate for the given cell, or None if it doesn't intersect. + /// + /// A candidate is marked terminal if: + /// - The region fully contains the cell, OR + /// - The cell is at max_level (cannot subdivide further) + fn new_candidate(&self, region: &R, cell: &S2Cell) -> Option { + if !region.may_intersect(cell) { + return None; + } + + let mut is_terminal = false; + let level = cell.level(); + + if level >= self.options.min_level { + // For exterior coverings (not interior): + if level + self.options.level_mod > self.options.max_level || region.contains_cell(cell) + { + is_terminal = true; + } + } + + let max_children = if is_terminal { + 0 + } else { + self.options.max_children() + }; + + Some(Candidate::new(*cell, is_terminal, max_children)) + } + + /// Expands children of a candidate by the given number of levels. + /// + /// Returns the number of terminal children. + fn expand_children( + &self, + region: &R, + candidate: &mut Candidate, + num_levels: usize, + ) -> usize { + let cell = candidate.cell; + self.expand_children_recursive(region, candidate, &cell, num_levels) + } + + /// Recursive helper for expand_children. + fn expand_children_recursive( + &self, + region: &R, + candidate: &mut Candidate, + cell: &S2Cell, + num_levels: usize, + ) -> usize { + let num_levels = num_levels - 1; + + let Some(children) = cell.subdivide() else { + return 0; + }; + + let mut num_terminals = 0; + + for child_cell in children { + if num_levels > 0 { + if region.may_intersect(&child_cell) { + num_terminals += + self.expand_children_recursive(region, candidate, &child_cell, num_levels); + } + continue; + } + + if let Some(child) = self.new_candidate(region, &child_cell) { + if child.is_terminal { + num_terminals += 1; + } + candidate.children.push(child); + candidate.num_children += 1; + } + } + + num_terminals + } + + // We negate the priority so that smaller absolute priorities are returned first. The + // heuristic is designed to refine the largest cells first, since those are where we have the + // largest potential gain. Among cells of the same size, we prefer the cells with the fewest + // children. Finally, among cells with equal numbers of children we prefer those with the + // smallest number of children that cannot be refined further. + fn compute_priority(&self, candidate: &Candidate, num_terminals: usize) -> i32 { + let shift = self.options.max_children_shift() as i32; + -(((((candidate.cell.level()) << shift) + candidate.num_children as i32) << shift) + + num_terminals as i32) + } + + /// If level > min_level(), then reduces "level" if necessary so that it also satisfies + /// level_mod(). Levels smaller than min_level() are not affected (since cells at these levels + /// are eventually expanded). + #[inline] + fn adjust_level(&self, level: i32) -> i32 { + if self.options.level_mod > 1 && level > self.options.min_level { + level - (level - self.options.min_level) % self.options.level_mod + } else { + level + } + } + + /// Ensures that all cells with level > min_level() also satisfy level_mod(), by replacing them + /// with an ancestor if necessary. Cell levels smaller than min_level() are not modified (see + /// AdjustLevel). The output is then normalized to ensure that no redundant cells are present. + fn adjust_cell_levels(&self, cells: &mut Vec) { + debug_assert!(cells.windows(2).all(|w| w[0] <= w[1])); // is_sorted + + if self.options.level_mod == 1 { + return; + } + + let mut out = 0; + for i in 0..cells.len() { + let mut id = cells[i]; + let level = id.level(); + let new_level = self.adjust_level(level); + + if new_level != level { + id = id.parent(new_level); + } + + // Skip if contained by previous output. + if out > 0 && cells[out - 1].contains(id) { + continue; + } + + // Remove any previous cells contained by this one. + while out > 0 && id.contains(cells[out - 1]) { + out -= 1; + } + + cells[out] = id; + out += 1; + } + + cells.truncate(out); + } + + /// Canonicalizes a covering to satisfy all constraints. + /// + /// This ensures cells are sorted, non-overlapping, and satisfy min_level, max_level, and + /// level_mod. + pub fn canonicalize_covering(&self, covering: &mut Vec) { + debug_assert!(self.options.min_level <= self.options.max_level); + + // If any cells are too small, or don't satisfy level_mod(), then replace them with + // ancestors. + if self.options.max_level < S2CellId::MAX_LEVEL || self.options.level_mod > 1 { + for id in covering.iter_mut() { + let level = id.level(); + let new_level = self.adjust_level(level.min(self.options.max_level)); + if new_level != level { + *id = id.parent(new_level); + } + } + } + + // Sort the cells and simplify them. + CellUnion::normalize_vec(covering); + + // Make sure that the covering satisfies min_level() and level_mod(), possibly at the + // expense of satisfying max_cells(). + if self.options.min_level > 0 || self.options.level_mod > 1 { + let normalized = std::mem::take(covering); + let cu = CellUnion::from_verbatim(normalized); + *covering = cu.denormalize(self.options.min_level, self.options.level_mod); + } + + // If there are too many cells and the covering is very large, use the S2RegionCoverer to + // compute a new covering. (This avoids possible O(n^2) behavior of the simpler algorithm + // below.) + let excess = covering.len() as i64 - self.options.max_cells as i64; + if excess <= 0 || self.is_canonical(covering) { + return; + } + + if excess * covering.len() as i64 > 10000 { + // For very large coverings, recompute using the full algorithm. + let cell_union = CellUnion::from_verbatim(std::mem::take(covering)); + self.get_covering_cells(&cell_union, covering); + } else { + // Repeatedly replace two adjacent cells by their lowest common ancestor. + self.reduce_to_max_cells_with_parent_check(covering); + } + debug_assert!(self.is_canonical(covering)); + } + + /// Reduces the covering to at most max_cells by merging adjacent cells. + /// + /// After merging, checks if all children of the parent are present and recursively replaces + /// them with the parent. + fn reduce_to_max_cells_with_parent_check(&self, covering: &mut Vec) { + while covering.len() > self.options.max_cells as usize { + // Find the pair of adjacent cells with the highest common ancestor + // at or above min_level. + let mut best_index: i32 = -1; + let mut best_level: i32 = -1; + + for i in 0..covering.len() - 1 { + let level = covering[i].common_ancestor_level(covering[i + 1]) as i32; + let level = self.adjust_level(level); + if level > best_level { + best_level = level; + best_index = i as i32; + } + } + + if best_level < self.options.min_level { + break; + } + + // Replace all cells contained by the new ancestor cell. + let mut id = covering[best_index as usize].parent(best_level); + self.replace_cells_with_ancestor(covering, id); + + // Now repeatedly check whether all children of the parent cell are present, in which + // case we can replace those cells with their parent. + while best_level > self.options.min_level { + best_level -= self.options.level_mod; + id = id.parent(best_level); + if !self.contains_all_children(covering, id) { + break; + } + self.replace_cells_with_ancestor(covering, id); + } + } + } + + /// Returns true if "covering" contains all children of "id" at level + /// (id.level() + options_.level_mod()). + fn contains_all_children(&self, covering: &[S2CellId], id: S2CellId) -> bool { + let level = id.level() + self.options.level_mod; + let mut it = covering.partition_point(|&c| c < id.range_min()); + let mut child = id.child_begin_at_level(level); + let end = id.child_end_at_level(level); + + while child != end { + if it >= covering.len() || covering[it] != child { + return false; + } + it += 1; + child = child.next(); + } + true + } + + /// Replaces all descendants of an ancestor in the covering with the ancestor. + fn replace_cells_with_ancestor(&self, covering: &mut Vec, id: S2CellId) { + let begin = covering.partition_point(|&c| c < id.range_min()); + let end = covering.partition_point(|&c| c <= id.range_max()); + + debug_assert!(begin != end); + covering[begin] = id; + if begin + 1 < end { + covering.drain(begin + 1..end); + } + } + + /// Returns true if the covering is canonical. + pub fn is_canonical(&self, covering: &[S2CellId]) -> bool { + debug_assert!(self.options.min_level <= self.options.max_level); + + let min_level = self.options.min_level; + let max_level = self.options.true_max_level(); + let level_mod = self.options.level_mod; + let too_many_cells = covering.len() > self.options.max_cells as usize; + let mut same_parent_count = 1; + let mut prev_id: Option = None; + + for &id in covering { + if !id.is_valid() { + return false; + } + + // Check that the S2CellId level is acceptable. + let level = id.level(); + if level < min_level || level > max_level { + return false; + } + if level_mod > 1 && (level - min_level) % level_mod != 0 { + return false; + } + + if let Some(prev) = prev_id { + // Check that cells are sorted and non-overlapping. + if prev.range_max() >= id.range_min() { + return false; + } + + // If there are too many cells, check that no pair of adjacent cells could be + // replaced by an ancestor. + if too_many_cells && id.common_ancestor_level(prev) as i32 >= min_level { + return false; + } + + // Check that there are no sequences of (4 ** level_mod) cells that all have the + // same parent (considering only multiples of "level_mod"). + let plevel = level - level_mod; + if plevel < min_level + || level != prev.level() + || id.parent(plevel) != prev.parent(plevel) + { + same_parent_count = 1; + } else { + same_parent_count += 1; + if same_parent_count == self.options.max_children() { + return false; + } + } + } + + prev_id = Some(id); + } + + true + } +} + +impl Default for RegionCoverer { + fn default() -> Self { + Self::with_defaults() + } +} + +#[cfg(test)] +#[path = "tests/region_coverer_tests.rs"] +mod tests; diff --git a/src/spatial/s1chord_angle.rs b/src/spatial/s1chord_angle.rs new file mode 100644 index 0000000000..c6432d340d --- /dev/null +++ b/src/spatial/s1chord_angle.rs @@ -0,0 +1,380 @@ +//! S1ChordAngle represents an angle as the squared chord length between two points on the unit +//! sphere. Chord angles are efficient for computing and comparing distances without trigonometry. +//! +//! Ported from s2/s1chord_angle.h (Eric Veach, Google). + +use std::f64::consts::PI; + +/// Maximum relative error when summing two S1ChordAngles together. Absolute error is length2() of +/// the summed S1ChordAngle times this value. See the definition of operator+() for the error +/// analysis. +pub const RELATIVE_SUM_ERROR: f64 = 2.02 * f64::EPSILON; + +const MAX_LENGTH2: f64 = 4.0; + +/// S1ChordAngle represents the angle subtended by a chord (i.e., the straight line segment +/// connecting two points on the unit sphere). Its representation makes it very efficient for +/// computing and comparing distances, but unlike S1Angle it is only capable of representing angles +/// between 0 and Pi radians. +#[derive(Clone, Copy, Debug)] +pub struct S1ChordAngle { + length2: f64, +} + +impl S1ChordAngle { + /// Return the zero chord angle. + pub const fn zero() -> Self { + Self { length2: 0.0 } + } + + /// Return a chord angle of 90 degrees (a "right angle"). + pub const fn right() -> Self { + Self { length2: 2.0 } + } + + /// Return a chord angle of 180 degrees (a "straight angle"). This is the maximum finite + /// chord angle. + pub const fn straight() -> Self { + Self { length2: 4.0 } + } + + /// Return a chord angle larger than any finite chord angle. The only valid operations on + /// Infinity() are comparisons, conversions, and Successor() / Predecessor(). + pub const fn infinity() -> Self { + Self { + length2: f64::INFINITY, + } + } + + /// Return a chord angle smaller than Zero(). The only valid operations on Negative() are + /// comparisons, conversions, and Successor() / Predecessor(). + pub const fn negative() -> Self { + Self { length2: -1.0 } + } + + /// Construct the S1ChordAngle corresponding to the distance between the two given points. The + /// points must be unit length. + pub fn from_points(x: &[f64; 3], y: &[f64; 3]) -> Self { + let dx = x[0] - y[0]; + let dy = x[1] - y[1]; + let dz = x[2] - y[2]; + let length2 = dx * dx + dy * dy + dz * dz; + // The squared distance may slightly exceed 4.0 due to roundoff errors. + Self { + length2: length2.min(MAX_LENGTH2), + } + } + + /// Construct an S1ChordAngle from the squared chord length. Note that the argument is + /// automatically clamped to a maximum of 4.0 to handle possible roundoff errors. The argument + /// must be non-negative. + pub fn from_length2(length2: f64) -> Self { + debug_assert!(length2 >= 0.0); + Self { + length2: length2.min(MAX_LENGTH2), + } + } + + /// Conversion from an angle in radians. Angles outside the range [0, Pi] are handled as + /// follows: Infinity is mapped to Infinity(), negative angles are mapped to Negative(), and + /// finite angles larger than Pi are mapped to Straight(). + pub fn from_radians(radians: f64) -> Self { + if radians < 0.0 { + return Self::negative(); + } + if !radians.is_finite() { + return Self::infinity(); + } + // The chord length is 2 * sin(angle / 2). + let length = 2.0 * (0.5 * radians.min(PI)).sin(); + Self { + length2: length * length, + } + } + + /// Convenience method implemented by converting from radians. + pub fn from_degrees(degrees: f64) -> Self { + Self::from_radians(degrees.to_radians()) + } + + /// This method uses the distance along the surface of the sphere as an upper bound on the + /// distance through the sphere's interior. Accurate to within 1% for distances up to about + /// 3100km on the Earth's surface. + pub fn fast_upper_bound_from_radians(radians: f64) -> Self { + Self::from_length2(radians * radians) + } + + /// The squared length of the chord. (Most clients will not need this.) + #[inline] + pub fn length2(&self) -> f64 { + self.length2 + } + + /// Converts to radians. Note that the conversion uses trigonometric functions and therefore + /// should be avoided in inner loops. + pub fn to_radians(&self) -> f64 { + if self.is_negative() { + return -1.0; + } + if self.is_infinity() { + return f64::INFINITY; + } + 2.0 * (0.5 * self.length2.sqrt()).asin() + } + + /// Convenience method implemented by calling to_radians() first. + pub fn to_degrees(&self) -> f64 { + self.to_radians().to_degrees() + } + + /// Returns true if this is the zero angle. + #[inline] + pub fn is_zero(&self) -> bool { + self.length2 == 0.0 + } + + /// Returns true if this is negative. + #[inline] + pub fn is_negative(&self) -> bool { + self.length2 < 0.0 + } + + /// Returns true if this is infinite. + #[inline] + pub fn is_infinity(&self) -> bool { + self.length2 == f64::INFINITY + } + + /// Negative or infinity. + #[inline] + pub fn is_special(&self) -> bool { + self.is_negative() || self.is_infinity() + } + + /// Return true if the internal representation is valid. Negative() and Infinity() are both + /// considered valid. + #[inline] + pub fn is_valid(&self) -> bool { + (self.length2 >= 0.0 && self.length2 <= MAX_LENGTH2) || self.is_special() + } + + /// Returns sin(a)**2, but computed more efficiently. + pub fn sin2(&self) -> f64 { + debug_assert!(!self.is_special()); + // Let "a" be the (non-squared) chord length, and let A be the + // corresponding half-angle (a = 2*sin(A)). The formula below can be + // derived from: + // sin(2*A) = 2 * sin(A) * cos(A) + // cos^2(A) = 1 - sin^2(A) + self.length2 * (1.0 - 0.25 * self.length2) + } + + /// Returns sin of the angle. + pub fn sin(&self) -> f64 { + self.sin2().sqrt() + } + + /// Returns cos of the angle. + pub fn cos(&self) -> f64 { + debug_assert!(!self.is_special()); + // cos(2*A) = cos^2(A) - sin^2(A) = 1 - 2*sin^2(A) + 1.0 - 0.5 * self.length2 + } + + /// Returns tan of the angle. + pub fn tan(&self) -> f64 { + self.sin() / self.cos() + } + + /// Returns the smallest representable S1ChordAngle larger than this object. This can be used + /// to convert a "<" comparison to a "<=" comparison. + /// + /// Note the following special cases: + /// Negative().Successor() == Zero() + /// Straight().Successor() == Infinity() + /// Infinity().Successor() == Infinity() + pub fn successor(&self) -> Self { + if self.length2 >= MAX_LENGTH2 { + return Self::infinity(); + } + if self.length2 < 0.0 { + return Self::zero(); + } + Self { + length2: next_after(self.length2, 10.0), + } + } + + /// Like Successor(), but returns the largest representable S1ChordAngle less than this object. + /// + /// Note the following special cases: + /// Infinity().Predecessor() == Straight() + /// Zero().Predecessor() == Negative() + /// Negative().Predecessor() == Negative() + pub fn predecessor(&self) -> Self { + if self.length2 <= 0.0 { + return Self::negative(); + } + if self.length2 > MAX_LENGTH2 { + return Self::straight(); + } + Self { + length2: next_after(self.length2, -10.0), + } + } + + /// Returns a new S1ChordAngle that has been adjusted by the given error bound (which can be + /// positive or negative). "error" should be the value returned by one of the error bound + /// methods below. If the angle is Negative() or Infinity(), don't change it. Otherwise clamp + /// it to the valid range. + pub fn plus_error(&self, error: f64) -> Self { + if self.is_special() { + return *self; + } + Self { + length2: (self.length2 + error).clamp(0.0, MAX_LENGTH2), + } + } + + /// Return the maximum error in length2() for the S1ChordAngle(x, y) constructor, assuming that + /// "x" and "y" are normalized to within the bounds guaranteed by S2Point::Normalize(). (The + /// error is defined with respect to the true distance after the points are projected to lie + /// exactly on the sphere.) + pub fn get_s2point_constructor_max_error(&self) -> f64 { + // There is a relative error of 2.5 * DBL_EPSILON when computing the + // squared distance, plus a relative error of 2 * DBL_EPSILON and an + // absolute error of (16 * DBL_EPSILON**2) because the lengths of the + // input points may differ from 1 by up to (2 * DBL_EPSILON) each. + 4.5 * f64::EPSILON * self.length2 + 16.0 * f64::EPSILON * f64::EPSILON + } +} + +impl PartialEq for S1ChordAngle { + fn eq(&self, other: &Self) -> bool { + self.length2 == other.length2 + } +} + +impl Eq for S1ChordAngle {} + +impl PartialOrd for S1ChordAngle { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for S1ChordAngle { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.length2.partial_cmp(&other.length2).unwrap() + } +} + +impl std::ops::Add for S1ChordAngle { + type Output = Self; + + // Note that this method is much more efficient than converting the chord angles to S1Angles + // and adding those. It requires only one square root plus a few additions and + // multiplications. + fn add(self, other: Self) -> Self { + debug_assert!(!self.is_special()); + debug_assert!(!other.is_special()); + + let a2 = self.length2; + let b2 = other.length2; + + // Optimization for the common case where "b" is an error tolerance parameter that happens + // to be set to zero. + if b2 == 0.0 { + return self; + } + + // Clamp the angle sum to at most 180 degrees. + if a2 + b2 >= MAX_LENGTH2 { + return Self::straight(); + } + + // Let "a" and "b" be the (non-squared) chord lengths, and let c = a+b. + // Let A, B, and C be the corresponding half-angles (a = 2*sin(A), etc). + // Then the formula below can be derived from c = 2 * sin(A+B) and the + // relationships sin(A+B) = sin(A)*cos(B) + sin(B)*cos(A) + // cos(X) = sqrt(1 - sin^2(X)) . + let x = a2 * (1.0 - 0.25 * b2); // is_valid() => non-negative + let y = b2 * (1.0 - 0.25 * a2); // is_valid() => non-negative + Self { + length2: (x + y + 2.0 * (x * y).sqrt()).min(MAX_LENGTH2), + } + } +} + +impl std::ops::Sub for S1ChordAngle { + type Output = Self; + + // See comments in operator+(). + fn sub(self, other: Self) -> Self { + debug_assert!(!self.is_special()); + debug_assert!(!other.is_special()); + + let a2 = self.length2; + let b2 = other.length2; + + if b2 == 0.0 { + return self; + } + if a2 <= b2 { + return Self::zero(); + } + + let x = a2 * (1.0 - 0.25 * b2); + let y = b2 * (1.0 - 0.25 * a2); + + // The calculation below is formulated differently (with two square roots rather than one) + // to avoid excessive cancellation error when two nearly equal values are subtracted. + let c = (x.sqrt() - y.sqrt()).max(0.0); + Self { length2: c * c } + } +} + +impl std::ops::AddAssign for S1ChordAngle { + fn add_assign(&mut self, other: Self) { + *self = *self + other; + } +} + +impl std::ops::SubAssign for S1ChordAngle { + fn sub_assign(&mut self, other: Self) { + *self = *self - other; + } +} + +impl std::fmt::Display for S1ChordAngle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}deg", self.to_degrees()) + } +} + +fn next_after(x: f64, y: f64) -> f64 { + if x == y { + return y; + } + if x.is_nan() || y.is_nan() { + return f64::NAN; + } + if x == 0.0 { + return if y > 0.0 { + f64::from_bits(1) + } else { + f64::from_bits(1u64 << 63 | 1) + }; + } + let bits = x.to_bits(); + let next_bits = if (x < y) == (x > 0.0) { + bits + 1 + } else { + bits - 1 + }; + f64::from_bits(next_bits) +} + +#[cfg(test)] +#[path = "tests/s1chord_angle_tests.rs"] +mod tests; diff --git a/src/spatial/s1interval.rs b/src/spatial/s1interval.rs new file mode 100644 index 0000000000..f5428a193f --- /dev/null +++ b/src/spatial/s1interval.rs @@ -0,0 +1,476 @@ +//! An S1Interval represents a closed interval on a unit circle (also known as a 1-dimensional +//! sphere). It is capable of representing the empty interval (containing no points), the full +//! interval (containing all points), and zero-length intervals (containing a single point). +//! +//! Points are represented by the angle they make with the positive x-axis in! the range [-Pi, Pi]. +//! An interval is represented by its lower and upper bounds (both inclusive, since the interval is +//! closed). The lower bound may be greater than the upper bound, in which case the interval is +//! "inverted" (i.e. it passes through the point (-1, 0)). +//! +//! Note that the point (-1, 0) has two valid representations, Pi and -Pi. The normalized +//! representation of this point internally is Pi, so that endpoints of normal intervals are in the +//! range (-Pi, Pi]. However, we take advantage of the point -Pi to construct two special +//! intervals: the Full() interval is [-Pi, Pi], and the Empty() interval is [Pi, -Pi]. +//! +//! S1Interval in s1interval.h and s1interval.cc + +use std::f64::consts::PI; + +use super::math::remainder; + +/// A closed interval on the unit circle. +/// +/// Points are represented by angles in the range [-PI, PI]. The interval is represented by its +/// lower and upper bounds (both inclusive). The lower bound may be greater than the upper bound, +/// in which case the interval is "inverted" (it passes through the point at PI/-PI). +/// +/// Special intervals: +/// - Empty: `[PI, -PI]` contains no points +/// - Full: `[-PI, PI]` contains all points +#[derive(Clone, Copy, Debug)] +pub struct S1Interval { + lo: f64, + hi: f64, +} + +impl S1Interval { + /// Creates a new interval with the given bounds. + /// + /// Both endpoints must be in the range [-PI, PI]. The value -PI is normalized to PI except for + /// the Full and Empty intervals. + pub fn new(lo: f64, hi: f64) -> Self { + let mut result = Self { lo, hi }; + if lo == -PI && hi != PI { + result.lo = PI; + } + if hi == -PI && lo != PI { + result.hi = PI; + } + debug_assert!(result.is_valid()); + result + } + + /// Internal constructor that assumes bounds are already normalized. + fn new_checked(lo: f64, hi: f64) -> Self { + let result = Self { lo, hi }; + debug_assert!(result.is_valid()); + result + } + + /// Returns the empty interval (contains no points). + pub fn empty() -> Self { + Self { lo: PI, hi: -PI } + } + + /// Returns the full interval (contains all points). + pub fn full() -> Self { + Self { lo: -PI, hi: PI } + } + + /// Creates an interval containing a single point. + pub fn from_point(mut p: f64) -> Self { + if p == -PI { + p = PI; + } + Self::new_checked(p, p) + } + + /// Creates the minimal interval containing two points. This is equivalent to starting with an + /// empty interval and calling `add_point()` twice, but more efficient. + pub fn from_point_pair(mut p1: f64, mut p2: f64) -> Self { + debug_assert!(p1.abs() <= PI); + debug_assert!(p2.abs() <= PI); + if p1 == -PI { + p1 = PI; + } + if p2 == -PI { + p2 = PI; + } + if positive_distance(p1, p2) <= PI { + Self::new_checked(p1, p2) + } else { + Self::new_checked(p2, p1) + } + } + + /// The low bound of the interval. + pub fn lo(&self) -> f64 { + self.lo + } + + /// The high bound of the interval. + pub fn hi(&self) -> f64 { + self.hi + } + + /// Returns true if the interval is valid. An interval is valid if neither bound exceeds PI in + /// absolute value, and the value -PI appears only in the Empty and Full intervals. + pub fn is_valid(&self) -> bool { + self.lo.abs() <= PI + && self.hi.abs() <= PI + && !(self.lo == -PI && self.hi != PI) + && !(self.hi == -PI && self.lo != PI) + } + + /// Returns true if the interval contains all points on the circle. + pub fn is_full(&self) -> bool { + self.lo == -PI && self.hi == PI + } + + /// Returns true if the interval is empty. + pub fn is_empty(&self) -> bool { + self.lo == PI && self.hi == -PI + } + + /// Returns true if lo > hi (the interval wraps around PI/-PI). + pub fn is_inverted(&self) -> bool { + self.lo > self.hi + } + + /// Returns the midpoint of the interval. For full and empty intervals, the result is + /// arbitrary. + pub fn center(&self) -> f64 { + let center = 0.5 * (self.lo + self.hi); + if !self.is_inverted() { + center + } else if center <= 0.0 { + center + PI + } else { + center - PI + } + } + + /// Returns the length of the interval. The length of an empty interval is negative. + pub fn length(&self) -> f64 { + let mut length = self.hi - self.lo; + if length >= 0.0 { + return length; + } + length += 2.0 * PI; + // Empty intervals have a negative length. + if length > 0.0 { + length + } else { + -1.0 + } + } + + /// Returns the complement of the interior of the interval. An interval and its complement have + /// the same boundary but do not share any interior values. + pub fn complement(&self) -> Self { + if self.lo == self.hi { + // Singleton. + Self::full() + } else { + // Handles empty and full correctly. + Self::new_checked(self.hi, self.lo) + } + } + + /// Returns the midpoint of the complement of the interval. For full and empty intervals, the + /// result is arbitrary. For a singleton interval, returns its antipodal point on S1. + pub fn complement_center(&self) -> f64 { + if self.lo != self.hi { + self.complement().center() + } else { + // Singleton. + if self.hi <= 0.0 { + self.hi + PI + } else { + self.hi - PI + } + } + } + + /// Returns true if the interval contains the given point. Return true if the interval (which + /// is closed) contains the point 'p'. Skips the normalization of 'p' from -Pi to Pi. + fn fast_contains(&self, p: f64) -> bool { + if self.is_inverted() { + (p >= self.lo || p <= self.hi) && !self.is_empty() + } else { + p >= self.lo && p <= self.hi + } + } + + /// Returns true if the interval contains the given point. The point must be in the range + /// [-PI, PI]. + pub fn contains_point(&self, mut p: f64) -> bool { + debug_assert!(p.abs() <= PI); + if p == -PI { + p = PI; + } + self.fast_contains(p) + } + + /// Returns true if the interior of the interval contains the point. + pub fn interior_contains_point(&self, mut p: f64) -> bool { + debug_assert!(p.abs() <= PI); + if p == -PI { + p = PI; + } + if self.is_inverted() { + p > self.lo || p < self.hi + } else { + (p > self.lo && p < self.hi) || self.is_full() + } + } + + /// Returns true if this interval contains the given interval. Works for empty, full, and + /// singleton intervals. + pub fn contains(&self, y: &S1Interval) -> bool { + if self.is_inverted() { + if y.is_inverted() { + y.lo >= self.lo && y.hi <= self.hi + } else { + (y.lo >= self.lo || y.hi <= self.hi) && !self.is_empty() + } + } else if y.is_inverted() { + self.is_full() || y.is_empty() + } else { + y.lo >= self.lo && y.hi <= self.hi + } + } + + /// Returns true if the interior of this interval contains the entire given interval. Note + /// that x.InteriorContains(x) is true only when x is the empty or full interval, and + /// x.InteriorContains(S1Interval(p,p)) is equivalent to x.InteriorContains(p). + pub fn interior_contains(&self, y: &S1Interval) -> bool { + if self.is_inverted() { + if !y.is_inverted() { + y.lo > self.lo || y.hi < self.hi + } else { + (y.lo > self.lo && y.hi < self.hi) || y.is_empty() + } + } else if y.is_inverted() { + self.is_full() || y.is_empty() + } else { + (y.lo > self.lo && y.hi < self.hi) || self.is_full() + } + } + + /// Return true if the two intervals contain any points in common. Note that the point +/-Pi + /// has two representations, so the intervals [-Pi,-3] and [2,Pi] intersect, for example. + pub fn intersects(&self, y: &S1Interval) -> bool { + if self.is_empty() || y.is_empty() { + return false; + } + if self.is_inverted() { + // Every non-empty inverted interval contains PI. + y.is_inverted() || y.lo <= self.hi || y.hi >= self.lo + } else if y.is_inverted() { + y.lo <= self.hi || y.hi >= self.lo + } else { + y.lo <= self.hi && y.hi >= self.lo + } + } + + /// Returns true if the interior of this interval intersects the given interval. Works for + /// empty, full, and singleton intervals. + pub fn interior_intersects(&self, y: &S1Interval) -> bool { + if self.is_empty() || y.is_empty() || self.lo == self.hi { + return false; + } + if self.is_inverted() { + y.is_inverted() || y.lo < self.hi || y.hi > self.lo + } else if y.is_inverted() { + y.lo < self.hi || y.hi > self.lo + } else { + (y.lo < self.hi && y.hi > self.lo) || self.is_full() + } + } + + /// Expands the interval to contain the given point. The point must be in the range [-PI, PI]. + pub fn add_point(&mut self, mut p: f64) { + debug_assert!(p.abs() <= PI); + if p == -PI { + p = PI; + } + + if self.fast_contains(p) { + return; + } + if self.is_empty() { + self.lo = p; + self.hi = p; + } else { + // Compute distance from p to each endpoint. + let dlo = positive_distance(p, self.lo); + let dhi = positive_distance(self.hi, p); + if dlo < dhi { + self.lo = p; + } else { + self.hi = p; + } + // Adding a point can never turn a non-full interval into a full one. + } + } + + /// Returns the closest point in the interval to the given point. The interval must be + /// non-empty. + pub fn project(&self, mut p: f64) -> f64 { + debug_assert!(!self.is_empty()); + debug_assert!(p.abs() <= PI); + if p == -PI { + p = PI; + } + if self.fast_contains(p) { + return p; + } + // Return the endpoint that is closer. + let dlo = positive_distance(p, self.lo); + let dhi = positive_distance(self.hi, p); + if dlo < dhi { + self.lo + } else { + self.hi + } + } + + /// Returns an interval expanded by the given margin on each side. If margin is negative, + /// shrinks instead. The result may be empty or full. Expansion of a full interval remains + /// full; expansion of an empty interval remains empty. + pub fn expanded(&self, margin: f64) -> S1Interval { + if margin >= 0.0 { + if self.is_empty() { + return *self; + } + // Check if this interval will be full after expansion. Allow for 1-bit rounding error. + if self.length() + 2.0 * margin + 2.0 * f64::EPSILON >= 2.0 * PI { + return S1Interval::full(); + } + } else { + if self.is_full() { + return *self; + } + // Check if this interval will be empty after contraction. + if self.length() + 2.0 * margin - 2.0 * f64::EPSILON <= 0.0 { + return S1Interval::empty(); + } + } + let mut new_lo = remainder(self.lo - margin, 2.0 * PI); + if new_lo <= -PI { + new_lo = PI; + } + let new_hi = remainder(self.hi + margin, 2.0 * PI); + S1Interval::new_checked(new_lo, new_hi) + } + + /// Returns the smallest interval containing both this and the given interval. + pub fn union(&self, y: &S1Interval) -> S1Interval { + if y.is_empty() { + return *self; + } + if self.fast_contains(y.lo) { + if self.fast_contains(y.hi) { + // Either this interval contains y, or the union is Full. + if self.contains(y) { + return *self; + } + return S1Interval::full(); + } + return S1Interval::new_checked(self.lo, y.hi); + } + if self.fast_contains(y.hi) { + return S1Interval::new_checked(y.lo, self.hi); + } + + // This interval contains neither endpoint of y. Either y contains all of this interval, or + // the two intervals are disjoint. + if self.is_empty() || y.fast_contains(self.lo) { + return *y; + } + + // Check which pair of endpoints are closer together. + let dlo = positive_distance(y.hi, self.lo); + let dhi = positive_distance(self.hi, y.lo); + if dlo < dhi { + S1Interval::new_checked(y.lo, self.hi) + } else { + S1Interval::new_checked(self.lo, y.hi) + } + } + + /// Returns the intersection of this interval with the given interval. Note that the region of + /// intersection may consist of two disjoint subintervals. In that case, this returns the + /// shorter of the two original intervals. + pub fn intersection(&self, y: &S1Interval) -> S1Interval { + if y.is_empty() { + return S1Interval::empty(); + } + if self.fast_contains(y.lo) { + if self.fast_contains(y.hi) { + // Either this interval contains y, or the intersection consists of two disjoint + // subintervals. Return the shorter original. + if y.length() < self.length() { + return *y; + } + return *self; + } + return S1Interval::new_checked(y.lo, self.hi); + } + if self.fast_contains(y.hi) { + return S1Interval::new_checked(self.lo, y.hi); + } + + // This interval contains neither endpoint of y. Either y contains all of this interval, or + // the two intervals are disjoint. + if y.fast_contains(self.lo) { + return *self; + } + debug_assert!(!self.intersects(y)); + S1Interval::empty() + } + + /// Return true if this interval can be transformed into the given interval by moving each + /// endpoint by at most "max_error" (and without the endpoints crossing, which would invert the + /// interval). Empty and full intervals are considered to start at an arbitrary point on the + /// unit circle, thus any interval with (length <= 2*max_error) matches the empty interval, and + /// any interval with (length >= 2*Pi - 2*max_error) matches the full interval. + pub fn approx_equals(&self, y: &S1Interval, max_error: f64) -> bool { + if self.is_empty() { + return y.length() <= 2.0 * max_error; + } + if y.is_empty() { + return self.length() <= 2.0 * max_error; + } + if self.is_full() { + return y.length() >= 2.0 * (PI - max_error); + } + if y.is_full() { + return self.length() >= 2.0 * (PI - max_error); + } + + // Check that moving endpoints does not invert the interval. + remainder(y.lo - self.lo, 2.0 * PI).abs() <= max_error + && remainder(y.hi - self.hi, 2.0 * PI).abs() <= max_error + && (self.length() - y.length()).abs() <= 2.0 * max_error + } +} + +impl PartialEq for S1Interval { + fn eq(&self, other: &Self) -> bool { + self.lo == other.lo && self.hi == other.hi + } +} + +impl std::fmt::Display for S1Interval { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[{}, {}]", self.lo, self.hi) + } +} + +/// Computes the distance from `a` to `b` in the range [0, 2*PI). This is more numerically stable +/// than `remainder(b - a - PI, 2*PI) + PI` for very small positive distances. +fn positive_distance(a: f64, b: f64) -> f64 { + let d = b - a; + if d >= 0.0 { + d + } else { + // Ensure that if b == PI and a == (-PI + eps), we return ~2*PI. + (b + PI) - (a - PI) + } +} + +#[cfg(test)] +#[path = "tests/s1interval_tests.rs"] +mod tests; diff --git a/src/spatial/s2cap.rs b/src/spatial/s2cap.rs new file mode 100644 index 0000000000..e97c845216 --- /dev/null +++ b/src/spatial/s2cap.rs @@ -0,0 +1,478 @@ +//! S2Cap represents a disc-shaped region defined by a center and radius. Technically this shape is +//! called a "spherical cap" (rather than disc) because it is not planar; the cap represents a +//! portion of the sphere that has been cut off by a plane. The boundary of the cap is the circle +//! defined by the intersection of the sphere and the plane. For containment purposes, the cap is +//! a closed set, i.e. it contains its boundary. +//! +//! The radius of the cap is measured along the surface of the sphere (rather than the +//! straight-line distance through the interior). Thus a cap of radius Pi/2 is a hemisphere, and a +//! cap of radius Pi covers the entire sphere. +//! +//! Ported from s2/s2cap.h (Eric Veach, Google). + +use std::f64::consts::{FRAC_PI_2, PI}; + +use super::latlng_rect::S2LatLngRect; +use super::math::{cross, dot, neg, norm2, remainder}; +use super::r1interval::R1Interval; +use super::region::Region; +use super::s1chord_angle::{S1ChordAngle, RELATIVE_SUM_ERROR}; +use super::s1interval::S1Interval; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; +use super::s2metrics::S2Metrics; + +/// S2Cap represents a disc-shaped region defined by a center and radius. +#[derive(Clone, Copy, Debug)] +pub struct S2Cap { + center: [f64; 3], + radius: S1ChordAngle, +} + +impl S2Cap { + /// Return an empty cap, i.e. a cap that contains no points. + pub fn empty() -> Self { + Self { + center: [1.0, 0.0, 0.0], + radius: S1ChordAngle::negative(), + } + } + + /// Return a full cap, i.e. a cap that contains all points. + pub fn full() -> Self { + Self { + center: [1.0, 0.0, 0.0], + radius: S1ChordAngle::straight(), + } + } + + /// Constructs a cap where the angle is expressed as an S1ChordAngle. A negative radius yields + /// an empty cap; a radius of 180 degrees or more yields a full cap. "center" should be unit + /// length. + pub fn new(center: [f64; 3], radius: S1ChordAngle) -> Self { + Self { center, radius } + } + + /// Constructs a cap with the given center and radius in radians. + pub fn from_center_angle(center: [f64; 3], radius_radians: f64) -> Self { + Self { + center, + radius: S1ChordAngle::from_radians(radius_radians.min(PI)), + } + } + + /// Convenience function that creates a cap containing a single point. + pub fn from_point(center: [f64; 3]) -> Self { + Self { + center, + radius: S1ChordAngle::zero(), + } + } + + /// Returns a cap with the given center and height (see comments above). A negative height + /// yields an empty cap; a height of 2 or more yields a full cap. "center" should be unit + /// length. + pub fn from_center_height(center: [f64; 3], height: f64) -> Self { + Self { + center, + radius: S1ChordAngle::from_length2(2.0 * height), + } + } + + /// Return a cap with the given center and surface area. Note that the area can also be + /// interpreted as the solid angle subtended by the cap (because the sphere has unit radius). A + /// negative area yields an empty cap; an area of 4*Pi or more yields a full cap. + pub fn from_center_area(center: [f64; 3], area: f64) -> Self { + Self { + center, + radius: S1ChordAngle::from_length2(area / PI), + } + } + + /// The center of the cap. + #[inline] + pub fn center(&self) -> [f64; 3] { + self.center + } + + /// The radius of the cap as an S1ChordAngle. + #[inline] + pub fn radius(&self) -> S1ChordAngle { + self.radius + } + + /// Returns the height of the cap, i.e. the distance from the center point to the cutoff plane. + #[inline] + pub fn height(&self) -> f64 { + 0.5 * self.radius.length2() + } + + /// Return the cap radius as an S1Angle. (Note that the cap angle is stored internally as an + /// S1ChordAngle, so this method requires a trigonometric operation and may yield a slightly + /// different result than the value passed to the (S2Point, S1Angle) constructor.) + pub fn get_radius_radians(&self) -> f64 { + self.radius.to_radians() + } + + /// Return the area of the cap. + pub fn get_area(&self) -> f64 { + 2.0 * PI * self.height().max(0.0) + } + + /// Return the true centroid of the cap multiplied by its surface area (see s2centroids.h for + /// details on centroids). The result lies on the ray from the origin through the cap's + /// center, but it is not unit length. Note that if you just want the "surface centroid", i.e. + /// the normalized result, then it is much simpler just to call center(). + /// + /// The reason for multiplying the result by the cap area is to make it easier to compute the + /// centroid of more complicated shapes. The centroid of a union of disjoint regions can be + /// computed simply by adding their GetCentroid() results. Caveat: for caps that contain a + /// single point (i.e., zero radius), this method always returns the origin (0, 0, 0). This is + /// because shapes with no area don't affect the centroid of a union whose total area is + /// positive. + pub fn get_centroid(&self) -> [f64; 3] { + if self.is_empty() { + return [0.0, 0.0, 0.0]; + } + // From symmetry, the centroid of the cap must be somewhere on the line from the origin to + // the center of the cap on the surface of the sphere. When a sphere is divided into slices + // of constant thickness by a set of parallel planes, all slices have the same surface + // area. This implies that the radial component of the centroid is simply the midpoint of + // the range of radial distances spanned by the cap. That is easily computed from the cap + // height. + let r = 1.0 - 0.5 * self.height(); + let scale = r * self.get_area(); + [ + scale * self.center[0], + scale * self.center[1], + scale * self.center[2], + ] + } + + /// We allow negative heights (to represent empty caps) but heights are normalized so that they + /// do not exceed 2. + pub fn is_valid(&self) -> bool { + let norm2 = self.center[0] * self.center[0] + + self.center[1] * self.center[1] + + self.center[2] * self.center[2]; + (norm2 - 1.0).abs() <= 5.0 * f64::EPSILON && self.radius.length2() <= 4.0 + } + + /// Return true if the cap is empty, i.e. it contains no points. + #[inline] + pub fn is_empty(&self) -> bool { + self.radius.is_negative() + } + + /// Return true if the cap is full, i.e. it contains all points. + #[inline] + pub fn is_full(&self) -> bool { + self.radius.length2() == 4.0 + } + + /// Return the complement of the interior of the cap. A cap and its complement have the same + /// boundary but do not share any interior points. The complement operator is not a bijection + /// because the complement of a singleton cap (containing a single point) is the same as the + /// complement of an empty cap. + pub fn complement(&self) -> Self { + // The complement of a full cap is an empty cap, not a singleton. Also make sure that the + // complement of an empty cap is full. + if self.is_full() { + return Self::empty(); + } + if self.is_empty() { + return Self::full(); + } + Self { + center: neg(self.center), + radius: S1ChordAngle::from_length2(4.0 - self.radius.length2()), + } + } + + /// Return true if and only if this cap contains the given other cap (in a set containment + /// sense, e.g. every cap contains the empty cap). + pub fn contains_cap(&self, other: &S2Cap) -> bool { + if self.is_full() || other.is_empty() { + return true; + } + self.radius >= S1ChordAngle::from_points(&self.center, &other.center) + other.radius + } + + /// Return true if and only if this cap intersects the given other cap, i.e. whether they have + /// any points in common. + pub fn intersects_cap(&self, other: &S2Cap) -> bool { + if self.is_empty() || other.is_empty() { + return false; + } + self.radius + other.radius >= S1ChordAngle::from_points(&self.center, &other.center) + } + + /// Return true if and only if the interior of this cap intersects the given other cap. (This + /// relationship is not symmetric, since only the interior of this cap is used.) + pub fn interior_intersects_cap(&self, other: &S2Cap) -> bool { + // Make sure this cap has an interior and the other cap is non-empty. + if self.radius.length2() <= 0.0 || other.is_empty() { + return false; + } + self.radius + other.radius > S1ChordAngle::from_points(&self.center, &other.center) + } + + /// The point "p" should be a unit-length vector. + #[inline] + pub fn contains_point(&self, p: &[f64; 3]) -> bool { + S1ChordAngle::from_points(&self.center, p) <= self.radius + } + + /// Return true if and only if the given point is contained in the interior of the cap (i.e. + /// the cap excluding its boundary). "p" should be a unit-length vector. + pub fn interior_contains_point(&self, p: &[f64; 3]) -> bool { + self.is_full() || S1ChordAngle::from_points(&self.center, p) < self.radius + } + + /// Increase the cap height if necessary to include the given point. If the cap is empty then + /// the center is set to the given point, but otherwise the center is not changed. "p" should + /// be a unit-length vector. + pub fn add_point(&mut self, p: &[f64; 3]) { + if self.is_empty() { + self.center = *p; + self.radius = S1ChordAngle::zero(); + } else { + let dist = S1ChordAngle::from_points(&self.center, p); + if dist > self.radius { + self.radius = dist; + } + } + } + + /// Increase the cap height if necessary to include "other". If the current cap is empty it is + /// set to the given other cap. + pub fn add_cap(&mut self, other: &S2Cap) { + if self.is_empty() { + *self = *other; + } else if !other.is_empty() { + // We round up the distance to ensure that the cap is actually contained. + let dist = S1ChordAngle::from_points(&self.center, &other.center) + other.radius; + let error = (2.0 * f64::EPSILON + RELATIVE_SUM_ERROR) * dist.length2(); + let dist_with_error = dist.plus_error(error); + if dist_with_error > self.radius { + self.radius = dist_with_error; + } + } + } + + /// Return a cap that contains all points within a given distance of this cap. Note that any + /// expansion of the empty cap is still empty. + pub fn expanded(&self, distance_radians: f64) -> Self { + debug_assert!(distance_radians >= 0.0); + if self.is_empty() { + return Self::empty(); + } + Self { + center: self.center, + radius: self.radius + S1ChordAngle::from_radians(distance_radians), + } + } + + // Return true if this cap intersects any point of 'cell' excluding its vertices (which are + // assumed to already have been checked). + fn intersects_cell(&self, cell: &S2Cell, vertices: &[[f64; 3]; 4]) -> bool { + // If the cap is a hemisphere or larger, the cell and the complement of the cap are both + // convex. Therefore since no vertex of the cell is contained, no other interior point of + // the cell is contained either. + if self.radius >= S1ChordAngle::right() { + return false; + } + + // We need to check for empty caps due to the center check just below. + if self.is_empty() { + return false; + } + + // Optimization: return true if the cell contains the cap center. (This allows half of the + // edge checks below to be skipped.) + if cell.contains_point(&self.center) { + return true; + } + + // At this point we know that the cell does not contain the cap center, and the cap does + // not contain any cell vertex. The only way that they can intersect is if the cap + // intersects the interior of some edge. + let sin2_angle = self.radius.sin2(); + for k in 0..4 { + let edge = cell.get_edge_raw(k); + let dotted = dot(&self.center, &edge); + if dotted > 0.0 { + // The center is in the interior half-space defined by the edge. We don't need to + // consider these edges, since if the cap intersects this edge then it also + // intersects the edge on the opposite side of the cell (because we know the center + // is not contained with the cell). + continue; + } + // The Norm2() factor is necessary because "edge" is not normalized. + let edge_norm2 = norm2(&edge); + if dotted * dotted > sin2_angle * edge_norm2 { + return false; + } + // Otherwise, the great circle containing this edge intersects the interior of the cap. + // We just need to check whether the point of closest approach occurs between the two + // edge endpoints. + let dir = cross(&edge, &self.center); + if dot(&dir, &vertices[k]) < 0.0 && dot(&dir, &vertices[(k + 1) & 3]) > 0.0 { + return true; + } + } + false + } + + /// Computes a covering of the S2Cap. In general the covering consists of at most 4 cells + /// except for very large caps, which may need up to 6 cells. The output is not sorted. + pub fn get_cell_union_bound(&self, cell_ids: &mut Vec) { + cell_ids.clear(); + + // Find the maximum level such that the cap contains at most one cell vertex and such that + // S2CellId::AppendVertexNeighbors() can be called. + let level = S2Metrics::get_level_for_min_width(self.get_radius_radians()) - 1; + + // If level < 0, then more than three face cells are required. + if level < 0 { + for face in 0..6 { + cell_ids.push(S2CellId::from_face(face)); + } + } else { + // The covering consists of the 4 cells at the given level that share the cell vertex + // that is closest to the cap center. + let center_cell = S2CellId::from_point(&self.center); + center_cell.append_vertex_neighbors(level, cell_ids); + } + } + + /// Computes the bounding latitude-longitude rectangle for the cap. + pub fn get_rect_bound_impl(&self) -> S2LatLngRect { + if self.is_empty() { + return S2LatLngRect::empty(); + } + + // Convert the center to a (lat, lng) pair, and compute the cap angle. + let center_lat = self.center[2].asin(); + let center_lng = self.center[1].atan2(self.center[0]); + let cap_angle = self.get_radius_radians(); + + let mut all_longitudes = false; + let mut lat_lo; + let mut lat_hi; + let mut lng_lo = -PI; + let mut lng_hi = PI; + + // Check whether cap includes the south pole. + lat_lo = center_lat - cap_angle; + if lat_lo <= -FRAC_PI_2 { + lat_lo = -FRAC_PI_2; + all_longitudes = true; + } + + // Check whether cap includes the north pole. + lat_hi = center_lat + cap_angle; + if lat_hi >= FRAC_PI_2 { + lat_hi = FRAC_PI_2; + all_longitudes = true; + } + + if !all_longitudes { + // Compute the range of longitudes covered by the cap. We use the law of sines for + // spherical triangles. Consider the triangle ABC where A is the north pole, B is the + // center of the cap, and C is the point of tangency between the cap boundary and a + // line of longitude. Then C is a right angle, and letting a,b,c denote the sides + // opposite A,B,C, we have sin(a)/sin(A) = sin(c)/sin(C), or sin(A) = sin(a)/sin(c). + // Here "a" is the cap angle, and "c" is the colatitude (90 degrees minus the + // latitude). This formula also works for negative latitudes. + // + // The formula for sin(a) follows from the relationship h = 1 - cos(a). + let sin_a = self.radius.sin(); + let sin_c = center_lat.cos(); + if sin_a <= sin_c { + let angle_a = (sin_a / sin_c).asin(); + lng_lo = remainder(center_lng - angle_a, 2.0 * PI); + lng_hi = remainder(center_lng + angle_a, 2.0 * PI); + } + } + + S2LatLngRect::new( + R1Interval::new(lat_lo, lat_hi), + S1Interval::new(lng_lo, lng_hi), + ) + } +} + +impl Region for S2Cap { + fn get_cap_bound(&self) -> S2Cap { + *self + } + + fn get_rect_bound(&self) -> S2LatLngRect { + self.get_rect_bound_impl() + } + + fn get_cell_union_bound(&self, cell_ids: &mut Vec) { + self.get_cell_union_bound(cell_ids); + } + + // If the cap does not contain all cell vertices, return false. We check the vertices before + // taking the Complement() because we can't accurately represent the complement of a very small + // cap (a height of 2-epsilon is rounded off to 2). + fn contains_cell(&self, cell: &S2Cell) -> bool { + let mut vertices = [[0.0; 3]; 4]; + for k in 0..4 { + vertices[k] = cell.get_vertex(k); + if !self.contains_point(&vertices[k]) { + return false; + } + } + // Otherwise, return true if the complement of the cap does not intersect the cell. (This + // test is slightly conservative, because technically we want + // Complement().InteriorIntersects() here.) + !self.complement().intersects_cell(cell, &vertices) + } + + fn may_intersect(&self, cell: &S2Cell) -> bool { + // If the cap contains any cell vertex, return true. + let mut vertices = [[0.0; 3]; 4]; + for k in 0..4 { + vertices[k] = cell.get_vertex(k); + if self.contains_point(&vertices[k]) { + return true; + } + } + self.intersects_cell(cell, &vertices) + } + + fn contains_point(&self, p: &[f64; 3]) -> bool { + self.contains_point(p) + } +} + +impl PartialEq for S2Cap { + fn eq(&self, other: &Self) -> bool { + (self.center == other.center && self.radius == other.radius) + || (self.is_empty() && other.is_empty()) + || (self.is_full() && other.is_full()) + } +} + +impl Eq for S2Cap {} + +impl std::fmt::Display for S2Cap { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "[Center=({:.6}, {:.6}, {:.6}), Radius={:.6}deg]", + self.center[0], + self.center[1], + self.center[2], + self.radius.to_degrees() + ) + } +} + +#[cfg(test)] +#[path = "tests/s2cap_tests.rs"] +mod tests; diff --git a/src/spatial/s2cell.rs b/src/spatial/s2cell.rs new file mode 100644 index 0000000000..ab4ecafa91 --- /dev/null +++ b/src/spatial/s2cell.rs @@ -0,0 +1,541 @@ +//! S2Cell: An S2Region object representing a cell. +//! +//! Unlike S2CellId, S2Cell supports efficient containment and intersection tests by caching the +//! cell's (u,v) bounds and providing vertex/edge access. +//! +//! S2Cell from s2cell.h and s2cell.cc. +use std::f64::consts::{FRAC_PI_2, FRAC_PI_4, PI}; + +use super::crossings::S2EdgeCrosser; +use super::latlng_rect::S2LatLngRect; +use super::math::{dot, ldexp, neg, normalize}; +use super::r1interval::R1Interval; +use super::r2rect::R2Rect; +use super::s1chord_angle::S1ChordAngle; +use super::s1interval::S1Interval; +use super::s2cap::S2Cap; +use super::s2cell_id::S2CellId; +use super::s2coords::{ + self, face_uv_to_xyz, face_xyz_to_uvw, get_u_axis, get_u_norm, get_v_axis, get_v_norm, + ij_to_st_min, st_to_uv, MAX_CELL_LEVEL, +}; +use super::s2edge_distances::update_min_distance; + +/// An S2Cell represents a cell in the S2 cell decomposition. +#[derive(Clone, Copy, Debug)] +pub struct S2Cell { + id: S2CellId, + face: u8, + level: u8, + orientation: u8, + uv: R2Rect, +} + +impl S2Cell { + /// Creates an S2Cell from an S2CellId. + pub fn new(id: S2CellId) -> Self { + let (face, i, j, orientation) = id.to_face_ij_orientation(); + let level = id.level(); + let uv = ij_level_to_bound_uv(i, j, level); + Self { + id, + face: face as u8, + level: level as u8, + orientation: orientation as u8, + uv, + } + } + + /// Creates an S2Cell for the given cube face (level 0). + pub fn from_face(face: i32) -> Self { + Self::new(S2CellId::from_face(face)) + } + + /// Creates an S2Cell containing the given point. + pub fn from_point(p: &[f64; 3]) -> Self { + Self::new(S2CellId::from_point(p)) + } + + /// Returns the cell ID. + #[inline] + pub fn id(&self) -> S2CellId { + self.id + } + + /// Returns the cube face (0-5). + #[inline] + pub fn face(&self) -> i32 { + self.face as i32 + } + + /// Returns the subdivision level (0-30). + #[inline] + pub fn level(&self) -> i32 { + self.level as i32 + } + + /// Returns the Hilbert curve orientation. + #[inline] + pub fn orientation(&self) -> i32 { + self.orientation as i32 + } + + /// Returns true if this is a leaf cell. + #[inline] + pub fn is_leaf(&self) -> bool { + self.level as i32 == MAX_CELL_LEVEL + } + + /// Returns the bounds in (u,v)-space. + #[inline] + pub fn get_bound_uv(&self) -> R2Rect { + self.uv + } + + /// Returns the size of the cell in (i,j) coordinates. + #[inline] + pub fn get_size_ij(&self) -> i32 { + S2CellId::size_ij_for_level(self.level as i32) + } + + /// Returns the size of the cell in (s,t) coordinates. + #[inline] + pub fn get_size_st(&self) -> f64 { + ij_to_st_min(self.get_size_ij()) + } + + /// Returns the k-th vertex (k = 0,1,2,3) in CCW order. + /// Lower left, lower right, upper right, upper left in UV plane. + #[inline] + pub fn get_vertex(&self, k: usize) -> [f64; 3] { + normalize(&self.get_vertex_raw(k)) + } + + /// Returns the k-th vertex without normalization. + #[inline] + pub fn get_vertex_raw(&self, k: usize) -> [f64; 3] { + let v = self.uv.get_vertex(k % 4); + face_uv_to_xyz(self.face as i32, v[0], v[1]) + } + + /// Returns the center of the cell. + #[inline] + pub fn get_center(&self) -> [f64; 3] { + normalize(&self.get_center_raw()) + } + + /// Returns the center without normalization. + #[inline] + pub fn get_center_raw(&self) -> [f64; 3] { + self.id.to_point_raw() + } + + /// Returns the inward-facing normal of edge k (k = 0,1,2,3). + /// Edge k connects vertex k to vertex k+1. + #[inline] + pub fn get_edge(&self, k: usize) -> [f64; 3] { + normalize(&self.get_edge_raw(k)) + } + + /// Returns the edge normal without normalization. + #[inline] + pub fn get_edge_raw(&self, k: usize) -> [f64; 3] { + match k & 3 { + 0 => get_v_norm(self.face as i32, self.uv[1].lo()), + 1 => get_u_norm(self.face as i32, self.uv[0].hi()), + 2 => neg(get_v_norm(self.face as i32, self.uv[1].hi())), + _ => neg(get_u_norm(self.face as i32, self.uv[0].lo())), + } + } + + /// Returns the UV coordinate of edge k (the constant coord along that edge). + pub fn get_uv_coord_of_edge(&self, k: usize) -> f64 { + match k & 3 { + 0 => self.uv[1].lo(), // Bottom: v = min + 1 => self.uv[0].hi(), // Right: u = max + 2 => self.uv[1].hi(), // Top: v = max + _ => self.uv[0].lo(), // Left: u = min + } + } + + /// Subdivides this cell into its four children. Returns None if this is a leaf cell. + pub fn subdivide(&self) -> Option<[S2Cell; 4]> { + if self.is_leaf() { + return None; + } + + // Compute the cell midpoint in uv-space. + let uv_mid = self.id.get_center_uv(); + let child_ids = self.id.children(); + + // Create children using the same optimization as S2. + let mut children = [S2Cell::new(child_ids[0]); 4]; + + for (pos, child_id) in child_ids.iter().enumerate() { + let ij = s2coords::pos_to_ij(self.orientation as i32, pos as i32); + let i = (ij >> 1) as usize; + let j = (ij & 1) as usize; + + children[pos] = S2Cell { + id: *child_id, + face: self.face, + level: self.level + 1, + orientation: (self.orientation as i32 ^ s2coords::pos_to_orientation(pos as i32)) + as u8, + uv: R2Rect::new( + R1Interval::new( + if i == 0 { self.uv[0].lo() } else { uv_mid[0] }, + if i == 0 { uv_mid[0] } else { self.uv[0].hi() }, + ), + R1Interval::new( + if j == 0 { self.uv[1].lo() } else { uv_mid[1] }, + if j == 0 { uv_mid[1] } else { self.uv[1].hi() }, + ), + ), + }; + } + + Some(children) + } + + /// Returns true if this cell contains the given point. Note that S2Cells are closed regions + /// (include their boundary). + pub fn contains_point(&self, p: &[f64; 3]) -> bool { + // Project point onto this face and check UV containment. + if let Some(uv) = s2coords::face_xyz_to_uv(self.face as i32, p) { + // Expand slightly to account for numerical error. + self.uv.expanded(f64::EPSILON).contains_point(uv) + } else { + false + } + } + + /// Returns true if this cell contains the given cell. + pub fn contains_cell(&self, other: &S2Cell) -> bool { + self.id.contains(other.id) + } + + /// Returns true if this cell may intersect the given cell. + pub fn may_intersect(&self, other: &S2Cell) -> bool { + self.id.intersects(other.id) + } + + /// Returns the average area of cells at the given level. + #[inline] + pub fn average_area(level: i32) -> f64 { + const AVG_AREA_BASE: f64 = 4.0 * PI / 6.0; // ~2.094 + ldexp(AVG_AREA_BASE, -2 * level) + } + + /// Returns the minimum distance from the cell to the given point. All + /// arguments should be unit length. Returns zero if the point is inside + /// the cell. + /// + /// Port of S2Cell::GetDistance(const S2Point&) in s2cell.cc. + pub fn get_distance_to_point(&self, target: &[f64; 3]) -> S1ChordAngle { + self.get_distance_internal(target, true) + } + + /// Returns the minimum distance from the cell boundary to the given point. + /// Returns zero if the point is on the boundary. + /// + /// Port of S2Cell::GetBoundaryDistance in s2cell.cc. + pub fn get_boundary_distance(&self, target: &[f64; 3]) -> S1ChordAngle { + self.get_distance_internal(target, false) + } + + /// Returns the minimum distance from the cell to the edge ab. Returns zero + /// if the edge intersects the cell or if either endpoint is inside. + /// + /// Port of S2Cell::GetDistance(const S2Point& a, const S2Point& b) in s2cell.cc. + pub fn get_distance_to_edge(&self, a: &[f64; 3], b: &[f64; 3]) -> S1ChordAngle { + // First, check the minimum distance to the edge endpoints A and B. + // (This also detects whether either endpoint is inside the cell.) + let mut min_dist = self.get_distance_to_point(a); + let d = self.get_distance_to_point(b); + if d < min_dist { + min_dist = d; + } + if min_dist == S1ChordAngle::zero() { + return min_dist; + } + + // Otherwise, check whether the edge crosses the cell boundary. + let v: [[f64; 3]; 4] = [ + self.get_vertex(0), + self.get_vertex(1), + self.get_vertex(2), + self.get_vertex(3), + ]; + let mut crosser = S2EdgeCrosser::new(a, b); + crosser.restart_at(&v[3]); + for i in 0..4 { + if crosser.crossing_sign(&v[i]) >= 0 { + return S1ChordAngle::zero(); + } + } + // Finally, check whether the minimum distance occurs between a cell vertex + // and the interior of the edge AB. (Some of this work is redundant, since + // it also checks the distance to the endpoints A and B again.) + // + // Note that we don't need to check the distance from the interior of AB to + // the interior of a cell edge, because the only way that this distance can + // be minimal is if the two edges cross (already checked above). + for i in 0..4 { + update_min_distance(&v[i], a, b, &mut min_dist); + } + min_dist + } + + // All calculations are done in the (u,v,w) coordinates of this cell's face. + // + // Port of S2Cell::GetDistanceInternal in s2cell.cc. + fn get_distance_internal(&self, target_xyz: &[f64; 3], to_interior: bool) -> S1ChordAngle { + let target = face_xyz_to_uvw(self.face as i32, target_xyz); + + // Compute dot products with all four upward or rightward-facing edge + // normals. "dirIJ" is the dot product for the edge corresponding to axis + // I, endpoint J. For example, dir01 is the right edge of the S2Cell + // (corresponding to the upper endpoint of the u-axis). + let dir00 = target[0] - target[2] * self.uv[0].lo(); + let dir01 = target[0] - target[2] * self.uv[0].hi(); + let dir10 = target[1] - target[2] * self.uv[1].lo(); + let dir11 = target[1] - target[2] * self.uv[1].hi(); + let mut inside = true; + if dir00 < 0.0 { + inside = false; // Target is to the left of the cell + if self.v_edge_is_closest(&target, 0) { + return edge_distance(-dir00, self.uv[0].lo()); + } + } + if dir01 > 0.0 { + inside = false; // Target is to the right of the cell + if self.v_edge_is_closest(&target, 1) { + return edge_distance(dir01, self.uv[0].hi()); + } + } + if dir10 < 0.0 { + inside = false; // Target is below the cell + if self.u_edge_is_closest(&target, 0) { + return edge_distance(-dir10, self.uv[1].lo()); + } + } + if dir11 > 0.0 { + inside = false; // Target is above the cell + if self.u_edge_is_closest(&target, 1) { + return edge_distance(dir11, self.uv[1].hi()); + } + } + if inside { + if to_interior { + return S1ChordAngle::zero(); + } + // Although you might think of S2Cells as rectangles, they are actually + // arbitrary quadrilaterals after they are projected onto the sphere. + // Therefore the simplest approach is just to find the minimum distance to + // any of the four edges. + return edge_distance(-dir00, self.uv[0].lo()) + .min(edge_distance(dir01, self.uv[0].hi())) + .min(edge_distance(-dir10, self.uv[1].lo())) + .min(edge_distance(dir11, self.uv[1].hi())); + } + // Otherwise, the closest point is one of the four cell vertices. Note that + // it is *not* trivial to narrow down the candidates based on the edge sign + // tests above, because (1) the edges don't meet at right angles and (2) + // there are points on the far side of the sphere that are both above *and* + // below the cell, etc. + self.vertex_chord_dist(&target, 0, 0) + .min(self.vertex_chord_dist(&target, 1, 0)) + .min(self.vertex_chord_dist(&target, 0, 1)) + .min(self.vertex_chord_dist(&target, 1, 1)) + } + + // Return the squared chord distance from point P to corner vertex (i,j). + // Port of S2Cell::VertexChordDist in s2cell.cc. + fn vertex_chord_dist(&self, p: &[f64; 3], i: usize, j: usize) -> S1ChordAngle { + let vertex = normalize(&[self.uv[0].get_bound(i), self.uv[1].get_bound(j), 1.0]); + S1ChordAngle::from_points(p, &vertex) + } + + // Given a point P and either the lower or upper edge of the S2Cell (specified + // by setting "v_end" to 0 or 1 respectively), return true if P is closer to + // the interior of that edge than it is to either endpoint. + // Port of S2Cell::UEdgeIsClosest in s2cell.cc. + fn u_edge_is_closest(&self, p: &[f64; 3], v_end: usize) -> bool { + let u0 = self.uv[0].lo(); + let u1 = self.uv[0].hi(); + let v = self.uv[1].get_bound(v_end); + // These are the normals to the planes that are perpendicular to the edge + // and pass through one of its two endpoints. + let dir0 = [v * v + 1.0, -u0 * v, -u0]; + let dir1 = [v * v + 1.0, -u1 * v, -u1]; + dot(p, &dir0) > 0.0 && dot(p, &dir1) < 0.0 + } + + // Given a point P and either the left or right edge of the S2Cell (specified + // by setting "u_end" to 0 or 1 respectively), return true if P is closer to + // the interior of that edge than it is to either endpoint. + // Port of S2Cell::VEdgeIsClosest in s2cell.cc. + fn v_edge_is_closest(&self, p: &[f64; 3], u_end: usize) -> bool { + let v0 = self.uv[1].lo(); + let v1 = self.uv[1].hi(); + let u = self.uv[0].get_bound(u_end); + let dir0 = [-u * v0, u * u + 1.0, -v0]; + let dir1 = [-u * v1, u * u + 1.0, -v1]; + dot(p, &dir0) > 0.0 && dot(p, &dir1) < 0.0 + } + + /// Returns the latitude of the cell vertex at (i, j) where i, j in {0, 1}. + fn get_latitude(&self, i: usize, j: usize) -> f64 { + let p = face_uv_to_xyz( + self.face as i32, + self.uv[0].get_bound(i), + self.uv[1].get_bound(j), + ); + (p[2] + 0.0).atan2((p[0] * p[0] + p[1] * p[1]).sqrt()) + } + + /// Returns the longitude of the cell vertex at (i, j). + fn get_longitude(&self, i: usize, j: usize) -> f64 { + let p = face_uv_to_xyz( + self.face as i32, + self.uv[0].get_bound(i), + self.uv[1].get_bound(j), + ); + (p[1] + 0.0).atan2(p[0] + 0.0) + } + + /// Returns the latitude-longitude bounding rectangle for this cell. + pub fn get_rect_bound(&self) -> S2LatLngRect { + if self.level > 0 { + let u = self.uv[0].lo() + self.uv[0].hi(); + let v = self.uv[1].lo() + self.uv[1].hi(); + + let u_axis = get_u_axis(self.face as i32); + let v_axis = get_v_axis(self.face as i32); + + let i = if u_axis[2] == 0.0 { + (u < 0.0) as usize + } else { + (u > 0.0) as usize + }; + let j = if v_axis[2] == 0.0 { + (v < 0.0) as usize + } else { + (v > 0.0) as usize + }; + + let lat = R1Interval::from_point_pair( + self.get_latitude(i, j), + self.get_latitude(1 - i, 1 - j), + ); + let lng = S1Interval::from_point_pair( + self.get_longitude(i, 1 - j), + self.get_longitude(1 - i, j), + ); + + S2LatLngRect::new(lat, lng) + .expanded(2.0 * f64::EPSILON, 2.0 * f64::EPSILON) + .polar_closure() + } else { + let pole_min_lat = (1.0_f64 / 3.0).sqrt().asin() - 0.5 * f64::EPSILON; + + let bound = match self.face { + 0 => S2LatLngRect::new( + R1Interval::new(-FRAC_PI_4, FRAC_PI_4), + S1Interval::new(-FRAC_PI_4, FRAC_PI_4), + ), + 1 => S2LatLngRect::new( + R1Interval::new(-FRAC_PI_4, FRAC_PI_4), + S1Interval::new(FRAC_PI_4, 3.0 * FRAC_PI_4), + ), + 2 => { + S2LatLngRect::new(R1Interval::new(pole_min_lat, FRAC_PI_2), S1Interval::full()) + } + 3 => S2LatLngRect::new( + R1Interval::new(-FRAC_PI_4, FRAC_PI_4), + S1Interval::new(3.0 * FRAC_PI_4, -3.0 * FRAC_PI_4), + ), + 4 => S2LatLngRect::new( + R1Interval::new(-FRAC_PI_4, FRAC_PI_4), + S1Interval::new(-3.0 * FRAC_PI_4, -FRAC_PI_4), + ), + _ => S2LatLngRect::new( + R1Interval::new(-FRAC_PI_2, -pole_min_lat), + S1Interval::full(), + ), + }; + + bound.expanded(f64::EPSILON, 0.0) + } + } + + /// Returns a bounding spherical cap for this cell. + pub fn get_cap_bound(&self) -> S2Cap { + let uv_center = self.uv.get_center(); + let center_raw = face_uv_to_xyz(self.face as i32, uv_center[0], uv_center[1]); + let center = normalize(¢er_raw); + + let mut cap = S2Cap::from_point(center); + for k in 0..4 { + cap.add_point(&self.get_vertex(k)); + } + cap + } +} + +// Given the dot product of a point P with the normal of a u- or v-edge at the +// given coordinate value, return the distance from P to that edge. +// Port of EdgeDistance in s2cell.cc. +fn edge_distance(dir_ij: f64, uv: f64) -> S1ChordAngle { + // Let P be the target point and let R be the closest point on the given + // edge AB. The desired distance PR can be expressed as PR^2 = PQ^2 + QR^2 + // where Q is the point P projected onto the plane through the great circle + // through AB. We can compute the distance PQ^2 perpendicular to the plane + // from "dirIJ" (the dot product of the target point P with the edge + // normal) and the squared length of the edge normal (1 + uv**2). + let pq2 = (dir_ij * dir_ij) / (1.0 + uv * uv); + + // We can compute the distance QR as (1 - OQ) where O is the sphere origin, + // and we can compute OQ^2 = 1 - PQ^2 using the Pythagorean theorem. + // (This calculation loses accuracy as angle POQ approaches Pi/2.) + let qr = 1.0 - (1.0 - pq2).sqrt(); + S1ChordAngle::from_length2(pq2 + qr * qr) +} + +/// Computes the (u,v) bounding rectangle for a cell at level `level` whose (i,j) coordinates are +/// `i` and `j`. +fn ij_level_to_bound_uv(i: i32, j: i32, level: i32) -> R2Rect { + let cell_size = S2CellId::size_ij_for_level(level); + + // Round down i and j to the nearest multiple of cell_size. + let i_lo = i & -cell_size; + let j_lo = j & -cell_size; + let i_hi = i_lo + cell_size; + let j_hi = j_lo + cell_size; + + R2Rect::new( + R1Interval::new(st_to_uv(ij_to_st_min(i_lo)), st_to_uv(ij_to_st_min(i_hi))), + R1Interval::new(st_to_uv(ij_to_st_min(j_lo)), st_to_uv(ij_to_st_min(j_hi))), + ) +} + +impl PartialEq for S2Cell { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } +} + +impl Eq for S2Cell {} + +impl PartialOrd for S2Cell { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for S2Cell { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.id.cmp(&other.id) + } +} diff --git a/src/spatial/s2cell_id.rs b/src/spatial/s2cell_id.rs new file mode 100644 index 0000000000..494ed82182 --- /dev/null +++ b/src/spatial/s2cell_id.rs @@ -0,0 +1,694 @@ +//! An S2CellId is a 64-bit unsigned integer that uniquely identifies a cell in the S2 cell +//! decomposition. It has the following format: +//! +//! ```text +//! id = [face][face_pos] +//! +//! face: a 3-bit number (range 0..5) encoding the cube face. +//! +//! face_pos: a 61-bit number encoding the position of the center of this +//! cell along the Hilbert curve over this face (see the Wiki +//! pages for details). +//! ``` +//! +//! Sequentially increasing cell ids follow a continuous space-filling curve over the entire +//! sphere. They have the following properties: +//! +//! - The id of a cell at level k consists of a 3-bit face number followed by k bit pairs that +//! recursively select one of the four children of each cell. The next bit is always 1, and all +//! other bits are 0. Therefore, the level of a cell is determined by the position of its +//! lowest-numbered bit that is turned on (for a cell at level k, this position is 2 * (kMaxLevel +//! - k).) +//! +//! - The id of a parent cell is at the midpoint of the range of ids spanned by its children (or by +//! its descendants at any level). +//! +//! Leaf cells are often used to represent points on the unit sphere, and this class provides +//! methods for converting directly between these two representations. For cells that represent 2D +//! regions rather than discrete point, it is better to use the S2Cell class. +//! +//! All methods require `is_valid()` to be true unless otherwise specified (although not all +//! methods enforce this). +//! +//! This class is intended to be copied by value as desired. It uses the default copy constructor +//! and assignment operator. +//! +//! s2cell_id.h and s2cell_id.cc +use super::math::normalize; +use super::s2coords::{ + face_si_ti_to_xyz, face_uv_to_xyz, si_ti_to_st, st_to_ij, st_to_uv, uv_to_st, xyz_to_face_uv, + MAX_CELL_LEVEL, NUM_FACES, POS_TO_IJ, POS_TO_ORIENTATION, +}; + +/// Number of bits used to encode the face. +const FACE_BITS: i32 = 3; + +/// Number of bits used for the position along the Hilbert curve. +const POS_BITS: i32 = 2 * MAX_CELL_LEVEL + 1; // 61 + +/// Maximum cell size (number of leaf cells along each axis of a face). +const MAX_SIZE: i32 = 1 << MAX_CELL_LEVEL; + +/// Hilbert curve orientation flags. +const SWAP_MASK: i32 = 0x01; +const INVERT_MASK: i32 = 0x02; + +/// Number of bits to process at once in lookup tables. +const LOOKUP_BITS: usize = 4; + +/// Lookup tables for Hilbert curve encoding/decoding. +/// Generated at compile time. +struct LookupTables { + /// (ij, orientation) -> (pos, new_orientation) + pos: [u16; 1 << (2 * LOOKUP_BITS + 2)], + /// (pos, orientation) -> (ij, new_orientation) + ij: [u16; 1 << (2 * LOOKUP_BITS + 2)], +} + +/// Generate lookup tables at compile time. +const fn make_lookup_tables() -> LookupTables { + let mut tables = LookupTables { + pos: [0; 1 << (2 * LOOKUP_BITS + 2)], + ij: [0; 1 << (2 * LOOKUP_BITS + 2)], + }; + + // Initialize for all 4 orientations + tables = init_lookup_cell(0, 0, 0, 0, 0, 0, tables); + tables = init_lookup_cell(0, 0, 0, SWAP_MASK, 0, SWAP_MASK, tables); + tables = init_lookup_cell(0, 0, 0, INVERT_MASK, 0, INVERT_MASK, tables); + tables = init_lookup_cell( + 0, + 0, + 0, + SWAP_MASK | INVERT_MASK, + 0, + SWAP_MASK | INVERT_MASK, + tables, + ); + + tables +} + +const fn init_lookup_cell( + level: usize, + i: usize, + j: usize, + orig_orientation: i32, + pos: usize, + orientation: i32, + mut tables: LookupTables, +) -> LookupTables { + if level == LOOKUP_BITS { + let ij = (i << LOOKUP_BITS) + j; + tables.pos[(ij << 2) + orig_orientation as usize] = + ((pos << 2) + orientation as usize) as u16; + tables.ij[(pos << 2) + orig_orientation as usize] = + ((ij << 2) + orientation as usize) as u16; + } else { + let next_level = level + 1; + let i2 = i << 1; + let j2 = j << 1; + let pos2 = pos << 2; + + // Process all 4 child positions + let r = POS_TO_IJ[orientation as usize]; + + tables = init_lookup_cell( + next_level, + i2 + ((r[0] >> 1) as usize), + j2 + ((r[0] & 1) as usize), + orig_orientation, + pos2, + orientation ^ POS_TO_ORIENTATION[0], + tables, + ); + tables = init_lookup_cell( + next_level, + i2 + ((r[1] >> 1) as usize), + j2 + ((r[1] & 1) as usize), + orig_orientation, + pos2 + 1, + orientation ^ POS_TO_ORIENTATION[1], + tables, + ); + tables = init_lookup_cell( + next_level, + i2 + ((r[2] >> 1) as usize), + j2 + ((r[2] & 1) as usize), + orig_orientation, + pos2 + 2, + orientation ^ POS_TO_ORIENTATION[2], + tables, + ); + tables = init_lookup_cell( + next_level, + i2 + ((r[3] >> 1) as usize), + j2 + ((r[3] & 1) as usize), + orig_orientation, + pos2 + 3, + orientation ^ POS_TO_ORIENTATION[3], + tables, + ); + } + tables +} + +/// The precomputed Hilbert curve lookup tables. +static LOOKUP: LookupTables = make_lookup_tables(); + +/// A 64-bit cell identifier. +/// +/// S2CellId uniquely identifies a cell in the S2 cell decomposition. The cell ID encodes the face +/// (0-5), position along the Hilbert curve, and subdivision level (0-30). +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)] +pub struct S2CellId(pub u64); + +impl S2CellId { + /// The maximum level (30) for S2CellIds. + pub const MAX_LEVEL: i32 = MAX_CELL_LEVEL; + + /// Creates an S2CellId from a raw 64-bit value. + #[inline] + pub const fn new(id: u64) -> Self { + Self(id) + } + + /// Returns the raw 64-bit value. + #[inline] + pub const fn id(&self) -> u64 { + self.0 + } + + /// Returns an invalid cell ID. + #[inline] + pub const fn none() -> Self { + Self(0) + } + + /// Returns a sentinel cell ID larger than all valid cell IDs. + #[inline] + pub const fn sentinel() -> Self { + Self(!0u64) + } + + /// Creates a cell ID for the given face (level 0). + #[inline] + pub fn from_face(face: i32) -> Self { + debug_assert!(face < NUM_FACES); + Self(((face as u64) << POS_BITS as u32) + Self::lsb_for_level(0)) + } + + /// Creates a cell ID from face, position, and level. + #[inline] + pub fn from_face_pos_level(face: i32, pos: u64, level: i32) -> Self { + debug_assert!(face < NUM_FACES); + debug_assert!(level <= MAX_CELL_LEVEL); + let cell = Self(((face as u64) << POS_BITS as u32) + (pos | 1)); + cell.parent(level) + } + + /// Creates a leaf cell from (face, i, j) coordinates. + pub fn from_face_ij(face: i32, i: i32, j: i32) -> Self { + debug_assert!(face < NUM_FACES); + debug_assert!((0..MAX_SIZE).contains(&i)); + debug_assert!((0..MAX_SIZE).contains(&j)); + + let mut n = (face as u64) << (POS_BITS as u32 - 1); + let mut bits = (face & SWAP_MASK) as u64; + + // Process 4 bits at a time using lookup tables + for k in (0..8).rev() { + let mask = (1 << LOOKUP_BITS) - 1; + bits += (((i >> (k * LOOKUP_BITS)) & mask) as u64) << (LOOKUP_BITS + 2); + bits += (((j >> (k * LOOKUP_BITS)) & mask) as u64) << 2; + bits = LOOKUP.pos[bits as usize] as u64; + n |= (bits >> 2) << (k * 2 * LOOKUP_BITS); + bits &= (SWAP_MASK | INVERT_MASK) as u64; + } + + Self(n * 2 + 1) + } + + /// Creates a leaf cell containing the given point. + pub fn from_point(p: &[f64; 3]) -> Self { + let (face, u, v) = xyz_to_face_uv(p); + let i = st_to_ij(uv_to_st(u)); + let j = st_to_ij(uv_to_st(v)); + Self::from_face_ij(face, i, j) + } + + /// Returns true if this is a valid cell ID. + #[inline] + pub fn is_valid(&self) -> bool { + self.face() < NUM_FACES && (self.lsb() & 0x1555555555555555u64) != 0 + } + + /// Returns the face (0-5) this cell belongs to. + #[inline] + pub fn face(&self) -> i32 { + (self.0 >> POS_BITS as u32) as i32 + } + + /// Returns the position along the Hilbert curve. + #[inline] + pub fn pos(&self) -> u64 { + self.0 & ((!0u64) >> FACE_BITS as u32) + } + + /// Returns the subdivision level (0-30). + #[inline] + pub fn level(&self) -> i32 { + debug_assert!(self.0 != 0); + MAX_CELL_LEVEL - (self.0.trailing_zeros() as i32 >> 1) + } + + /// Returns the lowest set bit (determines level). + #[inline] + pub fn lsb(&self) -> u64 { + self.0 & (!self.0).wrapping_add(1) + } + + /// Returns the LSB for cells at the given level. + #[inline] + pub const fn lsb_for_level(level: i32) -> u64 { + 1u64 << (2 * (MAX_CELL_LEVEL - level) as u32) + } + + /// Returns true if this is a leaf cell (level 30). + #[inline] + pub fn is_leaf(&self) -> bool { + (self.0 & 1) != 0 + } + + /// Returns true if this is a face cell (level 0). + #[inline] + pub fn is_face(&self) -> bool { + (self.0 & (Self::lsb_for_level(0) - 1)) == 0 + } + + /// Returns the size of this cell in (i,j) units. + #[inline] + pub fn size_ij(&self) -> i32 { + Self::size_ij_for_level(self.level()) + } + + /// Returns the size of cells at the given level. + #[inline] + pub const fn size_ij_for_level(level: i32) -> i32 { + 1 << (MAX_CELL_LEVEL - level) + } + + /// Returns the child position (0-3) within the parent. + #[inline] + pub fn child_position(&self) -> u8 { + self.child_position_at_level(self.level()) + } + + /// Returns the child position at the given level. + #[inline] + pub fn child_position_at_level(&self, level: i32) -> u8 { + debug_assert!(level >= 1 && level <= self.level()); + ((self.0 >> (2 * (MAX_CELL_LEVEL - level) as u32 + 1)) & 3) as u8 + } + + /// Returns the minimum cell ID contained by this cell. + #[inline] + pub fn range_min(&self) -> Self { + Self(self.0 - (self.lsb() - 1)) + } + + /// Returns the maximum cell ID contained by this cell. + #[inline] + pub fn range_max(&self) -> Self { + Self(self.0 + (self.lsb() - 1)) + } + + /// Returns true if this cell contains the other cell. + #[inline] + pub fn contains(&self, other: Self) -> bool { + debug_assert!(self.is_valid()); + debug_assert!(other.is_valid()); + other >= self.range_min() && other <= self.range_max() + } + + /// Returns true if this cell intersects the other cell. + #[inline] + pub fn intersects(&self, other: Self) -> bool { + debug_assert!(self.is_valid()); + debug_assert!(other.is_valid()); + other.range_min() <= self.range_max() && other.range_max() >= self.range_min() + } + + /// Returns the parent cell at the given level. + #[inline] + pub fn parent(&self, level: i32) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(level <= self.level()); + let new_lsb = Self::lsb_for_level(level); + Self((self.0 & (!new_lsb + 1)) | new_lsb) + } + + /// Returns the immediate parent cell. + #[inline] + pub fn immediate_parent(&self) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(!self.is_face()); + let new_lsb = self.lsb() << 2; + Self((self.0 & (!new_lsb + 1)) | new_lsb) + } + + /// Returns the child at the given position (0-3). + #[inline] + pub fn child(&self, position: i32) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(!self.is_leaf()); + debug_assert!(position < 4); + let new_lsb = self.lsb() >> 2; + let delta = (2 * position as i64 - 3) * new_lsb as i64; + Self((self.0 as i64 + delta) as u64) + } + + /// Returns the four children of this cell. + #[inline] + pub fn children(&self) -> [Self; 4] { + debug_assert!(self.is_valid()); + debug_assert!(!self.is_leaf()); + [self.child(0), self.child(1), self.child(2), self.child(3)] + } + + /// Returns the next cell at the same level along the Hilbert curve. + #[inline] + pub fn next(&self) -> Self { + Self(self.0 + (self.lsb() << 1)) + } + + /// Returns the previous cell at the same level along the Hilbert curve. + #[inline] + pub fn prev(&self) -> Self { + Self(self.0 - (self.lsb() << 1)) + } + + /// Returns the first cell at this level. + pub fn begin(level: i32) -> Self { + Self::from_face(0).child_begin_at_level(level) + } + + /// Returns one past the last cell at this level. + pub fn end(level: i32) -> Self { + Self::from_face(5).child_end_at_level(level) + } + + /// Returns the first child of this cell (at level + 1). + pub fn child_begin(&self) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(!self.is_leaf()); + let old_lsb = self.lsb(); + Self(self.0 - old_lsb + (old_lsb >> 2)) + } + + /// Returns one past the last child of this cell (at level + 1). + pub fn child_end(&self) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(!self.is_leaf()); + let old_lsb = self.lsb(); + Self(self.0 + old_lsb + (old_lsb >> 2)) + } + + /// Returns the first descendant at the given level. + pub fn child_begin_at_level(&self, level: i32) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(level >= self.level()); + debug_assert!(level <= MAX_CELL_LEVEL); + Self(self.0 - self.lsb() + Self::lsb_for_level(level)) + } + + /// Returns one past the last descendant at the given level. + pub fn child_end_at_level(&self, level: i32) -> Self { + debug_assert!(self.is_valid()); + debug_assert!(level >= self.level()); + debug_assert!(level <= MAX_CELL_LEVEL); + Self(self.0 + self.lsb() + Self::lsb_for_level(level)) + } + + /// Converts to (face, i, j) coordinates with optional orientation. + pub fn to_face_ij_orientation(&self) -> (i32, i32, i32, i32) { + let face = self.face(); + let mut bits = (face & SWAP_MASK) as u64; + let mut i = 0i32; + let mut j = 0i32; + + for k in (0..8).rev() { + let nbits = if k == 7 { + MAX_CELL_LEVEL as usize - 7 * LOOKUP_BITS + } else { + LOOKUP_BITS + }; + bits += ((self.0 >> (k * 2 * LOOKUP_BITS + 1)) & ((1 << (2 * nbits)) - 1)) << 2; + bits = LOOKUP.ij[bits as usize] as u64; + i += ((bits >> (LOOKUP_BITS + 2)) as i32) << (k * LOOKUP_BITS); + j += (((bits >> 2) & ((1 << LOOKUP_BITS) - 1)) as i32) << (k * LOOKUP_BITS); + bits &= (SWAP_MASK | INVERT_MASK) as u64; + } + + // Compute orientation based on LSB pattern + let orientation = if self.lsb() & 0x1111111111111110u64 != 0 { + (bits ^ SWAP_MASK as u64) as i32 + } else { + bits as i32 + }; + + (face, i, j, orientation) + } + + /// Returns the (face, si, ti) coordinates of the cell center. + pub fn get_center_si_ti(&self) -> (i32, u32, u32) { + let (face, i, j, _) = self.to_face_ij_orientation(); + + let delta = if self.is_leaf() { + 1 + } else if (i ^ ((self.0 >> 2) as i32)) & 1 != 0 { + 2 + } else { + 0 + }; + + let si = (2 * i + delta) as u32; + let ti = (2 * j + delta) as u32; + + (face, si, ti) + } + + /// Returns the center of this cell as a point. + pub fn to_point(&self) -> [f64; 3] { + normalize(&self.to_point_raw()) + } + + /// Returns the center of this cell as an unnormalized point. + pub fn to_point_raw(&self) -> [f64; 3] { + let (face, si, ti) = self.get_center_si_ti(); + face_si_ti_to_xyz(face, si, ti) + } + + /// Returns a debug string representation "f/dd...d". + pub fn to_debug_string(&self) -> String { + if !self.is_valid() { + return format!("Invalid: {:016x}", self.0); + } + let mut out = format!("{}/", self.face()); + for level in 1..=self.level() { + out.push(char::from_digit(self.child_position_at_level(level) as u32, 10).unwrap()); + } + out + } + + /// Returns a hex token representation. + pub fn to_token(&self) -> String { + if self.0 == 0 { + return "X".to_string(); + } + let num_zero_digits = self.0.trailing_zeros() / 4; + let id_shifted = self.0 >> (4 * num_zero_digits); + let width = 16 - num_zero_digits as usize; + format!("{:0>width$x}", id_shifted, width = width) + } + + /// Creates a cell ID from a hex token. + pub fn from_token(token: &str) -> Self { + if token.len() > 16 { + return Self::none(); + } + if token == "X" { + return Self(0); + } + + let mut id = 0u64; + for (i, c) in token.chars().enumerate() { + let d = match c { + '0'..='9' => c as u64 - '0' as u64, + 'a'..='f' => c as u64 - 'a' as u64 + 10, + 'A'..='F' => c as u64 - 'A' as u64 + 10, + _ => return Self::none(), + }; + id |= d << (60 - i * 4); + } + Self(id) + } + /// Returns the center of this cell in (u,v) coordinates. + pub fn get_center_uv(&self) -> [f64; 2] { + let (_, si, ti) = self.get_center_si_ti(); + [st_to_uv(si_ti_to_st(si)), st_to_uv(si_ti_to_st(ti))] + } + /// Appends the cells at `level` that share the closest vertex to this cell. + /// + /// Normally returns 4 cells, but returns only 3 if the vertex is one of the 8 cube vertices + /// (where only 3 faces meet). + /// + /// REQUIRES: level < self.level() (so we can determine which vertex is closest) + pub fn append_vertex_neighbors(&self, level: i32, output: &mut Vec) { + // "level" must be strictly less than this cell's level so that we can + // determine which vertex this cell is closest to. + debug_assert!(level < self.level()); + + let (face, i, j, _) = self.to_face_ij_orientation(); + + // Determine the i- and j-offsets to the closest neighboring cell in each direction. This + // involves looking at the next bit of "i" and "j" to determine which quadrant of + // self.parent(level) this cell lies in. + let halfsize = Self::size_ij_for_level(level + 1); + let size = halfsize << 1; + + let (ioffset, isame) = if (i & halfsize) != 0 { + // Cell is in the right half, so nearest vertex is to the right + (size, (i + size) < MAX_SIZE) + } else { + // Cell is in the left half, so nearest vertex is to the left + (-size, (i - size) >= 0) + }; + + let (joffset, jsame) = if (j & halfsize) != 0 { + (size, (j + size) < MAX_SIZE) + } else { + (-size, (j - size) >= 0) + }; + + // Add the parent cell (always present) + output.push(self.parent(level)); + + // Add the neighbor in the i-direction + output.push(Self::from_face_ij_same(face, i + ioffset, j, isame).parent(level)); + + // Add the neighbor in the j-direction + output.push(Self::from_face_ij_same(face, i, j + joffset, jsame).parent(level)); + + // Add the diagonal neighbor if both i- and j-neighbors are on the same face. If both are + // on different faces, this vertex is one of the 8 cube vertices and only has 3 neighbors + // (not 4). + if isame || jsame { + output.push( + Self::from_face_ij_same(face, i + ioffset, j + joffset, isame && jsame) + .parent(level), + ); + } + } + + /// Helper: creates a leaf cell from (face, i, j), wrapping to adjacent face if needed when + /// `same_face` is false. + fn from_face_ij_same(face: i32, i: i32, j: i32, same_face: bool) -> Self { + if same_face { + Self::from_face_ij(face, i, j) + } else { + Self::from_face_ij_wrap(face, i, j) + } + } + + /// Creates a leaf cell from (face, i, j), wrapping to the appropriate + /// adjacent face if (i, j) is outside [0, MAX_SIZE). + fn from_face_ij_wrap(face: i32, i: i32, j: i32) -> Self { + // Clamp i and j to just beyond the valid range to prevent overflow. + let i = i.clamp(-1, MAX_SIZE); + let j = j.clamp(-1, MAX_SIZE); + + // Convert to (u, v) using linear projection and clamp to barely outside the face to avoid + // numerical issues during reprojection. + const SCALE: f64 = 1.0 / MAX_SIZE as f64; + const LIMIT: f64 = 1.0 + f64::EPSILON; + + // The arithmetic below is designed to avoid 32-bit integer overflows. + let u = (SCALE * (2.0 * (i - MAX_SIZE / 2) as f64 + 1.0)).clamp(-LIMIT, LIMIT); + let v = (SCALE * (2.0 * (j - MAX_SIZE / 2) as f64 + 1.0)).clamp(-LIMIT, LIMIT); + + // Project onto sphere and find the containing face. + let p = face_uv_to_xyz(face, u, v); + let (new_face, new_u, new_v) = xyz_to_face_uv(&p); + + // Convert back to (i, j) on the new face using linear projection inverse. + let new_i = st_to_ij(0.5 * (new_u + 1.0)); + let new_j = st_to_ij(0.5 * (new_v + 1.0)); + + Self::from_face_ij(new_face, new_i, new_j) + } + + /// Returns the largest cell whose range_min is <= self and range_max limit. + /// + /// Used to iterate through a range of leaf cells efficiently: + /// ```ignore + /// let mut id = start.maximum_tile(limit); + /// while id != limit { + /// // process id + /// id = id.next().maximum_tile(limit); + /// } + /// ``` + pub fn maximum_tile(&self, limit: S2CellId) -> S2CellId { + let mut id = *self; + let start = id.range_min(); + + // If we're already at or past the limit, return limit. + if start >= limit.range_min() { + return limit; + } + + if id.range_max() >= limit { + // The cell is too large. Shrink it. Note that when generating coverings of S2CellId + // ranges, this loop usually executes only once. Also because id.range_min() < + // limit.range_min(), we will always exit the loop by the time we reach a leaf cell. + loop { + id = id.child(0); + if id.range_max() < limit { + break; + } + } + return id; + } + + // The cell may be too small. Grow it if necessary. Note that generally this loop only + // iterates once. + while !id.is_face() { + let parent = id.immediate_parent(); + if parent.range_min() != start || parent.range_max() >= limit { + break; + } + id = parent; + } + id + } + + /// Returns the level of the lowest common ancestor of this cell and other, or -1 if they are + /// on different faces. + pub fn common_ancestor_level(&self, other: S2CellId) -> i8 { + // The max() below is necessary for the case where one S2CellId is a descendant of the + // other. + let bits = (self.0 ^ other.0).max(self.lsb().max(other.lsb())); + + // Compute the position of the most significant bit, and then map the bit position as + // follows: + // {0} -> 30, {1,2} -> 29, {3,4} -> 28, ... , {59,60} -> 0, {61,62,63} -> -1. + let bit_width = 64 - bits.leading_zeros(); + ((61i32 - bit_width as i32).max(-1) >> 1) as i8 + } +} + +impl std::fmt::Display for S2CellId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.to_debug_string()) + } +} diff --git a/src/spatial/s2coords.rs b/src/spatial/s2coords.rs new file mode 100644 index 0000000000..939bbb8f56 --- /dev/null +++ b/src/spatial/s2coords.rs @@ -0,0 +1,301 @@ +//! S2 coordinate system transformations. +//! +//! Provides conversions between the various S2 coordinate systems: +//! (face, i, j), (face, s, t), (face, u, v), and (x, y, z). +//! +//! s2coords.h, s2coords.cc and s2coords_internal.h. + +/// Maximum cell level (30 levels of subdivision). +pub const MAX_CELL_LEVEL: i32 = 30; + +/// Maximum index of valid leaf cell coordinates. +pub const LIMIT_IJ: i32 = 1 << MAX_CELL_LEVEL; + +/// Maximum value of si/ti coordinates. +pub const MAX_SI_TI: u32 = 1 << (MAX_CELL_LEVEL + 1); + +/// Number of cube faces. +pub const NUM_FACES: i32 = 6; + +/// Convert an s- or t-value to the corresponding u- or v-value. This is a non-linear +/// transformation from [0,1] to [-1,1] that attempts to make the cell sizes more uniform. +#[inline] +pub fn st_to_uv(s: f64) -> f64 { + if s >= 0.5 { + (1.0 / 3.0) * (4.0 * s * s - 1.0) + } else { + (1.0 / 3.0) * (1.0 - 4.0 * (1.0 - s) * (1.0 - s)) + } +} + +/// The inverse of the st_to_uv transformation. Note that it is not always true that +/// uv_to_st(st_to_uv(x)) == x due to numerical errors. +#[inline] +pub fn uv_to_st(u: f64) -> f64 { + if u >= 0.0 { + 0.5 * (1.0 + 3.0 * u).sqrt() + } else { + 1.0 - 0.5 * (1.0 - 3.0 * u).sqrt() + } +} + +/// Convert the i- or j-index of a leaf cell to the minimum corresponding s- or t-value contained +/// by that cell. The argument must be in the range [0..2**30], i.e. up to one position beyond the +/// normal range of valid leaf cell indices. +#[inline] +pub fn ij_to_st_min(i: i32) -> f64 { + debug_assert!((0..=LIMIT_IJ).contains(&i)); + (1.0 / LIMIT_IJ as f64) * i as f64 +} + +/// Return the i- or j-index of the leaf cell containing the given s- or t-value. If the argument +/// is outside the range spanned by valid leaf cell indices, return the index of the closest valid +/// leaf cell (i.e., return values are clamped to the range of valid leaf cell indices). +#[inline] +pub fn st_to_ij(s: f64) -> i32 { + let ij = (LIMIT_IJ as f64 * s) as i32; + ij.clamp(0, LIMIT_IJ - 1) +} + +/// Convert si- or ti-coordinate to s- or t-value. +#[inline] +pub fn si_ti_to_st(si: u32) -> f64 { + debug_assert!(si <= MAX_SI_TI); + (1.0 / MAX_SI_TI as f64) * si as f64 +} + +/// Convert (face, u, v) coordinates to a direction vector (not necessarily unit length). +#[inline] +pub fn face_uv_to_xyz(face: i32, u: f64, v: f64) -> [f64; 3] { + match face { + 0 => [1.0, u, v], + 1 => [-u, 1.0, v], + 2 => [-u, -v, 1.0], + 3 => [-1.0, -v, -u], + 4 => [v, -1.0, -u], + _ => [v, u, -1.0], + } +} + +/// Return the face containing the given direction vector. (For points on the boundary between +/// faces, the result is arbitrary but repeatable.) +#[inline] +pub fn get_face(p: &[f64; 3]) -> i32 { + let ax = p[0].abs(); + let ay = p[1].abs(); + let az = p[2].abs(); + + // Strict > mirrors LargestAbsComponent() in vector.h. + let face = if ax > ay { + if ax > az { + 0 + } else { + 2 + } + } else if ay > az { + 1 + } else { + 2 + }; + + if p[face as usize] < 0.0 { + face + 3 + } else { + face + } +} + +/// Convert a direction vector (not necessarily unit length) to (face, u, v) coordinates. +#[inline] +pub fn xyz_to_face_uv(p: &[f64; 3]) -> (i32, f64, f64) { + let face = get_face(p); + let (u, v) = valid_face_xyz_to_uv(face, p); + (face, u, v) +} + +/// Given a *valid* face for the given point p (meaning that dot product of p with the face normal +/// is positive), return the corresponding u and v values (which may lie outside the range [-1,1]). +#[inline] +pub fn valid_face_xyz_to_uv(face: i32, p: &[f64; 3]) -> (f64, f64) { + match face { + 0 => (p[1] / p[0], p[2] / p[0]), + 1 => (-p[0] / p[1], p[2] / p[1]), + 2 => (-p[0] / p[2], -p[1] / p[2]), + 3 => (p[2] / p[0], p[1] / p[0]), + 4 => (p[2] / p[1], -p[0] / p[1]), + _ => (-p[1] / p[2], -p[0] / p[2]), + } +} + +/// Convert (face, si, ti) to XYZ direction vector. +#[inline] +pub fn face_si_ti_to_xyz(face: i32, si: u32, ti: u32) -> [f64; 3] { + let u = st_to_uv(si_ti_to_st(si)); + let v = st_to_uv(si_ti_to_st(ti)); + face_uv_to_xyz(face, u, v) +} + +/// IJ-to-position mapping (inverse of pos_to_ij). +pub const IJ_TO_POS: [[i32; 4]; 4] = [[0, 1, 3, 2], [0, 3, 1, 2], [2, 3, 1, 0], [2, 1, 3, 0]]; + +/// Converts (i,j) child indices to traversal position. +#[inline] +pub fn ij_to_pos(orientation: i32, i: i32, j: i32) -> i32 { + IJ_TO_POS[orientation as usize][((i << 1) + j) as usize] +} + +/// Face axis definitions for the S2 cube projection. +const FACE_U_AXES: [[f64; 3]; 6] = [ + [0.0, 1.0, 0.0], + [-1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], +]; + +const FACE_V_AXES: [[f64; 3]; 6] = [ + [0.0, 0.0, 1.0], + [0.0, 0.0, 1.0], + [0.0, -1.0, 0.0], + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, 0.0, 0.0], +]; + +/// Together with INVERT_MASK, defines a cell orientation. If true, the canonical traversal order +/// is flipped around the diagonal (i.e. i and j are swapped with each other). +pub const SWAP_MASK: i32 = 0x01; + +/// Together with SWAP_MASK, defines a cell orientation. If true, the traversal order is rotated by +/// 180 degrees (i.e. the bits of i and j are inverted, or equivalently, the axis directions are +/// reversed). +pub const INVERT_MASK: i32 = 0x02; + +/// Position-to-IJ mapping for each orientation. +pub const POS_TO_IJ: [[i32; 4]; 4] = [ + [0, 1, 3, 2], // canonical order + [0, 2, 3, 1], // axes swapped + [3, 2, 0, 1], // bits inverted + [3, 1, 0, 2], // swapped & inverted +]; + +/// Orientation modifier for each child position. +pub const POS_TO_ORIENTATION: [i32; 4] = [SWAP_MASK, 0, 0, INVERT_MASK | SWAP_MASK]; + +/// Returns the (i,j) indices for child at position `pos` with given orientation. +#[inline] +pub fn pos_to_ij(orientation: i32, pos: i32) -> i32 { + POS_TO_IJ[orientation as usize][pos as usize] +} + +/// Returns the orientation delta for child at position `pos`. +#[inline] +pub fn pos_to_orientation(pos: i32) -> i32 { + POS_TO_ORIENTATION[pos as usize] +} + +/// Return the right-handed normal (not necessarily unit length) for an edge in the direction of +/// the positive v-axis at the given u-value on the given face. (This vector is perpendicular to +/// the plane through the sphere origin that contains the given edge.) +#[inline] +pub fn get_u_norm(face: i32, u: f64) -> [f64; 3] { + match face { + 0 => [u, -1.0, 0.0], + 1 => [1.0, u, 0.0], + 2 => [1.0, 0.0, u], + 3 => [-u, 0.0, 1.0], + 4 => [0.0, -u, 1.0], + _ => [0.0, -1.0, -u], + } +} + +/// Return the right-handed normal (not necessarily unit length) for an edge in the direction of +/// the positive u-axis at the given v-value on the given face. +#[inline] +pub fn get_v_norm(face: i32, v: f64) -> [f64; 3] { + match face { + 0 => [-v, 0.0, 1.0], + 1 => [0.0, -v, 1.0], + 2 => [0.0, -1.0, -v], + 3 => [v, -1.0, 0.0], + 4 => [1.0, v, 0.0], + _ => [1.0, 0.0, v], + } +} + +/// If the dot product of p with the given face normal is positive, set the corresponding u and v +/// values (which may lie outside the range [-1,1]) and return true. Otherwise return false. +#[inline] +pub fn face_xyz_to_uv(face: i32, p: &[f64; 3]) -> Option<[f64; 2]> { + match face { + 0 => { + if p[0] > 0.0 { + Some([p[1] / p[0], p[2] / p[0]]) + } else { + None + } + } + 1 => { + if p[1] > 0.0 { + Some([-p[0] / p[1], p[2] / p[1]]) + } else { + None + } + } + 2 => { + if p[2] > 0.0 { + Some([-p[0] / p[2], -p[1] / p[2]]) + } else { + None + } + } + 3 => { + if p[0] < 0.0 { + Some([p[2] / p[0], p[1] / p[0]]) + } else { + None + } + } + 4 => { + if p[1] < 0.0 { + Some([p[2] / p[1], -p[0] / p[1]]) + } else { + None + } + } + _ => { + if p[2] < 0.0 { + Some([-p[1] / p[2], -p[0] / p[2]]) + } else { + None + } + } + } +} + +/// Transform the given point P to the (u,v,w) coordinate frame of the given face (where the w-axis +/// represents the face normal). +#[inline] +pub fn face_xyz_to_uvw(face: i32, p: &[f64; 3]) -> [f64; 3] { + match face { + 0 => [p[1], p[2], p[0]], + 1 => [-p[0], p[2], p[1]], + 2 => [-p[0], -p[1], p[2]], + 3 => [-p[2], -p[1], -p[0]], + 4 => [-p[2], p[0], -p[1]], + _ => [p[1], p[0], -p[2]], + } +} + +/// Return the u-axis for the given face. +#[inline] +pub fn get_u_axis(face: i32) -> [f64; 3] { + FACE_U_AXES[face as usize] +} + +/// Return the v-axis for the given face. +#[inline] +pub fn get_v_axis(face: i32) -> [f64; 3] { + FACE_V_AXES[face as usize] +} diff --git a/src/spatial/s2edge_clipping.rs b/src/spatial/s2edge_clipping.rs new file mode 100644 index 0000000000..be99748d24 --- /dev/null +++ b/src/spatial/s2edge_clipping.rs @@ -0,0 +1,342 @@ +//! Clipping geodesic edges to S2 cube faces and 2D edges against rectangles. +use super::crossings::robust_cross; +use super::math::{cross, dot, normalize}; +use super::r1interval::R1Interval; +use super::r2rect::R2Rect; +use super::s2coords::face_xyz_to_uvw; + +/// Maximum error in vertices returned by face clipping functions. +pub const FACE_CLIP_ERROR_UV_COORD: f64 = 9.0 * std::f64::consts::FRAC_1_SQRT_2 * f64::EPSILON; + +/// Maximum error in edge clipping functions. +pub const EDGE_CLIP_ERROR_UV_COORD: f64 = 2.25 * f64::EPSILON; + +/// Maximum error in IntersectsRect. +pub const INTERSECTS_RECT_ERROR_UV_DIST: f64 = 3.0 * std::f64::consts::SQRT_2 * f64::EPSILON; + +/// Represents an edge AB clipped to an S2 cube face. +#[derive(Clone, Copy, Debug)] +pub struct FaceSegment { + /// The cube face (0-5). + pub face: i32, + /// Endpoint A in (u,v) coordinates. + pub a: [f64; 2], + /// Endpoint B in (u,v) coordinates. + pub b: [f64; 2], +} + +/// Clips edge AB to the given cube face, returning the (u,v) coordinates. Returns None if the edge +/// does not intersect the face. +pub fn clip_to_face(a: &[f64; 3], b: &[f64; 3], face: i32) -> Option<([f64; 2], [f64; 2])> { + clip_to_padded_face(a, b, face, 0.0) +} + +/// Clips edge AB to a padded cube face. The face is expanded by `padding` in each direction. +pub fn clip_to_padded_face( + a_xyz: &[f64; 3], + b_xyz: &[f64; 3], + face: i32, + padding: f64, +) -> Option<([f64; 2], [f64; 2])> { + debug_assert!(padding >= 0.0); + + // Fast path: both endpoints on the same face. + let a_face = super::s2coords::get_face(a_xyz); + let b_face = super::s2coords::get_face(b_xyz); + + if a_face == face && b_face == face { + let (_, a_u, a_v) = super::s2coords::xyz_to_face_uv(a_xyz); + let (_, b_u, b_v) = super::s2coords::xyz_to_face_uv(b_xyz); + return Some(([a_u, a_v], [b_u, b_v])); + } + + // Convert everything into the (u,v,w) coordinates of the given face. Note + // that the cross product *must* be computed in the original (x,y,z) + // coordinate system because RobustCrossProd (unlike the mathematical cross + // product) can produce different results in different coordinate systems + // when one argument is a linear multiple of the other, due to the use of + // symbolic perturbations. + let n = face_xyz_to_uvw(face, &robust_cross(a_xyz, b_xyz)); + let a = face_xyz_to_uvw(face, a_xyz); + let b = face_xyz_to_uvw(face, b_xyz); + + // Scale u and v components to handle padding. + let scale_uv = 1.0 + padding; + let scaled_n = [scale_uv * n[0], scale_uv * n[1], n[2]]; + + // Check if the line intersects the (padded) face. + if !intersects_face(&scaled_n) { + return None; + } + + let n_norm = normalize(&n); + let a_tangent = cross(&n_norm, &a); + let b_tangent = cross(&b, &n_norm); + + // Clip both endpoints. + let neg_scaled_n = [-scaled_n[0], -scaled_n[1], -scaled_n[2]]; + let (a_uv, a_score) = clip_destination(&b, &a, &neg_scaled_n, &b_tangent, &a_tangent, scale_uv); + let (b_uv, b_score) = clip_destination(&a, &b, &scaled_n, &a_tangent, &b_tangent, scale_uv); + + if a_score + b_score < 3 { + Some((a_uv, b_uv)) + } else { + None + } +} + +// The three functions below all compare a sum (u + v) to a third value w. +// They are implemented in such a way that they produce an exact result even +// though all calculations are done with ordinary floating-point operations. +// Here are the principles on which these functions are based: +// +// A. If u + v < w in floating-point, then u + v < w in exact arithmetic. +// +// B. If u + v < w in exact arithmetic, then at least one of the following +// expressions is true in floating-point: +// u + v < w +// u < w - v +// v < w - u +// +// Proof: By rearranging terms and substituting ">" for "<", we can assume +// that all values are non-negative. Now clearly "w" is not the smallest +// value, so assume WLOG that "u" is the smallest. We want to show that +// u < w - v in floating-point. If v >= w/2, the calculation of w - v is +// exact since the result is smaller in magnitude than either input value, +// so the result holds. Otherwise we have u <= v < w/2 and w - v >= w/2 +// (even in floating point), so the result also holds. + +fn intersects_face(n: &[f64; 3]) -> bool { + let u = n[0].abs(); + let v = n[1].abs(); + let w = n[2].abs(); + // |Nu| + |Nv| >= |Nw| + (v >= w - u) && (u >= w - v) +} + +/// Returns true if the line intersects opposite edges of the face. +fn intersects_opposite_edges(n: &[f64; 3]) -> bool { + let u = n[0].abs(); + let v = n[1].abs(); + let w = n[2].abs(); + + // ||Nu| - |Nv|| >= |Nw| + if (u - v).abs() != w { + return (u - v).abs() >= w; + } + if u >= v { + u - w >= v + } else { + v - w >= u + } +} + +/// Returns the exit axis (0 for u, 1 for v). +fn get_exit_axis(n: &[f64; 3]) -> i32 { + debug_assert!(intersects_face(n)); + if intersects_opposite_edges(n) { + if n[0].abs() >= n[1].abs() { + 1 + } else { + 0 + } + } else { + let sign = |x: f64| x.is_sign_negative() as i32; + if (sign(n[0]) ^ sign(n[1]) ^ sign(n[2])) == 0 { + 1 + } else { + 0 + } + } +} + +/// Returns the exit point on the face for the given line. +fn get_exit_point(n: &[f64; 3], axis: i32) -> [f64; 2] { + if axis == 0 { + let u = if n[1] > 0.0 { 1.0 } else { -1.0 }; + [u, (-u * n[0] - n[2]) / n[1]] + } else { + let v = if n[0] < 0.0 { 1.0 } else { -1.0 }; + [(-v * n[1] - n[2]) / n[0], v] + } +} +/// Clips a destination point. +fn clip_destination( + a: &[f64; 3], + b: &[f64; 3], + scaled_n: &[f64; 3], + a_tangent: &[f64; 3], + b_tangent: &[f64; 3], + scale_uv: f64, +) -> ([f64; 2], i32) { + const MAX_SAFE_UV_COORD: f64 = 1.0 - FACE_CLIP_ERROR_UV_COORD; + + // Optimization: if B is within the safe region, use it directly. + if b[2] > 0.0 { + let uv = [b[0] / b[2], b[1] / b[2]]; + if uv[0].abs().max(uv[1].abs()) <= MAX_SAFE_UV_COORD { + return (uv, 0); + } + } + + // Find where line AB exits the face. + let exit_axis = get_exit_axis(scaled_n); + let exit_pt = get_exit_point(scaled_n, exit_axis); + let mut uv = [scale_uv * exit_pt[0], scale_uv * exit_pt[1]]; + let p = [uv[0], uv[1], 1.0]; + + // Determine if exit point is in the segment. + let mut score = 0; + let pa = [p[0] - a[0], p[1] - a[1], p[2] - a[2]]; + if dot(&pa, a_tangent) < 0.0 { + score = 2; // On wrong side of A + } else { + let pb = [p[0] - b[0], p[1] - b[1], p[2] - b[2]]; + if dot(&pb, b_tangent) < 0.0 { + score = 1; // On wrong side of B + } + } + + if score > 0 { + if b[2] <= 0.0 { + score = 3; // B cannot be projected + } else { + uv = [b[0] / b[2], b[1] / b[2]]; + } + } + + (uv, score) +} + +/// Returns true if edge AB intersects the rectangle. +pub fn intersects_rect(a: [f64; 2], b: [f64; 2], rect: &R2Rect) -> bool { + // Check bounding box intersection first. + let bound = R2Rect::from_point_pair(a, b); + if !rect.intersects(&bound) { + return false; + } + + // Check if all four vertices are on the same side of the line. + // n = ortho(b - a) = [-(b[1]-a[1]), b[0]-a[0]] + let n = [-(b[1] - a[1]), b[0] - a[0]]; + let i = if n[0] >= 0.0 { 1 } else { 0 }; + let j = if n[1] >= 0.0 { 1 } else { 0 }; + + let v_max = rect.get_vertex_ij(i, j); + let v_min = rect.get_vertex_ij(1 - i, 1 - j); + + // dot(n, v - a) + let max = n[0] * (v_max[0] - a[0]) + n[1] * (v_max[1] - a[1]); + let min = n[0] * (v_min[0] - a[0]) + n[1] * (v_min[1] - a[1]); + + (max >= 0.0) && (min <= 0.0) +} + +/// Clips edge AB to the rectangle, returning the clipped endpoints. Returns None if there is no +/// intersection. +pub fn clip_edge(a: [f64; 2], b: [f64; 2], clip: &R2Rect) -> Option<([f64; 2], [f64; 2])> { + let mut bound = R2Rect::from_point_pair(a, b); + if clip_edge_bound(a, b, clip, &mut bound) { + let ai = if a[0] > b[0] { 1 } else { 0 }; + let aj = if a[1] > b[1] { 1 } else { 0 }; + Some(( + bound.get_vertex_ij(ai, aj), + bound.get_vertex_ij(1 - ai, 1 - aj), + )) + } else { + None + } +} + +/// Updates the bounding rectangle of edge AB clipped to the given clip region. Returns false if AB +/// does not intersect the clip region. +pub fn clip_edge_bound(a: [f64; 2], b: [f64; 2], clip: &R2Rect, bound: &mut R2Rect) -> bool { + let diag = ((a[0] > b[0]) != (a[1] > b[1])) as usize; + + clip_bound_axis( + a[0], + b[0], + &mut bound.x, + a[1], + b[1], + &mut bound.y, + diag, + &clip[0], + ) && clip_bound_axis( + a[1], + b[1], + &mut bound.y, + a[0], + b[0], + &mut bound.x, + diag, + &clip[1], + ) +} + +/// Helper for clip_edge_bound: clips along one axis. +#[allow(clippy::too_many_arguments)] +fn clip_bound_axis( + a0: f64, + b0: f64, + bound0: &mut R1Interval, + a1: f64, + b1: f64, + bound1: &mut R1Interval, + diag: usize, + clip0: &R1Interval, +) -> bool { + if bound0.lo() < clip0.lo() { + if bound0.hi() < clip0.lo() { + return false; + } + *bound0 = R1Interval::new(clip0.lo(), bound0.hi()); + if !update_endpoint(bound1, diag, interpolate(clip0.lo(), a0, b0, a1, b1)) { + return false; + } + } + if bound0.hi() > clip0.hi() { + if bound0.lo() > clip0.hi() { + return false; + } + *bound0 = R1Interval::new(bound0.lo(), clip0.hi()); + if !update_endpoint(bound1, 1 - diag, interpolate(clip0.hi(), a0, b0, a1, b1)) { + return false; + } + } + true +} + +/// Updates an endpoint of an interval. +fn update_endpoint(bound: &mut R1Interval, end: usize, value: f64) -> bool { + if end == 0 { + if bound.hi() < value { + return false; + } + if bound.lo() < value { + *bound = R1Interval::new(value, bound.hi()); + } + } else { + if bound.lo() > value { + return false; + } + if bound.hi() > value { + *bound = R1Interval::new(bound.lo(), value); + } + } + true +} + +/// Interpolates a value along the line from (a0,a1) to (b0,b1). +fn interpolate(x: f64, a: f64, b: f64, a1: f64, b1: f64) -> f64 { + if a == b { + debug_assert!(x == a && a1 == b1); + return a1; + } + // Interpolate from the closer endpoint for better accuracy. + if (a - x).abs() <= (b - x).abs() { + a1 + (b1 - a1) * ((x - a) / (b - a)) + } else { + b1 + (a1 - b1) * ((x - b) / (a - b)) + } +} diff --git a/src/spatial/s2edge_distances.rs b/src/spatial/s2edge_distances.rs new file mode 100644 index 0000000000..1da8887ec1 --- /dev/null +++ b/src/spatial/s2edge_distances.rs @@ -0,0 +1,227 @@ +//! Functions for computing the distance from a point to an edge, and between pairs of edges, on +//! the unit sphere. +//! +//! Port of s2edge_distances.h and s2edge_distances.cc (Eric Veach, Google). + +use super::crossings::{robust_cross, S2EdgeCrosser}; +use super::math::{cross, dot, norm2, sub}; +use super::s1chord_angle::S1ChordAngle; + +/// Returns the minimum distance from x to any point on the edge ab as an S1ChordAngle. +/// All arguments should be unit length. The result is very accurate for small distances but may +/// have some numerical error if the distance is large (approximately pi/2 or greater). The case +/// a == b is handled correctly. +/// +/// Port of S2::GetDistance. +pub fn get_distance(x: &[f64; 3], a: &[f64; 3], b: &[f64; 3]) -> S1ChordAngle { + let mut min_dist = S1ChordAngle::from_length2(0.0); + always_update_min_distance::(x, a, b, &mut min_dist); + min_dist +} + +/// If the distance from x to the edge ab is less than min_dist, updates min_dist and returns +/// true. Otherwise returns false. The case a == b is handled correctly. +/// +/// Port of S2::UpdateMinDistance. +pub fn update_min_distance( + x: &[f64; 3], + a: &[f64; 3], + b: &[f64; 3], + min_dist: &mut S1ChordAngle, +) -> bool { + always_update_min_distance::(x, a, b, min_dist) +} + +/// Returns true if the distance from x to the edge ab is less than limit. +/// +/// Port of S2::IsDistanceLess. +pub fn is_distance_less(x: &[f64; 3], a: &[f64; 3], b: &[f64; 3], limit: S1ChordAngle) -> bool { + let mut limit = limit; + update_min_distance(x, a, b, &mut limit) +} + +/// If the minimum distance from x to ab is attained at an interior point of ab (i.e., not an +/// endpoint), and that distance is less than min_dist, then updates min_dist and returns true. +/// Otherwise returns false. +/// +/// Port of S2::UpdateMinInteriorDistance. +pub fn update_min_interior_distance( + x: &[f64; 3], + a: &[f64; 3], + b: &[f64; 3], + min_dist: &mut S1ChordAngle, +) -> bool { + let xa2 = norm2(&sub(x, a)); + let xb2 = norm2(&sub(x, b)); + always_update_min_interior_distance::(x, a, b, xa2, xb2, min_dist) +} + +/// Returns the maximum error in the result of update_min_distance (and associated functions), +/// assuming that all input points are normalized to within the bounds guaranteed by normalize(). +/// +/// Port of S2::GetUpdateMinDistanceMaxError. +pub fn get_update_min_distance_max_error(dist: S1ChordAngle) -> f64 { + f64::max( + get_update_min_interior_distance_max_error(dist), + dist.get_s2point_constructor_max_error(), + ) +} + +// Port of AlwaysUpdateMinInteriorDistance. +fn always_update_min_interior_distance( + x: &[f64; 3], + a: &[f64; 3], + b: &[f64; 3], + xa2: f64, + xb2: f64, + min_dist: &mut S1ChordAngle, +) -> bool { + // The closest point on AB could either be one of the two vertices (the + // "vertex case") or in the interior (the "interior case"). Let C = A x B. + // If X is in the spherical wedge extending from A to B around the axis + // through C, then we are in the interior case. Otherwise we are in the + // vertex case. + // + // Check whether we might be in the interior case. For this to be true, XAB + // and XBA must both be acute angles. Checking this condition exactly is + // expensive, so instead we consider the planar triangle ABX (which passes + // through the sphere's interior). The planar angles XAB and XBA are always + // less than the corresponding spherical angles, so if we are in the + // interior case then both of these angles must be acute. + // + // We check this by computing the squared edge lengths of the planar + // triangle ABX, and testing whether angles XAB and XBA are both acute using + // the law of cosines: + // + // | XA^2 - XB^2 | < AB^2 (*) + // + // There are two sources of error in the expression above (*). + let ab2 = norm2(&sub(a, b)); + let max_error = 4.75 * f64::EPSILON * (xa2 + xb2 + ab2) + 8.0 * f64::EPSILON * f64::EPSILON; + if (xa2 - xb2).abs() >= ab2 + max_error { + return false; + } + + // The minimum distance might be to a point on the edge interior. Let R + // be closest point to X that lies on the great circle through AB. Rather + // than computing the geodesic distance along the surface of the sphere, + // instead we compute the "chord length" through the sphere's interior. + let c = robust_cross(a, b); + let c2 = norm2(&c); + let x_dot_c = dot(x, &c); + let x_dot_c2 = x_dot_c * x_dot_c; + if !ALWAYS_UPDATE && x_dot_c2 > c2 * min_dist.length2() { + // The closest point on the great circle AB is too far away. + return false; + } + // Otherwise we do the exact, more expensive test for the interior case. + let cx = cross(&c, x); + if dot(&sub(a, x), &cx) >= 0.0 || dot(&sub(b, x), &cx) <= 0.0 { + return false; + } + // Compute the squared chord length XR^2 = XQ^2 + QR^2. + let qr = 1.0 - (norm2(&cx) / c2).sqrt(); + let dist2 = (x_dot_c2 / c2) + (qr * qr); + if !ALWAYS_UPDATE && dist2 >= min_dist.length2() { + return false; + } + *min_dist = S1ChordAngle::from_length2(dist2); + true +} + +// Port of AlwaysUpdateMinDistance. +fn always_update_min_distance( + x: &[f64; 3], + a: &[f64; 3], + b: &[f64; 3], + min_dist: &mut S1ChordAngle, +) -> bool { + let xa2 = norm2(&sub(x, a)); + let xb2 = norm2(&sub(x, b)); + if always_update_min_interior_distance::(x, a, b, xa2, xb2, min_dist) { + return true; // Minimum distance is attained along the edge interior. + } + // Otherwise the minimum distance is to one of the endpoints. + let dist2 = f64::min(xa2, xb2); + if !ALWAYS_UPDATE && dist2 >= min_dist.length2() { + return false; + } + *min_dist = S1ChordAngle::from_length2(dist2); + true +} + +/// Like update_min_distance(), but computes the minimum distance between the +/// given pair of edges. (If the two edges cross, the distance is zero.) +/// The cases a0 == a1 and b0 == b1 are handled correctly. +/// +/// Port of S2::UpdateEdgePairMinDistance in s2edge_distances.cc. +pub fn update_edge_pair_min_distance( + a0: &[f64; 3], + a1: &[f64; 3], + b0: &[f64; 3], + b1: &[f64; 3], + min_dist: &mut S1ChordAngle, +) -> bool { + if *min_dist == S1ChordAngle::zero() { + return false; + } + if crossing_sign(a0, a1, b0, b1) >= 0 { + *min_dist = S1ChordAngle::zero(); + return true; + } + // Otherwise, the minimum distance is achieved at an endpoint of at least + // one of the two edges. We use "|" rather than "||" below to ensure that + // all four possibilities are always checked. + // + // The calculation below computes each of the six vertex-vertex distances + // twice (this could be optimized). + let a = update_min_distance(a0, b0, b1, min_dist); + let b = update_min_distance(a1, b0, b1, min_dist); + let c = update_min_distance(b0, a0, a1, min_dist); + let d = update_min_distance(b1, a0, a1, min_dist); + a | b | c | d +} + +/// Returns true if the minimum distance between two edges is less than distance. +/// +/// Port of S2::IsEdgePairDistanceLess in s2edge_distances.cc. +pub fn is_edge_pair_distance_less( + a0: &[f64; 3], + a1: &[f64; 3], + b0: &[f64; 3], + b1: &[f64; 3], + distance: S1ChordAngle, +) -> bool { + // If the edges cross or share an endpoint, the minimum distance is zero. + if crossing_sign(a0, a1, b0, b1) >= 0 { + return distance != S1ChordAngle::zero(); + } + // Otherwise the minimum distance is achieved at an endpoint of at least + // one of the endpoints of the two edges. Written this way for short circuiting. + is_distance_less(a0, b0, b1, distance) + || is_distance_less(a1, b0, b1, distance) + || is_distance_less(b0, a0, a1, distance) + || is_distance_less(b1, a0, a1, distance) +} + +// Convenience wrapper matching S2::CrossingSign in s2edge_crossings.cc. +fn crossing_sign(a: &[f64; 3], b: &[f64; 3], c: &[f64; 3], d: &[f64; 3]) -> i32 { + let mut crosser = S2EdgeCrosser::new(a, b); + crosser.crossing_sign_two(c, d) +} + +// Port of GetUpdateMinInteriorDistanceMaxError. +fn get_update_min_interior_distance_max_error(dist: S1ChordAngle) -> f64 { + // If a point is more than 90 degrees from an edge, then the minimum + // distance is always to one of the endpoints, not to the edge interior. + if dist >= S1ChordAngle::right() { + return 0.0; + } + + let b = f64::min(1.0, 0.5 * dist.length2()); + let a = (b * (2.0 - b)).sqrt(); + ((2.5 + 2.0 * 3.0_f64.sqrt() + 8.5 * a) * a + + (2.0 + 2.0 * 3.0_f64.sqrt() / 3.0 + 6.5 * (1.0 - b)) * b + + (23.0 + 16.0 / 3.0_f64.sqrt()) * f64::EPSILON) + * f64::EPSILON +} diff --git a/src/spatial/s2metrics.rs b/src/spatial/s2metrics.rs new file mode 100644 index 0000000000..cd758b8288 --- /dev/null +++ b/src/spatial/s2metrics.rs @@ -0,0 +1,29 @@ +//! S2 cell metrics for determining appropriate cell levels. +//! +//! Ported from s2/s2metrics.h (Eric Veach, Google). + +/// S2 cell metrics for determining appropriate cell levels. +pub struct S2Metrics; + +impl S2Metrics { + // kMinWidth.deriv() for quadratic projection: 2 * sqrt(2) / 3 + const MIN_WIDTH_DERIV: f64 = 0.9428090415820634; + + /// Return the maximum level such that the metric is at least the given value, or 0 if there is + /// no such level. For example, kMinWidth.GetLevelForMinValue(0.1) returns the maximum level + /// such that all cells have a minimum width of 0.1 or larger. The return value is always a + /// valid level. + pub fn get_level_for_min_width(value: f64) -> i32 { + if value <= 0.0 { + return 30; + } + // This code is equivalent to computing a floating-point "level" value and rounding down. + let ratio = Self::MIN_WIDTH_DERIV / value; + let level = ratio.log2().floor() as i32; + level.clamp(0, 30) + } +} + +#[cfg(test)] +#[path = "tests/s2metrics_tests.rs"] +mod tests; diff --git a/src/spatial/s2padded_cell.rs b/src/spatial/s2padded_cell.rs new file mode 100644 index 0000000000..41dacb53d5 --- /dev/null +++ b/src/spatial/s2padded_cell.rs @@ -0,0 +1,255 @@ +//! S2PaddedCell represents an S2Cell whose (u,v)-range has been expanded on all sides by a given +//! amount of "padding". Unlike S2Cell, its methods and representation are optimized for clipping +//! edges against S2Cell boundaries to determine which cells are intersected by a given set of +//! edges. +use super::math::normalize; +use super::r1interval::R1Interval; +use super::r2rect::R2Rect; +use super::s2cell_id::S2CellId; +use super::s2coords::{ + self, face_si_ti_to_xyz, pos_to_ij, pos_to_orientation, si_ti_to_st, st_to_uv, INVERT_MASK, + MAX_CELL_LEVEL, SWAP_MASK, +}; + +/// An S2Cell whose (u,v)-range has been expanded on all sides by a given amount of "padding". +/// +/// Unlike S2Cell, its methods and representation are optimized for clipping edges against S2Cell +/// boundaries to determine which cells are intersected by a given set of edges. +#[derive(Clone, Debug)] +pub struct S2PaddedCell { + id: S2CellId, + padding: f64, + bound: R2Rect, + middle: R2Rect, + ij_lo: [i32; 2], + orientation: i32, + level: i32, +} + +fn middle(ij_lo: [i32; 2], level: i32, padding: f64) -> R2Rect { + let ij_size = S2CellId::size_ij_for_level(level); + let u = st_to_uv(si_ti_to_st((2 * ij_lo[0] + ij_size) as u32)); + let v = st_to_uv(si_ti_to_st((2 * ij_lo[1] + ij_size) as u32)); + R2Rect::new( + R1Interval::new(u - padding, u + padding), + R1Interval::new(v - padding, v + padding), + ) +} + +impl S2PaddedCell { + /// Construct an S2PaddedCell for the given cell id and padding. + pub fn new(id: S2CellId, padding: f64) -> Self { + let (bound, ij_lo, orientation, level) = if id.is_face() { + let limit = 1.0 + padding; + let bound = R2Rect::new( + R1Interval::new(-limit, limit), + R1Interval::new(-limit, limit), + ); + (bound, [0, 0], id.face() & 1, 0) + } else { + let (_, i, j, orientation) = id.to_face_ij_orientation(); + let level = id.level(); + let ij_size = S2CellId::size_ij_for_level(level); + let ij_lo = [i & -ij_size, j & -ij_size]; + let uv_bound = ij_level_to_bound_uv(i, j, level); + let bound = uv_bound.expanded(padding); + (bound, ij_lo, orientation, level) + }; + Self { + id, + padding, + bound, + middle: middle(ij_lo, level, padding), + ij_lo, + orientation, + level, + } + } + + /// Construct the child of "parent" with the given (i,j) index. + /// + /// The four child cells have indices of (0,0), (0,1), (1,0), (1,1), where the i and j indices + /// correspond to increasing u- and v-values respectively. + pub fn from_parent(parent: &S2PaddedCell, i: usize, j: usize) -> Self { + let pos = s2coords::ij_to_pos(parent.orientation, i as i32, j as i32); + let id = parent.id.child(pos); + let level = parent.level + 1; + let ij_size = S2CellId::size_ij_for_level(level); + + let ij_lo = [ + parent.ij_lo[0] + i as i32 * ij_size, + parent.ij_lo[1] + j as i32 * ij_size, + ]; + + let orientation = parent.orientation ^ pos_to_orientation(pos); + + // Compute child bound from parent bound and middle. + let parent_middle = parent.middle(); + let bound = R2Rect::new( + R1Interval::new( + if i == 0 { + parent.bound[0].lo() + } else { + parent_middle[0].lo() + }, + if i == 0 { + parent_middle[0].hi() + } else { + parent.bound[0].hi() + }, + ), + R1Interval::new( + if j == 0 { + parent.bound[1].lo() + } else { + parent_middle[1].lo() + }, + if j == 0 { + parent_middle[1].hi() + } else { + parent.bound[1].hi() + }, + ), + ); + Self { + id, + padding: parent.padding, + bound, + middle: middle(ij_lo, level, parent.padding), + ij_lo, + orientation, + level, + } + } + + /// Returns the cell ID. + #[inline] + pub fn id(&self) -> S2CellId { + self.id + } + + /// Returns the level. + #[inline] + pub fn level(&self) -> i32 { + self.level + } + + /// Returns the padded UV bounding box of this cell. + pub fn bound(&self) -> R2Rect { + self.bound + } + + /// Returns the "middle" rectangle that belongs to all four children. + pub fn middle(&self) -> R2Rect { + self.middle + } + + /// Returns the (i,j) coordinates for the child at the given traversal position. + #[inline] + pub fn get_child_ij(&self, pos: i32) -> (usize, usize) { + let ij = pos_to_ij(self.orientation, pos); + ((ij >> 1) as usize, (ij & 1) as usize) + } + + /// Returns the center of the cell. + pub fn get_center(&self) -> [f64; 3] { + let ij_size = S2CellId::size_ij_for_level(self.level); + let si = (2 * self.ij_lo[0] + ij_size) as u32; + let ti = (2 * self.ij_lo[1] + ij_size) as u32; + normalize(&face_si_ti_to_xyz(self.id.face(), si, ti)) + } + + /// Returns the vertex where the Hilbert curve enters this cell. + pub fn get_entry_vertex(&self) -> [f64; 3] { + let mut i = self.ij_lo[0] as u32; + let mut j = self.ij_lo[1] as u32; + + if self.orientation & INVERT_MASK != 0 { + let ij_size = S2CellId::size_ij_for_level(self.level) as u32; + i += ij_size; + j += ij_size; + } + + normalize(&face_si_ti_to_xyz(self.id.face(), 2 * i, 2 * j)) + } + + /// Returns the vertex where the Hilbert curve exits this cell. + pub fn get_exit_vertex(&self) -> [f64; 3] { + let mut i = self.ij_lo[0] as u32; + let mut j = self.ij_lo[1] as u32; + let ij_size = S2CellId::size_ij_for_level(self.level) as u32; + + if self.orientation == 0 || self.orientation == SWAP_MASK | INVERT_MASK { + i += ij_size; + } else { + j += ij_size; + } + + normalize(&face_si_ti_to_xyz(self.id.face(), 2 * i, 2 * j)) + } + + /// Returns the smallest cell that contains all descendants whose bounds + /// intersect the given rectangle. + /// + /// REQUIRES: self.bound.intersects(rect) + pub fn shrink_to_fit(&self, rect: &R2Rect) -> S2CellId { + debug_assert!(self.bound.intersects(rect)); + + // Quick rejection: if rect contains the center, no shrinking is possible. + let ij_size = S2CellId::size_ij_for_level(self.level); + if self.level == 0 { + if rect[0].contains_point(0.0) || rect[1].contains_point(0.0) { + return self.id; + } + } else { + let u = st_to_uv(si_ti_to_st((2 * self.ij_lo[0] + ij_size) as u32)); + let v = st_to_uv(si_ti_to_st((2 * self.ij_lo[1] + ij_size) as u32)); + if rect[0].contains_point(u) || rect[1].contains_point(v) { + return self.id; + } + } + + // Expand rect by padding to find the range of i,j coordinates. + let padded = rect.expanded(self.padding + 1.5 * f64::EPSILON); + + let mut ij_min = [0i32; 2]; + let mut ij_xor = [0i32; 2]; + + for d in 0..2 { + ij_min[d] = self.ij_lo[d].max(s2coords::st_to_ij(s2coords::uv_to_st(padded[d].lo()))); + let ij_max = (self.ij_lo[d] + ij_size - 1) + .min(s2coords::st_to_ij(s2coords::uv_to_st(padded[d].hi()))); + ij_xor[d] = ij_min[d] ^ ij_max; + } + + // Find the highest differing bit position. + let level_msb = ((ij_xor[0] | ij_xor[1]) << 1) + 1; + let level = MAX_CELL_LEVEL - (31 - level_msb.leading_zeros() as i32); + + if level <= self.level { + return self.id; + } + + S2CellId::from_face_ij(self.id.face(), ij_min[0], ij_min[1]).parent(level) + } +} + +/// Helper to compute (u,v) bounds from (i,j) coordinates and level. +fn ij_level_to_bound_uv(i: i32, j: i32, level: i32) -> R2Rect { + let cell_size = S2CellId::size_ij_for_level(level); + let i_lo = i & -cell_size; + let j_lo = j & -cell_size; + let i_hi = i_lo + cell_size; + let j_hi = j_lo + cell_size; + + R2Rect::new( + R1Interval::new( + st_to_uv(s2coords::ij_to_st_min(i_lo)), + st_to_uv(s2coords::ij_to_st_min(i_hi)), + ), + R1Interval::new( + st_to_uv(s2coords::ij_to_st_min(j_lo)), + st_to_uv(s2coords::ij_to_st_min(j_hi)), + ), + ) +} diff --git a/src/spatial/serializer.rs b/src/spatial/serializer.rs new file mode 100644 index 0000000000..9d4b7e267f --- /dev/null +++ b/src/spatial/serializer.rs @@ -0,0 +1,77 @@ +//! Spatial field serialization. +//! +//! Writes cell index and edge index data for spatial fields. Each field's cell index and edge +//! index are written as sub-sections of the composite files for the segment's SpatialCells and +//! SpatialEdges components. +use std::io; +use std::io::Write; + +use crate::directory::{CompositeWrite, WritePtr}; +use crate::schema::Field; +use crate::spatial::cell_index::{BuildOptions, IndexBuilder}; +use crate::spatial::edge_writer::EdgeWriter; +use crate::spatial::geometry_set::GeometrySet; +use crate::spatial::sphere::Sphere; + +/// Default skip interval for the edge index skip list directory. +pub(crate) const EDGE_SKIP_INTERVAL: u32 = 16; + +/// Serializes spatial field data into cell index and edge index files. +pub struct SpatialSerializer { + cells_write: CompositeWrite, + edges_write: CompositeWrite, +} + +impl SpatialSerializer { + /// Create the serializer from two write pointers (cells and edges). + pub fn new(cells_write: WritePtr, edges_write: WritePtr) -> io::Result { + Ok(SpatialSerializer { + cells_write: CompositeWrite::wrap(cells_write), + edges_write: CompositeWrite::wrap(edges_write), + }) + } + + /// Serialize one field's spatial data. Builds a CellIndex from the smashed GeometrySets + /// and streams edges through EdgeWriter. + pub fn serialize_field(&mut self, field: Field, sets: &[GeometrySet]) -> io::Result<()> { + if sets.is_empty() { + return Ok(()); + } + + // Build the cell index from smashed sets. + let builder = IndexBuilder::new(BuildOptions::default()); + let cell_index = builder.build_from_sets(sets); + + // Write the cell index. + let cells_out = self.cells_write.for_field(field); + cell_index.write(cells_out)?; + cells_out.flush()?; + + // Write the edge index. + let edges_out = self.edges_write.for_field(field); + { + let mut edge_writer = EdgeWriter::::new(edges_out, EDGE_SKIP_INTERVAL); + + for set in sets { + edge_writer.insert(set)?; + } + + edge_writer.finish()?; + } + edges_out.flush()?; + + Ok(()) + } + + /// Split into the two underlying CompositeWrites for direct access. Used by the merger. + pub fn into_composite_writes(self) -> (CompositeWrite, CompositeWrite) { + (self.cells_write, self.edges_write) + } + + /// Close both composite files. + pub fn close(self) -> io::Result<()> { + self.cells_write.close()?; + self.edges_write.close()?; + Ok(()) + } +} diff --git a/src/spatial/shape_index_region.rs b/src/spatial/shape_index_region.rs new file mode 100644 index 0000000000..f44588a5a2 --- /dev/null +++ b/src/spatial/shape_index_region.rs @@ -0,0 +1,430 @@ +//! ShapeIndexRegion: Query methods for shape-cell relationships. +//! +//! This module provides methods to determine whether shapes in a CellIndex intersect or contain a +//! given S2Cell target. Used to implement the Region trait over a CellIndex, enabling +//! RegionCoverer to produce coverings. +//! +//! Corresponds to S2ShapeIndexRegion in the C++ implementation. +//! From s2shape_index_region.h +use std::collections::HashMap; + +use super::cell_index::{CellIndex, CellRelation, ClippedShape, GeometryId}; +use super::cell_index_reader::CellIndexReader; +use super::cell_union::CellUnion; +use super::crossings::S2EdgeCrosser; +use super::edge_cache::EdgeCache; +use super::latlng_rect::S2LatLngRect; +use super::region::Region; +use super::s2cap::S2Cap; +use super::s2cell::S2Cell; +use super::s2cell_id::S2CellId; +use super::s2edge_clipping::{ + clip_to_padded_face, intersects_rect, FACE_CLIP_ERROR_UV_COORD, INTERSECTS_RECT_ERROR_UV_DIST, +}; +use super::sphere::Sphere; + +/// Provides vertex access for edge intersection tests. +/// +/// Implementors resolve a (geometry_id, edge_index) pair to the two endpoints of that edge. For +/// ingest-side indexes backed by EdgeCache, this reads from disk. For query-local indexes, this +/// indexes into an in-memory vertex array. +pub trait EdgeProvider { + /// Returns the two endpoints of the given edge. A point is an edge from itself to itself. + fn get_edge(&self, geometry_id: GeometryId, edge_idx: u32) -> ([f64; 3], [f64; 3]); + + /// Whether the geometry has an interior. Lines and points contain nothing. + fn is_closed(&self, geometry_id: GeometryId) -> bool; +} + +/// Minimal clipped shape data for indexed containment. Owned to avoid lifetime mismatches between +/// borrowed in-memory data and deserialized reader data. +pub struct ClippedInfo { + /// The cell containing the test point. + pub cell_id: S2CellId, + /// Whether the geometry contains the cell center. + pub contains_center: bool, + /// Edge indices clipped to this cell. + pub edge_indices: Vec, +} + +/// Abstracts cell lookup and edge resolution for indexed point-in-polygon. Takes &mut self because +/// the segment-side EdgeCache mutates its LRU cache. +pub trait ContainmentIndex { + /// Finds the cell containing the test point and returns the clipped shape data + /// for the given geometry, or None if the geometry is not present in that cell. + fn find_clipped(&mut self, geometry_id: GeometryId, point: &[f64; 3]) -> Option; + /// Resolves an edge index to its two endpoint vertices. + fn resolve_edge(&mut self, geometry_id: GeometryId, edge_idx: u32) -> ([f64; 3], [f64; 3]); +} + +/// Indexed point-in-polygon through either backend. Finds the cell containing the test point, +/// starts from contains_center, counts crossings of only the edges clipped to that cell. +pub fn contains_point( + index: &mut T, + geometry_id: GeometryId, + point: &[f64; 3], +) -> bool { + let info = match index.find_clipped(geometry_id, point) { + Some(c) => c, + None => return false, + }; + let mut inside = info.contains_center; + if !info.edge_indices.is_empty() { + let center = info.cell_id.to_point(); + let mut crosser = S2EdgeCrosser::new(¢er, point); + for &edge_idx in &info.edge_indices { + let (v0, v1) = index.resolve_edge(geometry_id, edge_idx); + inside ^= crosser.edge_or_vertex_crossing_two(&v0, &v1); + } + } + inside +} + +/// Segment containment index wrapping a CellIndexReader and an EdgeCache. +pub struct SegmentIndex<'a, 'b> { + /// The serialized cell index reader. + pub cell_reader: &'b CellIndexReader<'b>, + /// The edge cache for vertex resolution. + pub edge_cache: &'b mut EdgeCache<'a, Sphere>, +} + +impl ContainmentIndex for SegmentIndex<'_, '_> { + fn find_clipped(&mut self, geometry_id: GeometryId, point: &[f64; 3]) -> Option { + let cell_id = S2CellId::from_point(point); + let index_cell = self.cell_reader.find(cell_id)?; + let clipped = index_cell + .shapes + .iter() + .find(|s| s.geometry_id == geometry_id)?; + Some(ClippedInfo { + cell_id: index_cell.cell_id, + contains_center: clipped.contains_center, + edge_indices: clipped.edge_indices.clone(), + }) + } + + fn resolve_edge(&mut self, geometry_id: GeometryId, edge_idx: u32) -> ([f64; 3], [f64; 3]) { + self.edge_cache.get(geometry_id).edge(edge_idx) + } +} + +/// Maximum error in edge-cell intersection tests. +/// +/// This combines the error from face clipping and rectangle intersection to ensure we don't miss +/// any edges that might intersect the cell boundary. +pub const ANY_EDGE_INTERSECTS_MAX_ERROR: f64 = + FACE_CLIP_ERROR_UV_COORD + INTERSECTS_RECT_ERROR_UV_DIST; + +/// Returns true if any edge of the clipped shape intersects the target cell. +/// +/// The test uses UV-space bounds expanded by the maximum clipping error to avoid missing edges +/// that nearly touch the cell boundary. +/// +/// This method is conservative: it may return true for edges that are very close to (but don't +/// actually intersect) the cell boundary. The maximum error is less than 10 * DBL_EPSILON radians +/// (about 15 nanometers). +pub fn any_edge_intersects( + clipped: &ClippedShape, + target: &S2Cell, + edges: &E, +) -> bool { + let bound = target + .get_bound_uv() + .expanded(ANY_EDGE_INTERSECTS_MAX_ERROR); + let face = target.face(); + + for &edge_idx in &clipped.edge_indices { + let (v0, v1) = edges.get_edge(clipped.geometry_id, edge_idx); + + if let Some((p0, p1)) = clip_to_padded_face(&v0, &v1, face, ANY_EDGE_INTERSECTS_MAX_ERROR) { + if intersects_rect(p0, p1, &bound) { + return true; + } + } + } + + false +} + +/// Returns true if the given geometry in the index contains the point. Uses the cell index to +/// narrow the edge set — locates the cell containing the point, starts from contains_center, and +/// counts crossings of only the edges in that cell. A geometry without an interior contains +/// nothing. +/// +/// Port of S2ContainsPointQuery::ShapeContains. +pub fn index_contains_point( + index: &CellIndex, + edges: &E, + geometry_id: GeometryId, + point: &[f64; 3], +) -> bool { + if !edges.is_closed(geometry_id) { + return false; + } + let cell_id = S2CellId::from_point(point); + let index_cell = match index.find_cell(cell_id) { + Some(cell) => cell, + None => return false, + }; + let clipped = match index_cell.find_shape(geometry_id) { + Some(s) => s, + None => return false, + }; + clipped_shape_contains_point(index_cell.cell_id, clipped, point, |edge_idx| { + edges.get_edge(geometry_id, edge_idx) + }) +} + +/// Returns true if any geometry in the index that has an interior contains the point. Only the +/// geometries in the index cell containing the point can. +/// +/// Port of S2ContainsPointQuery::Contains. +pub fn any_shape_contains_point( + index: &CellIndex, + edges: &E, + point: &[f64; 3], +) -> bool { + let Some(index_cell) = index.find_cell(S2CellId::from_point(point)) else { + return false; + }; + index_cell.shapes.iter().any(|clipped| { + edges.is_closed(clipped.geometry_id) + && clipped_shape_contains_point(index_cell.cell_id, clipped, point, |edge_idx| { + edges.get_edge(clipped.geometry_id, edge_idx) + }) + }) +} + +/// Returns true if the given clipped shape contains the point. The cell_id identifies the index +/// cell the clipped shape belongs to — its center is the reference point for the edge crossing +/// test. The shape must have an interior. +pub(crate) fn clipped_shape_contains_point( + cell_id: S2CellId, + clipped: &ClippedShape, + point: &[f64; 3], + mut edge: impl FnMut(u32) -> ([f64; 3], [f64; 3]), +) -> bool { + let mut inside = clipped.contains_center; + if !clipped.edge_indices.is_empty() { + let cell_center = cell_id.to_point(); + let mut crosser = S2EdgeCrosser::new(&cell_center, point); + for &edge_idx in &clipped.edge_indices { + let (v0, v1) = edge(edge_idx); + inside ^= crosser.edge_or_vertex_crossing_two(&v0, &v1); + } + } + inside +} + +/// Locates a target cell relative to the CellIndex. +/// +/// Binary search on the sorted cells array. The three cases mirror S2's +/// ShapeIndexIterator::Locate. +fn locate(index: &CellIndex, target: S2CellId) -> CellRelation { + let cells = &index.cells; + if cells.is_empty() { + return CellRelation::Disjoint; + } + + // Binary search: find the first cell whose range_max >= target.range_min(). + let pos = cells.partition_point(|c| c.cell_id.range_max() < target.range_min()); + + if pos >= cells.len() { + return CellRelation::Disjoint; + } + + // If this cell contains the target, or equals it, it's Indexed. + if cells[pos].cell_id.contains(target) || cells[pos].cell_id == target { + return CellRelation::Indexed(pos); + } + + // If the target contains this cell, it's Subdivided. + if target.contains(cells[pos].cell_id) { + return CellRelation::Subdivided(pos); + } + + // The target's range starts after all index cells, or doesn't overlap. + if cells[pos].cell_id.range_min() > target.range_max() { + return CellRelation::Disjoint; + } + + CellRelation::Disjoint +} + +/// Visits all shapes that intersect an S2Cell target. For each shape, calls the visitor with the +/// geometry_id and whether the shape fully contains the target cell. The visitor returns true to +/// continue or false to terminate early. +pub fn visit_intersecting_shapes( + index: &CellIndex, + edges: &E, + target: &S2Cell, + mut visitor: F, +) -> bool +where + E: EdgeProvider, + F: FnMut(GeometryId, bool) -> bool, // (geometry_id, contains_target) -> continue +{ + match locate(index, target.id()) { + CellRelation::Disjoint => true, + + CellRelation::Subdivided(start) => { + // Target contains one or more index cells. + // A shape contains the target iff: + // - It appears in at least one cell + // - It has contains_center=true in ALL cells + // - It has NO edges in ANY cell + let mut shape_not_contains: HashMap = HashMap::new(); + let max = target.id().range_max(); + + let mut pos = start; + while pos < index.cells.len() && index.cells[pos].cell_id.range_min() <= max { + let cell = &index.cells[pos]; + for clipped in &cell.shapes { + let not_contains = shape_not_contains + .entry(clipped.geometry_id) + .or_insert(false); + + // Mark as not-containing if: + // - Has any edges (would intersect cell boundary), OR + // - Does not contain center (discontinuous interior) + *not_contains |= !clipped.edge_indices.is_empty() || !clipped.contains_center; + } + pos += 1; + } + + // Visit all shapes that appeared in any cell. + for (geometry_id, not_contains) in shape_not_contains { + let contains = !not_contains; + if !visitor(geometry_id, contains) { + return false; + } + } + true + } + + CellRelation::Indexed(pos) => { + // Target is contained by (or equals) an index cell. + let cell = &index.cells[pos]; + let index_cell_id = cell.cell_id; + let target_center = target.get_center(); + + for clipped in &cell.shapes { + let contains = if index_cell_id == target.id() { + // Exact match: shape contains target iff no edges AND contains_center + clipped.edge_indices.is_empty() && clipped.contains_center + } else { + // Target is smaller than index cell — need finer testing + if any_edge_intersects(clipped, target, edges) { + // Edge intersects target: shape intersects but doesn't contain + false + } else if !edges.is_closed(clipped.geometry_id) { + // A line or point with no edge in the target misses it. + continue; + } else { + // No edges intersect: check if shape contains target's center by + // ray-casting from the index cell's center (where we know contains_center) + // to the target center, counting crossings of only the edges in this cell. + if !clipped_shape_contains_point( + index_cell_id, + clipped, + &target_center, + |edge_idx| edges.get_edge(clipped.geometry_id, edge_idx), + ) { + continue; + } + true + } + }; + + if !visitor(clipped.geometry_id, contains) { + return false; + } + } + true + } + } +} + +/// Convenience method: Returns all geometry IDs that intersect the target. +pub fn get_intersecting_shapes( + index: &CellIndex, + edges: &E, + target: &S2Cell, +) -> Vec<(GeometryId, bool)> { + let mut results = Vec::new(); + visit_intersecting_shapes(index, edges, target, |geo_id, contains| { + results.push((geo_id, contains)); + true + }); + results +} + +/// Returns true if any shape in the index intersects the target cell. +/// +/// This is equivalent to S2ShapeIndexRegion::MayIntersect. +pub fn may_intersect(index: &CellIndex, edges: &E, target: &S2Cell) -> bool { + let mut intersects = false; + visit_intersecting_shapes(index, edges, target, |_geo_id, _contains| { + intersects = true; + false // Stop after first match + }); + intersects +} + +/// Returns true if any shape in the index contains the target cell. +/// +/// This is equivalent to S2ShapeIndexRegion::Contains(S2Cell). +pub fn any_shape_contains(index: &CellIndex, edges: &E, target: &S2Cell) -> bool { + let mut contains = false; + visit_intersecting_shapes(index, edges, target, |_geo_id, shape_contains| { + if shape_contains { + contains = true; + return false; // Stop after first containing shape + } + true // Continue looking + }); + contains +} + +/// Region adapter over a CellIndex. +/// +/// Wraps a CellIndex and an EdgeProvider to implement the Region trait, enabling RegionCoverer to +/// produce coverings of the indexed geometry. For the contains query, this wraps the query-local +/// CellIndex built from the query polygon. +pub struct CellIndexRegion<'a, E: EdgeProvider> { + index: &'a CellIndex, + edges: &'a E, +} + +impl<'a, E: EdgeProvider> CellIndexRegion<'a, E> { + /// Creates a region backed by a cell index and an edge provider. + pub fn new(index: &'a CellIndex, edges: &'a E) -> Self { + Self { index, edges } + } +} + +impl<'a, E: EdgeProvider> Region for CellIndexRegion<'a, E> { + fn get_cap_bound(&self) -> S2Cap { + let cell_ids: Vec = self.index.cells.iter().map(|c| c.cell_id).collect(); + let cu = CellUnion::from_cell_ids(cell_ids); + cu.get_cap_bound() + } + + fn get_rect_bound(&self) -> S2LatLngRect { + let cell_ids: Vec = self.index.cells.iter().map(|c| c.cell_id).collect(); + let cu = CellUnion::from_cell_ids(cell_ids); + cu.get_rect_bound() + } + + fn contains_cell(&self, cell: &S2Cell) -> bool { + any_shape_contains(self.index, self.edges, cell) + } + + fn may_intersect(&self, cell: &S2Cell) -> bool { + may_intersect(self.index, self.edges, cell) + } + + fn contains_point(&self, p: &[f64; 3]) -> bool { + any_shape_contains_point(self.index, self.edges, p) + } +} diff --git a/src/spatial/spatial_index_manager.rs b/src/spatial/spatial_index_manager.rs new file mode 100644 index 0000000000..7255874995 --- /dev/null +++ b/src/spatial/spatial_index_manager.rs @@ -0,0 +1,29 @@ +//! Registry for spatial index implemenations. +//! +//! `SpatialIndexManager` follows the same pattern as `TokenizerManager`. Users register +//! strategies by name at index creation time. The registered strategies are retrieved +//! at the writer, merger, and query integration points. +//! +//! Implemenations are monomorphic implementations behind a `dyn SpatialIndex` trait. The +//! vtable lookup occurs at field boundaries, not per-geometry, making the cost negligible. + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// HUSH +pub struct SpatialIndex {} + +/// Manages a map of index implementations keyed on a user specified index type name. +#[derive(Clone)] +pub struct SpatialIndexManager { + /// A map of spatial index implementations. + pub indices: Arc>>, +} + +impl Default for SpatialIndexManager { + fn default() -> Self { + Self { + indices: Arc::new(RwLock::new(HashMap::new())), + } + } +} diff --git a/src/spatial/sphere.rs b/src/spatial/sphere.rs new file mode 100644 index 0000000000..1daafc4476 --- /dev/null +++ b/src/spatial/sphere.rs @@ -0,0 +1,39 @@ +//! Implementation of a unit sphere surface. + +use crate::spatial::surface::Surface; + +/// The origin should not be a point that is commonly used in edge tests in order to avoid +/// triggering code to handle degenerate cases. (This rules out the north and south poles.) It +/// should also not be on the boundary of any low-level S2Cell for the same reason. +/// +/// The point chosen here is about 66km from the north pole towards the East Siberian Sea. See the +/// unittest for more details. It is written out explicitly using floating-point literals because +/// the optimizer doesn't seem willing to evaluate Normalize() at compile time. +pub const ORIGIN: [f64; 3] = [ + -0.00999946643502502, + 0.002592454260932412, + 0.999946643502502, + // x: -0.0099994664350250197, + // y: 0.0025924542609324121, + // z: 0.99994664350250195, +]; + +/// Implementation of a unit sphere surface. +pub struct Sphere; + +impl Surface for Sphere { + const DIMENSIONS: usize = 3; + type Point = [f64; 3]; + + fn project(lon: f64, lat: f64) -> [f64; 3] { + let lat = lat.to_radians(); + let lon = lon.to_radians(); + let cos_lat = lat.cos(); + [cos_lat * lon.cos(), cos_lat * lon.sin(), lat.sin()] + } + + #[inline] + fn origin_inside(ring: &[Self::Point]) -> bool { + crate::spatial::containment::compute_origin_inside(ring) + } +} diff --git a/src/spatial/surface.rs b/src/spatial/surface.rs new file mode 100644 index 0000000000..926f7ceb66 --- /dev/null +++ b/src/spatial/surface.rs @@ -0,0 +1,22 @@ +//! Geometry abstraction for spatial indexing. +//! +//! An abstraction one of two surfaces for geographic indices. For planar geometry, points are 2D +//! coordinates on a single face. For spherical geometry, points are 3D unit vectors on the six +//! faces of a cube projection. + +/// An abstraction one of two surfaces for geographic indices. For planar geometry, points are 2D +/// coordinates on a single face. For spherical geometry, points are 3D unit vectors on the six +/// faces of a cube projection. +pub trait Surface { + /// Number of f64 coordinates per vertex. 2 for Plane, 3 for Sphere. + const DIMENSIONS: usize; + + /// The point type of the surface. + type Point: Copy; + + /// Project a lon/lat coordinate onto this surface. + fn project(lon: f64, lat: f64) -> Self::Point; + + /// Whether the interior of a closed ring contains the surface's reference origin. + fn origin_inside(ring: &[Self::Point]) -> bool; +} diff --git a/src/spatial/tests/cell_union_tests.rs b/src/spatial/tests/cell_union_tests.rs new file mode 100644 index 0000000000..f4b4d91572 --- /dev/null +++ b/src/spatial/tests/cell_union_tests.rs @@ -0,0 +1,341 @@ +use super::*; + +#[test] +fn test_empty() { + let cu = CellUnion::new(); + assert!(cu.is_empty()); + assert_eq!(cu.len(), 0); + assert!(cu.is_valid()); + assert!(cu.is_normalized()); +} + +#[test] +fn test_whole_sphere() { + let cu = CellUnion::whole_sphere(); + assert_eq!(cu.len(), 6); + assert!(cu.is_normalized()); + // Should contain any point + assert!(cu.contains_point(&[1.0, 0.0, 0.0])); + assert!(cu.contains_point(&[0.0, 1.0, 0.0])); + assert!(cu.contains_point(&[0.0, 0.0, 1.0])); +} + +#[test] +fn test_normalize_removes_duplicates() { + let face0 = S2CellId::from_face(0); + let cu = CellUnion::from_cell_ids(vec![face0, face0, face0]); + assert_eq!(cu.len(), 1); + assert_eq!(cu.cell_id(0), face0); +} + +#[test] +fn test_normalize_removes_contained() { + let face0 = S2CellId::from_face(0); + let child = face0.child(0); + let grandchild = child.child(0); + + // Parent contains children, should keep only parent + let cu = CellUnion::from_cell_ids(vec![face0, child, grandchild]); + assert_eq!(cu.len(), 1); + assert_eq!(cu.cell_id(0), face0); +} + +#[test] +fn test_normalize_merges_siblings() { + let face0 = S2CellId::from_face(0); + let children = face0.children(); + + // Four siblings should merge to parent + let cu = CellUnion::from_cell_ids(vec![children[0], children[1], children[2], children[3]]); + assert_eq!(cu.len(), 1); + assert_eq!(cu.cell_id(0), face0); +} + +#[test] +fn test_normalize_partial_siblings() { + let face0 = S2CellId::from_face(0); + let children = face0.children(); + + // Three siblings should NOT merge + let cu = CellUnion::from_cell_ids(vec![children[0], children[1], children[2]]); + assert_eq!(cu.len(), 3); +} + +#[test] +fn test_are_siblings() { + let face0 = S2CellId::from_face(0); + let children = face0.children(); + + assert!(are_siblings( + children[0], + children[1], + children[2], + children[3] + )); + + // Not siblings if one is replaced + let face1 = S2CellId::from_face(1); + assert!(!are_siblings( + children[0], + children[1], + children[2], + face1.child(0) + )); +} + +#[test] +fn test_contains_cell() { + let face0 = S2CellId::from_face(0); + let cu = CellUnion::from_cell_ids(vec![face0]); + + // Should contain face0 and all its descendants + assert!(cu.contains(face0)); + assert!(cu.contains(face0.child(0))); + assert!(cu.contains(face0.child(0).child(0))); + + // Should not contain other faces + assert!(!cu.contains(S2CellId::from_face(1))); +} + +#[test] +fn test_intersects_cell() { + let face0 = S2CellId::from_face(0); + let child = face0.child(0); + let cu = CellUnion::from_cell_ids(vec![child]); + + // Should intersect with parent (child is contained by parent) + assert!(cu.intersects(face0)); + // Should intersect with self + assert!(cu.intersects(child)); + // Should not intersect with sibling + assert!(!cu.intersects(face0.child(1))); +} + +#[test] +fn test_from_min_max() { + // Create a small range + let face0 = S2CellId::from_face(0); + let min = face0.child_begin_at_level(S2CellId::MAX_LEVEL); + let max = face0 + .child(0) + .child_end_at_level(S2CellId::MAX_LEVEL) + .prev(); + + let cu = CellUnion::from_min_max(min, max); + assert!(cu.is_normalized()); + + // Should contain the range endpoints + assert!(cu.contains(min)); + assert!(cu.contains(max)); +} + +#[test] +fn test_from_begin_end_empty() { + let face0 = S2CellId::from_face(0); + let leaf = face0.child_begin_at_level(S2CellId::MAX_LEVEL); + + // Empty range + let cu = CellUnion::from_begin_end(leaf, leaf); + assert!(cu.is_empty()); +} + +#[test] +fn test_denormalize() { + let face0 = S2CellId::from_face(0); + let cu = CellUnion::from_cell_ids(vec![face0]); + + // Denormalize to level 2 + let cells = cu.denormalize(2, 1); + assert_eq!(cells.len(), 16); // 4^2 children at level 2 + + for cell in &cells { + assert_eq!(cell.level(), 2); + } +} + +#[test] +fn test_denormalize_level_mod() { + let level1_cell = S2CellId::from_face(0).child(0); + let cu = CellUnion::from_cell_ids(vec![level1_cell]); + + // Denormalize with min_level=0, level_mod=3 + // Level 1 cell should expand to level 3 (next level where (level 0) % 3 == 0) + let cells = cu.denormalize(0, 3); + for cell in &cells { + assert!(cell.level() % 3 == 0 || cell.level() == S2CellId::MAX_LEVEL); + } +} + +#[test] +fn test_iterator() { + let cu = CellUnion::from_cell_ids(vec![S2CellId::from_face(0), S2CellId::from_face(1)]); + + let mut count = 0; + for _id in &cu { + count += 1; + } + assert_eq!(count, 2); +} + +#[test] +fn test_from_iterator() { + let cells = vec![S2CellId::from_face(0), S2CellId::from_face(1)]; + let cu: CellUnion = cells.into_iter().collect(); + assert_eq!(cu.len(), 2); +} + +#[test] +fn test_maximum_tile() { + let face0 = S2CellId::from_face(0); + let begin = face0.child_begin_at_level(S2CellId::MAX_LEVEL); + let end = face0.child_end_at_level(S2CellId::MAX_LEVEL); + + // maximum_tile of first leaf cell with face0's end as limit shoul be face0 + let tile = begin.maximum_tile(end); + assert_eq!(tile, face0); + + // maximum_tile of cell with itself as limit returns limit + assert_eq!(face0.maximum_tile(face0), face0); +} + +#[test] +fn test_region_empty_cell_union() { + let cu = CellUnion::new(); + + // Empty cell union has empty cap bound + let cap = cu.get_cap_bound(); + assert!(cap.is_empty()); + + // Empty cell union has empty rect bound + let rect = cu.get_rect_bound(); + assert!(rect.is_empty()); + + // Empty cell union does not contain any cell + let cell = S2Cell::new(S2CellId::from_face(0)); + assert!(!cu.contains_cell(&cell)); + assert!(!cu.may_intersect(&cell)); + + // Empty cell union does not contain any point + assert!(!cu.contains_point(&[1.0, 0.0, 0.0])); +} + +#[test] +fn test_region_whole_sphere() { + let cu = CellUnion::whole_sphere(); + + // Whole sphere should have full cap bound + let cap = cu.get_cap_bound(); + assert!(cap.is_full()); + + // Whole sphere should have full rect bound + let rect = cu.get_rect_bound(); + assert!(rect.is_full()); + + // Whole sphere contains all cells + for face in 0..6 { + let cell = S2Cell::new(S2CellId::from_face(face)); + assert!(cu.contains_cell(&cell)); + assert!(cu.may_intersect(&cell)); + } + + // Whole sphere contains all points + assert!(cu.contains_point(&[1.0, 0.0, 0.0])); + assert!(cu.contains_point(&[0.0, 1.0, 0.0])); + assert!(cu.contains_point(&[0.0, 0.0, 1.0])); +} + +#[test] +fn test_region_single_face() { + let cu = CellUnion::from_cell_ids(vec![S2CellId::from_face(0)]); + + // Cap bound should contain face 0 + let cap = cu.get_cap_bound(); + assert!(!cap.is_empty()); + assert!(!cap.is_full()); + + // Should contain face 0 cell and children + let face0 = S2Cell::new(S2CellId::from_face(0)); + assert!(cu.contains_cell(&face0)); + assert!(cu.may_intersect(&face0)); + + let child = S2Cell::new(S2CellId::from_face(0).child(0)); + assert!(cu.contains_cell(&child)); + assert!(cu.may_intersect(&child)); + + // Should not contain other faces + let face1 = S2Cell::new(S2CellId::from_face(1)); + assert!(!cu.contains_cell(&face1)); + assert!(!cu.may_intersect(&face1)); +} + +#[test] +fn test_region_contains_cell_vs_may_intersect() { + let face0 = S2CellId::from_face(0); + let child = face0.child(0); + let cu = CellUnion::from_cell_ids(vec![child]); + + let child_cell = S2Cell::new(child); + let parent_cell = S2Cell::new(face0); + + // Cell union contains the child cell + assert!(cu.contains_cell(&child_cell)); + assert!(cu.may_intersect(&child_cell)); + + // Cell union does NOT contain the parent (parent is larger) + assert!(!cu.contains_cell(&parent_cell)); + // But it DOES intersect the parent + assert!(cu.may_intersect(&parent_cell)); +} + +#[test] +fn test_region_get_cell_union_bound() { + let cu = CellUnion::from_cell_ids(vec![S2CellId::from_face(0).child(0)]); + + let mut cell_ids = Vec::new(); + cu.get_cell_union_bound(&mut cell_ids); + + // Should return a small covering + assert!(!cell_ids.is_empty()); + assert!(cell_ids.len() <= 6); // At most face cells +} + +#[test] +fn test_region_cap_bound_contains_all_cells() { + // Create a cell union with multiple cells + let cu = CellUnion::from_cell_ids(vec![ + S2CellId::from_face(0).child(0), + S2CellId::from_face(0).child(1), + S2CellId::from_face(1).child(0), + ]); + + let cap = cu.get_cap_bound(); + + // Cap should contain all cells' centers + for &id in cu.cell_ids() { + let center = id.to_point(); + assert!( + cap.contains_point(¢er), + "Cap bound should contain cell center" + ); + } +} + +#[test] +fn test_region_rect_bound_contains_all_cells() { + let cu = CellUnion::from_cell_ids(vec![ + S2CellId::from_face(0).child(0), + S2CellId::from_face(0).child(1), + ]); + + let rect = cu.get_rect_bound(); + + // Rect should contain all cells' rect bounds + for &id in cu.cell_ids() { + let cell = S2Cell::new(id); + let cell_rect = cell.get_rect_bound(); + assert!( + rect.contains(&cell_rect), + "Rect bound should contain cell rect bound" + ); + } +} diff --git a/src/spatial/tests/latlng_rect_tests.rs b/src/spatial/tests/latlng_rect_tests.rs new file mode 100644 index 0000000000..e79a6b3583 --- /dev/null +++ b/src/spatial/tests/latlng_rect_tests.rs @@ -0,0 +1,222 @@ +use std::f64::consts::{FRAC_PI_2, FRAC_PI_4, PI}; + +use super::*; + +#[test] +fn test_region_trait_rect_contains_point() { + // Rectangle around prime meridian and equator + let rect = S2LatLngRect::new( + R1Interval::new(-FRAC_PI_4, FRAC_PI_4), + S1Interval::new(-FRAC_PI_4, FRAC_PI_4), + ); + + // Point on equator at prime meridian + let p = [1.0, 0.0, 0.0]; + assert!(rect.contains_point(&p)); + + // Point at north pole + let north_pole = [0.0, 0.0, 1.0]; + assert!(!rect.contains_point(&north_pole)); +} + +#[test] +fn test_region_trait_rect_contains_cell() { + // Full rectangle + let full = S2LatLngRect::full(); + + // Face 0 cell + let cell = S2Cell::from_face(0); + assert!(full.contains_cell(&cell)); + + // Empty rectangle contains nothing + let empty = S2LatLngRect::empty(); + assert!(!empty.contains_cell(&cell)); +} + +#[test] +fn test_region_trait_rect_may_intersect() { + // Rectangle in northern hemisphere + let rect = S2LatLngRect::new(R1Interval::new(0.0, FRAC_PI_2), S1Interval::new(-PI, PI)); + + // North pole face (face 2) + let north_cell = S2Cell::from_face(2); + assert!(rect.may_intersect(&north_cell)); + + // South pole face (face 5) + let south_cell = S2Cell::from_face(5); + assert!(!rect.may_intersect(&south_cell)); +} + +#[test] +fn test_region_trait_rect_get_cap_bound() { + // Small rectangle at equator + let rect = S2LatLngRect::new(R1Interval::new(-0.1, 0.1), S1Interval::new(-0.1, 0.1)); + + let cap = rect.get_cap_bound(); + assert!(!cap.is_empty()); + assert!(!cap.is_full()); + + // Cap should contain all rectangle vertices + for k in 0..4 { + let vertex = rect.get_vertex(k); + let cos_lat = vertex[1].cos(); + let p = [ + cos_lat * vertex[0].cos(), + cos_lat * vertex[0].sin(), + vertex[1].sin(), + ]; + assert!(cap.contains_point(&p)); + } +} + +#[test] +fn test_region_trait_rect_get_rect_bound() { + let rect = S2LatLngRect::new(R1Interval::new(-0.5, 0.5), S1Interval::new(-1.0, 1.0)); + + // Rect bound of a rect is itself + let bound = rect.get_rect_bound(); + assert_eq!(rect, bound); +} + +#[test] +fn test_region_trait_rect_get_cell_union_bound() { + // Small rectangle + let rect = S2LatLngRect::new(R1Interval::new(0.0, 0.1), S1Interval::new(0.0, 0.1)); + + let mut cell_ids = Vec::new(); + rect.get_cell_union_bound(&mut cell_ids); + + // Should return a small number of cells + assert!(!cell_ids.is_empty()); + assert!(cell_ids.len() <= 6); // At most 6 face cells +} + +#[test] +fn test_region_trait_empty_rect() { + let empty = S2LatLngRect::empty(); + + // Empty rect has empty cap bound + let cap = empty.get_cap_bound(); + assert!(cap.is_empty()); + + // Empty rect does not contain any point + assert!(!empty.contains_point(&[1.0, 0.0, 0.0])); + + // Empty rect does not contain any cell + let cell = S2Cell::from_face(0); + assert!(!empty.contains_cell(&cell)); + + // Empty rect does not intersect any cell + assert!(!empty.may_intersect(&cell)); +} + +#[test] +fn test_region_trait_full_rect() { + let full = S2LatLngRect::full(); + + // Full rect has full cap bound + let cap = full.get_cap_bound(); + assert!(cap.is_full()); + + // Full rect contains every point + assert!(full.contains_point(&[1.0, 0.0, 0.0])); + assert!(full.contains_point(&[0.0, 0.0, 1.0])); + + // Full rect contains every cell + for face in 0..6 { + let cell = S2Cell::from_face(face); + assert!(full.contains_cell(&cell)); + } +} + +/// A rectangle from "lat:lng" of its low and high corners, in degrees, as S2's tests write it. +fn rect(lo: (f64, f64), hi: (f64, f64)) -> S2LatLngRect { + S2LatLngRect::new( + R1Interval::new(lo.0.to_radians(), hi.0.to_radians()), + S1Interval::new(lo.1.to_radians(), hi.1.to_radians()), + ) +} + +fn point((lat, lng): (f64, f64)) -> [f64; 3] { + let (lat, lng) = (lat.to_radians(), lng.to_radians()); + [lat.cos() * lng.cos(), lat.cos() * lng.sin(), lat.sin()] +} + +fn crosses(rect: &S2LatLngRect, a: (f64, f64), b: (f64, f64)) -> bool { + rect.boundary_intersects(&point(a), &point(b)) +} + +#[test] +fn test_boundary_intersects_empty_and_full_rectangles() { + let (a, b) = (point((10.0, 20.0)), point((30.0, 40.0))); + assert!(!S2LatLngRect::empty().boundary_intersects(&a, &b)); + assert!(!S2LatLngRect::full().boundary_intersects(&a, &b)); +} + +#[test] +fn test_boundary_intersects_spherical_lune() { + // This rectangle only has two non-degenerate sides. + let rect = rect((-90.0, 100.0), (90.0, 120.0)); + assert!(!crosses(&rect, (60.0, 60.0), (90.0, 60.0))); + assert!(!crosses(&rect, (-60.0, 110.0), (60.0, 110.0))); + assert!(crosses(&rect, (-60.0, 95.0), (60.0, 110.0))); + assert!(crosses(&rect, (60.0, 115.0), (80.0, 125.0))); +} + +#[test] +fn test_boundary_intersects_hemispheres() { + // These rectangles have only one non-degenerate side. + let north = S2LatLngRect::new(R1Interval::new(0.0, FRAC_PI_2), S1Interval::full()); + assert!(!crosses(&north, (60.0, -180.0), (90.0, -180.0))); + assert!(!crosses(&north, (60.0, -170.0), (60.0, 170.0))); + assert!(crosses(&north, (-10.0, -180.0), (10.0, -180.0))); + let south = S2LatLngRect::new(R1Interval::new(-FRAC_PI_2, 0.0), S1Interval::full()); + assert!(!crosses(&south, (-90.0, -180.0), (-60.0, -180.0))); + assert!(!crosses(&south, (-60.0, -170.0), (-60.0, 170.0))); + assert!(crosses(&south, (-10.0, -180.0), (10.0, -180.0))); +} + +#[test] +fn test_boundary_intersects_rect_crossing_anti_meridian() { + let rect = rect((20.0, 170.0), (40.0, -170.0)); + assert!(rect.contains_point(&point((30.0, 180.0)))); + + // Check that crossings of all four sides are detected. + assert!(crosses(&rect, (25.0, 160.0), (25.0, 180.0))); + assert!(crosses(&rect, (25.0, -160.0), (25.0, -180.0))); + assert!(crosses(&rect, (15.0, 175.0), (30.0, 175.0))); + assert!(crosses(&rect, (45.0, 175.0), (30.0, 175.0))); + + // Check that the edges on the opposite side of the sphere but at the same latitude do not + // intersect the rectangle boundary. + assert!(!crosses(&rect, (25.0, -20.0), (25.0, 0.0))); + assert!(!crosses(&rect, (25.0, 20.0), (25.0, 0.0))); + assert!(!crosses(&rect, (15.0, -5.0), (30.0, -5.0))); + assert!(!crosses(&rect, (45.0, -5.0), (30.0, -5.0))); +} + +#[test] +fn test_an_edge_meets_a_parallel_twice_where_it_rises_above_it() { + let lng = S1Interval::new(0.0, 10f64.to_radians()); + let (a, b) = (point((9.99, 2.0)), point((9.99, 8.0))); + assert!(S2LatLngRect::intersects_lat_edge( + &a, + &b, + 10f64.to_radians(), + &lng + )); + assert!(!S2LatLngRect::intersects_lat_edge( + &a, + &b, + 10.1f64.to_radians(), + &lng + )); + // Where the edge meets the parallel, east of the longitudes asked about. + let west = S1Interval::new(0.0, 1f64.to_radians()); + assert!(!S2LatLngRect::intersects_lat_edge( + &a, + &b, + 10f64.to_radians(), + &west + )); +} diff --git a/src/spatial/tests/region_coverer_tests.rs b/src/spatial/tests/region_coverer_tests.rs new file mode 100644 index 0000000000..309fdadc6a --- /dev/null +++ b/src/spatial/tests/region_coverer_tests.rs @@ -0,0 +1,138 @@ +use super::super::s2cap::S2Cap; +use super::*; + +fn get_latlng_point(lat_degrees: f64, lng_degrees: f64) -> [f64; 3] { + let lat = lat_degrees.to_radians(); + let lng = lng_degrees.to_radians(); + let cos_lat = lat.cos(); + [cos_lat * lng.cos(), cos_lat * lng.sin(), lat.sin()] +} + +#[test] +fn test_covering_respects_max_level() { + // THIS TEST WAS FAILING - now should pass + let coverer = RegionCoverer::new(CovererOptions::new().max_level(5).max_cells(100)); + let center = get_latlng_point(0.0, 0.0); + let cap = S2Cap::from_center_angle(center, 0.001); // Tiny cap + + let covering = coverer.get_covering(&cap); + + // All cells should be at level <= 5 + for id in covering.cell_ids() { + assert!( + id.level() <= 5, + "Cell at level {} exceeds max_level 5", + id.level() + ); + } +} + +#[test] +fn test_get_initial_candidates_respects_max_level() { + // Direct test of get_initial_candidates behavior + let coverer = RegionCoverer::new(CovererOptions::new().max_level(3).max_cells(100)); + let center = get_latlng_point(0.0, 0.0); + let cap = S2Cap::from_center_angle(center, 0.0001); // Very tiny cap + + let mut cells = Vec::new(); + coverer.get_initial_candidates(&cap, &mut cells); + + // All initial candidates should be at level <= 3 + for id in &cells { + assert!( + id.level() <= 3, + "Initial candidate at level {} exceeds max_level 3", + id.level() + ); + } +} + +#[test] +fn test_canonicalize_respects_max_level() { + let coverer = RegionCoverer::new(CovererOptions::new().max_level(5).max_cells(10)); + + // Create a covering with cells too deep + let mut covering = vec![ + S2CellId::from_face(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0), // Level 6 + S2CellId::from_face(1) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0), // Level 7 + ]; + + coverer.canonicalize_covering(&mut covering); + + // All cells should now be at level <= 5 + for id in &covering { + assert!(id.level() <= 5); + } + assert!(coverer.is_canonical(&covering)); +} + +#[test] +fn test_fast_covering_respects_max_level() { + let coverer = RegionCoverer::new(CovererOptions::new().max_level(4).max_cells(10)); + let center = get_latlng_point(0.0, 0.0); + let cap = S2Cap::from_center_angle(center, 0.0001); // Very tiny cap + + let mut covering = Vec::new(); + coverer.get_fast_covering(&cap, &mut covering); + + // Fast covering should also respect max_level + for id in &covering { + assert!( + id.level() <= 4, + "Fast covering cell at level {} exceeds max_level 4", + id.level() + ); + } +} + +#[test] +fn test_is_canonical_rejects_too_deep_cells() { + let coverer = RegionCoverer::new(CovererOptions::new().max_level(5)); + + // A covering with a cell at level 6 should not be canonical + let bad_covering = vec![S2CellId::from_face(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0) + .child(0)]; + + assert!(!coverer.is_canonical(&bad_covering)); +} + +#[test] +fn test_coverings_refine_the_largest_cells_first() { + // Refining the finest cells first, rather than the coarsest as S2 does, leaves coverings + // more than twice as loose. + let coverer = RegionCoverer::new(CovererOptions::default()); + let mut ratios = Vec::new(); + for i in 0..200 { + let lat = (i as f64 * 37.0) % 160.0 - 80.0; + let lng = (i as f64 * 53.0) % 360.0 - 180.0; + let radius = 1e-4 * 1000f64.powf((i % 10) as f64 / 9.0); + let cap = S2Cap::from_center_angle(get_latlng_point(lat, lng), radius); + let covering = coverer.get_covering(&cap); + let area: f64 = covering + .cell_ids() + .iter() + .map(|id| S2Cell::average_area(id.level())) + .sum(); + ratios.push(area / cap.get_area()); + } + let mean = ratios.iter().sum::() / ratios.len() as f64; + assert!(mean < 3.0, "mean covering to cap area ratio {mean}"); +} diff --git a/src/spatial/tests/s1chord_angle_tests.rs b/src/spatial/tests/s1chord_angle_tests.rs new file mode 100644 index 0000000000..99fff561cd --- /dev/null +++ b/src/spatial/tests/s1chord_angle_tests.rs @@ -0,0 +1,28 @@ +use super::*; + +#[test] +fn test_from_points() { + let x = [1.0, 0.0, 0.0]; + let y = [0.0, 1.0, 0.0]; + let a = S1ChordAngle::from_points(&x, &y); + // 90 degrees = chord^2 of 2 + assert!((a.length2() - 2.0).abs() < 1e-14); + assert!((a.to_degrees() - 90.0).abs() < 1e-10); +} + +#[test] +fn test_sin_cos() { + let a = S1ChordAngle::from_radians(PI / 3.0); // 60 degrees + let expected_sin = (PI / 3.0).sin(); + let expected_cos = (PI / 3.0).cos(); + assert!((a.sin() - expected_sin).abs() < 1e-14); + assert!((a.cos() - expected_cos).abs() < 1e-14); +} + +#[test] +fn test_addition() { + let a = S1ChordAngle::from_degrees(30.0); + let b = S1ChordAngle::from_degrees(60.0); + let c = a + b; + assert!((c.to_degrees() - 90.0).abs() < 1e-10); +} diff --git a/src/spatial/tests/s1interval_tests.rs b/src/spatial/tests/s1interval_tests.rs new file mode 100644 index 0000000000..6dab108cba --- /dev/null +++ b/src/spatial/tests/s1interval_tests.rs @@ -0,0 +1,253 @@ +use super::*; + +const EPSILON: f64 = 1e-14; + +fn approx_eq(a: f64, b: f64) -> bool { + (a - b).abs() < EPSILON +} + +#[test] +fn test_empty_and_full() { + let empty = S1Interval::empty(); + let full = S1Interval::full(); + + assert!(empty.is_empty()); + assert!(!empty.is_full()); + assert!(empty.is_inverted()); + assert!(empty.length() < 0.0); + + assert!(full.is_full()); + assert!(!full.is_empty()); + assert!(!full.is_inverted()); + assert!(approx_eq(full.length(), 2.0 * PI)); +} + +#[test] +fn test_from_point() { + let i = S1Interval::from_point(1.0); + assert!(!i.is_empty()); + assert_eq!(i.lo(), 1.0); + assert_eq!(i.hi(), 1.0); + assert_eq!(i.length(), 0.0); + + // -PI should be normalized to PI + let i2 = S1Interval::from_point(-PI); + assert_eq!(i2.lo(), PI); + assert_eq!(i2.hi(), PI); +} + +#[test] +fn test_from_point_pair() { + // Normal interval + let i1 = S1Interval::from_point_pair(0.0, 1.0); + assert_eq!(i1.lo(), 0.0); + assert_eq!(i1.hi(), 1.0); + + // Inverted interval (crossing PI) + let i2 = S1Interval::from_point_pair(3.0, -3.0); + assert!(i2.is_inverted()); + assert_eq!(i2.lo(), 3.0); + assert_eq!(i2.hi(), -3.0); +} + +#[test] +fn test_inverted_interval() { + // An interval from 3 to -3 radians should wrap around + let i = S1Interval::new(3.0, -3.0); + assert!(i.is_inverted()); + assert!(!i.is_empty()); + assert!(!i.is_full()); + + // Length should be ~0.28 (6.28 - 6.0) + assert!(approx_eq(i.length(), 2.0 * PI - 6.0)); + + // Should contain PI (the wrap point) + assert!(i.contains_point(PI)); + assert!(i.contains_point(-PI)); + assert!(i.contains_point(3.1)); + assert!(i.contains_point(-3.1)); + + // Should not contain points in the gap + assert!(!i.contains_point(0.0)); + assert!(!i.contains_point(2.0)); +} + +#[test] +fn test_center() { + // Normal interval + let i1 = S1Interval::new(0.0, 2.0); + assert!(approx_eq(i1.center(), 1.0)); + + // Inverted interval spanning PI + let i2 = S1Interval::new(3.0, -3.0); + // Center should be at PI (or equivalent) + let c = i2.center(); + assert!(approx_eq(c.abs(), PI)); +} + +#[test] +fn test_contains_point() { + let normal = S1Interval::new(0.0, 2.0); + assert!(normal.contains_point(0.0)); + assert!(normal.contains_point(1.0)); + assert!(normal.contains_point(2.0)); + assert!(!normal.contains_point(-1.0)); + assert!(!normal.contains_point(3.0)); + + // Full contains everything + assert!(S1Interval::full().contains_point(0.0)); + assert!(S1Interval::full().contains_point(PI)); + assert!(S1Interval::full().contains_point(-PI)); + + // Empty contains nothing + assert!(!S1Interval::empty().contains_point(0.0)); +} + +#[test] +fn test_add_point() { + let mut i = S1Interval::empty(); + i.add_point(0.0); + assert_eq!(i.lo(), 0.0); + assert_eq!(i.hi(), 0.0); + + i.add_point(1.0); + assert_eq!(i.lo(), 0.0); + assert_eq!(i.hi(), 1.0); + + // Adding a point that would create a shorter inverted interval + i.add_point(-0.5); + assert_eq!(i.lo(), -0.5); + assert_eq!(i.hi(), 1.0); +} + +#[test] +fn test_add_point_wraparound() { + // Start near PI, add point near -PI + let mut i = S1Interval::from_point(3.0); + i.add_point(-3.0); + // Should create an inverted interval through PI + assert!(i.is_inverted()); + assert!(i.contains_point(PI)); +} + +#[test] +fn test_union() { + let a = S1Interval::new(0.0, 1.0); + let b = S1Interval::new(0.5, 2.0); + let u = a.union(&b); + assert_eq!(u.lo(), 0.0); + assert_eq!(u.hi(), 2.0); + + // Union with empty + assert_eq!(a.union(&S1Interval::empty()), a); + assert_eq!(S1Interval::empty().union(&a), a); + + // Union that becomes full + let c = S1Interval::new(-2.0, 2.0); + let d = S1Interval::new(2.0, -1.0); // inverted + let u2 = c.union(&d); + assert!(u2.is_full()); +} + +#[test] +fn test_intersection() { + let a = S1Interval::new(0.0, 2.0); + let b = S1Interval::new(1.0, 3.0); + let i = a.intersection(&b); + assert_eq!(i.lo(), 1.0); + assert_eq!(i.hi(), 2.0); + + // Disjoint intervals + let c = S1Interval::new(0.0, 1.0); + let d = S1Interval::new(2.0, 3.0); + assert!(c.intersection(&d).is_empty()); +} + +#[test] +fn test_complement() { + let i = S1Interval::new(0.0, 1.0); + let c = i.complement(); + assert_eq!(c.lo(), 1.0); + assert_eq!(c.hi(), 0.0); + assert!(c.is_inverted()); + assert!(c.contains_point(PI)); + assert!(!c.contains_point(0.5)); + + // Singleton complement is full + let s = S1Interval::from_point(0.0); + assert!(s.complement().is_full()); +} + +#[test] +fn test_expanded() { + let i = S1Interval::new(0.0, 1.0); + + // Expand + let e = i.expanded(0.5); + assert!(approx_eq(e.lo(), -0.5)); + assert!(approx_eq(e.hi(), 1.5)); + + // Contract + let c = i.expanded(-0.25); + assert!(approx_eq(c.lo(), 0.25)); + assert!(approx_eq(c.hi(), 0.75)); + + // Contract to empty + let e2 = i.expanded(-1.0); + assert!(e2.is_empty()); + + // Expand to full + let i2 = S1Interval::new(-2.0, 2.0); + let e3 = i2.expanded(2.0); + assert!(e3.is_full()); +} + +#[test] +fn test_positive_distance() { + assert!(approx_eq(positive_distance(0.0, 1.0), 1.0)); + assert!(approx_eq(positive_distance(1.0, 0.0), 2.0 * PI - 1.0)); + assert!(approx_eq(positive_distance(-1.0, 1.0), 2.0)); + assert!(approx_eq(positive_distance(1.0, -1.0), 2.0 * PI - 2.0)); +} + +#[test] +fn test_project() { + let i = S1Interval::new(0.0, 1.0); + assert_eq!(i.project(0.5), 0.5); + assert_eq!(i.project(0.0), 0.0); + assert_eq!(i.project(1.0), 1.0); + // Point outside, closer to lo + assert_eq!(i.project(-0.1), 0.0); + // Point outside, closer to hi + assert_eq!(i.project(1.1), 1.0); +} + +#[test] +fn test_approx_equals() { + let a = S1Interval::new(0.0, 1.0); + let b = S1Interval::new(0.0 + 1e-16, 1.0 - 1e-16); + assert!(a.approx_equals(&b, 1e-15)); + assert!(!a.approx_equals(&b, 1e-17)); + + // Empty intervals + let e1 = S1Interval::empty(); + let e2 = S1Interval::new(0.0, 1e-16); + assert!(e1.approx_equals(&e2, 1e-15)); +} + +#[test] +fn test_normalization_of_negative_pi() { + // -PI should become PI except for Full and Empty + let i = S1Interval::new(-PI, 1.0); + assert_eq!(i.lo(), PI); + assert_eq!(i.hi(), 1.0); + + let i2 = S1Interval::new(1.0, -PI); + assert_eq!(i2.lo(), 1.0); + assert_eq!(i2.hi(), PI); + + // Full is the exception + let f = S1Interval::full(); + assert_eq!(f.lo(), -PI); + assert_eq!(f.hi(), PI); +} diff --git a/src/spatial/tests/s2cap_tests.rs b/src/spatial/tests/s2cap_tests.rs new file mode 100644 index 0000000000..35671606f3 --- /dev/null +++ b/src/spatial/tests/s2cap_tests.rs @@ -0,0 +1,215 @@ +use std::f64::consts::{FRAC_PI_2, PI}; + +use super::*; +use crate::spatial::math::normalize; + +const EPS: f64 = 1e-14; + +fn get_latlng_point(lat_degrees: f64, lng_degrees: f64) -> [f64; 3] { + let lat = lat_degrees.to_radians(); + let lng = lng_degrees.to_radians(); + let cos_lat = lat.cos(); + [cos_lat * lng.cos(), cos_lat * lng.sin(), lat.sin()] +} + +#[test] +fn test_basic_properties() { + let empty = S2Cap::empty(); + let full = S2Cap::full(); + + assert!(empty.is_empty()); + assert!(!empty.is_full()); + assert!(empty.complement().is_full()); + + assert!(full.is_full()); + assert!(!full.is_empty()); + assert!(full.complement().is_empty()); + + assert!((full.height() - 2.0).abs() < EPS); + assert!((full.get_radius_radians() - PI).abs() < 1e-10); +} + +#[test] +fn test_singleton_cap() { + let xaxis = S2Cap::from_point([1.0, 0.0, 0.0]); + assert!(xaxis.contains_point(&[1.0, 0.0, 0.0])); + assert!(!xaxis.contains_point(&[1.0, 1e-20, 0.0])); + assert!(xaxis.get_radius_radians().abs() < EPS); + assert!(xaxis.height().abs() < EPS); + + // Complement of singleton is full + let xcomp = xaxis.complement(); + assert!(xcomp.is_full()); + assert!(xcomp.contains_point(&xaxis.center())); +} + +#[test] +fn test_hemisphere() { + let hemi = S2Cap::from_center_height(normalize(&[1.0, 0.0, 1.0]), 1.0); + assert!((hemi.complement().height() - 1.0).abs() < EPS); + assert!(hemi.contains_point(&[1.0, 0.0, 0.0])); + assert!(!hemi.complement().contains_point(&[1.0, 0.0, 0.0])); +} + +#[test] +fn test_cap_cap_containment() { + let empty = S2Cap::empty(); + let full = S2Cap::full(); + let xaxis = S2Cap::from_point([1.0, 0.0, 0.0]); + + assert!(empty.contains_cap(&empty)); + assert!(full.contains_cap(&empty)); + assert!(full.contains_cap(&full)); + assert!(!xaxis.contains_cap(&full)); + assert!(xaxis.contains_cap(&empty)); +} + +#[test] +fn test_cap_cap_intersection() { + let empty = S2Cap::empty(); + let full = S2Cap::full(); + let xaxis = S2Cap::from_point([1.0, 0.0, 0.0]); + + assert!(!empty.intersects_cap(&empty)); + assert!(full.intersects_cap(&full)); + assert!(!full.intersects_cap(&empty)); + assert!(full.intersects_cap(&xaxis)); +} + +#[test] +fn test_get_rect_bound_empty_full() { + assert!(S2Cap::empty().get_rect_bound().is_empty()); + assert!(S2Cap::full().get_rect_bound().is_full()); +} + +#[test] +fn test_get_rect_bound_south_pole() { + let cap = S2Cap::from_center_angle(get_latlng_point(-45.0, 57.0), 50.0_f64.to_radians()); + let rect = cap.get_rect_bound(); + + // Should include south pole + assert!((rect.lat_lo() - (-FRAC_PI_2)).abs() < 1e-10); + assert!(rect.lng().is_full()); +} + +#[test] +fn test_get_rect_bound_equator() { + let cap = S2Cap::from_center_angle(get_latlng_point(0.0, 50.0), 20.0_f64.to_radians()); + let rect = cap.get_rect_bound(); + + assert!((rect.lat_lo().to_degrees() - (-20.0)).abs() < 1e-10); + assert!((rect.lat_hi().to_degrees() - 20.0).abs() < 1e-10); + assert!((rect.lng_lo().to_degrees() - 30.0).abs() < 1e-10); + assert!((rect.lng_hi().to_degrees() - 70.0).abs() < 1e-10); +} + +#[test] +fn test_region_contains_cell() { + let full = S2Cap::full(); + let empty = S2Cap::empty(); + + let root_cell = S2Cell::from_face(0); + assert!(full.contains_cell(&root_cell)); + assert!(!empty.contains_cell(&root_cell)); +} + +#[test] +fn test_region_may_intersect() { + let full = S2Cap::full(); + let empty = S2Cap::empty(); + + let root_cell = S2Cell::from_face(0); + assert!(full.may_intersect(&root_cell)); + assert!(!empty.may_intersect(&root_cell)); +} + +#[test] +fn test_region_singleton_may_intersect() { + // A singleton cap at face center should intersect that face + for face in 0..6 { + let center = S2Cell::from_face(face).get_center(); + let singleton = S2Cap::from_point(center); + + let root_cell = S2Cell::from_face(face); + assert!(singleton.may_intersect(&root_cell)); + + // Should not intersect other faces + for other_face in 0..6 { + if other_face != face { + let other_cell = S2Cell::from_face(other_face); + assert!(!singleton.may_intersect(&other_cell)); + } + } + } +} + +#[test] +fn test_complement() { + let small = S2Cap::from_center_angle([1.0, 0.0, 0.0], 0.1); + let comp = small.complement(); + + // Complement should be large (almost full) + assert!((comp.get_radius_radians() - (PI - 0.1)).abs() < 1e-10); + + // Complement center should be opposite + assert!((comp.center()[0] - (-1.0)).abs() < EPS); + + // Double complement of empty is empty + assert!(S2Cap::empty().complement().complement().is_empty()); + + // Double complement of full is full + assert!(S2Cap::full().complement().complement().is_full()); +} + +#[test] +fn test_add_point() { + let mut cap = S2Cap::empty(); + cap.add_point(&[1.0, 0.0, 0.0]); + assert!(!cap.is_empty()); + assert!(cap.contains_point(&[1.0, 0.0, 0.0])); + + cap.add_point(&[0.0, 1.0, 0.0]); + assert!(cap.contains_point(&[1.0, 0.0, 0.0])); + assert!(cap.contains_point(&[0.0, 1.0, 0.0])); +} + +#[test] +fn test_expanded() { + let cap = S2Cap::from_point([1.0, 0.0, 0.0]); + let expanded = cap.expanded(0.1); + + assert!((expanded.get_radius_radians() - 0.1).abs() < 1e-10); + assert!(expanded.contains_point(&[1.0, 0.0, 0.0])); + + // Expanding empty cap stays empty + assert!(S2Cap::empty().expanded(0.1).is_empty()); +} + +#[test] +fn test_cell_intersection_hemisphere() { + // A hemisphere should contain 3 face cells and may-intersect more + let hemi = S2Cap::from_center_angle([0.0, 0.0, 1.0], FRAC_PI_2); + + let north_cell = S2Cell::from_face(2); // North pole face + assert!(hemi.may_intersect(&north_cell)); + + let south_cell = S2Cell::from_face(5); // South pole face + assert!(!hemi.may_intersect(&south_cell)); +} + +#[test] +fn test_cell_covering() { + // Small cap should return 3-4 cells + let small = S2Cap::from_center_angle([1.0, 0.0, 0.0], 0.01); + let mut cell_ids = Vec::new(); + small.get_cell_union_bound(&mut cell_ids); + + assert!(cell_ids.len() >= 3); + assert!(cell_ids.len() <= 4); + + // Large cap should return 6 face cells + let large = S2Cap::from_center_angle([1.0, 0.0, 0.0], 2.0); + cell_ids.clear(); + large.get_cell_union_bound(&mut cell_ids); + assert_eq!(cell_ids.len(), 6); +} diff --git a/src/spatial/tests/s2metrics_tests.rs b/src/spatial/tests/s2metrics_tests.rs new file mode 100644 index 0000000000..9b2ffd99df --- /dev/null +++ b/src/spatial/tests/s2metrics_tests.rs @@ -0,0 +1,14 @@ +use super::*; +#[test] +fn test_s2metrics_get_level_for_min_width() { + // Very small value -> level 30 + assert_eq!(S2Metrics::get_level_for_min_width(1e-10), 30); + + // Very large value -> level 0 or -1 + assert!(S2Metrics::get_level_for_min_width(2.0) <= 0); + + // Moderate value -> somewhere in between + let level = S2Metrics::get_level_for_min_width(0.01); + assert!(level > 0); + assert!(level < 30); +} diff --git a/src/spatial/within_query.rs b/src/spatial/within_query.rs new file mode 100644 index 0000000000..e0e3f00216 --- /dev/null +++ b/src/spatial/within_query.rs @@ -0,0 +1,82 @@ +//! Query for the geometries within a distance of a point. +//! +//! The closest edge query finds them too, nearest first, and keeps the distance of each. This +//! query only decides which they are, as a search for the area around the point: see +//! `area_query`. Interiors count: a polygon that contains the point is at distance zero. + +use common::BitSet; + +use super::area_query::{search_meeting, Area, Covering}; +use super::cell_index_reader::CellIndexReader; +use super::edge_cache::EdgeCache; +use super::region::Region; +use super::s1chord_angle::S1ChordAngle; +use super::s2cap::S2Cap; +use super::s2edge_distances::is_distance_less; +use super::sphere::Sphere; + +/// How much nearer than the distance, as a squared chord length, everything in a covering cell +/// must be for the cell to be taken whole. A geometry in an index cell may lie outside it by +/// the cell's padding. +const WITHIN_MARGIN: f64 = 1e-13; + +/// What lies nearer to a point than a distance. +struct Around { + center: [f64; 3], + distance: S1ChordAngle, +} + +impl Area for Around { + fn contains(&self, point: &[f64; 3]) -> bool { + is_distance_less(&self.center, point, point, self.distance) + } + + fn meets_edge(&self, a: &[f64; 3], b: &[f64; 3]) -> bool { + is_distance_less(&self.center, a, b, self.distance) + } + + fn inside(&self) -> &[f64; 3] { + &self.center + } +} + +/// Prepared query for the geometries nearer to a point than a distance, built once and applied +/// per segment. +pub struct WithinQuery { + around: Around, + /// None when nothing is nearer than the distance, which is then zero or less. + covering: Option, +} + +impl WithinQuery { + /// Build the query from a point, as a unit vector, and a distance. + pub fn new(center: [f64; 3], distance: S1ChordAngle) -> Self { + let covering = (distance > S1ChordAngle::zero()).then(|| { + // Everything is nearer than an infinite distance, which a cap cannot have. + if distance.is_infinity() { + return Covering::new(&S2Cap::full(), |_| true); + } + let length2 = (distance.length2() - WITHIN_MARGIN).max(0.0); + let nearer = S2Cap::new(center, S1ChordAngle::from_length2(length2)); + Covering::new(&S2Cap::new(center, distance), |cell| { + nearer.contains_cell(cell) + }) + }); + WithinQuery { + around: Around { center, distance }, + covering, + } + } + + /// Add the documents of one segment that match to `docs`. + pub fn search_segment( + &self, + cell_reader: &CellIndexReader, + edge_cache: &mut EdgeCache<'_, Sphere>, + docs: &mut BitSet, + ) { + if let Some(covering) = &self.covering { + search_meeting(&self.around, covering, cell_reader, edge_cache, docs); + } + } +} diff --git a/src/spatial/writer.rs b/src/spatial/writer.rs new file mode 100644 index 0000000000..c2670daa83 --- /dev/null +++ b/src/spatial/writer.rs @@ -0,0 +1,190 @@ +//! Spatial field writer. +//! +//! Accumulates geometry per field during indexing. At segment flush, builds cell indexes and +//! streams edge data through the SpatialSerializer. +use std::collections::HashMap; +use std::io; + +use crate::schema::Field; +use crate::spatial::geometry::Geometry; +use crate::spatial::geometry_set::{ + to_geometry_set_with_orientation, EdgeSet, GeometrySet, RingOrientation, +}; +use crate::spatial::plane::Plane; +use crate::spatial::serializer::SpatialSerializer; +use crate::spatial::sphere::Sphere; +use crate::DocId; + +/// Per-field accumulated documents with their geometries. +#[derive(Default)] +struct FieldData { + sets: Vec, +} + +/// Accumulates spatial geometry during indexing, then serializes at segment flush. +#[derive(Default)] +pub struct SpatialWriter { + data_by_field: HashMap, + /// The sum of `GeometrySet::mem_usage` over every accumulated set. The indexer asks for the + /// memory usage after every document, so it is kept current wherever a set is added. + sets_mem_usage: usize, +} + +impl SpatialWriter { + /// Add a geometry for a document and field, with polygon rings oriented as the field's are. + /// The values of a document's field make up one geometry set, so that a query matches the + /// document's geometry as a whole: contained only if all of it is. + pub fn add_geometry( + &mut self, + doc_id: DocId, + field: Field, + geometry: Geometry, + orientation: RingOrientation, + ) { + let data = self.data_by_field.entry(field).or_default(); + let projected = geometry.project::(); + let set = to_geometry_set_with_orientation(&projected, doc_id, orientation); + // An empty geometry, such as a polygon with no rings, adds nothing. + if set.members.is_empty() { + return; + } + match data.sets.last_mut() { + Some(last) if last.doc_id == doc_id => { + let member_size = std::mem::size_of::(); + // The incoming members keep their vertex and ring buffers. Only their outer + // buffer is discarded; account for the destination buffer's growth separately. + let added_buffers = set.mem_usage() - set.members.capacity() * member_size; + let previous_capacity = last.members.capacity(); + last.members.extend(set.members); + self.sets_mem_usage += + added_buffers + (last.members.capacity() - previous_capacity) * member_size; + } + _ => { + self.sets_mem_usage += set.mem_usage(); + data.sets.push(set); + } + } + } + + /// Memory usage estimate: the heap held by the sets, and the arrays that hold the sets + /// themselves, spare capacity included. + pub fn mem_usage(&self) -> usize { + self.sets_mem_usage + + self + .data_by_field + .values() + .map(|data| data.sets.capacity() * std::mem::size_of::()) + .sum::() + } + + /// Renumber the documents, as a segment sorted by a field is, and order each field's sets + /// by their new ids. + pub(crate) fn remap_doc_ids(&mut self, new_doc_id: impl Fn(DocId) -> DocId) { + for data in self.data_by_field.values_mut() { + for set in &mut data.sets { + set.doc_id = new_doc_id(set.doc_id); + } + data.sets.sort_by_key(|set| set.doc_id); + } + } + + /// Serialize all fields. + pub fn serialize(&mut self, mut serializer: SpatialSerializer) -> io::Result<()> { + for (field, data) in &self.data_by_field { + serializer.serialize_field(*field, &data.sets)?; + } + serializer.close()?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The memory usage recomputed from every set, spare capacity included. + fn summed_mem_usage(writer: &SpatialWriter) -> usize { + writer + .data_by_field + .values() + .map(|data| { + data.sets.capacity() * std::mem::size_of::() + + data.sets.iter().map(GeometrySet::mem_usage).sum::() + }) + .sum() + } + + #[test] + fn test_values_of_a_document_make_one_set() { + let mut writer = SpatialWriter::default(); + let field = Field::from_field_id(0); + writer.add_geometry( + 0, + field, + Geometry::Point([1.0, 1.0]), + RingOrientation::Smaller, + ); + writer.add_geometry( + 1, + field, + Geometry::Point([2.0, 2.0]), + RingOrientation::Smaller, + ); + writer.add_geometry( + 1, + field, + Geometry::MultiPoint(vec![[3.0, 3.0], [4.0, 4.0]]), + RingOrientation::Smaller, + ); + let sets = &writer.data_by_field[&field].sets; + let members: Vec<(u32, usize)> = sets.iter().map(|s| (s.doc_id, s.members.len())).collect(); + assert_eq!(members, [(0, 1), (1, 3)]); + assert_eq!(writer.mem_usage(), summed_mem_usage(&writer)); + } + + #[test] + fn test_mem_usage_counts_each_geometry_once() { + let mut writer = SpatialWriter::default(); + assert_eq!(writer.mem_usage(), 0); + let (points, lines) = (Field::from_field_id(0), Field::from_field_id(1)); + for doc in 0..1000 { + writer.add_geometry( + doc, + points, + Geometry::Point([19.94, 50.06]), + RingOrientation::Smaller, + ); + if doc % 3 == 0 { + let line = vec![[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]]; + writer.add_geometry( + doc, + lines, + Geometry::LineString(line), + RingOrientation::Smaller, + ); + } + assert_eq!(writer.mem_usage(), summed_mem_usage(&writer)); + } + } + + #[test] + fn test_mem_usage_for_multiple_values_includes_spare_capacity() { + let mut writer = SpatialWriter::default(); + let field = Field::from_field_id(0); + for doc in 0..2 { + for value in 0..257 { + writer.add_geometry( + doc, + field, + Geometry::Point([value as f64 % 180.0, 1.0]), + RingOrientation::Smaller, + ); + assert_eq!(writer.mem_usage(), summed_mem_usage(&writer)); + } + } + let sets = &writer.data_by_field[&field].sets; + assert_eq!(sets.len(), 2); + assert_eq!(sets[1].members.len(), 257); + assert!(sets[1].members.capacity() > sets[1].members.len()); + } +} diff --git a/src/spatial/xor.rs b/src/spatial/xor.rs new file mode 100644 index 0000000000..6e03be3b16 --- /dev/null +++ b/src/spatial/xor.rs @@ -0,0 +1,198 @@ +//! XOR delta compression for f64 polygon coordinates. +//! +//! Lossless compression for floating-point coordinates using XOR delta encoding on IEEE +//! 754 bit patterns with bit packing. Designed for per-geometry random access in the document +//! store, where each geometry compresses independently without requiring sequential decompression. +//! Additionally used to store per-geometry edge and point storage in the index referenced by +//! quadtree cells. +//! +//! Spatially local coordinates share most high-order bits. A municipal boundary spanning 1km has +//! consecutive vertices typically within 100-500 meters, meaning their f64 bit patterns share +//! 30-40 bits. XOR reveals these common bits as zeros, which bit-pack efficiently with minimal +//! bit-width. +//! +//! Unlike delta.rs which uses arithmetic deltas for integers, this module operates on f64 bit +//! patterns directly to preserve exact floating-point values. + +/// Compresses f64 coordinates using XOR delta encoding with automatic raw fallback. +/// +/// First pass computes XOR deltas between consecutive coordinates to find the maximum, which +/// determines the bit width needed. Second pass writes a header (1 byte bit width + 8 bytes first +/// coordinate) followed by bit-packed XOR deltas. If the compressed output would be larger than +/// raw storage (8 bytes per coordinate) or if the delta is too large to allow for a 7 bit +/// remainder, it falls back to raw encoding. +/// +/// Returns a byte vector that can be decompressed with `decompress_f64()` to recover exact +/// original values. +pub fn compress_f64(values: &[f64]) -> Vec { + // Subsequent code assumes a first element. + if values.is_empty() { + return Vec::new(); + } + + // Compute max XOR delta + let mut max_xor: u64 = 0; + let mut previous = values[0].to_bits(); + for &value in &values[1..] { + let bits = value.to_bits(); + let xor = bits ^ previous; + max_xor = max_xor.max(xor); + previous = bits; + } + + // Bits needed for delta. + let bits = if max_xor == 0 { + 0 + } else { + 64 - max_xor.leading_zeros() as u8 + }; + + // Determine if compression will indeed compress. + let count = values.len(); + let compressed_size = 9 + ((count - 1) * bits as usize).div_ceil(8); + let raw_size = count * 8; + + // Return the uncompressed array if we do not have enough room to buffer seven bits. Return the + // uncompressed array if compression buys us nothing. + if bits > 57 || compressed_size >= raw_size { + let mut output = Vec::with_capacity(raw_size); + for &value in values { + output.extend_from_slice(&value.to_le_bytes()); + } + return output; + } + + // Header: bit width (1 byte) + first value (8 bytes). + let mut output = Vec::with_capacity(compressed_size); + output.push(bits); + output.extend_from_slice(&values[0].to_le_bytes()); + + // Bit-packing of the XOR deltas. + let mut buffer: u64 = 0; + let mut buffer_bits: u32 = 0; + let mut previous = values[0].to_bits(); + + for &value in &values[1..] { + let current = value.to_bits(); + let xor = current ^ previous; + + // Shift buffer left, add new bits. + buffer = (buffer << bits) | xor; + buffer_bits += bits as u32; + + // Flush complete bytes. + while buffer_bits >= 8 { + buffer_bits -= 8; + output.push((buffer >> buffer_bits) as u8); + } + + previous = current; + } + + // Flush remaining bits left-aligned. + if buffer_bits > 0 { + output.push((buffer << (8 - buffer_bits)) as u8); + } + + output +} + +/// Decompresses f64 coordinates from XOR delta or raw encoding. +/// +/// If the byte length is equal to the uncompressed size, then the data is raw and copied directly. +/// Otherwise, reads the delta size from the first byte, the first coordinate from next 8 bytes, +/// then unpacks bit-packed XOR deltas reconstructing the original sequence. A value in the +/// sequence is derived by applying the XOR delta to the value preceding it. +/// +/// Returns exact f64 values that were passed to `compress_f64()`. +pub fn decompress_f64(bytes: &[u8], count: usize) -> Vec { + // Check for raw encoding, read raw values. + let mut values = Vec::with_capacity(count); + if bytes.len() == count * 8 { + for i in 0..count { + values.push(f64::from_le_bytes( + bytes[i * 8..(i + 1) * 8].try_into().unwrap(), + )); + } + return values; + } + + // Read header as unsigned integer for use with XOR. + let bits = bytes[0] as u32; + let first = u64::from_le_bytes(bytes[1..9].try_into().unwrap()); + + // Convert unisgned integer bits to float. + values.push(f64::from_bits(first)); + + // Unpack bit-packed XOR deltas + let mut buffer: u64 = 0; + let mut buffer_bits: u32 = 0; + let mut offset = 9; + let mut previous = first; + let mask = (1u64 << bits) - 1; + + for _ in 1..count { + // Refill buffer. + while buffer_bits < bits && offset < bytes.len() { + buffer = (buffer << 8) | bytes[offset] as u64; + buffer_bits += 8; + offset += 1; + } + + // Extract delta. + buffer_bits -= bits; + let xor = (buffer >> buffer_bits) & mask; + + let current = previous ^ xor; + values.push(f64::from_bits(current)); + previous = current; + } + + values +} + +#[cfg(test)] +mod test { + use super::*; + + // Compression ratio of 1.43 on small municipal polygons (12 vertices). Per FCBench (VLDB + // 2024), this exceeds the median lossless floating-point compression ratio of 1.16 and matches + // Gorilla/Chimp XOR-delta methods. Block-based methods (bitshuffle+zstd) achieve 1.4-1.5x on + // comparable low-entropy spatial data, but require sequential decompression. Nothing + // scientific about this observation, just a sanity check that we are, at first glance, on the + // right track. + #[test] + fn test_compress_spatial_locality() { + // Latitudes from a small town polygon. Hosmer, South Dakota. + let latitudes = vec![ + 45.577697, 45.571457, 45.571461, 45.571584, 45.571615, 45.57168, 45.57883, 45.586076, + 45.585926, 45.585953, 45.57873, 45.577697, + ]; + let bytes = compress_f64(&latitudes); + // Uncompressed is 96. + assert_eq!(bytes.len(), 67); + // Should decompress to exact original values + let decompressed = decompress_f64(&bytes, latitudes.len()); + assert_eq!(latitudes, decompressed); + } + #[test] + fn test_fallback_to_raw() { + // Random, widely scattered values - poor compression + let values = vec![ + 12345.6789, + -98765.4321, + 0.00001, + 999999.999, + -0.0, + std::f64::consts::PI, + std::f64::consts::E, + 42.0, + ]; + let bytes = compress_f64(&values); + // Should fall back to raw storage + assert_eq!(bytes.len(), values.len() * 8); + // Should still decompress correctly + let decompressed = decompress_f64(&bytes, values.len()); + assert_eq!(values, decompressed); + } +}