Оптимизируя sequences — или как мой код попал в kotlin

Введение

В начале этого года я провел небольшое исследование на тему измерения производительности sequences. Точнее я хотел бы, чтобы оно было небольшим, но в результате оно заняло 3-4 месяца моей активной работы.

В рамках этого исследования я измерял скорость преобразования коллекций через sequences в сравнении с обычными методами преобразования коллекций. Во время исследования, я обратил внимание на то, что некоторые функции sequences работают медленнее, чем могли бы.

Я предложил способы оптимизации этих функций и JetBrains приняли мои предложения и уже включили их в kotlin 2.0 (Issue в kotlin с моими оптимизациями)

Я подробно рассказываю о принципах работы каждой функции sequences в своей статье Измеряя sequences.

Здесь же, я хочу рассказать именно об оптимизациях и показать, как небольшие изменения в коде могут ускорять функции на 15-20%. Показать, насколько важно знать нюансы генерации kotlin в byte-code и как это влияет на скорость работы функций.

Оптимизация distinct

Я обратил внимание, что реализация функции distinct алгоритмически очень похожа на функцию filter. При этом функция distinct имеет существенный проигрыш по сравнению с коллекциями -15%, а функция filter выигрывает у коллекций примерно 3%-5% .

Результаты измерений filter

Результаты измерения distinct

Мне показалось это странным и я решил изучить код обеих функций, чтобы понять, за счет чего возникают такие потери в функции distinct.

Реализация distinct

Декоратор для distinct устроен довольно просто. Внутри он создает HashSet observed и накапливает в нем все возвращаемые элементы. Таким образом, если элемент уже есть в observed, то это значит, что он уже возвращался и дальше он будет пропускаться.

private class DistinctIterator<T, K>(
    private val source: Iterator<T>, 
    private val keySelector: (T) -> K
) : AbstractIterator<T>() {
    private val observed = HashSet<K>()

    override fun computeNext() {
        while (source.hasNext()) {
            val next = source.next()
            val key = keySelector(next)

            if (observed.add(key)) {
                setNext(next)
                return
            }
        }

        done()
    }
}

Если мы посмотрим реализацию кода коллекций, то мы увидим в коллекциях практически идентичный код. По сути, он делает тоже самое. Создает HashSet и накапливает в нем возвращенные элементы.

public inline fun <T, K> Iterable<T>.distinctBy(
    selector: (T) -> K
): List<T> {
    val set = HashSet<K>()
    val list = ArrayList<T>()
    for (e in this) {
        val key = selector(e)
        if (set.add(key))
            list.add(e)
    }
    return list
}

Таким образом понятно, что если реализация sequences и коллекций в целом совпадает, то потери где то в другом месте. Если мы посмотрим на объявление класса в sequences, то мы увидим какой то абстрактный класс AbstractIterator.

private class DistinctIterator<T, K>(
    private val source: Iterator<T>, 
    private val keySelector: (T) -> K
) : AbstractIterator<T>() {

  // ....
}

Именно в нем находится реализация всех базовых методов distinct декоратора. И если следовать логике, то и все наши потери должны находится в этом абстрактном классе.

Реализация класса AbstractIterator

По сути этот класс хранит в переменной state текущее состояние, вычислялся ли следующий элемент итератора или нет. Если он еще не вычислялся, то он вызывает функцию tryToComputeNext(), вычисляющую следующий элемент.

private enum class State {
    Ready,
    NotReady,
    Done,
    Failed
}

public abstract class AbstractIterator<T> : Iterator<T> {
    private var state = State.NotReady
    private var nextValue: T? = null

    override fun hasNext(): Boolean {
        require(state != State.Failed)
        return when (state) {
            State.Done -> false
            State.Ready -> true
            else -> tryToComputeNext()
        }
    }

    override fun next(): T {
        if (!hasNext()) throw NoSuchElementException()
        state = State.NotReady
        @Suppress("UNCHECKED_CAST")
        return nextValue as T
    }

    private fun tryToComputeNext(): Boolean {
        state = State.Failed
        computeNext()
        return state == State.Ready
    }

    /**
     * Computes the next item in the iterator.
     *
     * This callback method should call one of these two methods:
     *
     * * [setNext] with the next value of the iteration
     * * [done] to indicate there are no more elements
     *
     * Failure to call either method will result in the iteration terminating with a failed state
     */
    abstract protected fun computeNext(): Unit

    /**
     * Sets the next value in the iteration, called from the [computeNext] function
     */
    protected fun setNext(value: T): Unit {
        nextValue = value
        state = State.Ready
    }

