1use std::{
2 collections::{HashMap, HashSet, hash_map::Entry},
3 path::Path,
4 sync::Arc,
5};
6
7use anyhow::{Context, Result, bail};
8use serde::Serialize;
9use tantivy::{
10 Index, IndexReader, Score, Term,
11 collector::TopDocs,
12 columnar::Column,
13 query::{BooleanQuery, Occur, PhrasePrefixQuery, Query, QueryParser, TermQuery},
14 schema::{Field, IndexRecordOption},
15};
16
17use crate::{
18 pack::{PackReader, RecordId, RecordPointPrecision, RecordSummary},
19 record::{Layer, OsmObjectType},
20 spatial_index::SpatialIndexReader,
21 text_index::{TextIndexFields, normalize_index_text, open_text_index},
22};
23
24pub struct PackTextSearcher {
25 pack: Arc<PackReader>,
26 spatial: SpatialIndexReader,
27 index: Index,
28 reader: IndexReader,
29 fields: TextIndexFields,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct TextSearchOptions {
34 pub query: String,
35 pub limit: usize,
36 pub layer: Option<String>,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct TextAutocompleteOptions {
41 pub query: String,
42 pub limit: usize,
43 pub layer: Option<String>,
44}
45
46#[derive(Debug, Clone, Serialize, PartialEq)]
47pub struct TextSearchHit {
48 pub record_id: RecordId,
49 pub score: Score,
50 pub record: RecordSummary,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct AddressGeocodeOptions {
55 pub address: String,
56 pub locality: Option<String>,
57 pub region: Option<String>,
58 pub postcode: Option<String>,
59 pub limit: usize,
60 pub layer: Option<String>,
61}
62
63#[derive(Debug, Clone, Serialize, PartialEq)]
64pub struct AddressGeocodeHit {
65 pub query: String,
66 pub hit: TextSearchHit,
67}
68
69pub const DEFAULT_SEARCH_LIMIT: usize = 10;
70pub const MAX_AUTOCOMPLETE_LIMIT: usize = 20;
71const MIN_AUTOCOMPLETE_QUERY_CHARS: usize = 3;
72const AUTOCOMPLETE_PREFIX_MAX_EXPANSIONS: u32 = 1_024;
73const POSTCODE_AREA_RADIUS_M: f64 = 10_000.0;
77const POSTCODE_AREA_CHARS: usize = 3;
78const RANK_WINDOW_FACTOR: usize = 3;
80const MAX_RANK_WINDOW: usize = 256;
81
82impl PackTextSearcher {
83 pub fn open(pack_path: impl AsRef<Path>) -> Result<Self> {
84 Self::from_pack(Arc::new(PackReader::open(pack_path)?))
85 }
86
87 pub fn from_pack(pack: Arc<PackReader>) -> Result<Self> {
89 let index = open_text_index(pack.container())?;
90 let schema = index.schema();
91 let fields = TextIndexFields::from_schema(&schema)?;
92 let reader = index.reader().context("failed to open Tantivy reader")?;
93 let spatial = SpatialIndexReader::open(pack.container(), pack.records().clone())?;
94 Ok(Self {
95 spatial,
96 pack,
97 index,
98 reader,
99 fields,
100 })
101 }
102
103 pub fn search(&self, options: TextSearchOptions) -> Result<Vec<TextSearchHit>> {
104 let limit = effective_limit(options.limit);
105 if limit == 0 {
106 return Ok(Vec::new());
107 }
108
109 let query_text = options.query.trim();
110 if query_text.is_empty() {
111 bail!("search query cannot be empty");
112 }
113
114 let (_, hits) =
115 self.search_variants(query_text, options.layer.as_deref(), limit, |_| Ok(true))?;
116 Ok(hits)
117 }
118
119 pub fn geocode_address(
120 &self,
121 options: AddressGeocodeOptions,
122 ) -> Result<Option<AddressGeocodeHit>> {
123 let limit = effective_limit(options.limit);
124 if limit == 0 {
125 return Ok(None);
126 }
127
128 let address = options.address.trim();
129 if address.is_empty() {
130 return Ok(None);
131 }
132
133 for candidate in address_geocode_candidates(address, options.postcode.as_deref()) {
134 let (query, hits) =
135 self.search_variants(&candidate, options.layer.as_deref(), limit, |hit| {
136 Ok(hit.record.point.is_some()
137 && self.hit_matches_address_context(hit, &options)?)
138 })?;
139 if let Some(hit) = hits.into_iter().next() {
140 return Ok(Some(AddressGeocodeHit { query, hit }));
141 }
142 }
143
144 Ok(None)
145 }
146
147 pub fn autocomplete(&self, options: TextAutocompleteOptions) -> Result<Vec<TextSearchHit>> {
148 let limit = effective_autocomplete_limit(options.limit);
149 if limit == 0 {
150 return Ok(Vec::new());
151 }
152
153 let Some(query_text) = normalize_index_text(options.query.trim()) else {
154 return Ok(Vec::new());
155 };
156 if query_text
157 .chars()
158 .filter(|character| !character.is_whitespace())
159 .count()
160 < MIN_AUTOCOMPLETE_QUERY_CHARS
161 {
162 return Ok(Vec::new());
163 }
164
165 let Some(query) = self.build_autocomplete_query(&query_text, options.layer.as_deref())?
166 else {
167 return Ok(Vec::new());
168 };
169 self.ranked_hits(&query, &query_text, limit)
170 .with_context(|| format!("failed to autocomplete text index for {query_text:?}"))
171 }
172
173 fn build_query(&self, query_text: &str, layer: Option<&str>) -> Result<Box<dyn Query>> {
174 let mut query_parser = QueryParser::for_index(&self.index, self.search_fields());
175 query_parser.set_conjunction_by_default();
176 query_parser.set_field_boost(self.fields.label_text, 3.0);
177 query_parser.set_field_boost(self.fields.name_text, 2.5);
178 query_parser.set_field_boost(self.fields.address_number, 2.0);
179 query_parser.set_field_boost(self.fields.postcode_exact, 2.0);
180
181 let text_query = query_parser
182 .parse_query(query_text)
183 .with_context(|| format!("failed to parse search query {query_text:?}"))?;
184
185 let Some(layer) = layer.map(str::trim).filter(|layer| !layer.is_empty()) else {
186 return Ok(text_query);
187 };
188
189 let layer_query = TermQuery::new(
190 Term::from_field_text(self.fields.layer, layer),
191 IndexRecordOption::Basic,
192 );
193 Ok(Box::new(BooleanQuery::new(vec![
194 (Occur::Must, text_query),
195 (Occur::Must, Box::new(layer_query)),
196 ])))
197 }
198
199 fn search_variants(
204 &self,
205 query_text: &str,
206 layer: Option<&str>,
207 limit: usize,
208 mut accept: impl FnMut(&TextSearchHit) -> Result<bool>,
209 ) -> Result<(String, Vec<TextSearchHit>)> {
210 let mut last_query = query_text.trim().to_string();
211 let mut parse_error = None;
212 let mut parsed_any = false;
213 for variant in search_query_variants(query_text) {
214 last_query = variant.clone();
215 let query = match self.build_query(&variant, layer) {
216 Ok(query) => query,
217 Err(error) => {
218 parse_error = Some(error);
219 continue;
220 }
221 };
222 parsed_any = true;
223 let mut hits = Vec::new();
224 for hit in self
225 .ranked_hits(&query, &variant, limit)
226 .with_context(|| format!("failed to search text index for {variant:?}"))?
227 {
228 if accept(&hit)? {
229 hits.push(hit);
230 }
231 }
232 if !hits.is_empty() {
233 return Ok((variant, hits));
234 }
235 }
236
237 match parse_error {
238 Some(error) if !parsed_any => Err(error),
239 _ => Ok((last_query, Vec::new())),
240 }
241 }
242
243 fn hit_matches_address_context(
244 &self,
245 hit: &TextSearchHit,
246 options: &AddressGeocodeOptions,
247 ) -> Result<bool> {
248 let desired_region = normalized_for_match(options.region.as_deref());
249 let desired_locality = normalized_for_match(options.locality.as_deref());
250 let desired_postcode = normalized_postcode_for_match(options.postcode.as_deref());
251 if desired_region.is_none() && desired_locality.is_none() && desired_postcode.is_none() {
252 return Ok(true);
253 }
254
255 if let Some(desired) = &desired_postcode
256 && !self.hit_in_postcode(hit, desired)?
257 {
258 return Ok(false);
259 }
260
261 let Some(context) = self.pack.boundary_context(hit.record_id)? else {
262 return Ok(desired_region.is_none() && desired_locality.is_none());
263 };
264
265 let mut admin_labels = Vec::new();
266 {
267 let tuple = context.admin_context;
268 for (layer, record_id) in [
269 ("country", tuple.country_record_id),
270 ("region", tuple.region_record_id),
271 ("district", tuple.district_record_id),
272 ("locality", tuple.locality_record_id),
273 ("neighbourhood", tuple.neighbourhood_record_id),
274 ("place", tuple.place_record_id),
275 ] {
276 if let Some(record_id) = record_id
277 && let Some(record) = self.pack.context_record(record_id)?
278 {
279 admin_labels.push((
280 layer,
281 normalized_for_match(Some(&record.label)),
282 normalized_for_match(Some(&record.name)),
283 ));
284 }
285 }
286 }
287
288 if let Some(region) = desired_region
289 && !admin_labels.iter().any(|(layer, label, name)| {
290 *layer == "region"
291 && (label.as_deref() == Some(region.as_str())
292 || name.as_deref() == Some(region.as_str()))
293 })
294 {
295 return Ok(false);
296 }
297
298 if let Some(locality) = desired_locality {
299 let locality_layers = ["district", "locality", "neighbourhood", "place"];
300 if !admin_labels.iter().any(|(layer, label, name)| {
301 locality_layers.contains(layer)
302 && (label.as_deref() == Some(locality.as_str())
303 || name.as_deref() == Some(locality.as_str()))
304 }) {
305 return Ok(false);
306 }
307 }
308
309 Ok(true)
310 }
311
312 fn hit_in_postcode(&self, hit: &TextSearchHit, desired: &str) -> Result<bool> {
320 if let Some(postcode) = self.pack.records().postcode(hit.record_id)? {
321 return Ok(
322 normalized_postcode_for_match(Some(&postcode)).is_none_or(|postcode| {
323 postcode.starts_with(desired) || desired.starts_with(postcode.as_str())
324 }),
325 );
326 }
327 let Some(point) = hit.record.point else {
328 return Ok(true);
329 };
330 let area: String = desired.chars().take(POSTCODE_AREA_CHARS).collect();
331 let in_area = self.spatial.any_context_point(
332 point.lon,
333 point.lat,
334 Layer::Postcode,
335 POSTCODE_AREA_RADIUS_M,
336 |record_id| {
337 Ok(self
338 .pack
339 .records()
340 .postcode(record_id)?
341 .and_then(|postcode| normalized_postcode_for_match(Some(&postcode)))
342 .is_some_and(|postcode| {
343 postcode.starts_with(&area) || area.starts_with(postcode.as_str())
344 }))
345 },
346 )?;
347 Ok(in_area.unwrap_or(true))
349 }
350
351 fn search_fields(&self) -> Vec<tantivy::schema::Field> {
352 vec![
353 self.fields.label_text,
354 self.fields.name_text,
355 self.fields.content_text,
356 self.fields.address_number,
357 self.fields.postcode_exact,
358 ]
359 }
360
361 fn build_autocomplete_query(
362 &self,
363 query_text: &str,
364 layer: Option<&str>,
365 ) -> Result<Option<Box<dyn Query>>> {
366 let tokens = autocomplete_query_tokens(query_text);
367 if tokens.is_empty() {
368 return Ok(None);
369 }
370
371 let mut subqueries = Vec::new();
372
373 let subject_tokens = if tokens.len() > 1 && is_address_number_token(&tokens[0]) {
374 subqueries.push((
375 Occur::Must,
376 Box::new(TermQuery::new(
377 Term::from_field_text(self.fields.address_number, &tokens[0]),
378 IndexRecordOption::Basic,
379 )) as Box<dyn Query>,
380 ));
381 &tokens[1..]
382 } else {
383 tokens.as_slice()
384 };
385
386 if subject_tokens.is_empty() {
387 return Ok(None);
388 }
389
390 subqueries.push((
391 Occur::Must,
392 autocomplete_subject_query(self.fields.autocomplete_subject_text, subject_tokens),
393 ));
394
395 if let Some(layer) = layer.map(str::trim).filter(|layer| !layer.is_empty()) {
396 subqueries.push((
397 Occur::Must,
398 Box::new(TermQuery::new(
399 Term::from_field_text(self.fields.layer, layer),
400 IndexRecordOption::Basic,
401 )) as Box<dyn Query>,
402 ));
403 }
404
405 Ok(Some(Box::new(BooleanQuery::new(subqueries))))
406 }
407
408 fn ranked_hits(
412 &self,
413 query: &dyn Query,
414 query_text: &str,
415 limit: usize,
416 ) -> Result<Vec<TextSearchHit>> {
417 let searcher = self.reader.searcher();
418 let window = limit
419 .saturating_mul(RANK_WINDOW_FACTOR)
420 .min(MAX_RANK_WINDOW)
421 .max(limit);
422 let mut top_docs = searcher.search(query, &TopDocs::with_limit(window))?;
423 if let Some(&(cutoff, _)) = top_docs.get(limit.saturating_sub(1)) {
426 top_docs.retain(|(score, _)| *score >= cutoff);
427 }
428
429 let query_tokens = normalize_index_text(query_text)
430 .map(|text| {
431 text.split_whitespace()
432 .map(str::to_string)
433 .collect::<HashSet<_>>()
434 })
435 .unwrap_or_default();
436 let mut record_ids: HashMap<u32, Column<u64>> = HashMap::new();
437 let mut hits = Vec::with_capacity(top_docs.len());
438 for (score, doc_address) in top_docs {
439 let column = match record_ids.entry(doc_address.segment_ord) {
440 Entry::Occupied(entry) => entry.into_mut(),
441 Entry::Vacant(entry) => entry.insert(
442 searcher
443 .segment_reader(doc_address.segment_ord)
444 .fast_fields()
445 .u64("record_id")?,
446 ),
447 };
448 let record_id = column
449 .values_for_doc(doc_address.doc_id)
450 .next()
451 .context("text index hit is missing fast record_id")?;
452 let record = self.pack.record_summary(record_id)?;
453 hits.push((
454 TieBreak::new(&record, &query_tokens),
455 TextSearchHit {
456 record_id,
457 score,
458 record,
459 },
460 ));
461 }
462 hits.sort_by(|(left_key, left), (right_key, right)| {
463 right
464 .score
465 .total_cmp(&left.score)
466 .then_with(|| left_key.cmp(right_key))
467 });
468 hits.truncate(limit);
469 Ok(hits.into_iter().map(|(_, hit)| hit).collect())
470 }
471}
472
473#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
475struct TieBreak {
476 extra_label_tokens: usize,
479 precision: u8,
481 source: (u8, i64),
483 id: String,
484}
485
486impl TieBreak {
487 fn new(record: &RecordSummary, query_tokens: &HashSet<String>) -> Self {
488 let extra_label_tokens = normalize_index_text(&record.label)
489 .map(|label| {
490 label
491 .split_whitespace()
492 .filter(|token| !query_tokens.contains(*token))
493 .count()
494 })
495 .unwrap_or_default();
496 let precision = match record.point.map(|point| point.precision) {
497 Some(RecordPointPrecision::Point) => 0,
498 Some(RecordPointPrecision::Centroid) => 1,
499 Some(RecordPointPrecision::RepresentativePoint) => 2,
500 None => 3,
501 };
502 let source = match record.source.object_type {
503 Some(OsmObjectType::Node) => 0,
504 Some(OsmObjectType::Way) => 1,
505 Some(OsmObjectType::Relation) => 2,
506 None => 3,
507 };
508 Self {
509 extra_label_tokens,
510 precision,
511 source: (source, record.source.object_id.unwrap_or_default()),
512 id: record.id.clone(),
513 }
514 }
515}
516
517fn effective_limit(limit: usize) -> usize {
518 if limit == 0 {
519 DEFAULT_SEARCH_LIMIT
520 } else {
521 limit
522 }
523}
524
525fn effective_autocomplete_limit(limit: usize) -> usize {
526 let limit = effective_limit(limit);
527 limit.min(MAX_AUTOCOMPLETE_LIMIT)
528}
529
530fn address_geocode_candidates(address: &str, postcode: Option<&str>) -> Vec<String> {
531 let mut candidates = Vec::new();
532 if let Some(postcode) = postcode.and_then(normalize_index_text)
533 && let Some(address) = meaningful_address_query(address)
534 {
535 candidates.push(format!("{address} {postcode}"));
536 }
537 candidates.push(address.to_string());
538 unique_strings(candidates)
539}
540
541fn meaningful_address_query(address: &str) -> Option<String> {
542 let normalized = normalize_index_text(address)?;
543 let expanded = expand_address_abbreviations(&normalized);
544 strip_unit_terms(&expanded)
545}
546
547fn search_query_variants(query_text: &str) -> Vec<String> {
548 let mut variants = Vec::new();
549 if let Some(cleaned) = collapse_query(query_text) {
550 variants.push(cleaned);
551 }
552 if let Some(normalized) = normalize_index_text(query_text) {
553 variants.push(normalized.clone());
554 let expanded = expand_address_abbreviations(&normalized);
555 variants.push(expanded.clone());
556 if let Some(without_unit) = strip_unit_terms(&expanded) {
557 variants.push(without_unit);
558 }
559 }
560 unique_strings(variants)
561}
562
563fn collapse_query(value: &str) -> Option<String> {
564 let cleaned = value.split_whitespace().collect::<Vec<_>>().join(" ");
565 (!cleaned.is_empty()).then_some(cleaned)
566}
567
568fn expand_address_abbreviations(value: &str) -> String {
569 let tokens = value.split_whitespace().collect::<Vec<_>>();
570 let mut expanded = Vec::with_capacity(tokens.len());
571 for (index, token) in tokens.iter().enumerate() {
572 let replacement = match *token {
573 "ave" | "av" => Some("avenue"),
574 "blvd" => Some("boulevard"),
575 "cir" => Some("circle"),
576 "ct" | "crt" => Some("court"),
577 "cres" => Some("crescent"),
578 "dr" => Some("drive"),
579 "hwy" => Some("highway"),
580 "ln" => Some("lane"),
581 "pkwy" => Some("parkway"),
582 "pl" => Some("place"),
583 "rd" => Some("road"),
584 "sq" => Some("square"),
585 "st" if index > 0 && !is_numeric_token(tokens[index - 1]) => Some("street"),
586 "ter" | "terr" => Some("terrace"),
587 "trl" | "tr" => Some("trail"),
588 "wy" => Some("way"),
589 "e" if previous_token_is_street_type(&expanded) => Some("east"),
590 "n" if previous_token_is_street_type(&expanded) => Some("north"),
591 "s" if previous_token_is_street_type(&expanded) => Some("south"),
592 "w" if previous_token_is_street_type(&expanded) => Some("west"),
593 _ => None,
594 };
595 expanded.push(replacement.unwrap_or(*token));
596 }
597 expanded.join(" ")
598}
599
600fn previous_token_is_street_type(tokens: &[&str]) -> bool {
601 tokens.last().is_some_and(|token| {
602 matches!(
603 *token,
604 "avenue"
605 | "boulevard"
606 | "circle"
607 | "court"
608 | "crescent"
609 | "drive"
610 | "highway"
611 | "lane"
612 | "parkway"
613 | "place"
614 | "road"
615 | "square"
616 | "street"
617 | "terrace"
618 | "trail"
619 | "way"
620 )
621 })
622}
623
624fn strip_unit_terms(value: &str) -> Option<String> {
625 let tokens = value.split_whitespace().collect::<Vec<_>>();
626 let mut stripped = Vec::with_capacity(tokens.len());
627 let mut index = 0;
628 while index < tokens.len() {
629 if is_unit_designator(tokens[index]) {
630 index += 1;
631 if index < tokens.len() && is_unit_value_token(tokens[index]) {
632 index += 1;
633 }
634 continue;
635 }
636 stripped.push(tokens[index]);
637 index += 1;
638 }
639
640 if stripped.len() > 2 && is_numeric_token(stripped[0]) && is_numeric_token(stripped[1]) {
641 stripped.remove(0);
642 }
643
644 let stripped = stripped.join(" ");
645 (!stripped.is_empty()).then_some(stripped)
646}
647
648fn is_unit_designator(token: &str) -> bool {
649 matches!(
650 token,
651 "apt"
652 | "apartment"
653 | "bldg"
654 | "building"
655 | "dept"
656 | "department"
657 | "fl"
658 | "floor"
659 | "rm"
660 | "room"
661 | "ste"
662 | "suite"
663 | "unit"
664 )
665}
666
667fn is_unit_value_token(token: &str) -> bool {
668 token.chars().any(|character| character.is_ascii_digit())
669}
670
671fn is_numeric_token(token: &str) -> bool {
672 token.chars().all(|character| character.is_ascii_digit())
673}
674
675fn normalized_for_match(value: Option<&str>) -> Option<String> {
676 normalize_index_text(value?.trim())
677}
678
679fn normalized_postcode_for_match(value: Option<&str>) -> Option<String> {
680 normalized_for_match(value).map(|value| value.split_whitespace().collect())
681}
682
683fn unique_strings(values: Vec<String>) -> Vec<String> {
684 let mut unique = Vec::new();
685 for value in values {
686 if !unique.contains(&value) {
687 unique.push(value);
688 }
689 }
690 unique
691}
692
693fn autocomplete_query_tokens(query_text: &str) -> Vec<String> {
694 query_text
695 .split_whitespace()
696 .map(str::to_string)
697 .collect::<Vec<_>>()
698}
699
700fn autocomplete_subject_query(field: Field, tokens: &[String]) -> Box<dyn Query> {
701 let terms = tokens
702 .iter()
703 .map(|token| Term::from_field_text(field, token))
704 .collect::<Vec<_>>();
705 let mut query = PhrasePrefixQuery::new(terms);
706 query.set_max_expansions(AUTOCOMPLETE_PREFIX_MAX_EXPANSIONS);
707 Box::new(query)
708}
709
710fn is_address_number_token(token: &str) -> bool {
711 token
712 .chars()
713 .next()
714 .is_some_and(|character| character.is_ascii_digit())
715}
716
717#[cfg(test)]
718mod tests {
719 use std::collections::BTreeMap;
720
721 use crate::{
722 pack::PackWriter,
723 record::{
724 AddressComponents, AddressRecord, DerivedSourceProvenance, LocationPrecision,
725 OsmObjectType, PostcodeRecord, SourceProvenance, StreetRecord, point_geometry,
726 },
727 };
728
729 use super::*;
730
731 #[test]
732 fn search_query_variants_expand_common_address_abbreviations() {
733 let variants = search_query_variants("33 Princess St Suite 170");
734 assert!(variants.contains(&"33 princess street suite 170".to_string()));
735 assert!(variants.contains(&"33 princess street".to_string()));
736 }
737
738 #[test]
739 fn search_query_variants_strip_leading_unit_numbers() {
740 let variants = search_query_variants("306-1333 Sheppard Ave E");
741 assert!(variants.contains(&"306 1333 sheppard avenue east".to_string()));
742 assert!(variants.contains(&"1333 sheppard avenue east".to_string()));
743 }
744
745 #[test]
746 fn search_query_variants_do_not_treat_initial_saint_as_street() {
747 let variants = search_query_variants("St Clair Ave W");
748 assert!(variants.contains(&"st clair avenue west".to_string()));
749 assert!(!variants.contains(&"street clair avenue west".to_string()));
750 }
751
752 #[test]
753 fn equal_scores_prefer_the_closest_label_then_points_then_lower_ids() {
754 let temp_dir = temp_pack_path("search-ties");
755 let _ = std::fs::remove_dir_all(&temp_dir);
756
757 let mut writer = PackWriter::create(&temp_dir).expect("writer");
758 let mut prince =
759 address_record("osm:node:1", "", "342", "Prince Albert Street", None, None);
760 prince.geometry = point_geometry(-75.6, 45.4);
761 let mut way = address_record("osm:node:9", "", "342", "Albert Street", None, None);
762 way.source.object_type = OsmObjectType::Way;
763 for record in [
764 prince,
765 way,
766 address_record("osm:node:7", "", "342", "Albert Street", None, None),
767 address_record("osm:node:5", "", "342", "Albert Street", None, None),
768 ] {
769 writer.write(&record.into(), None).expect("write");
770 }
771 writer.finish().expect("finish");
772
773 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
774 let hits = searcher
775 .search(TextSearchOptions {
776 query: "342 Albert Street".to_string(),
777 limit: 3,
778 layer: None,
779 })
780 .expect("search");
781 let ids = hits
782 .iter()
783 .map(|hit| hit.record.id.as_str())
784 .collect::<Vec<_>>();
785 assert_eq!(ids, vec!["osm:node:5", "osm:node:7", "osm:way:9"]);
786
787 let _ = std::fs::remove_dir_all(temp_dir);
788 }
789
790 fn geocode(
791 searcher: &PackTextSearcher,
792 address: &str,
793 locality: Option<&str>,
794 postcode: Option<&str>,
795 ) -> Option<String> {
796 searcher
797 .geocode_address(AddressGeocodeOptions {
798 address: address.to_string(),
799 locality: locality.map(str::to_string),
800 region: None,
801 postcode: postcode.map(str::to_string),
802 limit: 10,
803 layer: None,
804 })
805 .expect("geocode")
806 .map(|hit| hit.hit.record.id)
807 }
808
809 #[test]
810 fn geocode_rejects_hits_in_a_different_postcode() {
811 let temp_dir = temp_pack_path("geocode-postcode");
812 let _ = std::fs::remove_dir_all(&temp_dir);
813 let mut writer = PackWriter::create(&temp_dir).expect("writer");
814 writer
815 .write(
816 &address_record("osm:node:1", "", "10", "King Street", None, Some("M5V 1A1"))
817 .into(),
818 None,
819 )
820 .expect("write");
821 writer.finish().expect("finish");
822 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
823
824 assert_eq!(
825 geocode(&searcher, "10 King Street", None, Some("K7L 1A1")),
826 None
827 );
828 assert_eq!(
829 geocode(&searcher, "10 King Street", None, Some("m5v1a1")).as_deref(),
830 Some("osm:node:1")
831 );
832 assert_eq!(
834 geocode(&searcher, "10 King Street", None, Some("M5V")).as_deref(),
835 Some("osm:node:1")
836 );
837 let _ = std::fs::remove_dir_all(temp_dir);
838 }
839
840 #[test]
841 fn geocode_infers_the_postcode_area_of_hits_without_one() {
842 let temp_dir = temp_pack_path("geocode-inferred-postcode");
843 let _ = std::fs::remove_dir_all(&temp_dir);
844 let mut writer = PackWriter::create(&temp_dir).expect("writer");
845 let mut cobden = address_record("osm:node:1", "", "44", "Gould Street", None, None);
847 cobden.geometry = point_geometry(-76.88, 45.63);
848 writer.write(&cobden.into(), None).expect("cobden");
849 writer
850 .write(
851 &PostcodeRecord {
852 postcode: "K0J 1K0".to_string(),
853 geometry: point_geometry(-76.881, 45.631),
854 source: DerivedSourceProvenance::osm_address_records(5),
855 }
856 .into(),
857 None,
858 )
859 .expect("postcode");
860 let mut remote = address_record("osm:node:2", "", "7", "Far Road", None, None);
862 remote.geometry = point_geometry(-90.0, 50.0);
863 writer.write(&remote.into(), None).expect("remote");
864 writer.finish().expect("finish");
865 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
866
867 assert_eq!(
868 geocode(&searcher, "44 Gould Street", None, Some("M5V 1A1")),
869 None
870 );
871 assert_eq!(
872 geocode(&searcher, "44 Gould Street", None, Some("K0J 1K0")).as_deref(),
873 Some("osm:node:1")
874 );
875 assert_eq!(
876 geocode(&searcher, "7 Far Road", None, Some("M5V 1A1")).as_deref(),
877 Some("osm:node:2")
878 );
879 let _ = std::fs::remove_dir_all(temp_dir);
880 }
881
882 #[test]
883 fn geocode_accepts_hits_near_a_busier_neighbouring_postcode_area() {
884 let temp_dir = temp_pack_path("geocode-postcode-border");
885 let _ = std::fs::remove_dir_all(&temp_dir);
886 let mut writer = PackWriter::create(&temp_dir).expect("writer");
887 let mut border = address_record("osm:node:1", "", "5", "Line Road", None, None);
888 border.geometry = point_geometry(-79.0, 43.0);
889 writer.write(&border.into(), None).expect("address");
890 let mut postcode = |code: String, lon: f64| {
891 writer
892 .write(
893 &PostcodeRecord {
894 postcode: code,
895 geometry: point_geometry(lon, 43.0),
896 source: DerivedSourceProvenance::osm_address_records(1),
897 }
898 .into(),
899 None,
900 )
901 .expect("postcode");
902 };
903 for index in 0..20 {
906 postcode(
907 format!("B2B {index}A{index}"),
908 -79.0 + 0.0001 * index as f64,
909 );
910 }
911 postcode("A1A 1A1".to_string(), -78.9);
912 writer.finish().expect("finish");
913 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
914
915 assert_eq!(
916 geocode(&searcher, "5 Line Road", None, Some("A1A 2B2")).as_deref(),
917 Some("osm:node:1")
918 );
919 assert_eq!(
920 geocode(&searcher, "5 Line Road", None, Some("C3C 3C3")),
921 None
922 );
923 let _ = std::fs::remove_dir_all(temp_dir);
924 }
925
926 #[test]
927 fn geocode_tries_expanded_variants_when_earlier_hits_fail_the_context() {
928 use crate::{
929 context::AdminContextTuple,
930 pack::RecordContext,
931 record::{PlaceLayer, PlaceRecord, Record},
932 };
933
934 let temp_dir = temp_pack_path("geocode-variants");
935 let _ = std::fs::remove_dir_all(&temp_dir);
936 let mut writer = PackWriter::create(&temp_dir).expect("writer");
937 let mut locality = |name: &str, object_id: i64| {
938 writer
939 .write(
940 &Record::Place(
941 PlaceLayer::Locality,
942 PlaceRecord {
943 name: name.to_string(),
944 place_type: "city".to_string(),
945 geometry: point_geometry(-79.0, 43.0),
946 source: SourceProvenance::osm(OsmObjectType::Node, object_id),
947 },
948 ),
949 None,
950 )
951 .expect("locality")
952 };
953 let hamilton = locality("Hamilton", 100);
954 let toronto = locality("Toronto", 101);
955 let in_locality = |record_id| {
956 Some(RecordContext {
957 admin_context: AdminContextTuple {
958 locality_record_id: Some(record_id),
959 ..AdminContextTuple::default()
960 },
961 flags: 0,
962 })
963 };
964 writer
965 .write(
966 &address_record("osm:node:1", "", "10", "King St W", None, None).into(),
967 in_locality(hamilton),
968 )
969 .expect("hamilton address");
970 writer
971 .write(
972 &address_record("osm:node:2", "", "10", "King Street West", None, None).into(),
973 in_locality(toronto),
974 )
975 .expect("toronto address");
976 writer.finish().expect("finish");
977 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
978
979 assert_eq!(
980 geocode(&searcher, "10 King St W", Some("Toronto"), None).as_deref(),
981 Some("osm:node:2")
982 );
983 let _ = std::fs::remove_dir_all(temp_dir);
984 }
985
986 #[test]
987 fn search_reports_queries_that_cannot_be_parsed() {
988 let temp_dir = temp_pack_path("search-parse-error");
989 let _ = std::fs::remove_dir_all(&temp_dir);
990 let mut writer = PackWriter::create(&temp_dir).expect("writer");
991 writer
992 .write(
993 &address_record("osm:node:1", "", "10", "King Street", None, None).into(),
994 None,
995 )
996 .expect("write");
997 writer.finish().expect("finish");
998 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
999
1000 let error = searcher
1001 .search(TextSearchOptions {
1002 query: "(".to_string(),
1003 limit: 5,
1004 layer: None,
1005 })
1006 .expect_err("unparseable query");
1007 assert!(format!("{error:#}").contains("failed to parse search query"));
1008 let _ = std::fs::remove_dir_all(temp_dir);
1009 }
1010
1011 #[test]
1012 fn searches_and_hydrates_records_from_pack() {
1013 let temp_dir = temp_pack_path("search-hydrates");
1014 let _ = std::fs::remove_dir_all(&temp_dir);
1015
1016 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1017 writer
1018 .write(
1019 &address_record(
1020 "osm:node:1",
1021 "10 King Street, Toronto",
1022 "10",
1023 "King Street",
1024 Some("Toronto"),
1025 Some("M5V 1A1"),
1026 )
1027 .into(),
1028 None,
1029 )
1030 .expect("write king address");
1031 writer
1032 .write(
1033 &address_record(
1034 "osm:node:2",
1035 "20 Queen Street, Toronto",
1036 "20",
1037 "Queen Street",
1038 Some("Toronto"),
1039 Some("M5V 1A1"),
1040 )
1041 .into(),
1042 None,
1043 )
1044 .expect("write queen address");
1045 writer.finish().expect("finish");
1046
1047 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1048 let hits = searcher
1049 .search(TextSearchOptions {
1050 query: "King Street Toronto".to_string(),
1051 limit: 5,
1052 layer: None,
1053 })
1054 .expect("search");
1055
1056 assert!(!hits.is_empty());
1057 assert_eq!(hits[0].record_id, 0);
1058 assert_eq!(hits[0].record.id, "osm:node:1");
1059 assert_eq!(hits[0].record.label, "10 King Street, Toronto, M5V 1A1");
1060
1061 let _ = std::fs::remove_dir_all(temp_dir);
1062 }
1063
1064 #[test]
1065 fn filters_hits_by_layer() {
1066 let temp_dir = temp_pack_path("search-layer");
1067 let _ = std::fs::remove_dir_all(&temp_dir);
1068
1069 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1070 writer
1071 .write(
1072 &address_record(
1073 "osm:node:1",
1074 "10 King Street, Toronto",
1075 "10",
1076 "King Street",
1077 Some("Toronto"),
1078 None,
1079 )
1080 .into(),
1081 None,
1082 )
1083 .expect("write address");
1084 writer
1085 .write(&street_record("osm:way:9", "King Street").into(), None)
1086 .expect("write street");
1087 writer.finish().expect("finish");
1088
1089 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1090 let hits = searcher
1091 .search(TextSearchOptions {
1092 query: "King Street".to_string(),
1093 limit: 10,
1094 layer: Some("street".to_string()),
1095 })
1096 .expect("search");
1097
1098 assert_eq!(hits.len(), 1);
1099 assert_eq!(hits[0].record_id, 1);
1100 assert_eq!(hits[0].record.layer, "street");
1101
1102 let _ = std::fs::remove_dir_all(temp_dir);
1103 }
1104
1105 #[test]
1106 fn searches_postcode_text() {
1107 let temp_dir = temp_pack_path("search-postcode");
1108 let _ = std::fs::remove_dir_all(&temp_dir);
1109
1110 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1111 writer
1112 .write(
1113 &PostcodeRecord {
1114 postcode: "M5V".to_string(),
1115 geometry: point_geometry(-79.4, 43.6),
1116 source: DerivedSourceProvenance::osm_address_records(2),
1117 }
1118 .into(),
1119 None,
1120 )
1121 .expect("write postcode");
1122 writer.finish().expect("finish");
1123
1124 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1125 let hits = searcher
1126 .search(TextSearchOptions {
1127 query: "M5V".to_string(),
1128 limit: 5,
1129 layer: Some("postcode".to_string()),
1130 })
1131 .expect("search");
1132
1133 assert_eq!(hits.len(), 1);
1134 assert_eq!(hits[0].record_id, 0);
1135 assert_eq!(hits[0].record.layer, "postcode");
1136
1137 let _ = std::fs::remove_dir_all(temp_dir);
1138 }
1139
1140 #[test]
1141 fn autocompletes_prefixes_and_hydrates_records_from_pack() {
1142 let temp_dir = temp_pack_path("autocomplete-prefix");
1143 let _ = std::fs::remove_dir_all(&temp_dir);
1144
1145 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1146 writer
1147 .write(
1148 &address_record(
1149 "osm:node:1",
1150 "10 King Street, Toronto",
1151 "10",
1152 "King Street",
1153 Some("Toronto"),
1154 Some("M5V 1A1"),
1155 )
1156 .into(),
1157 None,
1158 )
1159 .expect("write king address");
1160 writer
1161 .write(
1162 &address_record(
1163 "osm:node:2",
1164 "20 Queen Street, Toronto",
1165 "20",
1166 "Queen Street",
1167 Some("Toronto"),
1168 None,
1169 )
1170 .into(),
1171 None,
1172 )
1173 .expect("write queen address");
1174 writer.finish().expect("finish");
1175
1176 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1177 let hits = searcher
1178 .autocomplete(TextAutocompleteOptions {
1179 query: "kin".to_string(),
1180 limit: 5,
1181 layer: None,
1182 })
1183 .expect("autocomplete");
1184
1185 assert!(!hits.is_empty());
1186 assert_eq!(hits[0].record_id, 0);
1187 assert_eq!(hits[0].record.label, "10 King Street, Toronto, M5V 1A1");
1188
1189 let _ = std::fs::remove_dir_all(temp_dir);
1190 }
1191
1192 #[test]
1193 fn autocompletes_multi_token_prefixes() {
1194 let temp_dir = temp_pack_path("autocomplete-multi-token");
1195 let _ = std::fs::remove_dir_all(&temp_dir);
1196
1197 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1198 writer
1199 .write(
1200 &address_record(
1201 "osm:node:1",
1202 "10 King Street, Toronto",
1203 "10",
1204 "King Street",
1205 Some("Toronto"),
1206 None,
1207 )
1208 .into(),
1209 None,
1210 )
1211 .expect("write king address");
1212 writer.finish().expect("finish");
1213
1214 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1215 let hits = searcher
1216 .autocomplete(TextAutocompleteOptions {
1217 query: "king st".to_string(),
1218 limit: 5,
1219 layer: None,
1220 })
1221 .expect("autocomplete");
1222
1223 assert_eq!(hits.len(), 1);
1224 assert_eq!(hits[0].record_id, 0);
1225
1226 let _ = std::fs::remove_dir_all(temp_dir);
1227 }
1228
1229 #[test]
1230 fn autocompletes_with_layer_filter() {
1231 let temp_dir = temp_pack_path("autocomplete-layer");
1232 let _ = std::fs::remove_dir_all(&temp_dir);
1233
1234 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1235 writer
1236 .write(
1237 &address_record(
1238 "osm:node:1",
1239 "10 King Street, Toronto",
1240 "10",
1241 "King Street",
1242 Some("Toronto"),
1243 None,
1244 )
1245 .into(),
1246 None,
1247 )
1248 .expect("write address");
1249 writer
1250 .write(&street_record("osm:way:9", "King Street").into(), None)
1251 .expect("write street");
1252 writer.finish().expect("finish");
1253
1254 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1255 let hits = searcher
1256 .autocomplete(TextAutocompleteOptions {
1257 query: "kin".to_string(),
1258 limit: 10,
1259 layer: Some("street".to_string()),
1260 })
1261 .expect("autocomplete");
1262
1263 assert_eq!(hits.len(), 1);
1264 assert_eq!(hits[0].record_id, 1);
1265 assert_eq!(hits[0].record.layer, "street");
1266
1267 let _ = std::fs::remove_dir_all(temp_dir);
1268 }
1269
1270 #[test]
1271 fn autocompletes_postcode_but_not_standalone_house_number_prefixes() {
1272 let temp_dir = temp_pack_path("autocomplete-postcode-house-number");
1273 let _ = std::fs::remove_dir_all(&temp_dir);
1274
1275 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1276 writer
1277 .write(
1278 &address_record(
1279 "osm:node:1",
1280 "221 Baker Street, London, NW1",
1281 "221",
1282 "Baker Street",
1283 Some("London"),
1284 Some("NW1 6XE"),
1285 )
1286 .into(),
1287 None,
1288 )
1289 .expect("write baker address");
1290 writer.finish().expect("finish");
1291
1292 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1293 let postcode_hits = searcher
1294 .autocomplete(TextAutocompleteOptions {
1295 query: "nw16".to_string(),
1296 limit: 5,
1297 layer: None,
1298 })
1299 .expect("postcode autocomplete");
1300 let number_hits = searcher
1301 .autocomplete(TextAutocompleteOptions {
1302 query: "221".to_string(),
1303 limit: 5,
1304 layer: None,
1305 })
1306 .expect("number autocomplete");
1307
1308 assert_eq!(postcode_hits.len(), 1);
1309 assert!(number_hits.is_empty());
1310 assert_eq!(postcode_hits[0].record_id, 0);
1311
1312 let _ = std::fs::remove_dir_all(temp_dir);
1313 }
1314
1315 #[test]
1316 fn autocompletes_house_number_with_street_prefix() {
1317 let temp_dir = temp_pack_path("autocomplete-number-street-prefix");
1318 let _ = std::fs::remove_dir_all(&temp_dir);
1319
1320 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1321 writer
1322 .write(
1323 &address_record(
1324 "osm:node:1",
1325 "221 Baker Street, London, NW1",
1326 "221",
1327 "Baker Street",
1328 Some("London"),
1329 Some("NW1 6XE"),
1330 )
1331 .into(),
1332 None,
1333 )
1334 .expect("write baker address");
1335 writer.finish().expect("finish");
1336
1337 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1338 let hits = searcher
1339 .autocomplete(TextAutocompleteOptions {
1340 query: "221 bak".to_string(),
1341 limit: 5,
1342 layer: None,
1343 })
1344 .expect("number plus street autocomplete");
1345
1346 assert_eq!(hits.len(), 1);
1347 assert_eq!(hits[0].record_id, 0);
1348
1349 let _ = std::fs::remove_dir_all(temp_dir);
1350 }
1351
1352 #[test]
1353 fn autocomplete_ignores_blank_and_short_queries() {
1354 let temp_dir = temp_pack_path("autocomplete-short");
1355 let _ = std::fs::remove_dir_all(&temp_dir);
1356
1357 let mut writer = PackWriter::create(&temp_dir).expect("writer");
1358 writer
1359 .write(
1360 &address_record(
1361 "osm:node:1",
1362 "10 King Street, Toronto",
1363 "10",
1364 "King Street",
1365 Some("Toronto"),
1366 None,
1367 )
1368 .into(),
1369 None,
1370 )
1371 .expect("write address");
1372 writer.finish().expect("finish");
1373
1374 let searcher = PackTextSearcher::open(&temp_dir).expect("searcher");
1375
1376 assert!(
1377 searcher
1378 .autocomplete(TextAutocompleteOptions {
1379 query: " ".to_string(),
1380 limit: 5,
1381 layer: None,
1382 })
1383 .expect("blank autocomplete")
1384 .is_empty()
1385 );
1386 assert!(
1387 searcher
1388 .autocomplete(TextAutocompleteOptions {
1389 query: "k".to_string(),
1390 limit: 5,
1391 layer: None,
1392 })
1393 .expect("single character autocomplete")
1394 .is_empty()
1395 );
1396 assert!(
1397 searcher
1398 .autocomplete(TextAutocompleteOptions {
1399 query: "ki".to_string(),
1400 limit: 5,
1401 layer: None,
1402 })
1403 .expect("two character autocomplete")
1404 .is_empty()
1405 );
1406
1407 let _ = std::fs::remove_dir_all(temp_dir);
1408 }
1409
1410 fn temp_pack_path(name: &str) -> std::path::PathBuf {
1411 std::env::temp_dir().join(format!("open-geocode-{name}-{}", std::process::id()))
1412 }
1413
1414 fn address_record(
1415 id: &str,
1416 _label: &str,
1417 number: &str,
1418 street: &str,
1419 locality: Option<&str>,
1420 postcode: Option<&str>,
1421 ) -> AddressRecord {
1422 let object_id = id
1423 .strip_prefix("osm:node:")
1424 .and_then(|value| value.parse::<i64>().ok())
1425 .unwrap_or(1);
1426 AddressRecord {
1427 address: AddressComponents {
1428 number: number.to_string(),
1429 street: Some(street.to_string()),
1430 place: None,
1431 unit: None,
1432 locality: locality.map(str::to_string),
1433 region: None,
1434 postcode: postcode.map(str::to_string),
1435 country: None,
1436 },
1437 geometry: point_geometry(-79.0, 43.0),
1438 location_precision: LocationPrecision::Point,
1439 source: SourceProvenance {
1440 dataset: "osm".to_string(),
1441 object_type: OsmObjectType::Node,
1442 object_id,
1443 tags: Some(BTreeMap::new()),
1444 },
1445 }
1446 }
1447
1448 fn street_record(id: &str, label: &str) -> StreetRecord {
1449 let object_id = id
1450 .strip_prefix("osm:way:")
1451 .and_then(|value| value.parse::<i64>().ok())
1452 .unwrap_or(9);
1453 StreetRecord {
1454 name: label.to_string(),
1455 geometry: point_geometry(-79.0, 43.0),
1456 representative_point: [-79.0, 43.0],
1457 source: SourceProvenance {
1458 dataset: "osm".to_string(),
1459 object_type: OsmObjectType::Way,
1460 object_id,
1461 tags: Some(BTreeMap::new()),
1462 },
1463 }
1464 }
1465}