1+ package org.utbot.fuzzer
2+
3+ fun <T > trieOf (vararg values : Iterable <T >): Trie <T , T > = IdentityTrie <T >().apply {
4+ values.forEach(this ::add)
5+ }
6+
7+ fun stringTrieOf (vararg values : String ): StringTrie = StringTrie ().apply {
8+ values.forEach(this ::add)
9+ }
10+
11+ class StringTrie : IdentityTrie <Char >() {
12+ fun add (string : String ) = super .add(string.toCharArray().asIterable())
13+ fun removeCompletely (string : String ) = super .removeCompletely(string.toCharArray().asIterable())
14+ fun remove (string : String ) = super .remove(string.toCharArray().asIterable())
15+ operator fun get (string : String ) = super .get(string.toCharArray().asIterable())
16+ fun collect () = asSequence().map { String (it.toCharArray()) }.toSet()
17+ }
18+
19+ open class IdentityTrie <T > : Trie <T , T >({ it })
20+
21+ /* *
22+ * Implementation of a trie for any iterable values.
23+ */
24+ open class Trie <T , K >(
25+ private val keyExtractor : (T ) -> K
26+ ) : Iterable<List<T>> {
27+
28+ private val roots = HashMap <K , NodeImpl <T , K >>()
29+ private val implementations = HashMap <Node <T >, NodeImpl <T , K >>()
30+
31+ /* *
32+ * Adds value into a trie.
33+ *
34+ * If value already exists then do nothing except increasing internal counter of added values.
35+ * The counter can be returned by [Node.count].
36+ *
37+ * @return corresponding [Node] of the last element in the `values`
38+ */
39+ fun add (values : Iterable <T >): Node <T > {
40+ val root = try { values.first() } catch (e: NoSuchElementException ) { error(" Empty list are not allowed"
41+ var key = keyExtractor(root)
42+ var node = roots.computeIfAbsent(key) { NodeImpl (root, null ) }
43+ values.asSequence().drop(1 ).forEach { value ->
44+ key = keyExtractor(value)
45+ node = node.children.computeIfAbsent(key) { NodeImpl (value, node) }
46+ }
47+ node.count++
48+ implementations[node] = node
49+ return node
50+ }
51+
52+ /* *
53+ * Decreases node counter value or removes the value completely if `counter == 1`.
54+ *
55+ * Use [removeCompletely] to remove the value from the trie regardless of counter value.
56+ *
57+ * @return removed node if value exists.
58+ */
59+ fun remove (values : Iterable <T >): Node <T >? {
60+ val node = findImpl(values) ? : return null
61+ return when {
62+ node.count == 1 -> removeCompletely(values)
63+ node.count > 1 -> node.apply { count-- }
64+ else -> throw IllegalStateException (" count should be 1 or greater"
65+ }
66+ }
67+
68+ /* *
69+ * Removes value from a trie.
70+ *
71+ * The value is removed completely from the trie. Thus, the next code is true:
72+ *
73+ * ```
74+ * trie.remove(someValue)
75+ * trie.get(someValue) == null
76+ * ```
77+ *
78+ * Use [remove] to decrease counter value instead of removal.
79+ *
80+ * @return removed node if value exists
81+ */
82+ fun removeCompletely (values : Iterable <T >): Node <T >? {
83+ val node = findImpl(values) ? : return null
84+ if (node.count > 0 && node.children.isEmpty()) {
85+ var n: NodeImpl <T , K >? = node
86+ while (n != null ) {
87+ val key = keyExtractor(n.data)
88+ n = n.parent
89+ if (n == null ) {
90+ val removed = roots.remove(key)
91+ check(removed != null )
92+ } else {
93+ val removed = n.children.remove(key)
94+ check(removed != null )
95+ if (n.count != 0 ) {
96+ break
97+ }
98+ }
99+ }
100+ }
101+ return if (node.count > 0 ) {
102+ node.count = 0
103+ implementations.remove(node)
104+ node
105+ } else {
106+ null
107+ }
108+ }
109+
110+ operator fun get (values : Iterable <T >): Node <T >? {
111+ return findImpl(values)
112+ }
113+
114+ operator fun get (node : Node <T >): List <T >? {
115+ return implementations[node]?.let (this ::buildValue)
116+ }
117+
118+ private fun findImpl (values : Iterable <T >): NodeImpl <T , K >? {
119+ val root = try { values.first() } catch (e: NoSuchElementException ) { return null }
120+ var key = keyExtractor(root)
121+ var node = roots[key] ? : return null
122+ values.asSequence().drop(1 ).forEach { value ->
123+ key = keyExtractor(value)
124+ node = node.children[key] ? : return null
125+ }
126+ return node.takeIf { it.count > 0 }
127+ }
128+
129+ override fun iterator (): Iterator <List <T >> {
130+ return iterator {
131+ roots.values.forEach { node ->
132+ traverseImpl(node)
133+ }
134+ }
135+ }
136+
137+ private suspend fun SequenceScope<List<T>>.traverseImpl (node : NodeImpl <T , K >) {
138+ val stack = ArrayDeque <NodeImpl <T , K >>()
139+ stack.addLast(node)
140+ while (stack.isNotEmpty()) {
141+ val n = stack.removeLast()
142+ if (n.count > 0 ) {
143+ yield (buildValue(n))
144+ }
145+ n.children.values.forEach(stack::addLast)
146+ }
147+ }
148+
149+ private fun buildValue (node : NodeImpl <T , K >): List <T > {
150+ return generateSequence(node) { it.parent }.map { it.data }.toList().asReversed()
151+ }
152+
153+ interface Node <T > {
154+ val data: T
155+ val count: Int
156+ }
157+
158+ /* *
159+ * Trie node
160+ *
161+ * @param data data to be stored
162+ * @param parent reference to the previous element of the value
163+ * @param count number of value insertions
164+ * @param children list of children mapped by their key
165+ */
166+ private class NodeImpl <T , K >(
167+ override val data : T ,
168+ val parent : NodeImpl <T , K >? ,
169+ override var count : Int = 0 ,
170+ val children : MutableMap <K , NodeImpl <T , K >> = HashMap (),
171+ ) : Node<T>
172+ }
0 commit comments