Kotlin под капотом: как избавится от рекурсии

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

Tailrec и хвостовая рекурсия

Хвостовая рекурсия - это частный случай рекурсии, при котором рекурсивный вызов является последней операцией перед возвратом из функции. Если упрощенно, то функция должна вернуть либо конечный результат, либо вернуть результат своего рекурсивного вызова.

Данный пример вычисления последовательности Фибоначчи очень хорошо демонстрирует принцип хвостовой рекурсии 

tailrec fun fibonacci(n: Int, a: BigInteger, b: BigInteger): BigInteger {
  return if (n == 0) b else fibonacci(n - 1, a + b, a)
}

Хвостовая рекурсия очень хорошо оптимизируется через итерацию в плоском цикле. Давайте посмотрим декомпилированный код, чтобы понять как tailrec оптимизирует нашу рекурсию

@NotNull
public final BigInteger fibonacci(
    int n, @NotNull BigInteger a, @NotNull BigInteger b
) {
    while(true) {
        if (n == 0) {
            return b;
        }

        n = n - 1;
        BigInteger var10001 = a.add(b);
        b = a;
        a = var10001;
    }
}

Как видите, все превратилось в обычный цикл и в коде не осталось никакого следа нашей рекурсии. Но это будет работать строго для случаев простой хвостовой рекурсии.

Рекурсия в деревьях

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

А вот ее простейшее решение через рекурсию.

fun ViewGroup.findViewRecursion(predicate: (View) -> Boolean): List<View> {
    val accumulator = mutableListOf<View>()
    if (predicate(this)) {
        accumulator.add(this)
    }
    children.forEach { child ->
        when {
            child is ViewGroup -> accumulator.addAll(
                child.findViewRecursion(predicate)
            )
            predicate(child) -> accumulator.add(child)
        }
    }
    return accumulator
}

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

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

fun ViewGroup.findViewRecursionOpt(predicate: (View) -> Boolean): List<View> {
    val accumulator = mutableListOf<View>()
    this.internalFindView(predicate, accumulator)
    return accumulator
 }

private fun ViewGroup.internalFindView(
    predicate: (View) -> Boolean,
    accumulator: MutableList<View> = mutableListOf()
) {
    if (predicate(this)) {
        accumulator.add(this)
    }
    children.forEach { child ->
        when {
            child is ViewGroup -> {
                child.internalFindView(predicate, accumulator)
            }
            predicate(child) -> accumulator.add(child)
        }
    }
}

Первое, что я попробовал сделать, это добавить ключевое слово tailrec для этой рекурсивной функции. Я был уверен, что это не сработает и я получу ошибку компиляции, но kotlin выдал мне всего лишь warning "A function is marked as tail-recursive but no tail calls are found". Лично я даже не сразу его заметил и многие разработчики могут посчитать, что удачно избавились от проблем рекурсии. 

tailrec fun ViewGroup.findViewRecursion(predicate: (View) -> Boolean): List<View> {
    val accumulator = mutableListOf<View>()
    if (predicate(this)) {
        accumulator.add(this)
    }
    children.forEach { child ->
        when {
            child is ViewGroup -> accumulator.addAll(
                child.findViewRecursion(predicate)
            )
            predicate(child) -> accumulator.add(child)
        }
    }
    return accumulator
}

Добавление tailrec к функции без хвостовой рекурсии скомпилируется, но не даст вам никакой оптимизации. Максимум что вы получите - это warning компилятора.

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

public final List findViewRecursion(ViewGroup $this, Function1 predicate) {
    List accumulator = new ArrayList<View>();
    if ((Boolean)predicate.invoke($this)) {
        accumulator.add($this);
    }

    Iterator iterator = ViewGroupKt.getChildren($this).iterator();

    while(iterator.hasNext()) {
        View child = (View)iterator.next();
        if (child instanceof ViewGroup) {
            accumulator.addAll(findViewRecursion(child, predicate));
        } else if ((Boolean)predicate.invoke(child)) {
            accumulator.add(child);
        }
    }

    return accumulator;
}

Честно говоря, я был обескуражен и пожалуй создам Issue в JetBrains, потому что на мой взгляд это состояние ошибки, а не warning. Но возможно у них будет какое то разумное объяснение.

Но как же нам избавится от рекурсии в этом достаточно типовом способе обхода дерева?

Стандартный способ избавления от рекурсии через очередь 

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

Лучше всего этот принцип оптимизации демонстрирует сам код.

