Flatten 2D vector

Vector2DwithState.kt

data class State(
    val outerIndex: Int,
    val innerIndex: Int
)

class Vector2D(private val vec: Array<IntArray>) {
    var outerIndex = 0
    var innerIndex = 0

    fun next(): Int {
        while (true) {
            val currInner = vec[outerIndex]
        	if (innerIndex < currInner.size) {
            	return currInner[innerIndex++]
        	} else {
            	outerIndex++
            	innerIndex = 0
        	}
        }

    }

    fun hasNext(): Boolean {
        var o = outerIndex
        var i = innerIndex
        while (o < vec.size) {
            val currInner = vec[o]
        	if (i < currInner.size) {
            	return true
        	} else {
            	o++
            	i = 0
        	}
        }
        return false
    }
    
    
    fun getState() = State(outerIndex, innerIndex)

    fun setState(state: State) {
        outerIndex = state.outerIndex
        innerIndex = state.innerIndex
    }
}

fun `test one single array surrounded by two empty arrays`() {
    val testVector = Vector2D(arrayOf(intArrayOf(),intArrayOf(1, 3), intArrayOf()))
    assert(testVector.hasNext())
    assert(testVector.next() == 1)
    assert(testVector.hasNext())
    assert(testVector.hasNext()) // test doing multiple hasNext()s will not break
    assert(testVector.next() == 3)
    assert(!testVector.hasNext())
}

fun `test one single array`() {
    val testVector = Vector2D(arrayOf(intArrayOf(1, 3)))
    assert(testVector.hasNext())
    assert(testVector.next() == 1)
    assert(testVector.hasNext())
    assert(testVector.hasNext())
    assert(testVector.next() == 3)
    assert(!testVector.hasNext())
}

fun `test two arrays`() {
    val testVector = Vector2D(arrayOf(intArrayOf(1),intArrayOf(3)))
    assert(testVector.hasNext())
    assert(testVector.next() == 1)
    assert(testVector.hasNext())
    assert(testVector.hasNext())
    assert(testVector.next() == 3)
    assert(!testVector.hasNext())
}

fun `test two arrays with greater lengths`() {
    val testVector = Vector2D(arrayOf(intArrayOf(1,3),intArrayOf(),intArrayOf(4)))
    assert(testVector.hasNext())
    assert(testVector.next() == 1)
    assert(testVector.hasNext())
    assert(testVector.hasNext())
    assert(testVector.next() == 3)
    assert(testVector.hasNext())
    assert(testVector.next() == 4)
    assert(!testVector.hasNext())
}

fun `test empty array`() {
    val testVector = Vector2D(arrayOf(intArrayOf()))
    assert(!testVector.hasNext())
}

fun `test empty arrays`() {
    val testVector = Vector2D(arrayOf(intArrayOf(),intArrayOf(),intArrayOf()))
    assert(!testVector.hasNext())
}

fun `test two arrays with greater lengths with state`() {
    val testVector = Vector2D(arrayOf(intArrayOf(1,3),intArrayOf(),intArrayOf(4)))
    assert(testVector.hasNext())
    val state = testVector.getState()
    assert(testVector.next() == 1)
    assert(testVector.hasNext())
    assert(testVector.hasNext())
    assert(testVector.next() == 3)
    assert(testVector.hasNext())
    assert(testVector.next() == 4)
    assert(!testVector.hasNext())
    testVector.setState(state)
    assert(testVector.next() == 1)
}

fun main() {
    `test one single array surrounded by two empty arrays`()
    `test one single array`()
    `test two arrays`()
    `test two arrays with greater lengths`()
    `test empty array`()
    `test empty arrays`()
    `test two arrays with greater lengths with state`()
}