    /**
     * Sets the state to done so that the iteration terminates.
     */
    protected fun done() {
        state = State.Done
    }
}

Кажется, что код написан очень грамотно и на первый взгляд проблем здесь быть не должно. Но где же они тогда?

Декоратор sequences для функции fliter реализован очень похожим образом, однако, он не дает такого проигрыша по сравнению с коллекциями. Я начал сравнивать код двух декораторов и искать существенные отличия.

Реализации функции filter

Как видите, принцип реализации фильтра очень похож. Тут тоже есть состояние вычисления следующего элемента и оно также хранится в поле nextState

Также есть метод вычисления следующего элемента calcNext(), который является полной аналогией метода tryToComputeNext() в функции distinct

internal class FilteringSequence<T>(
    private val sequence: Sequence<T>,
    private val sendWhen: Boolean = true,
    private val predicate: (T) -> Boolean
) : Sequence<T> {

    override fun iterator(): Iterator<T> = object : Iterator<T> {
        val iterator = sequence.iterator()
        // -1 for unknown, 0 for done, 1 for continue
        var nextState: Int = -1 
        var nextItem: T? = null

        private fun calcNext() {
            while (iterator.hasNext()) {
                val item = iterator.next()
                if (predicate(item) == sendWhen) {
                    nextItem = item
                    nextState = 1
                    return
                }
            }
            nextState = 0
        }

        override fun next(): T {
            if (nextState == -1)
                calcNext()
            if (nextState == 0)
                throw NoSuchElementException()
            val result = nextItem
            nextItem = null
            nextState = -1
            @Suppress("UNCHECKED_CAST")
            return result as T
        }

        override fun hasNext(): Boolean {
            if (nextState == -1)
                calcNext()
            return nextState == 1
        }
    }
}

Кажется, что за исключением мелких деталей реализации, эти две функции очень похожи по принципам своей работы.

Так где же потери функции distinct?

Сначала, в силу своей некомпетенции, я подозревал абстрактный класс и вызовы виртуальных методов. Но нет, если убрать абстрактный класс и перенести все в один общий класс, то функция не начинает работать быстрее.

Единственное, чем еще существенно отличаются эти две функции - это Enum.

В случае filter мы храним состояние в виде обычного числа Int

        // -1 for unknown, 0 for done, 1 for continue
        var nextState: Int = -1 

А в случае distinct мы используем для этого типизированный Enum

private enum class State {
    Ready,
    NotReady,
    Done,
    Failed
}

    private var state = State.NotReady

Но не может же использование Enum класса давать 15% потери производительности.

Но других вариантов из анализа кода я больше не видел, поэтому решил спуститься на уровень ниже и перейти к анализу byte-code обеих функций.

Если мы посмотрим byte-code проверки переменной nextState для функции filter, то не увидим там ничего лишнего. Просто загрузка переменных в стек, проверка их значения и передача управления на вызов соответствующего блока кода.

Но если посмотреть byte-code проверки переменной state в функции distinct, то мы увидим нечто интересное.

Это код проверки стейта в функции distinct

        return when (state) {
            State.Done -> false
            State.Ready -> true
            else -> tryToComputeNext()
        }

А это byte-code, который ей соответсвует

#1  private final hasNext()Z
#2   L0
#3    LINENUMBER 20 L0
#4    ALOAD 0
#5    GETFIELD AbstractIterator.state : Lcom/AbstractIterator$State;
#6    GETSTATIC AbstractIterator$WhenMappings.$EnumSwitchMapping$0 : [I
#7    SWAP
#8    INVOKEVIRTUAL AbstractIterator$State.ordinal ()I
#9    IALOAD
#10   L1
#11    TABLESWITCH
#12      1: L2
#13      2: L3
#14      default: L4
#15   L2

Не думаю что многие умеют читать byte-code (по крайней мере я не умел), поэтому подробно прокомментирую что здесь происходит буквально построчно.

#5 - загружает в локальный стек указатель на переменную state
#6 - загружает в локальный стек массив ordinal значений enum класса
#7 - меняет две этих переменные местами в локальном стеке
#8 - получает ordinal значение переменной state
#9 - IALOAD ищет в массиве ordinal значений enum класса индекс ordinal значения переменной state
#11 - найденный индекс передается на вход оператору TABLESWITCH и дальше происходит переход по нужной ветке кода.

В итоге получается, что при использовании оператора when в сочетании с enum, мы получаем при каждом сравнении лишнюю загрузку массива всех значений enum в локальный стек.

