1package cache
2
3import (
4 "sort"
5 "strings"
6 "unicode"
7)
8
9// NewTrie creates a new empty trie.
10func NewTrie() *Trie {
11 return &Trie{
12 Root: &TrieNode{
13 Children: make(map[rune]*TrieNode),
14 Entries: make([]uint32, 0),
15 },
16 }
17}
18
19// Insert adds an entry to the trie for all its prefixes.
20func (t *Trie) Insert(text string, entryID uint32) {
21 text = normalizeText(text)
22
23 // Insert all meaningful prefixes
24 words := strings.Fields(text)
25 for _, word := range words {
26 if len(word) < 2 { // Skip single characters
27 continue
28 }
29 t.insertWord(word, entryID)
30
31 // Also insert domain parts for URLs
32 if strings.Contains(word, ".") {
33 parts := strings.Split(word, ".")
34 for _, part := range parts {
35 if len(part) >= 2 {
36 t.insertWord(part, entryID)
37 }
38 }
39 }
40 }
41}
42
43// insertWord inserts a single word and its prefixes into the trie.
44func (t *Trie) insertWord(word string, entryID uint32) {
45 node := t.Root
46
47 // Insert the full word
48 for _, char := range word {
49 if node.Children[char] == nil {
50 node.Children[char] = &TrieNode{
51 Children: make(map[rune]*TrieNode),
52 Entries: make([]uint32, 0),
53 }
54 }
55 node = node.Children[char]
56 }
57
58 // Add the entry ID to all nodes in the path
59 current := t.Root
60 for _, char := range word {
61 current = current.Children[char]
62 if !containsUint32(current.Entries, entryID) {
63 current.Entries = append(current.Entries, entryID)
64 }
65 }
66 node.IsEnd = true
67}
68
69// Search finds all entries that match the given prefix.
70func (t *Trie) Search(prefix string) []uint32 {
71 if len(prefix) == 0 {
72 return nil
73 }
74
75 prefix = normalizeText(prefix)
76 node := t.Root
77
78 // Navigate to the prefix node
79 for _, char := range prefix {
80 if node.Children[char] == nil {
81 return nil // Prefix not found
82 }
83 node = node.Children[char]
84 }
85
86 // Return all entries at this node
87 result := make([]uint32, len(node.Entries))
88 copy(result, node.Entries)
89 return result
90}
91
92// PrefixSearch finds all entries that start with any word in the query.
93func (t *Trie) PrefixSearch(query string) []uint32 {
94 query = normalizeText(query)
95 words := strings.Fields(query)
96
97 entrySet := make(map[uint32]struct{})
98
99 for _, word := range words {
100 if len(word) < 2 {
101 continue
102 }
103
104 entries := t.Search(word)
105 for _, entryID := range entries {
106 entrySet[entryID] = struct{}{}
107 }
108 }
109
110 // Convert set back to slice
111 result := make([]uint32, 0, len(entrySet))
112 for entryID := range entrySet {
113 result = append(result, entryID)
114 }
115
116 return result
117}
118
119// buildTrigramIndex creates a trigram index from the cache entries.
120func (c *DmenuFuzzyCache) buildTrigramIndex() {
121 c.trigramIndex = make(map[string][]uint32)
122
123 for i, entry := range c.entries {
124 entryID := uint32(i)
125
126 // Build trigrams from URL and title
127 text := normalizeText(entry.URL + " " + entry.Title)
128 trigrams := extractTrigrams(text)
129
130 for _, trigram := range trigrams {
131 if !containsUint32(c.trigramIndex[trigram], entryID) {
132 c.trigramIndex[trigram] = append(c.trigramIndex[trigram], entryID)
133 }
134 }
135 }
136}
137
138// buildPrefixTrie creates a prefix trie from the cache entries.
139func (c *DmenuFuzzyCache) buildPrefixTrie() {
140 c.prefixTrie = NewTrie()
141
142 for i, entry := range c.entries {
143 entryID := uint32(i)
144
145 // Add URL and title to trie
146 c.prefixTrie.Insert(entry.URL, entryID)
147 if entry.Title != "" {
148 c.prefixTrie.Insert(entry.Title, entryID)
149 }
150 }
151}
152
153// extractTrigrams extracts all 3-character sequences from text.
154func extractTrigrams(text string) []string {
155 if len(text) < 3 {
156 return nil
157 }
158
159 trigrams := make([]string, 0, len(text)-2)
160 runes := []rune(text)
161
162 for i := 0; i <= len(runes)-3; i++ {
163 trigram := string(runes[i : i+3])
164 trigrams = append(trigrams, trigram)
165 }
166
167 return trigrams
168}
169
170// normalizeText normalizes text for consistent matching.
171func normalizeText(text string) string {
172 // Convert to lowercase and remove extra whitespace
173 text = strings.ToLower(text)
174
175 // Remove protocol prefixes for URLs
176 text = strings.TrimPrefix(text, "https://")
177 text = strings.TrimPrefix(text, "http://")
178 text = strings.TrimPrefix(text, "www.")
179
180 // Replace non-alphanumeric characters with spaces
181 var result strings.Builder
182 result.Grow(len(text))
183
184 for _, r := range text {
185 if unicode.IsLetter(r) || unicode.IsDigit(r) {
186 result.WriteRune(r)
187 } else {
188 result.WriteRune(' ')
189 }
190 }
191
192 // Normalize whitespace
193 return strings.Join(strings.Fields(result.String()), " ")
194}
195
196// containsUint32 checks if a slice contains a uint32 value.
197func containsUint32(slice []uint32, value uint32) bool {
198 for _, v := range slice {
199 if v == value {
200 return true
201 }
202 }
203 return false
204}
205
206// intersectUint32 finds the intersection of two uint32 slices.
207func intersectUint32(a, b []uint32) []uint32 {
208 if len(a) == 0 || len(b) == 0 {
209 return nil
210 }
211
212 // Use the smaller slice for the map to save memory
213 if len(a) > len(b) {
214 a, b = b, a
215 }
216
217 // Create a set from the smaller slice
218 set := make(map[uint32]struct{}, len(a))
219 for _, v := range a {
220 set[v] = struct{}{}
221 }
222
223 // Find intersection
224 result := make([]uint32, 0, len(a))
225 for _, v := range b {
226 if _, exists := set[v]; exists {
227 result = append(result, v)
228 delete(set, v) // Avoid duplicates
229 }
230 }
231
232 return result
233}
234
235// unionUint32 finds the union of multiple uint32 slices.
236func unionUint32(slices ...[]uint32) []uint32 {
237 if len(slices) == 0 {
238 return nil
239 }
240
241 // Estimate capacity
242 capacity := 0
243 for _, slice := range slices {
244 capacity += len(slice)
245 }
246
247 set := make(map[uint32]struct{}, capacity)
248 for _, slice := range slices {
249 for _, v := range slice {
250 set[v] = struct{}{}
251 }
252 }
253
254 result := make([]uint32, 0, len(set))
255 for v := range set {
256 result = append(result, v)
257 }
258
259 return result
260}
261
262// getTrigramCandidates finds entry candidates using trigram matching.
263func (c *DmenuFuzzyCache) getTrigramCandidates(query string) []uint32 {
264 query = normalizeText(query)
265 trigrams := extractTrigrams(query)
266
267 if len(trigrams) == 0 {
268 return nil
269 }
270
271 // Get candidates from first trigram
272 candidates := c.trigramIndex[trigrams[0]]
273 if len(candidates) == 0 {
274 return nil
275 }
276
277 // Intersect with other trigrams for more accuracy
278 for i := 1; i < len(trigrams) && len(candidates) > 0; i++ {
279 nextCandidates := c.trigramIndex[trigrams[i]]
280 if len(nextCandidates) == 0 {
281 continue // Skip empty trigrams
282 }
283 candidates = intersectUint32(candidates, nextCandidates)
284 }
285
286 return candidates
287}
288
289// getPrefixCandidates finds entry candidates using prefix matching.
290func (c *DmenuFuzzyCache) getPrefixCandidates(query string) []uint32 {
291 if c.prefixTrie == nil {
292 return nil
293 }
294
295 return c.prefixTrie.PrefixSearch(query)
296}
297
298// getAllCandidates combines trigram and prefix matching for comprehensive candidate selection.
299func (c *DmenuFuzzyCache) getAllCandidates(query string) []uint32 {
300 trigramCandidates := c.getTrigramCandidates(query)
301 prefixCandidates := c.getPrefixCandidates(query)
302
303 if len(trigramCandidates) == 0 && len(prefixCandidates) == 0 {
304 return nil
305 }
306
307 if len(trigramCandidates) == 0 {
308 return prefixCandidates
309 }
310
311 if len(prefixCandidates) == 0 {
312 return trigramCandidates
313 }
314
315 // Combine both candidate sets
316 return unionUint32(trigramCandidates, prefixCandidates)
317}
318
319// buildSortedIndex creates a pre-sorted index of entries by score for O(1) top entries retrieval.
320func (c *DmenuFuzzyCache) buildSortedIndex() {
321 // Create index array
322 c.sortedIndex = make([]uint32, len(c.entries))
323 for i := range c.sortedIndex {
324 c.sortedIndex[i] = uint32(i)
325 }
326
327 // Sort indices by entry scores (highest first)
328 sort.Slice(c.sortedIndex, func(i, j int) bool {
329 return c.entries[c.sortedIndex[i]].Score > c.entries[c.sortedIndex[j]].Score
330 })
331}