LongBitTree.kt

typealias LArray1 = Array<LongArray?>
typealias LArray2 = Array<LArray1?>
typealias LArray3 = Array<LArray2?>
typealias LArray4 = Array<LArray3?>
typealias LArray5 = Array<LArray4?>
typealias LArray6 = Array<LArray5?>

@Suppress("ConvertTwoComparisonsToRangeCheck")
class LongBitTree : Iterable<Long> {
    @Suppress("NOTHING_TO_INLINE")
    companion object {
        private const val BITS_PER_ENTRY = 10
        private const val BITS_PER_ENTRY_2 = 20
        private const val BITS_PER_FIRST_ENTRY = 16

        private const val ADDRESS_BITS_PER_WORD = 6
        private const val BITS_PER_WORD = 1 shl ADDRESS_BITS_PER_WORD
        private const val WORD_MASK = -1L

        private const val ELEMENTS_PER_ENTRY = 1 shl BITS_PER_ENTRY
        private const val WORDS_PER_ENTRY = ((ELEMENTS_PER_ENTRY - 1) shr ADDRESS_BITS_PER_WORD) + 1
        private const val MASK = (ELEMENTS_PER_ENTRY - 1).toLong()

        private const val ELEMENTS_PER_FIRST_ENTRY = 1 shl BITS_PER_FIRST_ENTRY
        private const val WORDS_PER_FIRST_ENTRY = ((ELEMENTS_PER_FIRST_ENTRY - 1) shr ADDRESS_BITS_PER_WORD) + 1
        private const val MASK_FIRST = (ELEMENTS_PER_FIRST_ENTRY - 1).toLong()

        fun LongArray.nextSetBit(index: Int): Int {
            if (index < 0) {
                throw IndexOutOfBoundsException("negative index " + index)
            }

            var u = intWordIndex(index)
            if (u >= size) {
                return -1
            }

            var word = this[u] and (WORD_MASK shl index)

            while (true) {
                if (word != 0L) {
                    return u * BITS_PER_WORD + java.lang.Long.numberOfTrailingZeros(word)
                }
                if (++u == size) {
                    return -1
                }
                word = this[u]
            }
        }

        private inline fun getEntryIndex(value: Long): Long {
            return value shr BITS_PER_ENTRY
        }

        private inline fun intWordIndex(bitIndex: Int): Int {
            return bitIndex shr ADDRESS_BITS_PER_WORD
        }

        private inline fun LongArray.hasBit(index: Int): Boolean {
            val wordIndex = intWordIndex(index)
            return (this[wordIndex] and (1L shl index)) != 0L
        }

        private inline fun LongArray.clearBit(key: Int) {
            val wordIndex = intWordIndex(key)
            this[wordIndex] = this[wordIndex] and ((1L shl key).inv())
        }

        private inline fun LongArray.setBit(key: Int) {
            val wordIndex = intWordIndex(key)
            this[wordIndex] = this[wordIndex] or (1L shl key)
        }
    }

    var size = 0

    val gen0 = LongArray(WORDS_PER_FIRST_ENTRY)
    val gen1: LArray1 = arrayOfNulls(ELEMENTS_PER_ENTRY)
    val gen2: LArray2 = arrayOfNulls(ELEMENTS_PER_ENTRY)
    //val gen3: LArray3 = arrayOfNulls(ELEMENTS_PER_ENTRY)
    //val gen4: LArray4 = arrayOfNulls(ELEMENTS_PER_ENTRY)
    //val gen5: LArray5 = arrayOfNulls(ELEMENTS_PER_ENTRY)

    val mask1 = LongArray(WORDS_PER_ENTRY)
    val mask2_1 = LongArray(WORDS_PER_ENTRY)
    val mask2_2: LArray1 = arrayOfNulls(ELEMENTS_PER_ENTRY)

    //val mask3 = arrayOfNulls<LArray1?>(ELEMENTS_PER_ENTRY)
    //val mask4 = arrayOfNulls<LArray2?>(ELEMENTS_PER_ENTRY)

    fun count(): Int = size

    fun add(key: Long) {
        if (0 <= key && key < ELEMENTS_PER_FIRST_ENTRY) {
            gen0.setBit(key.toInt())
        } else {
            val parent: LArray1
            val mask: LongArray
            var index0 = getEntryIndex(key)
            if (index0 >= ELEMENTS_PER_ENTRY) {
                val index1 = getEntryIndex(index0)
                if (index1 >= ELEMENTS_PER_ENTRY) {
                    throw UnsupportedOperationException()
                } else {
                    val index = index1.toInt()
                    parent = gen2[index] ?: arrayOfNulls<LongArray?>(ELEMENTS_PER_ENTRY).apply {
                        gen2[index] = this
                    }
                    mask = mask2_2[index] ?: LongArray(WORDS_PER_ENTRY).apply {
                        mask2_2[index] = this
                    }
                    mask2_1.setBit(index)
                }
                index0 = index0 and MASK
            } else {
                parent = gen1
                mask = mask1
            }
            val index = index0.toInt()
            val container = parent[index] ?: LongArray(WORDS_PER_ENTRY).apply {
                parent[index] = this
                mask.setBit(index)
            }
            container.setBit((key and MASK).toInt())
        }
        size++
    }