При этом, если мы перепишем сравнение через ordinal значения, то можем избежать этой ненужной и затратной операции

        return when (state.ordinal) {
            State.Done.ordinal -> false
            State.Ready.ordinal -> true
            else -> tryToComputeNext()
        }

По сути такая проверка значения enum через поиск в массиве ordinal значений enum класса может быть актуальна только в одном случае. Когда код проверки значения может быть скомпилирован отдельно от кода объявления enum класса. Кажется что в случае с android это невозможно.

Чуть позже, после того, как я докопался до сути проблемы, я нашел старую статью Jake Wharton, в которой он описывает похожую проблему

Jake Wharton - R8 Optimization: Enum Switch Maps, 2019

В этой статье он обещал исправить это еще в 2019 году, но видимо не успел.

Результаты оптимизации distinct

Я попробовал переписать функцию distinct с учетом нюансов обработки enum. Я убрал использование enum и использую для хранения состояния обычные Int константы.

В результате это ускорило работу функции примерно на 10%-15%, что на мой взгляд довольно существенно

Результаты замеров оригинальной и оптимизированной функции distinct

private class OptimizedDistinctIterator<T, K>(
    private val source: Iterator<T>, 
    private val keySelector: (T) -> K
) : Iterator<T>{
    private val observed = HashSet<K>()
    // { UNDEFINED_STATE, HAS_NEXT_ITEM, HAS_FINISHED }
    private var nextState: Int = UNDEFINED_STATE
    private var nextItem: T? = null

    override fun hasNext(): Boolean {
        if (nextState == UNDEFINED_STATE)
            calcNext()
        return nextState == HAS_NEXT_ITEM
    }

    override fun next(): T {
        if (nextState == UNDEFINED_STATE)
            calcNext()
        if (nextState == HAS_FINISHED)
            throw NoSuchElementException()
        nextState = UNDEFINED_STATE
        return nextItem as T
    }

    private fun calcNext() {
        while (source.hasNext()) {
            val next = source.next()
            val key = keySelector(next)

            if (observed.add(key)) {
                nextItem = next
                nextState = HAS_NEXT_ITEM // found next item
                return
            }
        }
        nextState = HAS_FINISHED // end of iterator
    }
}

Оптимизация flatten

Когда я проводил свое исследование, то обнаружил, что функция flatten проигрывает коллекциям почти в 2 раза. Эта функция позволяет развернуть список списков в один большой, линейный список. Это довольно часто используется в различных преобразованиях данных.

Результаты измерений flatten

Проигрыш в два раза, это слишком много.

По сути это ставит крест на использовании sequences для преобразования коллекций, если в нашем преобразовании встречается хотя бы одно преобразование через flatten

Что еще обидней, flatten является базовым декоратором в sequences и на его основе сделаны также и другие функции. Например функция plus

public operator fun <T> Sequence<T>.plus(elements: Iterable<T>): Sequence<T> {
    return sequenceOf(this, elements.asSequence()).flatten()
}

Мне захотелось это исправить и вдохновленный успехом с оптимизацией distinct, я решил попробовать оптимизировать и flatten

Реализация flatten

Декоратор для flatten довольно сложен для понимания. В целом принцип его работы следующий:

На каждый вызов метода hasNext(), вызывается функция ensureItemIterator(), которая вычисляет итератор вложенного списка для текущего элемента и записывает его значение в поле itemIterator.

При вызове метода next() он просто вызывает itemIterator.next() и таким образом последовательно перебираются все элементы и список списков разворачивается в линейный список.

internal class FlatteningSequence<T, R, E>
constructor(
    private val sequence: Sequence<T>,
    private val transformer: (T) -> R,
    private val iterator: (R) -> Iterator<E>
) : Sequence<E> {
  
    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
        }
    }
}

На первый взгляд, класс написан сложно, но если присмотреться, то понимаешь, что он работает достаточно оптимально и зацепится особо не за что.

Но потом я обратил внимание на то, что поле itemIterator объявлено как nullable тип

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

А это значит, что при обращении к этому полю, kotlin каждый раз будет добавлять проверку на то, что значение поля не Null. А это дополнительная операция чтения.

Я пометил в коде звездочками все такие места, где мы имеем лишнюю проверку на Null значение.

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

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

Я устранил этот недочет через создание синглтона EmptyIterator, который реализует логику, которая должна быть у ItemIterator в случае Null значения.

