PriorityOrderedSet.kt

import java.util.*
import kotlin.NoSuchElementException


class PriorityOrderedSet<T>
private constructor(private val defaultPriority: Priority,
                    private val itemToNodeMap: MutableMap<T, Node<T>>) :
        MutableSet<T> by itemToNodeMap.keys {

    constructor(defaultPriority: Priority = Priority.Normal) : this(defaultPriority, mutableMapOf<T, Node<T>>())

    val head: T? get() = _head?.value

    private val priorityToNodeMap = mutableMapOf<Priority, Node<T>?>()
    private val _head: Node<T>?
        @Synchronized
        get() {
            for (p in Priority.values())
                return priorityToNodeMap[p] ?: continue

            return null
        }
    private var _tail: Node<T>? = null

    @Suppress("NOTHING_TO_INLINE")
    inline fun getPriority(element: T): Priority =
            getPriorityOrNull(element) ?: throw NoSuchElementException()

    fun getPriorityOrNull(element: T): Priority? = itemToNodeMap[element]?.priority

    @Synchronized
    fun pop(): T {
        val res = _head?.value ?: throw EmptyStackException()
        remove(res)
        return res
    }

    @Synchronized
    fun addWithPriority(element: T, priority: Priority, toTheBottom: Boolean = priority != Priority.High): Boolean {
        val exist = itemToNodeMap.containsKey(element)
        if (exist) return false

        itemToNodeMap[element] =
                if (toTheBottom) addToTheBottom(element, priority)
                else addToTheTop(element, priority)

        return true
    }

    @Synchronized
    fun addAfter(afterElement: T, newElement: T): Boolean {
        if (itemToNodeMap.containsKey(newElement)) return false

        itemToNodeMap[newElement] = itemToNodeMap.getValue(afterElement).addAfter(newElement)

        return true
    }

    @Synchronized
    fun addBefore(beforeElement: T, newElement: T): Boolean {
        if (itemToNodeMap.containsKey(newElement)) return false

        itemToNodeMap[newElement] = itemToNodeMap.getValue(beforeElement).addBefore(newElement)

        return true
    }

    private fun addToTheTop(element: T, priority: Priority): Node<T> {
        fun getNextNodeRecursively(priority: Priority): Node<T>? {
            return priorityToNodeMap[priority]
                    ?: getNextNodeRecursively(priority.next ?: return null)
        }

        val nextNode = getNextNodeRecursively(priority)
        val prevNode = if (priority != Priority.High) nextNode?.prev ?: _tail else null
        val newNode = Node(element, priority, nextNode, prevNode)

        prevNode?.next = newNode
        nextNode?.prev = newNode

        priorityToNodeMap[priority] = newNode

        _tail = _tail ?: newNode

        return newNode
    }

    private fun addToTheBottom(element: T, priority: Priority): Node<T> {
        fun getNextNodeRecursively(priority: Priority): Node<T>? = run {
            val nextPriority = priority.next ?: return null
            priorityToNodeMap[nextPriority] ?: getNextNodeRecursively(nextPriority)
        }

        val nextNode = getNextNodeRecursively(priority)
        val prevNode = if (nextNode != null) nextNode.prev else _tail
        val newNode = Node(element, priority, nextNode, prevNode)

        prevNode?.next = newNode
        nextNode?.prev = newNode

        if (priorityToNodeMap[priority] == null)
            priorityToNodeMap[priority] = newNode

        if (nextNode == null)
            _tail = newNode

        return newNode
    }

    private fun Node<T>.addBefore(element: T): Node<T> {
        val nextNode = this
        val prevNode = nextNode.prev
        val newNode = Node(element, priority, nextNode, prevNode)

        prevNode?.next = newNode
        nextNode.prev = newNode

        if (priorityToNodeMap[priority] == nextNode)
            priorityToNodeMap[priority] = newNode

        return newNode
    }

    private fun Node<T>.addAfter(element: T): Node<T> {
        val prevNode = this
        val nextNode = prevNode.next
        val newNode = Node(element, priority, nextNode, prevNode)

        prevNode.next = newNode
        nextNode?.prev = newNode

        if (nextNode == null)
            _tail = newNode

        return newNode
    }

    override fun iterator(): MutableIterator<T> = Iterator()

    override fun add(element: T): Boolean =
            addWithPriority(element, defaultPriority)

    override fun addAll(elements: Collection<T>): Boolean =
            elements.fold(false) { acc, it -> acc || add(it) }

    @Synchronized
    override fun clear() {
        itemToNodeMap.clear()
        priorityToNodeMap.clear()
        _tail = null
    }

    @Synchronized
    override fun remove(element: T): Boolean {
        val node = itemToNodeMap.remove(element) ?: return false
        val nextNode = node.next
        val prevNode = node.prev
        val headOfPriority = priorityToNodeMap.entries.find { it.value == node }?.key

        prevNode?.next = nextNode
        nextNode?.prev = prevNode

        if (nextNode == null) _tail = prevNode

        if (headOfPriority != null) {
            priorityToNodeMap[headOfPriority] = nextNode?.takeIf { it !in priorityToNodeMap.values }
        }

        return true
    }

    override fun removeAll(elements: Collection<T>): Boolean =
            elements.fold(false) { acc, it -> acc || remove(it) }

    override fun retainAll(elements: Collection<T>): Boolean =
            removeAll(subtract(elements))

    private data class Node<T>(val value: T,
                               val priority: Priority,
                               var next: Node<T>?,
                               var prev: Node<T>?)

    private inner class Iterator : MutableIterator<T> {

        private var cursor = _head

        override fun hasNext(): Boolean = cursor != null

        override fun next(): T {
            val res = cursor?.value ?: throw NoSuchElementException()
            cursor = cursor?.next
            return res
        }

        override fun remove() {
            val toRemove = cursor?.value ?: throw NoSuchElementException()
            cursor = cursor?.next
            remove(toRemove)
        }
    }
}

enum class Priority {
    High, Normal, Low;

    val next: Priority? get() = values().getOrNull(ordinal + 1)
}