源码精读 TNG-Hooks - 再看 useState 实现

Loongerz
1,332 阅读6分钟

之前写过一篇博客 深入解析类 React's Hooks 的实现原理 大约是至少 2 年前了。

最近翻看自己写的博客，发觉写的不是很清晰，自己都被绕得有点迷糊。于是开始再次找出 TNG-Hooks 源码进行学习。尝试再以更系统、更高的视角去分析。

这次我们尝试简化整个博客的思路，仅实现 useState，并分析其思路。

设个题目

先设定一个目标，我们期望下面代码示例最终运行打印如下：

function _clickExpanded() {
  let [count, setCount] = useState(0);
  let [expanded, setExpanded] = useState(false);

  setCount(count + 1);
  if (!expanded) {
    setExpanded(true);
    renderUsername('David');
  } else {
    setExpanded(false);
    renderUsername('David is a King');
  }

  console.log('count', count);
  console.log('expanded', expanded);
}

function _renderUsername(username) {
  // using the `useState(..)` hook
  var [activated, setActivated] = useState(false);

  // do update
  // usernameElem.innerHTML = username;

  // only run this code the first time
  if (!activated) {
    setActivated(true);
    // set listener
    // usernameElem.addEventListener("click",onClickUsername,false);
  }

  console.log('activated', activated);
}

// ************ 

function TNG(fn) {
  // ...
  return fn;
}

function useState(initialVal) {
  // ...
  return [initialVal, () => initialVal];
}

// ************

const clickExpanded = TNG(_clickExpanded);
const renderUsername = TNG(_renderUsername);

clickExpanded();
/**
activated false
count 0
expanded false
 */
clickExpanded();
/**
activated:  true
count 1
expanded true
 */
clickExpanded();
/**
activated:  true
count 2
expanded false
 */

大家感兴趣可以先去尝试实现下，再回来看剩下的博客。

题干分析

读题干，有两个函数 _clickExpanded、_renderUserName 被 TNG 处理后，返回的函数 clickExpanded、renderUserName 拥有了状态记忆 的能力。

  • 在第 1 次调用时，可以使用初始默认状态。并设置下次状态。
  • 在第 >1 次调用时，可以获取上次状态。并设置下次状态。

那么我们此时需要实现的就是两个函数：

  • TNG 函数：用以给目标函数设定上下文
  • useState：实现目标状态的 get/set

当然，除此之外，可以自己根据需要实现任何辅助函数、变量均可。最终满足调用条件即可。（Done is better than perfect，写完后可以再进行优化

实现

通过阅读源码思想，简化实现 useState 如下。大家也可以先跳过这段源码，先看下面一章节的分析。

// ************
let buckets = new WeakMap();
let tngStack = [];

// 获取 tng 栈顶函数对应的状态桶
// 1. tng 栈可能为空，此时直接返回空即可
// 2. 栈顶函数可能未初始化其状态桶，此时初始化即可
function getCurrentBucket() {
  if (tngStack.length > 0) {
    let tngf = tngStack[tngStack.length - 1];
    let bucket;

    if (!buckets.has(tngf)) {
      bucket = {
        nextStateSlotIdx: 0,
        stateSlots: [],
      };
      buckets.set(tngf, bucket);
    }

    return buckets.get(tngf);
  }
}

// 将无状态函数 fn 转为有状态函数 tngf
// 一个很复杂的闭包
function TNG(fn) {
  return function tngf(...args) {
    // tngf 它被执行的时候，此时将它入栈，符合直觉
    // 潜台词：它未出栈时，currentBucket 一直都是它的 bucket
    tngStack.push(tngf);

    // 该 tngf 第 1 次执行时，其实就是在初始化的新 bucket
    // 该 tngf 第 >1 次执行时，其实就是已经有状态的老 bucket
    let bucket = getCurrentBucket();

    // 将其重置为 0 的原因是
    // 同一个 tngf 重复调用时，进行重置该指针
    // 使得每次 useState 的对应到的 slot 均是一致
    // 这就是为什么 useState 必须位置稳定，不能使用 if else 来条件式出现
    bucket.nextStateSlotIdx = 0;

    try {
      return fn.apply(this, args);
    } finally {
      // 将该 tngf 出栈，此时再获取的 currentBucket 就是当前栈顶的另一个 tngf
      tngStack.pop();
    }
  };
}