// Empty iterator for cause when we haven't next element
private object EmptyIterator: Iterator<Nothing> {
    override fun hasNext(): Boolean = false
    override fun next(): Nothing = throw NoSuchElementException()
}

Это позволило мне убрать все проверки на Null при обращении к полю itemIterator

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

        override fun next(): E {
            if (!ensureItemIterator())
                throw NoSuchElementException()
          return itemIterator.next() // было itemIterator!!.next()
        }

        private fun ensureItemIterator(): Boolean {
          if (itemIterator.hasNext() == false) // было itemIterator?.hasNext()
                itemIterator = null
      

Затем я обратил внимание на то, что при каждом вызове next() и hasNext() мы всегда безусловно вызываем достаточно сложную функцию ensureItemIterator()

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

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

Так как с итераторами мы всегда работаем парой методов hasNext() + next(), то кажется расточительным вызывать два раза функцию ensureItemIterator() для обработки каждого элемента. Ее результат вполне можно закешировать в поле, хранящем состояние ее вычисления.

Я ввел поле state и запоминаю его при каждом вычислении ensureItemIterator(), а сбрасываю при вызове метода next(), когда мы переходим к следующему элементу. Таким образом функция ensureItemIterator() вызывается теперь только один раз на пару вызовов hasNext() + next()

Все это позволило мне значительно оптимизировать время выполнения функции flatten.

Результаты оптимизации flatten

Отказ от nullable переменной и введение состояния вычисления для ensureItemIterator существенно ускорило работу функции. Выигрыш 35%-37%, что на мой взгляд довольно существенно.

Результаты замеров оригинальной и оптимизированной функции flatten

// Empty iterator for cause when we haven't next element
private object EmptyIterator: Iterator<Nothing> {
    override fun hasNext(): Boolean = false
    override fun next(): Nothing = throw NoSuchElementException()
}

internal class FlatteningSequence<T, R, E>
constructor(
    private val sequence: Sequence<T>,
    private val transformer: (T) -> R,
    private val iterator: (R) -> Iterator<E>
) : Sequence<E> {

    override fun iterator(): Iterator<E> = object : Iterator<E> {
        private val iterator = sequence.iterator()
        private var itemIterator: Iterator<E> = EmptyIterator 
        
        // { UNDEFINED_STATE, HAS_NEXT_ITEM, HAS_FINISHED }
        private var state: Int = UNDEFINED_STATE 

        override fun next(): E {
            if (state == UNDEFINED_STATE) { 
                ensureItemIterator()
            }
            state = UNDEFINED_STATE
            return itemIterator.next()
        }

        override fun hasNext(): Boolean {
            return when (state) { 
                HAS_NEXT_ITEM -> true
                HAS_FINISHED -> false
                else -> ensureItemIterator()
            }
        }

        private fun ensureItemIterator(): Boolean {
            if (itemIterator.hasNext()) {
                state = HAS_NEXT_ITEM
                return true
            } else {
                while (iterator.hasNext()) {
                    val nextItemIterator = iterator(transformer(iterator.next()))
                    if (nextItemIterator.hasNext()) {
                        itemIterator = nextItemIterator
                        state = HAS_NEXT_ITEM
                        return true
                    }
                }
                state = HAS_FINISHED
                itemIterator = EmptyIterator
                return false
            }
        }
    }
}

Отправка оптимизаций в JetBrains

Сначала я боялся, что передача моих предложений по оптимизации sequences будет сложным и мучительным процессом. Все таки JetBrains - это глобальная компания и мне было сложно представить, что они могут заинтересоваться моими предложениями. Но все оказалось намного проще, чем я думал.

В JetBrains есть открытый youtrack, в котором кто угодно может создать свою Issue и описать проблему или предложение.

Довольно долго моя issue с предложениями по оптимизации висела без движения. Ее почти сразу разметили нужными тегами и отнесли к библиотекам kotlin. А затем, в течение нескольких месяцев, по ней не было никаких изменений.

Но спустя 3 или 4 месяца, до нее все таки добрался разработчик и буквально через неделю она уже была влита в основную ветку kotlin 2.0.

И теперь частичка меня, мои идеи, есть в kotlin. И от осознания этого меня "прет" до сих пор.

Если ты обнаружил проблему или видишь какое то хорошее решение, то не надо боятся создавать Issue в глобальные продукты.

Эти продукты делают такие же люди как ты и вполне возможно, что именно тебе пришла в голову гениальная идея, которая поможет сделать продукт лучше.

Если вам понравилась моя статья, проголосуйте за мою статью на medium "Measuring sequences". Это поможет продвинуть ее на международной площадке.