news.shamcode.ru | Отделяем стек от рекурсии

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

Постановка проблемы

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

Наивная реализация

const ackermann1 = (m, n) => {
  if (m === 0) {
    return n + 1;
  }

  if (n === 0) {
    return ackermann1(m - 1, 1);
  }

  return ackermann1(m - 1, ackermann1(m, n - 1));
};

Довольно быстро мы упрёмся в ограничения стека. При m = 3 и n = 12, функция падает с ошибкой. Обычно, в таком случае рекурсию заменяют на цикл, и вместо стека вызовов используют обычный стек. В некоторых задачах (DFS, обходы деревьев) код не становится сильно сложнее. В общем же случае, приходится "рвать" функцию, сохранять состояние явно, после чего восстанавливать его, из-за чего программа на высокоуровневом языке становится похожа на программу на языке ассемблера (или на конечный автомат) и сильно страдает читаемость кода.

Цикл со стеком

const ackermann2 = (m, n) => {
  const stack = []; // Стековые кадры

  let state = "initial"; // Instruction Pointer
  let value = undefined; // Возвращаемое значение рекурсивного вызова

  const call = (params, returnState) => {
    // Добавляем кадр и переходим в начало функции
    stack.push({ params, returnState });
    state = "initial";
  };

  const ret = (val) => {
    // Удаляем кадр и переходим к месту возврата
    const frame = stack.pop();
    value = val;
    state = frame.returnState;
  };

  call([m, n]);

  while (stack.length !== 0) {
    const frame = stack[stack.length - 1];
    const {
      params: [m, n],
    } = frame;

    if (state === "initial") {
      if (m === 0) {
        ret(n + 1);
        continue;
      }

      if (n === 0) {
        call([m - 1, 1], "recursive call n=0");
        continue;
      }

      call([m, n - 1], "general recursive call 1");
    } else if (state === "recursive call n=0") {
      ret(value);
    } else if (state === "general recursive call 1") {
      call([m - 1, value], "general recursive call 2");
    } else if (state === "general recursive call 2") {
      ret(value);
    }
  }

  return value;
};

Было бы здорово совместить плюсы обоих подходов, и получить читаемость как у наивной версии, а стек как у версии с циклом

Решение

Во многих современных языках есть способ прервать исполнение функции сохранив при этом её состояние. Я говорю, конечно, об асинхронных функциях и генераторах. В старых версиях javascript, когда уже появились генераторы, но ещё не появился async-await, можно было написать функцию превращающую генератор в асинхронную функцию. Код далее будет использовать схожие идеи, но по отношению к рекурсивным функциям.

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

Реализация на генераторах

const ackermann3 = recursive(function* (m, n) {
  if (m === 0) {
    return n + 1;
  }

  if (n === 0) {
    return yield [m - 1, 1];
  }

  return yield [m - 1, yield [m, n - 1]];
});

Осталось написать функцию recursive связывающую вычисления.

Функция recursive

const recursive = (generator) => (...params) => {
  let generatorObj = generator(...params);

  let request = undefined; // Рекурсивный запрос от генератора
  let result = undefined; // Ответ на запрос

  const stack = [];

  while (true) {
    // Получаем IteratorResult 
    const message = generatorObj.next(result); 
    
    if (message.done) {
      // Если генератор завершил работу - передаём
      // значение выше по стеку или завершаемся, если 
      // выше никого нет  
      if (stack.length === 0) {
        return message.value;
      }

      result = message.value;
      generatorObj = stack.pop();
    } else {
      // В противном случае - сохраняем состояние в стек, 
      // создаём новый генератор и дальше работаем с ним

      result = undefined;
      request = message.value;

      stack.push(generatorObj);
      generatorObj = generator(...request);
    }
  }
};

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

Сравнение

Весь код будет запускаться на node v20.16.0 LTS. Первым делом узнаем сколько мы потеряли в скорости. Для этого сделаем сто замеров для каждой версии при параметрах m = 3, n = 10.

Код бенчмарка

const benchmark = (fn, params) => {
  try {
    const start = performance.now();
    fn(...params);
    const end = performance.now();
    return end - start;
  } catch (e) {
    return undefined;
  }
};

const params = [3, 10];
let fns = [
  ["ackermann1", ackermann1],
  ["ackermann2", ackermann2],
  ["ackermann3", ackermann3],
];

for (let [name, fn] of fns) {
  const samples = 100;
  let data = [];

  for (let i = 0; i < samples; i++) {
    data.push(benchmark(fn, params));
  }

  const total = data.reduce((acc, time) => acc + time, 0);
  console.log(`${name} mean: ${total/samples}ms`);
}

Результаты:

ackermann1 mean: 178.13901175999723ms
ackermann2 mean: 916.1059493000008ms
ackermann3 mean: 2552.887184250003ms

Получаем, что наша версия в 2.7 раз медленнее версии с явным стеком и в 14.3 раза медленнее наивной реализации.

Далее узнаем насколько больше стал стек вызовов. Для этого будем использовать другие функции:

const count1 = (n) => {
  if (n === 0) {
    return 0;
  }

  return 1 + count1(n - 1);
};

const count2 = recursive(function* (n) {
  if (n === 0) {
    return 0;
  }

  return 1 + (yield [n - 1]);
});

На моей машине count1 запускается на значениях до 10468 и начинает ломаться на 10469. count2 запускается на значениях до 29_229_797 и ломается на 29_229_798. Таким образом на стандартных настройках стек стал больше в 2792 раза.

Ну вот и всё. Спасибо за внимание