Skip to main content

open_geocode/
search.rs

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;
73/// For a hit that states no postcode: postcode records within this radius are
74/// compared with the requested postcode on their leading characters (a
75/// Canadian FSA, a US sectional centre, the start of a UK outward code).
76const POSTCODE_AREA_RADIUS_M: f64 = 10_000.0;
77const POSTCODE_AREA_CHARS: usize = 3;
78/// Hits examined per requested hit, so ties at the cutoff can be ordered.
79const 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    /// Open the text index using an existing shared pack reader.
88    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    /// Try the query as written, then normalized and expanded rewrites, and
200    /// return the first variant with hits that `accept` keeps. A variant whose
201    /// hits are all rejected falls through to the next one. Fails only when no
202    /// variant could be parsed at all.
203    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    /// Whether a hit can be in the requested (normalized) postcode.
313    ///
314    /// A hit that states a postcode must agree with it; either may be a prefix
315    /// of the other, so "M5V" matches "M5V 1A1". A hit without one is judged by
316    /// the postcode areas around it: if there are some nearby and none shares
317    /// the requested postcode's leading characters, the hit is somewhere else.
318    /// With no postcode data nearby there is nothing to contradict the row.
319    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        // No postcode data nearby means nothing contradicts the row.
348        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    /// The best `limit` hits. Hits whose score ties the cutoff are ordered by
409    /// how closely their label matches the query, then by precision, then by
410    /// source id, instead of by where records happen to sit in the Pack.
411    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        // Only hits that can still make the cut need hydrating: everything above
424        // the cutoff score plus whatever ties it.
425        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/// Secondary order for hits with equal scores.
474#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
475struct TieBreak {
476    /// Label words the query did not ask for: "342 Albert Street" beats
477    /// "342 Prince Albert Street" for the query "342 albert street".
478    extra_label_tokens: usize,
479    /// Exact points before centroids before line midpoints.
480    precision: u8,
481    /// Nodes before ways before relations, then the lower id.
482    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        // A forward sortation area matches the full postcode it starts.
833        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        // An address without addr:postcode, inside postcode area K0J.
846        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        // An address far from every postcode record: its area is unknown.
861        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        // Twenty centres of the neighbouring area right next to the address,
904        // and the address's own area 8 km away.
905        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}