Sequences: Implementation of the flatten

How flatten decorator works under hood

This is a quite complex decorator, so let’s break it down step by step.

The decorator takes the sequence parameter as input, from which it obtains the original iterator, iterator. This iterator represents the main list. The itemIterator variable holds a reference to the iterator of the nested list for the current element of the main list.

    override fun iterator(): Iterator<E> = object : Iterator<E> {
        val iterator = sequence.iterator()
        var itemIterator: Iterator<E>? = null

When hasNext() is called, the decorator invokes the internal function ensureItemIterator(), which computes the itemIterator variable and assigns it a reference to the iterator of the nested list for the current element. It also returns whether there is a next element in the nested list.

        override fun hasNext(): Boolean {
            return ensureItemIterator()
        }

        override fun next(): E {
            if (!ensureItemIterator())
                throw NoSuchElementException()
            return itemIterator!!.next()
        }

Most of the main work is done in the ensureItemIterator() function. It computes the iterator for the current element of the main list.

        private fun ensureItemIterator(): Boolean {
            if (itemIterator?.hasNext() == false)
                itemIterator = null

            while (itemIterator == null) {
                if (!iterator.hasNext()) {
                    return false
                } else {
                    val element = iterator.next()
                    val nextItemIterator = iterator(transformer(element))
                    if (nextItemIterator.hasNext()) {
                        itemIterator = nextItemIterator
                        return true
                    }
                }
            }
            return true
        }

Source code of the flatten decorator

    override fun iterator(): Iterator<E> = object : Iterator<E> {
        val iterator = sequence.iterator()
        var itemIterator: Iterator<E>? = null

        override fun next(): E {
            if (!ensureItemIterator())
                throw NoSuchElementException()
            return itemIterator!!.next()
        }

        override fun hasNext(): Boolean {
            return ensureItemIterator()
        }

        private fun ensureItemIterator(): Boolean {
            if (itemIterator?.hasNext() == false)
                itemIterator = null

            while (itemIterator == null) {
                if (!iterator.hasNext()) {
                    return false
                } else {
                    val element = iterator.next()
                    val nextItemIterator = iterator(transformer(element))
                    if (nextItemIterator.hasNext()) {
                        itemIterator = nextItemIterator
                        return true
                    }
                }
            }
            return true
        }
    }

You can learn more from my sequences performance study “Measuring sequences”