Separating the stack from recursion

In this article I will tell you how you can use generators to modify recursion so that it uses a heap instead of a stack and at the same time is almost no different from regular recursion.

Statement of the problem

Let's say we have a recursive function that is difficult to express with a regular loop. For example, I'll take the Ackermann function

Naive implementation

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));
};

We'll quickly run into stack limitations. m = 3 And n = 12the function crashes with an error. Usually, in such a case, recursion is replaced with a cycle, and a regular stack is used instead of a call stack. In some tasks (DFS, tree traversals), the code does not become much more complicated. In the general case, you have to “tear” the function, save the state explicitly, and then restore it, which is why a program in a high-level language becomes similar to a program in assembly language (or a finite state machine) and the readability of the code suffers greatly.

Cycle with stack

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;
};

It would be great to combine the advantages of both approaches and get the readability of the naive version, and the stack of the version with a loop.

Solution

Many modern languages have a way to interrupt the execution of a function while preserving its state. I am talking, of course, about asynchronous functions and generators. In older versions of javascript, when generators had already appeared, but had not yet appeared async–awaityou could write function which turns the generator into an asynchronous function. The code below will use similar ideas, but with respect to recursive functions.

For this we will use generators, and if necessary to make a recursive call, we will use yield. The “recursive scheduler” will save the state, make a recursive call, and then restore the function, passing it the result of the calculation.

Implementation on generators

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]];
});

All that remains is to write the function recursive linking calculations.

recursive function

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);
    }
  }
};

Now you can write recursive code without worrying about call stack limitations.

Comparison

All code will run on node v20.16.0 LTS. First of all, we will find out how much we have lost in speed. To do this, we will make a hundred measurements for each version with the parameters m = 3, n = 10.

Benchmark code

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`);
}

Results:

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

We get that our version is 2.7 times slower than the version with an explicit stack and 14.3 times slower than the naive implementation.

Next, we'll find out how much larger the call stack has become. For this, we'll use other functions:

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]);
});

On my car count1 starts at values up to 10468 and starts to break down 10469. count2 starts at values up to 29_229_797 and breaks down 29_229_798. Thus, on standard settings the stack became larger in 2792 times.

Well, that's all. Thanks for your attention.