前言
本文灵感以及代码实现参考自源码：TNG-Hooks。此库的开发者也是大名鼎鼎的 You-Dont-Know-JS 的作者 **Kyle Simpson。**所以本文也可以叫做 TNG-Hooks 源码精读 :-)。
在 Github README 中提到 TNG-Hooks 的灵感起源于 React's Hooks，它提供类似于 useState、useReducer、useEffect 等钩子函数给到普通函数用于状态以及副作用的管理。
注意，TNG-Hooks 的目标对象是普通函数，并不依赖 React 无状态组件（不过实现原理都是大致相似的）。因此，无需了解 React 也可以学习它的源码实现。而且源码行数也只有约 300 行，可谓短小精悍，代码实现也相当优雅。源码十分适合阅读学习。
无论 Hooks 是作用于独立的普通函数，还是我们常见的 React 无状态组件，原理基本类似。所以本文会更倾向于分析无 React 依赖的 Hooks 实现，然后再回到 React 中进行讨论。
话不多说，本文正式开始！
Hooks 从 0 到 1 指南
基本原理概述
在开始之前，先思考一个问题：
怎样让一个函数记住状态？
当然一般情况下，函数不会记住状态。假如借助一些全局变量是否可以记住状态呢？（当然污染全局的做法不推荐）
于是编写实验代码如下，希望可以通过全局变量 val 记住 count 值：
let val: unknown;
function rememberAndUpdateState(initialVal: any) {
const updateVal = (v: any) => {
val = v;
};
updateVal(initialVal);
return [val as any, updateVal];
}
function hit() {
var [count, setCount] = rememberAndUpdateState(0);
count++;
setCount(count);
console.log(`Hit count: ${count}`);
}
满心欢喜的运行：
hit(); // Hit count: 1
hit(); // Hit count: 1
hit(); // Hit count: 1
结果完全不符合预期。到底错在哪呢？可以发现每次运行 hit 函数，rememberAndUpdateState 内部的逻辑也都全部运行。因此 updateVal(initialVal) 这一句也会每次重置 count 为初始值。
所以我们还需要思考让 updateVal(initialVal) 只在第一次运行。我们尝试增加 funcIsFirstTimeCall 数组，并将 hit 函数引用传入，如果 hit 函数还未曾出现过（也就是第一次执行 rememberAndUpdateState），则缓存函数用以后续判断，并执行 updateVal(initialVal) 。代码如下：
let val: unknown;
const funcIsFirstTimeCall = [];
function rememberAndUpdateState(initialVal: any) {
const updateVal = (v: any) => {
val = v;
};
// 获取调用 rememberAndUpdateState 的函数，也即是 hit
const host = arguments.callee.caller;
// 只有在第一次才会调用初始操作
if (!funcIsFirstTimeCall.includes(host)) {
funcIsFirstTimeCall.push(host);
updateVal(initialVal);
}
return [val as any, updateVal];
}
function hit() {
var [count, setCount] = rememberAndUpdateState(0);
count++;
setCount(count);
console.log(`Hit count: ${count}`);
}
成功运行！
hit(); // Hit count: 1
hit(); // Hit count: 2
hit(); // Hit count: 3
于是 hit 函数就简单的记住了 count 值，rememberAndUpdateState 也就是一个 useState 的最简单实现。不过，即使如此实现了，这种做法也不够优雅，因为它将不需要的细节暴露给了开发者，同时没有扩展性可言。
那如何实现一个扩展性良好的 rememberAndUpdateState （也即是 useState）呢？我们总结上述 demo，得出两条结论如下：
- 需要在函数之外通过其他对象变量记住状态；
- 需要在函数之外标记调用次数，标记区别依据可以是函数本身；
两条结论其实归纳起来，就是让函数可以有上下文。如下：
function contextWrapper() {
let val: unknown;
const funcIsFirstTimeCall = [];
function rememberAndUpdateState(initialVal: any) {
// ...
}
function hit() {
// ...
}
return hit;
}
const hitWithContext = contextWrapper();
hitWithContext(); // Hit count: 1
hitWithContext(); // Hit count: 2
hitWithContext(); // Hit count: 3
咦？这不就是闭包么？这不就是高阶函数么？本质上还是通过闭包实现了一个上下文。这也是 Hooks 为什么要运行在特定的上下文中，才会发挥作用。（比如 React Hooks 只能运行在 React 函数组件中）
上下文处理
那么？如何把上述实现搞得高级点、灵活点、优雅点。我们期望的使用方式如下：
function hit() {
const [count, setCount] = useState(0);
const newCount = count + 1;
setCount(newCount);
console.log(`Hit count: ${newCount}`);
}
const hitWithContext = createHC(hit);
hitWithContext(); // Hit count: 1
hitWithContext(); // Hit count: 2
hitWithContext(); // Hit count: 3
其中 createHC 就是给 hit 创建上下文的函数（你可以理解 React 中在运行时中也对函数组件也做了类似的事情）。
我们尽管可以照猫画虎地实现，但是还是有些别扭。而且 arguments.callee 也是不推荐使用的方法。
我们知道，JavaScript 语言特性之一就是单线程。也就意味着任意时间只存在一个函数执行，函数的执行顺序在 调用栈（Call Stack）中保存。函数运行前入栈，运行后出栈。简单示例如下：
在这种调用栈中，最大的好处是我们可以知道哪个函数正在执行！就如上文中使用 arguments.callee.caller 获取当前函数的调用者。如下：
那么，我们是不是也可以通过在运行时中手动实现这个机制，从而获取到当前运行 hook 的是哪个函数呢？
基于此想法，编写实验代码：
// 使用 WeakMap 记录 func 和 bucket 的映射关系
// bucket 用来记录上下文状态
const buckets = new WeakMap();
// 使用堆栈跟踪当前运行的函数
const runtimeStack = [];
function createHC(func: Function) {
// 返回高阶函数封装上下文处理操作
return function HOFWithContext(...args: any) {
runtimeStack.push(func);
if (!buckets.has(func)) {
buckets.set(func, {
stateSlot: [],
});
}
try {
return func.apply(this, args);
} finally {
runtimeStack.pop();
}
};
}
function useState(initialVal: any) {
// 此时的 caller 即为当前运行 useState 的函数
const caller = runtimeStack[runtimeStack.length - 1];
const bucket = buckets.get(caller);
// 如果找不到当前 bucket，证明 useState 被错误使用在无上下文环境
if (!bucket) {
throw new Error(
'useState() only valid inside an Articulated Function or a Custom Hook.'
);
}
// 只有在首次执行初始化操作
if (bucket.stateSlot.length === 0) {
const slot = [
typeof initialVal == 'function' ? initialVal() : initialVal,
function updateSlot(v: unknown) {
slot[0] = v;
},
];
bucket.stateSlot = slot;
}
return bucket.stateSlot;
}
function hit() {
const [count, setCount] = useState(0);
const newCount = count + 1;
setCount(newCount);
console.log(`Hit count: ${newCount}`);
}
const hitWithContext = createHC(hit);
hitWithContext(); // Hit count: 1
hitWithContext(); // Hit count: 2
hitWithContext(); // Hit count: 3
runtimeStack 代表了函数的先后执行关系，而基于 WeakMap 的哈希表映射了 runtimeStack 中的函数与对应的上下文关系。
当然，上述实现代码通用性并不好，我们可以扩展一下 bucket 对象用以支持更丰富的上下文。大家可以大致阅读一下下述完整实现代码。
const buckets = new WeakMap<Function, IBucket>();
const runtimeStack: Function[] = [];
function getCurrentBucket() {
if (runtimeStack.length > 0) {
let bucket: IBucket;
const func = runtimeStack[runtimeStack.length - 1];
// 不存在则新建 bucket
if (!buckets.has(func)) {
bucket = {
nextStateSlotIdx: 0,
nextEffectIdx: 0,
nextMemoizationIdx: 0,
stateSlots: [],
effects: [],
cleanups: [],
memoizations: [],
};
buckets.set(func, bucket);
}
return buckets.get(func);
}
return null;
}
export function createHC(func: Function) {
function HOFWithContext(...args: any) {
runtimeStack.push(func);
const bucket = getCurrentBucket();
// e.g. 运行 hit 函数重置 hooks 的索引，因此 hooks 的执行依赖顺序。
bucket.nextStateSlotIdx = 0;
bucket.nextEffectIdx = 0;
bucket.nextMemoizationIdx = 0;
try {
return func.apply(this, args);
} finally {
try {
// 执行副作用
runEffects(bucket);
} finally {
runtimeStack.pop();
}
}
// useEffect 钩子函数会依赖此函数运行副作用，这里先暂不介绍
function runEffects(bucket: IBucket) {
for (let [idx, [effect, guards]] of bucket.effects.entries()) {
try {
if (typeof effect === 'function') {
effect();
}
} finally {
bucket.effects[idx][0] = undefined;
}
}
}
}
return HOFWithContext;
}
确保理解了上下文的实现原理后，后续所有的 hooks 都依赖上述代码进行开发。
实现 useReducer
我们先来了解 useReducer 的实现机制。reduce 是函数式编程的概念，在前端世界里也有 Array.prototype.reduce() 的数组工具方法，其他编程语言也基本都有类似 reduce 的说法，当然可能叫做 fold。了解了 reduce 的概念后，我们再来看 useReducer 对此概念的应用。
export function useReducer(
reducerFn: Function,
initialVal: any,
...initialReduction: any
) {
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useReducer() only valid inside an Articulated Function or a Custom Hook.'
);
}
if (!(bucket.nextStateSlotIdx in bucket.stateSlots)) {
const slot: StateSlot = [
typeof initialVal == 'function' ? initialVal() : initialVal,
function updateSlot(v: unknown) {
slot[0] = reducerFn(slot[0], v);
},
];
bucket.stateSlots[bucket.nextStateSlotIdx] = slot;
if (initialReduction.length > 0) {
bucket.stateSlots[bucket.nextStateSlotIdx][1](initialReduction[0]);
}
}
return [...bucket.stateSlots[bucket.nextStateSlotIdx++]];
}
上述代码中，在 getCurrentBucket 中先去通过函数堆栈获取当前调用者的上下文 bucket 对象，若无则代表 useReducer 被错误使用，立即抛错阻止代码执行。后文中所有钩子都有同样的处理，之后不再赘述。
我们看到 bucket 对象有 nextStateSlotIdx 和 stateSlots，这是因为一个函数可能会调用多个 useReducer，比如一个函数内调用三次 useReducer （不考虑 useState）则会存在三个 state slot 如下。
声明 slot 数组的时候，在数组中的第二项里使用传入的 reducerFn 对值进行计算。使用方式如下：
function hit(amount = 1) {
const [count, setCount] = useReducer(function reducer(accumulator: number, currentValue: number) {
return accumulator * currentValue;
}, 10);
setCount(amount);
console.log(`Hit Count: ${count}`);
}
const hitWithContext = createHC(hit);
hitWithContext(2); // Hit count: 10
hitWithContext(4); // Hit count: 20
hitWithContext(8); // Hit count: 80
请注意，这里每次 console log 打印出来的都是上次计算的值。因为 count 导出时，当前的 setCount 还未执行。
实现 useState
了解了 useReducer 的实现。我们会发现 useState 能干的事情，useReducer 也全都可以干。因此势必底层实现也可以用后者实现前者，如下：
export function useState(initialVal: any) {
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useState() only valid inside an Articulated Function or a Custom Hook.'
);
}
return useReducer(function reducer(preVal: any, vOrFn: any) {
return typeof vOrFn == 'function' ? vOrFn(preVal) : vOrFn;
}, initialVal);
}
发现 useReducer 中传入一个恒等的 reducer 即可实现 useState。假设 vOrFn 并非函数的话，简化来看就是：
useReducer(function reducer(preVal: any, vOrFn: any) {
return vOrFn;
}, initialVal);
实现 useRef
useRef 背后也是使用 useState 实现的。请看实现：
export function useRef(initialVal: any) {
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useRef() only valid inside an Articulated Function or a Custom Hook.'
);
}
const [ref] = useState({ current: initialVal });
return ref;
}
实现 useMemo
useMemo 的实践一般在于通过缓存减少重复计算量，从而提高应用性能。实现如下：
export function useMemo(func: Function, guards?: Array<any>) {
let realGuards: Array<any>;
if (guards && guards.length > 0) {
realGuards = guards;
} else {
realGuards = [func];
}
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useMemo() only valid inside an Articulated Function or a Custom Hook.'
);
}
if (!(bucket.nextMemoizationIdx in bucket.memoizations)) {
bucket.memoizations[bucket.nextMemoizationIdx] = [];
}
const memoization = bucket.memoizations[bucket.nextMemoizationIdx];
if (guardsChanged(memoization[1], realGuards)) {
try {
memoization[0] = func();
} finally {
memoization[1] = realGuards;
}
}
bucket.nextMemoizationIdx++;
return memoization[0];
}
先收集当前 guards 依赖，从 memoizations 中取得上次的依赖，通过 guardsChange 算法判断是否依赖更新，若无更新则使用上次计算的值，若更新了则重新计算。
memoization 数组格式为 [value, guards]。须知第一次初始化依赖时，此时没有上次的依赖，但又产生了新依赖，所以必然会调用一次计算 memoization[0] = func()。
至于 guardsChange 算法实现如下：
function guardsChanged(guards1: any, guards2: any): boolean {
if (guards1 === undefined || guards2 === undefined) {
return true;
}
if (guards1.length !== guards2.length) {
return true;
}
for (let [idx, guard] of guards1.entries()) {
if (!Object.is(guard, guards2[idx])) {
return true;
}
}
return false;
}
实现 useCallback
和 useRef -> useState -> useReducer 一样，useCallback 也可以使用 useMemo 实现，说白了其实不过是 Hooks 语法糖。
export function useCallback(func: Function, guards?: Array<any>) {
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useCallback() only valid inside an Articulated Function or a Custom Hook.'
);
}
return useMemo(function callback() {
return func;
}, guards);
}
可以理解为在 useMemo 传入了一个高阶函数用以返回 func，从而保留了函数本身，而不是函数调用后的返回值。不得不再次感叹高阶函数的威力。
实现 useEffect
useEffect 是我们很常用的 Hook 函数。这里需要留意的是 useEffect 会在当前函数执行完后再执行（参加上下文处理中的 runEffects）。同时，假如存在上一次 cleanup 函数（即为清理副作用的函数），则会优先执行，然后再执行当前这次的副作用函数。
代码实现如下：
export function useEffect(func: Function, guards?: Array<any>) {
const bucket = getCurrentBucket();
if (!bucket) {
throw new Error(
'useEffect() only valid inside an Articulated Function or a Custom Hook.'
);
}
if (!(bucket.nextEffectIdx in bucket.effects)) {
bucket.effects[bucket.nextEffectIdx] = [undefined, undefined];
}
const effectIdx = bucket.nextEffectIdx;
const effect = bucket.effects[effectIdx];
if (guardsChanged(effect[1], guards)) {
effect[0] = function effect() {
if (typeof bucket.cleanups[effectIdx] === 'function') {
try {
bucket.cleanups[effectIdx]();
} finally {
bucket.cleanups[effectIdx] = undefined;
}
}
const ret = func();
if (typeof ret === 'function') {
bucket.cleanups[effectIdx] = ret;
}
};
effect[1] = guards;
}
bucket.nextEffectIdx++;
}
回到 React
到这里，基本的 Hook 实现原理已经介绍完毕。
但是，如何将上述的 Hook 逻辑实现在一个 UI 框架中，这是一个值得思考的问题。思想当然都是一致的，但是实现方式绝不止一种，可能数据结构发生了变化，比如使用 Current-Owner 而非上文中的 runtimeStack 记录函数与上下文的映射信息，还有如何在适当时机触发函数组件更新等等，这些本质上和 Hooks 的原理无关了，而是一个 UI 框架的设计问题。
以上，感谢阅读。
