2020-05-03 14:36:19 -04:00

369 lines
11 KiB
Kotlin

package exh.util
import android.util.SparseArray
import java.util.AbstractMap
import java.util.LinkedList
class NakedTrieNode<T>(val key: Int, var parent: NakedTrieNode<T>?) {
val children = SparseArray<NakedTrieNode<T>>(1)
var hasData: Boolean = false
var data: T? = null
// Walks in ascending order
// Consumer should return true to continue walking, false to stop walking
inline fun walk(prefix: String, consumer: (String, T) -> Boolean, leavesOnly: Boolean) {
// Special case root
if (hasData && (!leavesOnly || children.size() <= 0)) {
if (!consumer(prefix, data!! as T)) return
}
val stack = LinkedList<Pair<String, NakedTrieNode<T>>>()
SparseArrayValueCollection(children, true).forEach {
stack += prefix + it.key.toChar() to it
}
while (!stack.isEmpty()) {
val (key, bottom) = stack.removeLast()
SparseArrayValueCollection(bottom.children, true).forEach {
stack += key + it.key.toChar() to it
}
if (bottom.hasData && (!leavesOnly || bottom.children.size() <= 0)) {
if (!consumer(key, bottom.data!! as T)) return
}
}
}
fun getAsNode(key: String): NakedTrieNode<T>? {
var current = this
for (c in key) {
current = current.children.get(c.toInt()) ?: return null
if (!current.hasData) return null
}
return current
}
}
/**
* Fast, memory efficient and flexible trie implementation with implementation details exposed
*/
class NakedTrie<T> : MutableMap<String, T> {
/**
* Returns the number of key/value pairs in the map.
*/
override var size: Int = 0
private set
/**
* Returns `true` if the map is empty (contains no elements), `false` otherwise.
*/
override fun isEmpty() = size <= 0
/**
* Removes all elements from this map.
*/
override fun clear() {
root.children.clear()
root.hasData = false
root.data = null
size = 0
}
val root = NakedTrieNode<T>(-1, null)
private var version: Long = 0
override fun put(key: String, value: T): T? {
// Traverse to node location in tree, making parent nodes if required
var current = root
for (c in key) {
val castedC = c.toInt()
var node = current.children.get(castedC)
if (node == null) {
node = NakedTrieNode(castedC, current)
current.children.put(castedC, node)
}
current = node
}
// Add data to node or replace existing data
val previous = if (current.hasData) {
current.data
} else {
current.hasData = true
size++
null
}
current.data = value
version++
return previous
}
override fun get(key: String): T? {
val current = getAsNode(key) ?: return null
return if (current.hasData) current.data else null
}
fun getAsNode(key: String): NakedTrieNode<T>? {
return root.getAsNode(key)
}
override fun containsKey(key: String): Boolean {
var current = root
for (c in key) {
current = current.children.get(c.toInt()) ?: return false
if (!current.hasData) return false
}
return current.hasData
}
/**
* Removes the specified key and its corresponding value from this map.
*
* @return the previous value associated with the key, or `null` if the key was not present in the map.
*/
override fun remove(key: String): T? {
// Traverse node tree while keeping track of the nodes we have visited
val nodeStack = LinkedList<NakedTrieNode<T>>()
for (c in key) {
val bottomOfStack = nodeStack.last
val current = bottomOfStack.children.get(c.toInt()) ?: return null
if (!current.hasData) return null
nodeStack.add(bottomOfStack)
}
// Mark node as having no data
val bottomOfStack = nodeStack.last
bottomOfStack.hasData = false
val oldData = bottomOfStack.data
bottomOfStack.data = null // Clear data field for GC
// Remove nodes that we visited that are useless
for (curBottom in nodeStack.descendingIterator()) {
val parent = curBottom.parent ?: break
if (!curBottom.hasData && curBottom.children.size() <= 0) {
// No data or child nodes, this node is useless, discard
parent.children.remove(curBottom.key)
} else break
}
version++
size--
return oldData
}
/**
* Updates this map with key/value pairs from the specified map [from].
*/
override fun putAll(from: Map<out String, T>) {
// No way to optimize this so yeah...
from.forEach { (s, u) ->
put(s, u)
}
}
// Walks in ascending order
// Consumer should return true to continue walking, false to stop walking
inline fun walk(consumer: (String, T) -> Boolean) {
walk(consumer, false)
}
// Walks in ascending order
// Consumer should return true to continue walking, false to stop walking
inline fun walk(consumer: (String, T) -> Boolean, leavesOnly: Boolean) {
root.walk("", consumer, leavesOnly)
}
fun getOrPut(key: String, producer: () -> T): T {
// Traverse to node location in tree, making parent nodes if required
var current = root
for (c in key) {
val castedC = c.toInt()
var node = current.children.get(castedC)
if (node == null) {
node = NakedTrieNode(castedC, current)
current.children.put(castedC, node)
}
current = node
}
// Add data to node or replace existing data
if (!current.hasData) {
current.hasData = true
current.data = producer()
size++
version++
}
return current.data!!
}
// Includes root
fun subMap(prefix: String, leavesOnly: Boolean = false): Map<String, T> {
val node = getAsNode(prefix) ?: return emptyMap()
return object : Map<String, T> {
/**
* Returns a read-only [Set] of all key/value pairs in this map.
*/
override val entries: Set<Map.Entry<String, T>>
get() {
val out = mutableSetOf<Map.Entry<String, T>>()
node.walk(
"",
{ k, v ->
out.add(AbstractMap.SimpleImmutableEntry(k, v))
true
},
leavesOnly
)
return out
}
/**
* Returns a read-only [Set] of all keys in this map.
*/
override val keys: Set<String>
get() {
val out = mutableSetOf<String>()
node.walk(
"",
{ k, _ ->
out.add(k)
true
},
leavesOnly
)
return out
}
/**
* Returns the number of key/value pairs in the map.
*/
override val size: Int get() {
var s = 0
node.walk("", { _, _ -> s++; true }, leavesOnly)
return s
}
/**
* Returns a read-only [Collection] of all values in this map. Note that this collection may contain duplicate values.
*/
override val values: Collection<T>
get() {
val out = mutableSetOf<T>()
node.walk(
"",
{ _, v ->
out.add(v)
true
},
leavesOnly
)
return out
}
/**
* Returns `true` if the map contains the specified [key].
*/
override fun containsKey(key: String): Boolean {
if (!key.startsWith(prefix)) return false
val childNode = node.getAsNode(key.removePrefix(prefix)) ?: return false
return childNode.hasData && (!leavesOnly || childNode.children.size() <= 0)
}
/**
* Returns `true` if the map maps one or more keys to the specified [value].
*/
override fun containsValue(value: T): Boolean {
node.walk(
"",
{ _, v ->
if (v == value) return true
true
},
leavesOnly
)
return false
}
/**
* Returns the value corresponding to the given [key], or `null` if such a key is not present in the map.
*/
override fun get(key: String): T? {
if (!key.startsWith(prefix)) return null
val childNode = node.getAsNode(key.removePrefix(prefix)) ?: return null
if (!childNode.hasData || (leavesOnly && childNode.children.size() > 0)) return null
return childNode.data
}
/**
* Returns `true` if the map is empty (contains no elements), `false` otherwise.
*/
override fun isEmpty(): Boolean {
if (node.children.size() <= 0 && !root.hasData) return true
if (!leavesOnly) return false
node.walk("", { _, _ -> return false }, leavesOnly)
return true
}
}
}
// Slow methods below
/**
* Returns `true` if the map maps one or more keys to the specified [value].
*/
override fun containsValue(value: T): Boolean {
walk { _, t ->
if (t == value) {
return true
}
true
}
return false
}
/**
* Returns a [MutableSet] of all key/value pairs in this map.
*/
override val entries: MutableSet<MutableMap.MutableEntry<String, T>>
get() = FakeMutableSet.fromSet(
mutableSetOf<MutableMap.MutableEntry<String, T>>().apply {
walk { k, v ->
this += FakeMutableEntry.fromPair(k, v)
true
}
}
)
/**
* Returns a [MutableSet] of all keys in this map.
*/
override val keys: MutableSet<String>
get() = FakeMutableSet.fromSet(
mutableSetOf<String>().apply {
walk { k, _ ->
this += k
true
}
}
)
/**
* Returns a [MutableCollection] of all values in this map. Note that this collection may contain duplicate values.
*/
override val values: MutableCollection<T>
get() = FakeMutableCollection.fromCollection(
mutableListOf<T>().apply {
walk { _, v ->
this += v
true
}
}
)
}