All Articles

Lazy Evaluation in JavaScript With Generators

When I was learning Haskell last year, I was fascinated by how Haskell handles calculations that are awkward to implement in other languages, say, JavaScript. Take a look at this:

sum (takeWhile (<10000) (filter odd (map (^2) [1..])))

You just can’t do it in an expressive way like this in JavaScript. So I wondered if I could somehow hack the JavaScript language to support lazy evaluation. I admit that I was not innovative enough to get this done by myself, so I googled. To my delight, someone had already done that! I found this article Lazy Evaluation in JavaScript with Generators, Map, Filter, and Reduce.

This article only handles infinite iterables. That’s not enough, so what I want to do in this post is to extend the original implementation to support any iterables.

First, I’ll adapt the original code to do the calculation I mentioned in the beginning.

const Lazy = iterator => {
  const next = iterator.next.bind(iterator)

  /* map overrides the original `next` function,
   * so that each `next` call will first override
   * the inner value of the iteratee
   * with the result of specified calculation.
   */
  const map = f => {
    const modifiedNext = () => {
      const item = next()
      const mappedValue = f(item.value)
      return {
        value: mappedValue,
        done: item.done,
      }
    }
    const newIter = { ...iterator, next: modifiedNext }
    return Lazy(newIter)
  }

  /* filter overrides the original `next` function,
   * so that only the `next` call that satisfies
   * specified predicate yields a iteratee.
   */
  const filter = predicate => {
    const modifiedNext = () => {
      while (true) {
        const item = next()
        if (predicate(item.value)) {
          return item
        }
      }
    }
    const newIter = { ...iterator, next: modifiedNext }
    return Lazy(newIter)
  }

  /* takeWhile terminates the calculation and
   * pulls out values until specified predicated is not met.
   */
  const takeWhile = predicate => {
    const result = []
    let value = next().value
    while (predicate(value)) {
      result.push(value)
      value = next().value
    }
    return result
  }

  /* I use Object.freeze to ensure data safety.
   * To learn more, check out this article:
   * https://medium.freecodecamp.org/elegant-patterns-in-modern-javascript-ice-factory-4161859a0eee
   */
  return Object.freeze({
    map,
    filter,
    takeWhile,
    next,
  })
}

const numbers = function*() {
  let i = 1
  while (true) {
    yield i++
  }
}

/* In my opinion, the following code is much
 * readable than its equivalent in Haskell
 */
Lazy(numbers())
  .map(x => x ** 2)
  .filter(x => x % 2 === 1)
  .takeWhile(x => x < 10000)
  .reduce((x, y) => x + y)
// => 16650

Mission completed. However, I feel that we can do more. The Lazy function only handles infinite iterables, what if we want to handle any iterables? Can we do that? Turns out, with the hard work done, we can easily extend the Lazy function to achieve our new goal.

Let’s first limit the kind of types the Lazy function accepts. To make it easy, I’ll only handle generator functions and iterables. We need a way to check if a function is a generator. That’s not an easy task to do by yourself, so I suggest looking for help. I found one that works.

const Lazy = source => {
  if (
    typeof source[Symbol.iterator] !== 'function' &&
    !isGeneratorFunction(source)
  )
    throw new Error(
      'The source input must be an iterable or a generator function'
    )

  const iterator = isGeneratorFunction(source)
    ? source()
    : source[Symbol.iterator]()

  const _lazy = it => {
    const next = it.next.bind(it)

    const map = f => {
      const modifiedNext = () => {
        const item = next()
        const mappedValue = f(item.value)
        return {
          value: mappedValue,
          done: item.done,
        }
      }
      const newIter = { ...it, next: modifiedNext }
      return _lazy(newIter)
    }

    const filter = predicate => {
      const modifiedNext = () => {
        let item = next()
        /* Notice the difference from the previous
         * implementation,
         * We can't use an infinite loop here
         */
        while (!item.done) {
          if (predicate(item.value)) {
            return item
          }
          item = next()
        }
        /* Also notice that once the loop is done,
         * we need to explicitly notify the iterator
         */
        return { value: null, done: true }
      }
      const newIter = { ...it, next: modifiedNext }
      return _lazy(newIter)
    }

    const takeWhile = predicate => {
      const result = []
      let value = next().value
      while (predicate(value)) {
        result.push(value)
        value = next().value
      }
      return result
    }

    const take = n => {
      const values = []
      let item = next()
      for (let i = 0; i < n; i++) {
        /* If the length of the iterable is
         * smaller than n, we need to break the loop early
         */
        if (item.done) break
        values.push(item.value)
        item = next()
      }

      return values
    }

    /* `zipWith` zips an iterable with another
     * iterable that is also wrapped in the `Lazy` function
     * with provided function fn
     */
    const zipWith = (fn, other) => {
      const modifiedNext = () => {
        const first = next()
        const second = other.next()
        /* If any one of the two iterables finishes,
         * we notify the iterator explicitly to terminate
         */
        if (first.done || second.done) {
          return { value: null, done: true }
        }
        return {
          value: fn(first.value, second.value),
          done: false,
        }
      }
      const newIter = { ...it, next: modifiedNext }
      return _lazy(newIter)
    }

    return Object.freeze({
      map,
      filter,
      takeWhile,
      next,
      take,
      zipWith,
    })
  }
  return _lazy(iterator)
}

Let’s test it:

Lazy([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
  .map(x => x * 2)
  .filter(x => x % 3 === 0)
  .take(10) // => [6, 12, 18]

Lazy([1, 2, 3, 4, 5, 6, 7, 8])
  .zipWith((x, y) => x + y, Lazy([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]))
  .map(x => x + 1)
  .take(10) // => ​​​​​[ 3, 5, 7, 9, 11, 13, 15, 17 ]​​​​​

Caution: The code I’ve presented in this post is just a proof of concept. You can certainly extend it further and even write a library based on it, but don’t use it in production yet!