Блог им. dmitrys99

Оптимизация хвостовой рекурсии в трансляторе

О чем это вообще?

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

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

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

Что сложного то?

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

Перейдем к ~~кошкам~~ примеру.

Есть функция for():

//  ML-подобный язык
for (init, predicate, fn) {
    if (predicate(init))
        for (fn(init), predicate, fn)
    else init
}

Транслятор переведет ее в Java или Javascript один в один, никакой магии (приведен код на Javascript):


// Javascript
function for__(init, predicate, fn) {
        if (predicate(init)) {
            return for__(fn(init), predicate, fn);
        } else {
            return init;
        }
    }
}

Она даже синтаксически будет выглядеть практически также.

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

Поэтому мы получим рекурсивный вызов. Ну и следом во время выполнения получим исчерпание стека вызовов.

Что делать? Спасите, помогите!

Чтобы исправить ситуацию и как-то заставить это работать, разработана техника, позволяющая экономить место в стеке вызовов или избавиться от рекурсивных вызовов вообще.

Мы модифицируем функции таким образом, чтобы изменить рекурсивный вызов на цикл. Это возможно, но, к сожалению, заменить можно не все рекурсивные вызовы. Так, чтобы вызов функции можно было заменить на итерацию цикла, нужно соблюсти несколько условий:

вызов функции должен быть последним в последовательности операций;
это должен быть именно рекурсивное обращение к функции, без каких-либо дополнительный операций (например, в вычислении факториала, куда уж без него!, выражение n*fac(n-1) нельзя автоматически заменить на циклический вызов);
вызов функции не может быть вложен в замыкание, т.е. в случае с for() конструкция for(..., \x->for()..., ...) не будет заменена на циклическую.

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

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

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

Я покажу два решения, которые успешно применяются в трансляторах.

Решение первое

Функция предварительно готовится с помощью вспомогательной операции:

// Javascript
function TCO(fn, fn_name) {
    var top_args;
    window[fn_name] = function() {
        var result, old_top_args = top_args;
        top_args = arguments;
        while (top_args !== null) {
            var cur_args = top_args;
            top_args = null;
            result = fn.apply(null, cur_args);
        }
        top_args = old_top_args;
        return result;
    };
    window['tc_' + fn_name] = function() {
        top_args = arguments;
    };
}

TCO((init, predicate, fn) => {
    if (predicate(init)) {
        return tc_for__(fn(init), predicate, fn);
    } else {
        return init;
    }
}, 'for__');

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

Метод работает безупречно, правда медленно. Виной тому использование метода не самого быстрого метода fn.apply(). В Java используется рефлексия, которая тоже дает падение производительности.

Этих недостатков лишен второй метод.

Второй метод

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

// Javascript
function for__(init, predicate, fn) {
    T: while (true) {
        if (predicate(init)) {
            let $a_ = fn(init);
            let $b_ = predicate;
            let $c_ = fn;
            init = $a_;
            predicate = $b_;
            fn = $c_;
            continue T
        } else {
            return init;
        }
    }
}

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

Очевидную оптимизацию с неизменяемыми аргументами (см. $b_ = predicate;) оставлю пытливому читателю для самостоятельного изучения.

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

Неожиданная проблема

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

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

// ML-подобный язык
main() {
    recursiveFn(1,  \res -> println("Внешний вызов. Результат = " + int2str(res)));
}

recursiveFn(counter : int, onOK : (int) -> void) -> void {
	println("Шапка recursiveFn. Counter = " + int2str(counter));

	if (counter > 1) {
		println("Первая ветка recursiveFn. Counter = " + int2str(counter));
		onOK(counter)
	} else {
		println("Вторая ветка recursiveFn. Counter = " + int2str(counter));
		recursiveFn(counter + 1, \newCount -> {
			println("Внутренний вызов. newCounter = " + int2str(newCount));
			onOK(newCount);
		});
	}
}

Ключевой момент — вызов функции onOK после println("Внутренний вызов. newCounter = " + int2str(newCount)).

Если мы, что называется, в лоб, преобразуем функцию в Javascript, результатом будет бесконечная рекурсия.

// Javascript
function main() {
    return recursiveFn(1, ((res) => {
        return println(("Внешний вызов. Результат = " + int2str(res)));
    }));
}

function recursiveFn(counter, onOK) {
    T: while (true) {
        println(("Шапка recursiveFn. Counter = " + int2str(counter)));
        if ((counter > 1)) {
            println(("Первая ветка recursiveFn. Counter = " + int2str(counter)));
            return onOK(counter);
        } else {
            println(("Вторая ветка recursiveFn. Counter = " + int2str(counter)));

            let $a_ = ((counter + 1) | 0);
            let $b_ = ((newCount) => {
                println(("Внутренний вызов. newCounter = " + int2str(newCount)));
                return onOK(newCount);
            });
            counter = $a_;
            onOK = $b_;
            continue T
        }
    }
}

Неважно что будет стоять после строки println(("Внутренний вызов. newCounter = " + int2str(newCount))), onOK или $b_. Мы получим бесконечную рекурсию из-за того, что замыкание использует внешнее значение функционального параметра onOK.

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

// Javascript
function main() {
    return recursiveFn(1, ((res) => {
        return println(("Внешний вызов. Результат = " + int2str(res)));
    }));
}

function recursiveFn(counter, onOK) {
    T: while (true) {
        println(("Шапка recursiveFn. Counter = " + int2str(counter)));
        if ((counter > 1)) {
            println(("Первая ветка recursiveFn. Counter = " + int2str(counter)));
            return onOK(counter);
        } else {
            println(("Вторая ветка recursiveFn. Counter = " + int2str(counter)));

            let $_ta = onOK;
            let $a_ = ((counter + 1) | 0);
            let $b_ = ((newCount) => {
                println(("Внутренний вызов. newCounter = " + int2str(newCount)));
                return $_ta(newCount);
            });
            counter = $a_;
            onOK = $b_;
            continue T
        }
    }
}

Внимание следует обратить на переменную $_ta, которая уже внутри замыкания используется вместо onOK.

При трансляции мы модифицируем лексическое окружение для генерации кода замыкания. После генерации тела замыкания лексическое окружение восстанавливается, чтобы не поломать нижележащий код. Это делается тем же способом, что и генерация кода для рабочих серверов (в неотладочном режиме).

Этот метод трансляции работает корректно и уже используется в боевых системах.

И что, вот так вот и работает?

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

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