补充一个替代 for 循环的新姿势

本文英文版发表在Lei's Blog

我也没想到我还在这个问题上死磕……

最近复习以前的知识点，发现之前钻的不够深，然后继续开了下脑洞，然后就有了今天要写的内容。

可能有不熟悉背景的朋友，这里我简单介绍下。我之前在掘金写了一篇很引战的文章 《如何在 JS 代码中消灭 for 循环》 。这个标题严谨推敲一下是不合适的，但已经成历史了，我就不改了。后来我又解释了为什么要在 特定场合下 避免 for 循环，这里就不赘述了。接下来我要写的代码也包含大量 for 循环，这和我之前的解释并不冲突。

去年学习 Haskell 的时候发现 Haskell 的一种写法很简洁和优雅，很难用 JS 来实现。那段代码长这样：

sum (takeWhile (<10000) (filter odd (map (^2) [1..])))
复制代码

我当时就想在 JS 里面怎么用高阶函数实现这段代码。当然我是想不出来的，然后我 Google 搜了一下，还真有人写了。 Lazy Evaluation in JavaScript with Generators, Map, Filter, and Reduce 。

我基于这篇文章，加了一个 takeWhile 方法，然后用 JS 翻译了上面那段 Haskell 代码。然后我写了一篇文章，逐块解释代码，发表在掘金（后来删了）。前时间我把这篇文章重新分享出来，然后就有人质疑我从哪抄的代码。

我确实一开始没有标明出处，只有认了。同时别人的监督也促使我反省我为什么不基于别人的成果扩展些自己的东西？

我接下来要基于原代码从两个方向扩展。

一是将 class 改写成工厂函数。这里不想再引起争议，我不是提倡不要用 class。在 Node 开发中，为了性能是要用 class 的。这里折腾改写出于两个原因：

1. 通过改写彻底弄懂代码。
2. 个人偏好上我更喜欢工厂函数。扩展性和安全性上，工厂函数更优（欢迎证明我是错的）。

二是将原文 Lazy 函数改写成接受任何 iterable（原文只接受无限 iterator），并加上常用的列表操作高阶函数。这一步完成后，我们其实就能用 Lazy 来实现大数据集遍历了。我之前解决这个问题是用的 Transducer。惰性求值和 transduce 原理不一样，但目的有重合。

我首先用工厂函数改写原版：

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

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

  const filter = predicate => {
    const modifiedNext = () => {
      while (true) {
        const item = next()
        if (predicate(item.value)) {
          return item
        }
      }
    }
    const newIter = { ...iterator, 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
  }

  return Object.freeze({
    map,
    filter,
    takeWhile,
    next,
  })
}

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

Lazy(numbers())
  .map(x => x ** 2)
  .filter(x => x % 2 === 1)
  .takeWhile(x => x < 10000)
  .reduce((x, y) => x + y)
// => 16650
复制代码

以上代码在我博客中有逐步解释。

接着扩展 Lazy 函数，让它接受任何 Iterable 和自定义 Generator：

/*
 * JS 中判断 generator 有些麻烦，我用了这个库：
 * https://github.com/ljharb/is-generator-function
 */

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()
        /*
         * 注意这里的改动，这里为了处理有限数据集，
         * 不能再用死循环了
         */
        while (!item.done) {
          if (predicate(item.value)) {
            return item
          }
          item = next()
        }
        /*
         * 如果上面的循环完成，还没匹配值，
         * 通知 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++) {
        /*
         * 如果数据集长度比 n 要小，
         * 遍历提前完成，要及时 break 循环
         */
        if (item.done) break
        values.push(item.value)
        item = next()
      }

      return values
    }

    /*
     * ZipWith 把两个包在 Lazy 函数中的 Iterable
     * 按指定回调函数拼接起来
     */
    const zipWith = (fn, other) => {
      const modifiedNext = () => {
        const first = next()
        const second = other.next()
        /*
         * 只要有一个 Iterable 遍历完成，
         * 告诉外层 Iterator 结束
         */
        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)
}
复制代码

来试验一下：

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 ]​​​​​
复制代码

注意本文只是提供了一个 proof of concept, 性能问题和其它潜在的问题我还没验证过，所以不建议在生产中直接使用本文代码。

另外，有一个叫 lazy.js 的库提供了类似的功能。感兴趣的朋友可以参考 lazy.js 的 API，基于本文代码继续扩展。比如，给 Lazy 加上异步迭代的功能就可以很简单做到。欢迎折腾反馈。