    fun remove(key: Long) {
        if (0 <= key && key < ELEMENTS_PER_FIRST_ENTRY) {
            gen0.clearBit(key.toInt())
        } else {
            val parent: Array<LongArray?>
            var index0 = getEntryIndex(key)
            if (index0 >= ELEMENTS_PER_ENTRY) {
                val index1 = getEntryIndex(index0)
                if (index1 >= ELEMENTS_PER_ENTRY) {
                    throw UnsupportedOperationException()
                } else {
                    val index = index1.toInt()
                    parent = gen2[index] ?: return
                }
                index0 = index0 and MASK
            } else {
                parent = gen1
            }

            val container = parent[index0.toInt()] ?: return
            container.clearBit((key and MASK).toInt())
        }
        size--
    }

    fun contains(key: Long): Boolean {
        if (0 <= key && key < ELEMENTS_PER_FIRST_ENTRY) {
            return gen0.hasBit(key.toInt())
        } else {
            val index0 = getEntryIndex(key)
            if (index0 >= ELEMENTS_PER_ENTRY) {
                val index1 = getEntryIndex(index0)
                if (index1 >= ELEMENTS_PER_ENTRY) {
                    throw UnsupportedOperationException()
                } else {
                    val index = index1.toInt()
                    if (!mask2_1.hasBit(index)) return false

                    val deeperIndex = (index0 and MASK).toInt()
                    if (!mask2_2[index]!!.hasBit(deeperIndex)) return false
                    return gen2[index]!![deeperIndex]!!.hasBit((key and MASK).toInt())
                }
            } else {
                val index = index0.toInt()
                if (!mask1.hasBit(index)) return false
                return gen1[index]!!.hasBit((key and MASK).toInt())
            }
        }
    }

    override fun iterator(): Iterator<Long> {
        return object : Iterator<Long> {
            var base = 0L
            var entry = gen0
            var next: Long = entry.nextSetBit(0).toLong().let {
                if (it < 0) {
                    fetchNextEntry(MASK_FIRST)
                } else {
                    it
                }
            }

            override fun hasNext(): Boolean {
                return next != -1L
            }

            override fun next(): Long {
                val result = next + base
                val index = next.toInt()
                val nextBit = entry.nextSetBit(index + 1).toLong()
                next = if (nextBit < 0) {
                    fetchNextEntry(result)
                } else {
                    nextBit
                }
                return result
            }

            fun fetchNextEntry(key: Long): Long {
                val index0 = getEntryIndex(key)
                if (index0 < ELEMENTS_PER_ENTRY) {
                    val nextIndex = mask1.nextSetBit(index0.toInt() + 1)
                    if (nextIndex > 0) {
                        base = (nextIndex shl BITS_PER_ENTRY).toLong()
                        entry = gen1[nextIndex]!!
                        return entry.nextSetBit(0).toLong()
                    }
                    val index1 = mask2_1.nextSetBit(0)
                    if (index1 < 0) {
                        return -1
                    } else {
                        val entryIndex = mask2_2[index1]!!.nextSetBit(0)
                        base = (index1.toLong() shl BITS_PER_ENTRY_2) + (entryIndex.toLong() shl BITS_PER_ENTRY)
                        entry = gen2[index1]!![entryIndex]!!
                        return entry.nextSetBit(0).toLong()
                    }
                } else {
                    val index1 = getEntryIndex(index0)
                    if (index1 >= ELEMENTS_PER_ENTRY) {
                        throw UnsupportedOperationException()
                    } else {
                        val startingIndex = index1.toInt()
                        var index = mask2_1.nextSetBit(startingIndex) // greater or equal
                        if (index >= 0) {
                            val startingIndex0 = (index0 and MASK).toInt()
                            var deeperIndex = startingIndex0
                            val keyIndex = key and MASK
                            val nextKeyIndex = if (keyIndex == MASK) {
                                deeperIndex++
                                0
                            } else {
                                keyIndex.toInt() + 1
                            }
                            var entryIndex = mask2_2[index]!!.nextSetBit(deeperIndex)
                            var nextEntry = gen2[index]!![entryIndex]!!
                            var nextBit = nextEntry.nextSetBit(nextKeyIndex).toLong()
                            if (nextBit < 0) {
                                if (entryIndex == startingIndex0) {
                                    entryIndex = mask2_2[index]!!.nextSetBit(deeperIndex + 1)
                                    if (entryIndex >= 0) {
                                        nextEntry = gen2[index]!![entryIndex]!!
                                        nextBit = nextEntry.nextSetBit(0).toLong()
                                    }
                                }
                                if (nextBit < 0 && index == startingIndex) {
                                    index = mask2_1.nextSetBit(startingIndex + 1)
                                    if (index < 0) {
                                        return -1
                                    }
                                    entryIndex = mask2_2[index]!!.nextSetBit(0)
                                    nextEntry = gen2[index]!![entryIndex]!!
                                    nextBit = nextEntry.nextSetBit(0).toLong()
                                }
                            }
                            base = (index.toLong() shl BITS_PER_ENTRY_2) + (entryIndex.toLong() shl BITS_PER_ENTRY)
                            entry = nextEntry
                            return nextBit
                        } else {
                            return -1
                        }
                    }
                }
            }
        }
    }
}