Improve country selection search

submodules/CountrySelectionUI/Sources/AuthorizationSequenceCountrySelectionControllerNode.swift

@@ -2,6 +2,7 @@ import Foundation
 import UIKit
 import AsyncDisplayKit
 import Display
+import Postbox
 import TelegramCore
 import SyncCore
 import TelegramPresentationData
@@ -73,6 +74,75 @@ func localizedCountryNamesAndCodes(strings: PresentationStrings) -> [((String, S
     return result
 }
 
+private func stringTokens(_ string: String) -> [ValueBoxKey] {
+    let nsString = string.replacingOccurrences(of: ".", with: "").folding(options: .diacriticInsensitive, locale: .current).lowercased() as NSString
+    
+    let flag = UInt(kCFStringTokenizerUnitWord)
+    let tokenizer = CFStringTokenizerCreate(kCFAllocatorDefault, nsString, CFRangeMake(0, nsString.length), flag, CFLocaleCopyCurrent())
+    var tokenType = CFStringTokenizerAdvanceToNextToken(tokenizer)
+    var tokens: [ValueBoxKey] = []
+    
+    var addedTokens = Set<ValueBoxKey>()
+    while tokenType != [] {
+        let currentTokenRange = CFStringTokenizerGetCurrentTokenRange(tokenizer)
+        
+        if currentTokenRange.location >= 0 && currentTokenRange.length != 0 {
+            let token = ValueBoxKey(length: currentTokenRange.length * 2)
+            nsString.getCharacters(token.memory.assumingMemoryBound(to: unichar.self), range: NSMakeRange(currentTokenRange.location, currentTokenRange.length))
+            if !addedTokens.contains(token) {
+                tokens.append(token)
+                addedTokens.insert(token)
+            }
+        }
+        tokenType = CFStringTokenizerAdvanceToNextToken(tokenizer)
+    }
+    
+    return tokens
+}
+
+private func matchStringTokens(_ tokens: [ValueBoxKey], with other: [ValueBoxKey]) -> Bool {
+    if other.isEmpty {
+        return false
+    } else if other.count == 1 {
+        let otherToken = other[0]
+        for token in tokens {
+            if otherToken.isPrefix(to: token) {
+                return true
+            }
+        }
+    } else {
+        for otherToken in other {
+            var found = false
+            for token in tokens {
+                if otherToken.isPrefix(to: token) {
+                    found = true
+                    break
+                }
+            }
+            if !found {
+                return false
+            }
+        }
+        return true
+    }
+    return false
+}
+
+private func searchCountries(items: [((String, String), String, Int)], query: String) -> [((String, String), String, Int)] {
+    let queryTokens = stringTokens(query.lowercased())
+    
+    var result: [((String, String), String, Int)] = []
+    for item in items {
+        let string = "\(item.0) \(item.1)"
+        let tokens = stringTokens(string)
+        if matchStringTokens(tokens, with: queryTokens) {
+            result.append(item)
+        }
+    }
+    
+    return result
+}
+
 final class AuthorizationSequenceCountrySelectionControllerNode: ASDisplayNode, UITableViewDelegate, UITableViewDataSource {
     let itemSelected: (((String, String), String, Int)) -> Void
     
@@ -88,6 +158,7 @@ final class AuthorizationSequenceCountrySelectionControllerNode: ASDisplayNode,
     private let sectionTitles: [String]
     
     private var searchResults: [((String, String), String, Int)] = []
+    private let countryNamesAndCodes: [((String, String), String, Int)]
     
     init(theme: PresentationTheme, strings: PresentationStrings, displayCodes: Bool, itemSelected: @escaping (((String, String), String, Int)) -> Void) {
         self.theme = theme
@@ -107,6 +178,7 @@ final class AuthorizationSequenceCountrySelectionControllerNode: ASDisplayNode,
         }
         
         let countryNamesAndCodes = localizedCountryNamesAndCodes(strings: strings)
+        self.countryNamesAndCodes = countryNamesAndCodes
         
         var sections: [(String, [((String, String), String, Int)])] = []
         for (names, id, code) in countryNamesAndCodes.sorted(by: { lhs, rhs in
@@ -176,18 +248,7 @@ final class AuthorizationSequenceCountrySelectionControllerNode: ASDisplayNode,
             self.searchTableView.reloadData()
             self.searchTableView.isHidden = true
         } else {
-            let normalizedQuery = query.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
-            
-            var results: [((String, String), String, Int)] = []
-            for (_, items) in self.sections {
-                for item in items {
-                    if item.0.0.lowercased().hasPrefix(normalizedQuery) || item.0.1.lowercased().hasPrefix(normalizedQuery) {
-                        results.append(item)
-                    }
-                }
-            }
-            
-            self.searchResults = results
+            self.searchResults = searchCountries(items: self.countryNamesAndCodes, query: query)
             self.searchTableView.isHidden = false
             self.searchTableView.reloadData()
         }