Оптимизируя sequences — или как мой код попал в kotlin25.10.2023 08:15

Введение

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

Я представил его в своем докладе на весеннем Мобиус 2023 «Измеряя sequences».

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

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

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

Issue в kotlin с моими оптимизациями

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

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

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

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

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

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

Мне показалось это странным и я полез под капот изучать код

Реализация distinct

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

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

    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  Iterable.distinctBy(
    selector: (T) -> K
): List {
    val set = HashSet()
    val list = ArrayList()
    for (e in this) {
        val key = selector(e)
        if (set.add(key))
            list.add(e)
    }
    return list
}

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

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

  // ....
}

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

Давайте посмотрим как устроен это абстрактный класс AbstractIterator

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

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

public abstract class AbstractIterator : Iterator {
    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(
    private val sequence: Sequence,
    private val sendWhen: Boolean = true,
    private val predicate: (T) -> Boolean
) : Sequence {

    override fun iterator(): Iterator = object : Iterator {
        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 мы получаем при каждом сравнении лишнюю загрузку массива значений в локальный стек.

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

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

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

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

R8 Optimization: Enum Switch Maps, 2019

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

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

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

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

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

Полный код оптимизированной фунции distinct

private class OptimizedDistinctIterator(
    private val source: Iterator, 
    private val keySelector: (T) -> K
) : Iterator{
    private val observed = HashSet()
    // { 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  Sequence.plus(elements: Iterable): Sequence {
    return sequenceOf(this, elements.asSequence()).flatten()
}

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

Реализация flatten

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

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

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

internal class FlatteningSequence
constructor(
    private val sequence: Sequence,
    private val transformer: (T) -> R,
    private val iterator: (R) -> Iterator
) : Sequence {
  
    override fun iterator(): Iterator = object : Iterator {
        val iterator = sequence.iterator()
        var itemIterator: Iterator? = 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 = object : Iterator {
        val iterator = sequence.iterator()
        var itemIterator: Iterator? = 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 {
    override fun hasNext(): Boolean = false
    override fun next(): Nothing = throw NoSuchElementException()
}

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

    override fun iterator(): Iterator = object : Iterator {
        val iterator = sequence.iterator()
        var itemIterator: Iterator = 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

Полный код оптимизированной функции flatten

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

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

    override fun iterator(): Iterator = object : Iterator {
        private val iterator = sequence.iterator()
        private var itemIterator: Iterator = 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 есть открытый youtrack, в котором кто угодно может создать свою Issue и описать проблему или предложение.

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

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

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

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

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

© Habrahabr.ru