src/nominatim_db/tokenizer/token_analysis/simple_trie.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2025 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Simple dict-based implementation of a trie structure.
   9 """
  10 from typing import TypeVar, Generic, Tuple, Optional, List, Dict
  11 from collections import defaultdict
  12
  13 T = TypeVar('T')
  14
  15
  16 class SimpleTrie(Generic[T]):
  17     """ A simple read-only trie structure.
  18         This structure supports examply one lookup operation,
  19         which is longest-prefix lookup.
  20     """
  21
  22     def __init__(self, data: Optional[List[Tuple[str, T]]] = None) -> None:
  23         self._tree: Dict[str, 'SimpleTrie[T]'] = defaultdict(SimpleTrie[T])
  24         self._value: Optional[T] = None
  25         self._prefix = ''
  26
  27         if data:
  28             for key, value in data:
  29                 self._add(key, 0, value)
  30
  31             self._make_compact()
  32
  33     def _add(self, word: str, pos: int, value: T) -> None:
  34         """ (Internal) Add a sub-word to the trie.
  35             The word is added from index 'pos'. If the sub-word to add
  36             is empty, then the trie saves the given value.
  37         """
  38         if pos < len(word):
  39             self._tree[word[pos]]._add(word, pos + 1, value)
  40         else:
  41             self._value = value
  42
  43     def _make_compact(self) -> None:
  44         """ (Internal) Compress tree where there is exactly one subtree
  45             and no value.
  46
  47             Compression works recursively starting at the leaf.
  48         """
  49         for t in self._tree.values():
  50             t._make_compact()
  51
  52         if len(self._tree) == 1 and self._value is None:
  53             assert not self._prefix
  54             for k, v in self._tree.items():
  55                 self._prefix = k + v._prefix
  56                 self._tree = v._tree
  57                 self._value = v._value
  58
  59     def longest_prefix(self, word: str, start: int = 0) -> Tuple[Optional[T], int]:
  60         """ Return the longest prefix match for the given word starting at
  61             the position 'start'.
  62
  63             The function returns a tuple with the value for the longest match and
  64             the position of the word after the match. If no match was found at
  65             all, the function returns (None, start).
  66         """
  67         cur = self
  68         pos = start
  69         result: Tuple[Optional[T], int] = None, start
  70
  71         while True:
  72             if cur._prefix:
  73                 if not word.startswith(cur._prefix, pos):
  74                     return result
  75                 pos += len(cur._prefix)
  76
  77             if cur._value:
  78                 result = cur._value, pos
  79
  80             if pos >= len(word) or word[pos] not in cur._tree:
  81                 return result
  82
  83             cur = cur._tree[word[pos]]
  84             pos += 1