1import * as Equal from "../Equal.ts"
2import { format } from "../Formatter.ts"
3import { dual } from "../Function.ts"
4import * as Hash from "../Hash.ts"
5import { NodeInspectSymbol, toJson } from "../Inspectable.ts"
6import * as Option from "../Option.ts"
7import type * as Ordering from "../Ordering.ts"
8import { pipeArguments } from "../Pipeable.ts"
9import { hasProperty } from "../Predicate.ts"
10import * as Result from "../Result.ts"
11import type * as TR from "../Trie.ts"
12import type { NoInfer } from "../Types.ts"
13
14/** @internal */
15export const TrieTypeId = "~effect/Trie"
16
17type TraversalMap<K, V, A> = (k: K, v: V) => A
18
19type TraversalFilter<K, V> = (k: K, v: V) => boolean
20
21/** @internal */
22export interface TrieImpl<in out V> extends TR.Trie<V> {
23 readonly _root: Node<V> | undefined
24 readonly _count: number
25}
26
27const trieVariance = {
28 /* c8 ignore next */
29 _Value: (_: never) => _
30}
31
32const trieSeed = Hash.string(TrieTypeId)
33
34const TrieProto: TR.Trie<unknown> = {
35 [TrieTypeId]: trieVariance,
36 [Symbol.iterator]<V>(this: TrieImpl<V>): Iterator<[string, V]> {
37 return new TrieIterator(this, (k, v) => [k, v], () => true)
38 },
39 [Hash.symbol](this: TR.Trie<unknown>): number {
40 let hash = trieSeed
41 for (const item of this) {
42 hash ^= Hash.combine(Hash.hash(item[0]), Hash.hash(item[1]))
43 }
44 return Hash.optimize(hash)
45 },
46 [Equal.symbol]<V>(this: TrieImpl<V>, that: unknown): boolean {
47 if (isTrie(that) && size(this) === size(that)) {
48 const entries = Array.from(that)
49 return Array.from(this).every((itemSelf, i) => {
50 const itemThat = entries[i]
51 return Equal.equals(itemSelf[0], itemThat[0]) && Equal.equals(itemSelf[1], itemThat[1])
52 })
53 }
54 return false
55 },
56 toString() {
57 return `Trie(${format(Array.from(this))})`
58 },
59 toJSON() {
60 return {
61 _id: "Trie",
62 values: toJson(Array.from(this))
63 }
64 },
65 [NodeInspectSymbol]() {
66 return this.toJSON()
67 },
68 pipe() {
69 return pipeArguments(this, arguments)
70 }
71}
72
73const makeImpl = <V>(root: Node<V> | undefined): TrieImpl<V> => {
74 const trie = Object.create(TrieProto)
75 trie._root = root
76 trie._count = root?.count ?? 0
77 return trie
78}
79
80class TrieIterator<in out V, out T> implements IterableIterator<T> {
81 stack: Array<[Node<V>, string, boolean]> = []
82 readonly trie: TrieImpl<V>
83 readonly f: TraversalMap<string, V, T>
84 readonly filter: TraversalFilter<string, V>
85
86 constructor(
87 trie: TrieImpl<V>,
88 f: TraversalMap<string, V, T>,
89 filter: TraversalFilter<string, V>
90 ) {
91 this.trie = trie
92 this.f = f
93 this.filter = filter
94 const root = trie._root !== undefined ? trie._root : undefined
95 if (root !== undefined) {
96 this.stack.push([root, "", false])
97 }
98 }
99
100 next(): IteratorResult<T> {
101 while (this.stack.length > 0) {
102 const [node, keyString, isAdded] = this.stack.pop()!
103
104 if (isAdded) {
105 const value = node.value
106 if (value !== undefined) {
107 const key = keyString + node.key
108 if (this.filter(key, value.value)) {
109 return { done: false, value: this.f(key, value.value) }
110 }
111 }
112 } else {
113 this.addToStack(node, keyString)
114 }
115 }
116
117 return { done: true, value: undefined }
118 }
119
120 addToStack(node: Node<V>, keyString: string) {
121 if (node.right !== undefined) {
122 this.stack.push([node.right, keyString, false])
123 }
124 if (node.mid !== undefined) {
125 this.stack.push([node.mid, keyString + node.key, false])
126 }
127 this.stack.push([node, keyString, true])
128 if (node.left !== undefined) {
129 this.stack.push([node.left, keyString, false])
130 }
131 }
132
133 [Symbol.iterator](): IterableIterator<T> {
134 return new TrieIterator(this.trie, this.f, this.filter)
135 }
136}
137
138/** @internal */
139export const isTrie: {
140 <V>(u: Iterable<readonly [string, V]>): u is TR.Trie<V>
141 (u: unknown): u is TR.Trie<unknown>
142} = (u: unknown): u is TR.Trie<unknown> => hasProperty(u, TrieTypeId)
143
144/** @internal */
145export const empty = <V = never>(): TR.Trie<V> => makeImpl<V>(undefined)
146
147/** @internal */
148export const fromIterable = <V>(entries: Iterable<readonly [string, V]>) => {
149 let trie = empty<V>()
150 for (const [key, value] of entries) {
151 trie = insert(trie, key, value)
152 }
153 return trie
154}
155
156/** @internal */
157export const make = <Entries extends Array<readonly [string, any]>>(...entries: Entries): TR.Trie<
158 Entries[number] extends readonly [any, infer V] ? V : never
159> => {
160 return fromIterable(entries)
161}
162
163/** @internal */
164export const insert = dual<
165 <V>(key: string, value: V) => (self: TR.Trie<V>) => TR.Trie<V>,
166 <V>(self: TR.Trie<V>, key: string, value: V) => TR.Trie<V>
167>(3, <V>(self: TR.Trie<V>, key: string, value: V) => {
168 if (key.length === 0) return self
169
170 // -1:left | 0:mid | 1:right
171 const dStack: Array<Ordering.Ordering> = []
172 const nStack: Array<Node<V>> = []
173 let n: Node<V> = (self as TrieImpl<V>)._root ?? {
174 key: key[0],
175 count: 0
176 }
177 let count = n.count + 1
178 let cIndex = 0
179
180 while (cIndex < key.length) {
181 const c = key[cIndex]
182 nStack.push(n)
183 if (c > n.key) {
184 dStack.push(1)
185 if (n.right === undefined) {
186 n = { key: c, count }
187 } else {
188 n = n.right
189 }
190 } else if (c < n.key) {
191 dStack.push(-1)
192 if (n.left === undefined) {
193 n = { key: c, count }
194 } else {
195 n = n.left
196 }
197 } else {
198 if (cIndex === key.length - 1) {
199 if (n.value !== undefined) {
200 count -= 1
201 }
202 nStack[nStack.length - 1] = {
203 key: n.key,
204 count,
205 value: { value },
206 left: n.left,
207 mid: n.mid,
208 right: n.right
209 }
210 } else if (n.mid === undefined) {
211 dStack.push(0)
212 n = { key: key[cIndex + 1], count }
213 } else {
214 dStack.push(0)
215 n = n.mid
216 }
217
218 cIndex += 1
219 }
220 }
221
222 // Rebuild path to leaf node (Path-copying immutability)
223 for (let s = nStack.length - 2; s >= 0; --s) {
224 const n2 = nStack[s]
225 const d = dStack[s]
226 if (d === -1) {
227 // left
228 nStack[s] = {
229 key: n2.key,
230 count,
231 value: n2.value,
232 left: nStack[s + 1],
233 mid: n2.mid,
234 right: n2.right
235 }
236 } else if (d === 1) {
237 // right
238 nStack[s] = {
239 key: n2.key,
240 count,
241 value: n2.value,
242 left: n2.left,
243 mid: n2.mid,
244 right: nStack[s + 1]
245 }
246 } else {
247 // mid
248 nStack[s] = {
249 key: n2.key,
250 count,
251 value: n2.value,
252 left: n2.left,
253 mid: nStack[s + 1],
254 right: n2.right
255 }
256 }
257 }
258
259 nStack[0].count = count
260 return makeImpl(nStack[0])
261})
262
263/** @internal */
264export const size = <V>(self: TR.Trie<V>): number => (self as TrieImpl<V>)._root?.count ?? 0
265
266/** @internal */
267export const isEmpty = <V>(self: TR.Trie<V>): boolean => size(self) === 0
268
269/** @internal */
270export const keys = <V>(self: TR.Trie<V>): IterableIterator<string> =>
271 new TrieIterator(self as TrieImpl<V>, (key) => key, () => true)
272
273/** @internal */
274export const values = <V>(self: TR.Trie<V>): IterableIterator<V> =>
275 new TrieIterator(self as TrieImpl<V>, (_, value) => value, () => true)
276
277/** @internal */
278export const entries = <V>(self: TR.Trie<V>): IterableIterator<[string, V]> =>
279 new TrieIterator(self as TrieImpl<V>, (key, value) => [key, value], () => true)
280
281/** @internal */
282export const reduce = dual<
283 <Z, V>(
284 zero: Z,
285 f: (accumulator: Z, value: V, key: string) => Z
286 ) => (self: TR.Trie<V>) => Z,
287 <Z, V>(self: TR.Trie<V>, zero: Z, f: (accumulator: Z, value: V, key: string) => Z) => Z
288>(3, (self, zero, f) => {
289 let accumulator = zero
290 for (const entry of self) {
291 accumulator = f(accumulator, entry[1], entry[0])
292 }
293 return accumulator
294})
295
296/** @internal */
297export const map = dual<
298 <A, V>(f: (value: V, key: string) => A) => (self: TR.Trie<V>) => TR.Trie<A>,
299 <V, A>(self: TR.Trie<V>, f: (value: V, key: string) => A) => TR.Trie<A>
300>(2, (self, f) =>
301 reduce(
302 self,
303 empty(),
304 (trie, value, key) => insert(trie, key, f(value, key))
305 ))
306
307/** @internal */
308export const filter: {
309 <A, B extends A>(f: (a: NoInfer<A>, k: string) => a is B): (self: TR.Trie<A>) => TR.Trie<B>
310 <A>(f: (a: NoInfer<A>, k: string) => boolean): (self: TR.Trie<A>) => TR.Trie<A>
311 <A, B extends A>(self: TR.Trie<A>, f: (a: A, k: string) => a is B): TR.Trie<B>
312 <A>(self: TR.Trie<A>, f: (a: A, k: string) => boolean): TR.Trie<A>
313} = dual(
314 2,
315 <A>(self: TR.Trie<A>, f: (a: A, k: string) => boolean): TR.Trie<A> =>
316 reduce(
317 self,
318 empty(),
319 (trie, value, key) => f(value, key) ? insert(trie, key, value) : trie
320 )
321)
322
323/** @internal */
324export const filterMap = dual<
325 <A, B, X>(
326 f: (input: A, key: string) => Result.Result<B, X>
327 ) => (self: TR.Trie<A>) => TR.Trie<B>,
328 <A, B, X>(self: TR.Trie<A>, f: (input: A, key: string) => Result.Result<B, X>) => TR.Trie<B>
329>(2, <A, B, X>(self: TR.Trie<A>, f: (input: A, key: string) => Result.Result<B, X>): TR.Trie<B> =>
330 reduce(
331 self,
332 empty<B>(),
333 (trie, value, key) => {
334 const result = f(value, key)
335 return Result.isSuccess(result) ? insert(trie, key, result.success) : trie
336 }
337 ))
338
339/** @internal */
340export const compact = <A>(self: TR.Trie<Option.Option<A>>) =>
341 reduce(
342 self,
343 empty<A>(),
344 (trie, option, key) => Option.isSome(option) ? insert(trie, key, option.value) : trie
345 )
346
347/** @internal */
348export const forEach = dual<
349 <V>(f: (value: V, key: string) => void) => (self: TR.Trie<V>) => void,
350 <V>(self: TR.Trie<V>, f: (value: V, key: string) => void) => void
351>(2, (self, f) => reduce(self, void 0 as void, (_, value, key) => f(value, key)))
352
353/** @internal */
354export const keysWithPrefix = dual<
355 (prefix: string) => <V>(self: TR.Trie<V>) => IterableIterator<string>,
356 <V>(self: TR.Trie<V>, prefix: string) => IterableIterator<string>
357>(
358 2,
359 <V>(self: TR.Trie<V>, prefix: string): IterableIterator<string> =>
360 new TrieIterator(self as TrieImpl<V>, (key) => key, (key) => key.startsWith(prefix))
361)
362
363/** @internal */
364export const valuesWithPrefix = dual<
365 (prefix: string) => <V>(self: TR.Trie<V>) => IterableIterator<V>,
366 <V>(self: TR.Trie<V>, prefix: string) => IterableIterator<V>
367>(
368 2,
369 <V>(self: TR.Trie<V>, prefix: string): IterableIterator<V> =>
370 new TrieIterator(self as TrieImpl<V>, (_, value) => value, (key) => key.startsWith(prefix))
371)
372
373/** @internal */
374export const entriesWithPrefix = dual<
375 (prefix: string) => <V>(self: TR.Trie<V>) => IterableIterator<[string, V]>,
376 <V>(self: TR.Trie<V>, prefix: string) => IterableIterator<[string, V]>
377>(
378 2,
379 <V>(self: TR.Trie<V>, prefix: string): IterableIterator<[string, V]> =>
380 new TrieIterator(self as TrieImpl<V>, (key, value) => [key, value], (key) => key.startsWith(prefix))
381)
382
383/** @internal */
384export const toEntriesWithPrefix = dual<
385 (prefix: string) => <V>(self: TR.Trie<V>) => Array<[string, V]>,
386 <V>(self: TR.Trie<V>, prefix: string) => Array<[string, V]>
387>(
388 2,
389 <V>(self: TR.Trie<V>, prefix: string): Array<[string, V]> => Array.from(entriesWithPrefix(self, prefix))
390)
391
392/** @internal */
393export const get = dual<
394 (key: string) => <V>(self: TR.Trie<V>) => Option.Option<V>,
395 <V>(self: TR.Trie<V>, key: string) => Option.Option<V>
396>(
397 2,
398 <V>(self: TR.Trie<V>, key: string) => {
399 let n: Node<V> | undefined = (self as TrieImpl<V>)._root
400 if (n === undefined || key.length === 0) return Option.none()
401 let cIndex = 0
402 while (cIndex < key.length) {
403 const c = key[cIndex]
404 if (c > n.key) {
405 if (n.right === undefined) {
406 return Option.none()
407 } else {
408 n = n.right
409 }
410 } else if (c < n.key) {
411 if (n.left === undefined) {
412 return Option.none()
413 } else {
414 n = n.left
415 }
416 } else {
417 if (cIndex === key.length - 1) {
418 return n.value === undefined ? Option.none() : Option.some(n.value.value)
419 } else {
420 if (n.mid === undefined) {
421 return Option.none()
422 } else {
423 n = n.mid
424 cIndex += 1
425 }
426 }
427 }
428 }
429 return Option.none()
430 }
431)
432
433/** @internal */
434export const has = dual<
435 (key: string) => <V>(self: TR.Trie<V>) => boolean,
436 <V>(self: TR.Trie<V>, key: string) => boolean
437>(2, (self, key) => Option.isSome(get(self, key)))
438
439/** @internal */
440export const getUnsafe = dual<
441 (key: string) => <V>(self: TR.Trie<V>) => V,
442 <V>(self: TR.Trie<V>, key: string) => V
443>(2, (self, key) => {
444 const element = get(self, key)
445 if (Option.isNone(element)) {
446 throw new Error("Expected trie to contain key")
447 }
448 return element.value
449})
450
451/** @internal */
452export const remove = dual<
453 (key: string) => <V>(self: TR.Trie<V>) => TR.Trie<V>,
454 <V>(self: TR.Trie<V>, key: string) => TR.Trie<V>
455>(
456 2,
457 <V>(self: TR.Trie<V>, key: string) => {
458 let n: Node<V> | undefined = (self as TrieImpl<V>)._root
459 if (n === undefined || key.length === 0) return self
460
461 const count = n.count - 1
462 // -1:left | 0:mid | 1:right
463 const dStack: Array<Ordering.Ordering> = []
464 const nStack: Array<Node<V>> = []
465
466 let cIndex = 0
467 while (cIndex < key.length) {
468 const c = key[cIndex]
469 if (c > n.key) {
470 if (n.right === undefined) {
471 return self
472 } else {
473 nStack.push(n)
474 dStack.push(1)
475 n = n.right
476 }
477 } else if (c < n.key) {
478 if (n.left === undefined) {
479 return self
480 } else {
481 nStack.push(n)
482 dStack.push(-1)
483 n = n.left
484 }
485 } else {
486 if (cIndex === key.length - 1) {
487 if (n.value !== undefined) {
488 nStack.push(n)
489 dStack.push(0)
490 cIndex += 1
491 } else {
492 return self
493 }
494 } else {
495 if (n.mid === undefined) {
496 return self
497 } else {
498 nStack.push(n)
499 dStack.push(0)
500 n = n.mid
501 cIndex += 1
502 }
503 }
504 }
505 }
506
507 const removeNode = nStack[nStack.length - 1]
508 nStack[nStack.length - 1] = {
509 key: removeNode.key,
510 count,
511 left: removeNode.left,
512 mid: removeNode.mid,
513 right: removeNode.right
514 }
515
516 // Rebuild path to leaf node (Path-copying immutability)
517 for (let s = nStack.length - 2; s >= 0; --s) {
518 const n2 = nStack[s]
519 const d = dStack[s]
520 const child = nStack[s + 1]
521 const nc =
522 child.value === undefined && child.left === undefined && child.mid === undefined && child.right === undefined
523 ? undefined
524 : child
525 if (d === -1) {
526 // left
527 nStack[s] = {
528 key: n2.key,
529 count,
530 value: n2.value,
531 left: nc,
532 mid: n2.mid,
533 right: n2.right
534 }
535 } else if (d === 1) {
536 // right
537 nStack[s] = {
538 key: n2.key,
539 count,
540 value: n2.value,
541 left: n2.left,
542 mid: n2.mid,
543 right: nc
544 }
545 } else {
546 // mid
547 nStack[s] = {
548 key: n2.key,
549 count,
550 value: n2.value,
551 left: n2.left,
552 mid: nc,
553 right: n2.right
554 }
555 }
556 }
557
558 nStack[0].count = count
559 return makeImpl(nStack[0])
560 }
561)
562
563/** @internal */
564export const removeMany = dual<
565 (keys: Iterable<string>) => <V>(self: TR.Trie<V>) => TR.Trie<V>,
566 <V>(self: TR.Trie<V>, keys: Iterable<string>) => TR.Trie<V>
567>(2, (self, keys) => {
568 let trie = self
569 for (const key of keys) {
570 trie = remove(key)(trie)
571 }
572 return trie
573})
574
575/** @internal */
576export const insertMany = dual<
577 <V>(iter: Iterable<[string, V]>) => (self: TR.Trie<V>) => TR.Trie<V>,
578 <V>(self: TR.Trie<V>, iter: Iterable<[string, V]>) => TR.Trie<V>
579>(2, (self, iter) => {
580 let trie = self
581 for (const [key, value] of iter) {
582 trie = insert(key, value)(trie)
583 }
584 return trie
585})
586
587/** @internal */
588export const modify = dual<
589 <V>(key: string, f: (v: V) => V) => (self: TR.Trie<V>) => TR.Trie<V>,
590 <V>(self: TR.Trie<V>, key: string, f: (v: V) => V) => TR.Trie<V>
591>(
592 3,
593 <V>(self: TR.Trie<V>, key: string, f: (v: V) => V): TR.Trie<V> => {
594 let n: Node<V> | undefined = (self as TrieImpl<V>)._root
595 if (n === undefined || key.length === 0) return self
596
597 // -1:left | 0:mid | 1:right
598 const dStack: Array<Ordering.Ordering> = []
599 const nStack: Array<Node<V>> = []
600
601 let cIndex = 0
602 while (cIndex < key.length) {
603 const c = key[cIndex]
604 if (c > n.key) {
605 if (n.right === undefined) {
606 return self
607 } else {
608 nStack.push(n)
609 dStack.push(1)
610 n = n.right
611 }
612 } else if (c < n.key) {
613 if (n.left === undefined) {
614 return self
615 } else {
616 nStack.push(n)
617 dStack.push(-1)
618 n = n.left
619 }
620 } else {
621 if (cIndex === key.length - 1) {
622 if (n.value !== undefined) {
623 nStack.push(n)
624 dStack.push(0)
625 cIndex += 1
626 } else {
627 return self
628 }
629 } else {
630 if (n.mid === undefined) {
631 return self
632 } else {
633 nStack.push(n)
634 dStack.push(0)
635 n = n.mid
636 cIndex += 1
637 }
638 }
639 }
640 }
641
642 const updateNode = nStack[nStack.length - 1]
643 if (updateNode.value === undefined) {
644 return self
645 }
646
647 nStack[nStack.length - 1] = {
648 key: updateNode.key,
649 count: updateNode.count,
650 value: { value: f(updateNode.value.value) }, // Update
651 left: updateNode.left,
652 mid: updateNode.mid,
653 right: updateNode.right
654 }
655
656 // Rebuild path to leaf node (Path-copying immutability)
657 for (let s = nStack.length - 2; s >= 0; --s) {
658 const n2 = nStack[s]
659 const d = dStack[s]
660 const child = nStack[s + 1]
661 if (d === -1) {
662 // left
663 nStack[s] = {
664 key: n2.key,
665 count: n2.count,
666 value: n2.value,
667 left: child,
668 mid: n2.mid,
669 right: n2.right
670 }
671 } else if (d === 1) {
672 // right
673 nStack[s] = {
674 key: n2.key,
675 count: n2.count,
676 value: n2.value,
677 left: n2.left,
678 mid: n2.mid,
679 right: child
680 }
681 } else {
682 // mid
683 nStack[s] = {
684 key: n2.key,
685 count: n2.count,
686 value: n2.value,
687 left: n2.left,
688 mid: child,
689 right: n2.right
690 }
691 }
692 }
693
694 return makeImpl(nStack[0])
695 }
696)
697
698/** @internal */
699export const longestPrefixOf = dual<
700 (key: string) => <V>(self: TR.Trie<V>) => Option.Option<[string, V]>,
701 <V>(self: TR.Trie<V>, key: string) => Option.Option<[string, V]>
702>(
703 2,
704 <V>(self: TR.Trie<V>, key: string): Option.Option<[string, V]> => {
705 let n: Node<V> | undefined = (self as TrieImpl<V>)._root
706 if (n === undefined || key.length === 0) return Option.none()
707 let longestPrefixNode: Option.Option<[string, V]> = Option.none()
708 let cIndex = 0
709 while (cIndex < key.length) {
710 const c = key[cIndex]
711 if (c > n.key) {
712 if (n.right === undefined) {
713 break
714 } else {
715 n = n.right
716 }
717 } else if (c < n.key) {
718 if (n.left === undefined) {
719 break
720 } else {
721 n = n.left
722 }
723 } else {
724 if (n.value !== undefined) {
725 longestPrefixNode = Option.some([key.slice(0, cIndex + 1), n.value.value])
726 }
727 if (n.mid === undefined) {
728 break
729 } else {
730 n = n.mid
731 cIndex += 1
732 }
733 }
734 }
735
736 return longestPrefixNode
737 }
738)
739
740interface Node<V> {
741 key: string
742 count: number
743 value?: { readonly value: V } | undefined
744 left?: Node<V> | undefined
745 mid?: Node<V> | undefined
746 right?: Node<V> | undefined
747}