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

Уже давно определённые вещи из мира функционального программирования всё сильнее проникают в нефункциональные языки. Может показаться странным, что в Java смогли интегрировать лямбда-выражения, а вот оптимизацию хвостовой рекурсии (преобразование рекурсии в эквивалентный цикл) до сих пор не сделали, ведь она гораздо проще. Почему её нет?
Попробуем разобраться с причинами и посмотрим, что можно сделать своими руками.

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

static int add(int x, int y) {
    if (y == 0) return x;
    return add(x ^ y, (x & y) << 1);
}

Это рекурсивный способ сложения 2-х целых чисел. Он подходит под определение хвостовой рекурсии: за каждым рекурсивным вызовом непосредственно следует операция return. Оптимизация заключается в том, чтобы при рекурсивном вызове не создавать новый кадр стэка, а переиспользовать текущий. Для этого нужно новые параметры положить на место старых и выполнить безусловный переход на первую инструкцию метода. В псевдокоде это будет выглядеть так:

static int add(int x, int y) {
    start: {
        if (y == 0) return x;
        (x, y) = (x ^ y, (x & y) << 1);
        goto start;
    }
}

Далее возможны различные стилистические вариации, но в целом код после ручной оптимизации станет таким (можно и красивее, многословность тут лишь для ясности):

static int add(int x, int y) {
    while (true) {
        if (y == 0) return x;
        int _x = x ^ y;
        int _y = (x & y) << 1;
        x = _x;
        y = _y;
    }
}

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

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

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

• полиморфизм:
  для того, чтобы осуществить GOTO, нужно точно знать, что по факту должен вызываться тот же метод. Это требование выполнено для:
  — статических методов;
  — приватных методов;
  — методов final классов;
  — явных вызовов invokespecial. Данный вариант не может быть создан компилятором из исходного Java кода, так как invokespecial доступен только для вызова методов суперкласса.
• исключения:
  если вдруг рекурсивный вызов происходит внутри try блока, то мы не можем просто так взять и передать управление инструкции вне этого самого блока. Вот пример:

  static int f(boolean shouldThrow) {
      if (shouldThrow) throw new RuntimeException();
      try {
          f(!shouldThrow);
      } catch (Exception e) {
      }
      return -1;
  }
  

  Вызов f(false) должен возвратить -1, но если сделать GOTO в месте рекурсивного вызова, то мы получим RuntimeException, а это явно отличается от того, что должно происходить при корректной оптимизации.

Есть минимум 2 проверенных способа модифицировать байткод Java классов:

• препроцессор — встраивается в компилятор, изменения байткода попадут в class файлы;
• инструментатор — встраивается в ClassLoader, изменения будут видны только в рантайме.

Я выбрал второй вариант и написал простенький Java Agent, оптимизирующий хвостовую рекурсию. Он сможет провести оптимизацию только при выше описанных условиях:

• static method / final method / final class;
• рекурсивные вызовы находятся вне try блоков.

Для нетерпеливых ссылка на github: github.com/ibessonov/java-tailrec-agent.
Там настроенный IDEA проект, в котором есть примеры, чтобы с ними поиграться. А для терпеливых — небольшое пояснение на примере.
Проверка модификаторов доступа в отдельных пояснениях не нуждается в силу очевидности её реализации. Поэтому опустим её и перейдём к рассмотрению типичного метода и проследим, что с ним происходит:

static int gcd(int n, int m) {
    try {
        if (m == 0) return n;
    } catch (Throwable t) {
        // do nothing
    }
    return gcd(m, n % m);
}

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

static gcd(II)I
    TRYCATCHBLOCK TryBlockStart TryBlockEnd CatchBlockStart java/lang/Throwable
TryBlockStart:
    ILOAD 1
    IFNE ElseBlock
    ILOAD 0
TryBlockEnd:
    IRETURN
CatchBlockStart:
    ASTORE 2 // сохранение исключения во временную переменную - дефолтный пустой обработчик
ElseBlock:
    ILOAD 1
    ILOAD 0
    ILOAD 1
    IREM
    INVOKESTATIC Main.gcd(II)I
    IRETURN

Каждый метод содержит в себе массив описаний try блоков. У каждого описания есть 4 составляющих: метка начала блока, метка конца блока, метка обработчика исключения и дескриптор класса исключения. Эта информация позволяет нам однозначно определить инструкции, находящиеся внутри try блоков, и не производить для них оптимизацию.
Далее нужно найти все инструкции INVOKE* с дескриптором метода, совпадающим с самим методом (в данном случае ищется дескриптор gcd(II)I метода класса Main), за которыми непосредственно находится инструкция типа RETURN.
Найденный INVOKESTATIC надо преобразовать из вызова в безусловный переход. Известно, что в момент вызова на стэке находятся все параметры. Для статических методов всё просто, нужно сохранить эти параметры обратно в локальные переменные и сделать безусловный переход в самое начало:

static gcd(II)I
    TRYCATCHBLOCK TryBlockStart TryBlockEnd CatchBlockStart java/lang/Throwable
StartLabel:
TryBlockStart:
    ILOAD 1
    IFNE ElseBlock
    ILOAD 0
TryBlockEnd:
    IRETURN
CatchBlockStart:
    ASTORE 2
ElseBlock:
    ILOAD 1
    ILOAD 0
    ILOAD 1
    IREM

    ISTORE 1
    ISTORE 0
    GOTO StartLabel

Для нестатических методов встаёт одна интересная проблема: метод вызывается на объекте, который, вообще говоря, не обязан совпадать с this. Технически возможно найти в байткоде место вычисления этого объекта и проводить оптимизацию только тогда, когда это вычисление равно ALOAD 0.
Я же поступил лениво и уже вычисленное значение нужного класса сохранил вместо текущего this (сделав ASTORE 0). Как ни странно, такой подход работает, но я, в силу недостаточности знаний, не могу гарантировать, что JIT в этой ситуации будет вести себя корректно. Хотел бы узнать ответ в комментариях — есть ли тут какие-нибудь риски.

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

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

Автор: ibessonov

Источник