Оптимизируя 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». Это поможет продвинуть ее на международной площадке.