fun ViewGroup.findViewQueue(predicate: (View) -> Boolean): List<View> {
    val accumulator = mutableListOf<View>()
    val queue: Queue<View> = LinkedList()

    // add self as a first element of queue
    queue.add(this) 
    while (queue.isNotEmpty()) {
        // get and remove next item from queue
        val view = queue.poll() 

        if (predicate(view)) {
            accumulator.add(view)
        }

        // add to queue all childs for current view
        if (view is ViewGroup) { 
            view.children.forEach { queue.add(it) }
        }
    }

    return accumulator
}

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

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

Способ избавления от рекурсии через итератор 

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

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

/**
 * Iterator for iterate by abstract tree
 * @param getChildIterator function which return child iterator for current item
 * if current item has child return child iterator for item
 * else return null
 *
 * Example for View hierarchy
 * TreeIterator<View>(root) { (it as? ViewGroup)?.children?.iterator() }
 *
 */
class TreeIterator<T>(
    root: T,
    private val getChildIterator: (T) -> Iterator<T>?
) : Iterator<T> {
    private val stack = mutableListOf<Iterator<T>>()

    private var item: T = root
    private var iterator: Iterator<T> = getChildIterator(root) ?: EmptyIterator()

    override fun hasNext(): Boolean {
        return iterator.hasNext()
    }

    override fun next(): T {
        item = iterator.next()
        prepareNextIterator()
        return item
    }

   /**
     * calculate next iterator
     * if current item has child then get child iterator and save current iterator to stack
     * else if current iterator hasn't more elements then restore parent iterator from stack
     */
    private fun prepareNextIterator() {
        val childIterator = getChildIterator(item)
        if (childIterator != null && childIterator.hasNext()) {
            stack.add(iterator)
            iterator = childIterator
        } else {
            while (!iterator.hasNext() && stack.isNotEmpty()) {
                iterator = stack.last()
                stack.removeLast()
            }
        }
    }
}

Реализация нашего поиска View с помощью данного итератора становится тривиальной задачей. И такой итератор легко встраивать в любые ленивые цепочки обработки данных через sequence.

fun ViewGroup.findViewTreeIterator(predicate: (View) -> Boolean): Sequence<View> {
    val treeIterator = TreeIterator<View>(this) { view -> 
        (view as? ViewGroup)?.children?.iterator() 
    }
    return treeIterator.asSequence()
        .filter { predicate(it) }
}

Способ избавления от рекурсии через sequence.yield 

Но существует и другой интересный способ избавления от рекурсии, о котором хочется здесь рассказать. В sequence существует механизм ленивых операторов yield, о которых я раньше не знал. Собственно, я обнаружил их, когда писал эту статью и изучал код стандартных расширений androidx.core.view для ViewGroup. 

public val ViewGroup.descendants: Sequence<View>
    get() = sequence {
        forEach { child ->
            yield(child)
            if (child is ViewGroup) {
                yieldAll(child.descendants)
            }
        }
    }

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

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

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

На самом деле я еще не до конца разобрался с тем, как вызов yield работает под капотом. Он использует принцип, очень похожий на принцип прерывания suspend функций в корутинах. Возможно чуть позже я напишу об этом большую статью

И наконец, сама наша функция поиска View, которая в случае с yeld тоже стала тривиальной.

fun ViewGroup.findViewYield(predicate: (View) -> Boolean): Sequence<View> {
    return sequenceOf(this)
        .plus(this.descendants)
        .filter { predicate(it) }
}

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

Мне стало интересно, какой из способов отказа от рекурсии даст наибольший выигрыш в производительности.

Я создал иерархию View с разной глубиной вложенности и написал тесты, которые производят поиск в этой иерархии всеми рассмотренными способами. Для замеров я использовал библиотеку kotlinx.benchmark. 

Признаться честно, результаты меня удивили.

Как видите, обычная рекурсия с копированием списков дала ошибку уже на 3000 уровней вложенности. Оптимизированная рекурсия с аккумулятором дала ошибку на глубине в 5000 уровней.

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

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

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

Возможно, я займусь изучением и оптимизацией yield в sequence и когда нибудь напишу об этом отдельную большую статью. 

Выводы

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

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

И еще один важный момент. Любые оптимизации необходимо замерять, так как очень часто бывает, что логичный и оптимизированный код (по мнению разработчика) начинает работать медленнее простого и не оптимизированного. То что кажется нам логичным и оптимальным не всегда логично и оптимально для компилятора.

Пожалуй, я создам issue в google, чтобы заменить реализацию обхода иерархии View в функции ViewGroup.descendants на мое решение с ленивым итератором. Это стандартная функция и ее используют многие разработчики, не догадываясь, что производительность обработки иерархии через эту функцию падает в сотни раз.