titles/search.go

   1 package titles
   2
   3 import (
   4         "regexp"
   5         "strings"
   6 )
   7
   8 // Exact string search, returns first N matches (nondeterministic)
   9 func (db *TitlesDatabase) ExactSearchN(s string, n int) (matches SearchMatches) {
  10         return db.doSearchN(func(k string) bool { return s == k }, n)
  11 }
  12
  13 // Exact string search (nondeterministic order)
  14 func (db *TitlesDatabase) ExactSearchAll(s string) (matches SearchMatches) {
  15         return db.ExactSearchN(s, -1)
  16 }
  17
  18 // Exact string search, returns first match (nondeterministic)
  19 func (db *TitlesDatabase) ExactSearch(s string) (m SearchMatch) {
  20         return firstMatch(db.ExactSearchN(s, 1))
  21 }
  22
  23 // String search with case folding, returns first N matches (nondeterministic)
  24 func (db *TitlesDatabase) ExactSearchFoldN(s string, n int) (matches SearchMatches) {
  25         return db.doSearchN(func(k string) bool { return strings.EqualFold(k, s) }, n)
  26 }
  27
  28 // String search with case folding (nondeterministic order)
  29 func (db *TitlesDatabase) ExactSearchFoldAll(s string) (matches SearchMatches) {
  30         return db.ExactSearchFoldN(s, -1)
  31 }
  32
  33 // String search with case folding, returns first match (nondeterministic)
  34 func (db *TitlesDatabase) ExactSearchFold(s string) (m SearchMatch) {
  35         return firstMatch(db.ExactSearchFoldN(s, 1))
  36 }
  37
  38 // Regular expression search, returns first N matches (nondeterministic)
  39 func (db *TitlesDatabase) RegexpSearchN(re *regexp.Regexp, n int) (matches SearchMatches) {
  40         return db.doSearchN(func(k string) bool { return re.MatchString(k) }, n)
  41 }
  42
  43 // Regular expression search (nondeterministic order)
  44 func (db *TitlesDatabase) RegexpSearchAll(re *regexp.Regexp) (matches SearchMatches) {
  45         return db.RegexpSearchN(re, -1)
  46 }
  47
  48 // Regular expression search, returns first match (nondeterministic)
  49 func (db *TitlesDatabase) RegexpSearch(re *regexp.Regexp) (m SearchMatch) {
  50         return firstMatch(db.RegexpSearchN(re, 1))
  51 }
  52
  53 // Prefix exact string search, returns first N matches (nondeterministic)
  54 func (db *TitlesDatabase) PrefixSearchN(s string, n int) (matches SearchMatches) {
  55         return db.doSearchN(func(k string) bool { return strings.HasPrefix(k, s) }, n)
  56 }
  57
  58 // Prefix exact string search (nondeterministic order)
  59 func (db *TitlesDatabase) PrefixSearchAll(s string) (matches SearchMatches) {
  60         return db.PrefixSearchN(s, -1)
  61 }
  62
  63 // Prefix exact string search, returns first match (nondeterministic)
  64 func (db *TitlesDatabase) PrefixSearch(s string) (m SearchMatch) {
  65         return firstMatch(db.PrefixSearchN(s, 1))
  66 }
  67
  68 // Suffix exact string search, returns first N matches (nondeterministic)
  69 func (db *TitlesDatabase) SuffixSearchN(s string, n int) (matches SearchMatches) {
  70         return db.doSearchN(func(k string) bool { return strings.HasSuffix(k, s) }, n)
  71 }
  72
  73 // Suffix exact string search (nondeterministic order)
  74 func (db *TitlesDatabase) SuffixSearchAll(s string) (matches SearchMatches) {
  75         return db.SuffixSearchN(s, -1)
  76 }
  77
  78 // Suffix exact string search, returns first match (nondeterministic)
  79 func (db *TitlesDatabase) SuffixSearch(s string) (m SearchMatch) {
  80         return firstMatch(db.SuffixSearchN(s, 1))
  81 }
  82
  83 // Prefix string search with case folding, returns first N matches (nondeterministic)
  84 func (db *TitlesDatabase) PrefixSearchFoldN(s string, n int) (matches SearchMatches) {
  85         s = strings.ToLower(s)
  86         return db.doSearchN(func(k string) bool { return strings.HasPrefix(strings.ToLower(k), s) }, n)
  87 }
  88
  89 // Prefix string search with case folding (nondeterministic order)
  90 func (db *TitlesDatabase) PrefixSearchFoldAll(s string) (matches SearchMatches) {
  91         return db.PrefixSearchFoldN(s, -1)
  92 }
  93
  94 // Prefix string search with case folding, returns first match (nondeterministic)
  95 func (db *TitlesDatabase) PrefixSearchFold(s string) (m SearchMatch) {
  96         return firstMatch(db.PrefixSearchFoldN(s, 1))
  97 }
  98
  99 // Suffix string search with case folding, returns first N matches (nondeterministic)
 100 func (db *TitlesDatabase) SuffixSearchFoldN(s string, n int) (matches SearchMatches) {
 101         s = strings.ToLower(s)
 102         return db.doSearchN(func(k string) bool { return strings.HasSuffix(strings.ToLower(k), s) }, n)
 103 }
 104
 105 // Suffix string search with case folding (nondeterministic order)
 106 func (db *TitlesDatabase) SuffixSearchFoldAll(s string) (matches SearchMatches) {
 107         return db.SuffixSearchFoldN(s, -1)
 108 }
 109
 110 // Suffix string search with case folding, returns first match (nondeterministic)
 111 func (db *TitlesDatabase) SuffixSearchFold(s string) (m SearchMatch) {
 112         return firstMatch(db.SuffixSearchFoldN(s, 1))
 113 }
 114
 115 // \b doesn't consider the boundary between e.g. '.' and ' ' in ". "
 116 // to be a word boundary, but . may be significant in a title
 117 const wordBound = ` `
 118
 119 // Fuzzy string search with algorithm similar to the official Chii[AR] IRC bot.
 120 //
 121 // First attempts an exact search. Otherwise, uses strings.Fields to split the string
 122 // into words and tries, in order, the following alternate matches:
 123 //
 124 // * Initial words (prefix, but ending at word boundary)
 125 //
 126 // * Final words (suffix, but starting at word boundary)
 127 //
 128 // * Infix words
 129 //
 130 // * Prefix
 131 //
 132 // * Suffix
 133 //
 134 // * Initial words in the given order, but with possible words between them
 135 //
 136 // * Final words in the given order
 137 //
 138 // * Infix words in the given order
 139 //
 140 // * Initial strings in the given order, but with other possible strings between them
 141 //
 142 // * Final strings in the given order
 143 //
 144 // * Any match with strings in the given order
 145 //
 146 // Failing all those cases, the search returns a nil ResultSet.
 147 func (db *TitlesDatabase) FuzzySearch(s string) (rs ResultSet) {
 148         // whole title
 149         if matches := db.ExactSearchAll(s); len(matches) > 0 {
 150                 // log.Printf("best case: %q", s)
 151                 return matches.ToResultSet(db)
 152         }
 153
 154         // all regexes are guaranteed to compile:
 155         // the user-supplied token already went through regexp.QuoteMeta
 156         // all other tokens are hardcoded, so a compilation failure is reason for panic
 157
 158         words := strings.Fields(regexp.QuoteMeta(s))
 159         q := strings.Join(words, `.*`)
 160
 161         candidates := db.RegexpSearchAll(regexp.MustCompile(q)).ToResultSet(db)
 162         if len(candidates) == 0 {
 163                 // log.Printf("no results: %q", s)
 164                 return nil
 165         }
 166         q = strings.Join(words, ` `)
 167
 168         // initial words (prefix, but ending at word boundary)
 169         re := regexp.MustCompile(`\A` + q + wordBound)
 170         reCmp := func(k string) bool { return re.MatchString(k) }
 171         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 172                 // log.Printf("1st case: %q", s)
 173                 return
 174         }
 175
 176         // final words (suffix, but starting at a word boundary)
 177         re = regexp.MustCompile(wordBound + q + `\z`)
 178         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 179                 // log.Printf("2nd case: %q", s)
 180                 return
 181         }
 182
 183         // infix words
 184         re = regexp.MustCompile(wordBound + q + wordBound)
 185         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 186                 // log.Printf("3rd case: %q", s)
 187                 return
 188         }
 189
 190         // initial substring
 191         if rs = candidates.FilterByTitles(
 192                 func(k string) bool {
 193                         return strings.HasPrefix(k, s)
 194                 }); len(rs) > 0 {
 195                 // log.Printf("4th case: %q", s)
 196                 return
 197         }
 198
 199         // terminal substring
 200         if rs = candidates.FilterByTitles(
 201                 func(k string) bool {
 202                         return strings.HasSuffix(k, s)
 203                 }); len(rs) > 0 {
 204                 // log.Printf("5th case: %q", s)
 205                 return
 206         }
 207
 208         // words in that order, but with possible words between them...
 209         q = strings.Join(words, ` +(?:[^ ]+ +)*`)
 210
 211         // ... initial ...
 212         re = regexp.MustCompile(`\A` + q + wordBound)
 213         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 214                 // log.Printf("6th case: %q", s)
 215                 return
 216         }
 217
 218         // ... then final ...
 219         re = regexp.MustCompile(wordBound + q + `\z`)
 220         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 221                 // log.Printf("7th case: %q", s)
 222                 return
 223         }
 224
 225         // ... then anywhere
 226         re = regexp.MustCompile(wordBound + q + wordBound)
 227         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 228                 // log.Printf("8th case: %q", s)
 229                 return
 230         }
 231
 232         // then it's that, but with any or no characters between the input words...
 233         q = strings.Join(words, `.*`)
 234
 235         // and the same priority order as for the substring case
 236         // initial
 237         re = regexp.MustCompile(`\A` + q)
 238         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 239                 // log.Printf("9th case: %q", s)
 240                 return
 241         }
 242
 243         // final
 244         re = regexp.MustCompile(q + `\z`)
 245         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 246                 // log.Printf("10th case: %q", s)
 247                 return
 248         }
 249
 250         // no result better than the inital candidates
 251         // log.Printf("worst case: %q", s)
 252         return candidates
 253 }
 254
 255 // Version with case folding of FuzzySearch.
 256 //
 257 // See the FuzzySearch documentation for details.
 258 func (db *TitlesDatabase) FuzzySearchFold(s string) (rs ResultSet) {
 259         // whole title
 260         if matches := db.ExactSearchFoldAll(s); len(matches) > 0 {
 261                 return matches.ToResultSet(db)
 262         }
 263
 264         words := strings.Fields(`(?i:` + regexp.QuoteMeta(s) + `)`)
 265         q := strings.Join(words, `.*`)
 266
 267         candidates := db.RegexpSearchAll(regexp.MustCompile(q)).ToResultSet(db)
 268         if len(candidates) == 0 {
 269                 // log.Printf("no results: %q", s)
 270                 return nil
 271         }
 272         q = strings.Join(words, `\s+`)
 273
 274         // initial words (prefix, but ending at word boundary)
 275         re := regexp.MustCompile(`\A` + q + wordBound)
 276         reCmp := func(k string) bool { return re.MatchString(k) }
 277         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 278                 // log.Printf("1st case: %q", s)
 279                 return
 280         }
 281
 282         // final words (suffix, but starting at a word boundary)
 283         re = regexp.MustCompile(wordBound + q + `\z`)
 284         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 285                 // log.Printf("2nd case: %q", s)
 286                 return
 287         }
 288
 289         // infix words
 290         re = regexp.MustCompile(wordBound + q + wordBound)
 291         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 292                 // log.Printf("3rd case: %q", s)
 293                 return
 294         }
 295
 296         // initial substring
 297         ls := strings.ToLower(s)
 298         if rs = candidates.FilterByTitles(
 299                 func(k string) bool {
 300                         return strings.HasPrefix(strings.ToLower(k), ls)
 301                 }); len(rs) > 0 {
 302                 // log.Printf("4th case: %q", s)
 303                 return
 304         }
 305
 306         // terminal substring
 307         if rs = candidates.FilterByTitles(
 308                 func(k string) bool {
 309                         return strings.HasSuffix(strings.ToLower(k), ls)
 310                 }); len(rs) > 0 {
 311                 // log.Printf("5th case: %q", s)
 312                 return
 313         }
 314
 315         // words in that order, but with possible words between them...
 316         q = strings.Join(words, `\s+(?:\S+\s+)*`)
 317
 318         // ... initial ...
 319         re = regexp.MustCompile(`\A` + q + wordBound)
 320         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 321                 // log.Printf("6th case: %q", s)
 322                 return
 323         }
 324
 325         // ... then final ...
 326         re = regexp.MustCompile(wordBound + q + `\z`)
 327         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 328                 // log.Printf("7th case: %q", s)
 329                 return
 330         }
 331
 332         // ... then anywhere
 333         re = regexp.MustCompile(wordBound + q + wordBound)
 334         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 335                 // log.Printf("8th case: %q", s)
 336                 return
 337         }
 338
 339         // then it's that, but with any or no characters between the input words...
 340         q = strings.Join(words, `.*`)
 341
 342         // and the same priority order as for the substring case
 343         // initial
 344         re = regexp.MustCompile(`\A` + q)
 345         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 346                 // log.Printf("9th case: %q", s)
 347                 return
 348         }
 349
 350         // final
 351         re = regexp.MustCompile(q + `\z`)
 352         if rs = candidates.FilterByTitles(reCmp); len(rs) > 0 {
 353                 // log.Printf("10th case: %q", s)
 354                 return
 355         }
 356
 357         // no result better than the inital candidates
 358         // log.Printf("worst case: %q", s)
 359         return candidates
 360 }