function useState(initialVal) {
  let bucket = getCurrentBucket();

  if (bucket) {
    // 在第 1 次 tngf 调用时，相当于 useState 每出现，都会先走初始化 slot 的步骤
    // 在第 >1 次 tngf 调用时，stateSlots 均已经出现，所以直接按顺序递增指针即可。
    if (bucket.nextStateSlotIdx >= bucket.stateSlots.length) {
      const slot = [
        initialVal,
        function update(v) {
          slot[0] = v;
        },
      ];

      bucket.stateSlots[bucket.nextStateSlotIdx] = slot;
    }

    // 导出该 slot, 并指向下个 slot（未存在）
    return [...bucket.stateSlots[bucket.nextStateSlotIdx++]];
  } else {
    throw new Error(
      'useState() only valid inside an Articulated Function or a Custom Hook.'
    );
  }
}
// ************

分析

我们试着从 Demo 调用去分析，我们调用了三次 clickExpanded，根据内部实现，我们分析下在单线程下的函数调用栈是怎么样的。（我们不关心其他函数）

tngf 调用栈

执行第一次 clickExpanded 如下（其中 tngf1 是经过 TNG 封装的 clickExpanded，tngf2 是经过 TNG 封装的 renderUsername）：

image.png

那么三次调用的调用栈变化情况是怎样的呢？

image.png

调用栈周而复始，起于平凡，踏浪而行，最终功成身就后，归于平静，风度非凡。

那么有趣的是，这个逻辑维护在浏览器中，对于 js 代码逻辑来说是不可见的。我们想追踪它唯有通过手动维护一个栈的方式去管理。

在上述代码中，其实看到 tngStack 相关逻辑就是手动维护的栈。简化逻辑见下：

function TNG(fn) {
    return function tngf(...args) {
        tngStack.push(tngf);
        fn.apply(this, args);
        tngStack.pop();
    }
}

那维护这么一个调用栈的好处有哪些？

  1. 可以知道调用关系链是怎样的？比如 tngf1 -> tngf2
  2. 某一时刻可以通过看栈顶，知道代码中正在运行的函数。（依赖单线程）

当然，维护所有函数的调用栈是复杂的，在我们只关心 TNG 封装的函数调用情况下的话，就简单很多。

依赖建立

除此之外，还需要有一个 map 来记录不同的 tngf 有哪些依赖。为了让 key 可以是函数对象，并解决内存泄露问题，我们通过 WeakMap 建立依赖表。

一个 tngf 都让它提一个 bucket 桶（让每个 tngf 都很有“体统” ...希望没有太冷哈哈哈）

image.png

这个桶里可以放任何东西，目前我们只放 stateSlots 和 nextStateSlotIdx 这两个重要物件。

那么到此为止，我们可以做到一个事情，就是根据 tngStack 取栈顶，结合这个 buckets 的依赖表，可以立马获取这个 tngf 的对应桶。

这就是 getCurrentBucket 的逻辑。当然它里面也默认做了个逻辑，就是这个 tngf 第一次报到的时候，别让它这么惨，都会给默认分配一个空荡荡的桶。第二次再出现的时候，则不会再分配桶，直接取它已有的桶，做到人手一个，公平公正公开透明。

stateSlot 顺序

我们知道，一个函数里，可能会使用多个 useState。那么每个 useState 其实都会对应到一个 slot。

image.png

在每次 tngf 调用的时候，nextStateSlotIdx 指针会重置，其中：

  • 在第 1 次 tngf 调用时，相当于 useState 每出现，都会先走初始化 slot 的步骤
    • 会给每个 useState 初始化分配一个 slot。
    • nextStateSlotIdx 作为一个指针，会指向下一个待分配的 slot。
  • 在第 >1 次 tngf 调用时，stateSlots 均已分配，所以直接按顺序递增 nextStateSlotIdx 指针逐个取 slot 即可。

这里暗含了一个条件约定，就是 tngf 里所有的 useState 的调用顺序总是需要稳定出现，否则会出现 useState 找不到 slot 或者找到错误的 slot，从而引来异常 bug。

其实，stateSlots 也是一个 map。相当于每个 useState 通过出现的索引顺序， 认领一个 slot。见上图。

  • useState(A) 通过 nextStateSlotIdx 0 永远对应 slot 1
  • useState(B) 通过 nextStateSlotIdx 1 永远对应 slot 2
  • useState(C) 通过 nextStateSlotIdx 2 永远对应 slot 3

最终运行，结果符合预期。

小结

对于我来说，初读 TNG Hooks 源码，让我最惊艳的一个点就是，js 维护目标函数的调用栈，这种思维是很少见的，相当于用模拟了 js 引擎的能力，让我印象深刻。

时隔两年再读源码，还是颇有收益。因此再写篇博客重新分享。

感谢阅读。

avatar
文章
阅读
粉丝