46 lines
1.2 KiB
JavaScript
46 lines
1.2 KiB
JavaScript
import { useLayoutUpdateEffect } from "@rc-component/util/es/hooks/useLayoutEffect";
|
|
import { useRef, useState } from 'react';
|
|
|
|
/**
|
|
* Help to merge callback with `useLayoutEffect`.
|
|
* One time will only trigger once.
|
|
*/
|
|
export default function useUpdate(callback) {
|
|
const [count, setCount] = useState(0);
|
|
const effectRef = useRef(0);
|
|
const callbackRef = useRef();
|
|
callbackRef.current = callback;
|
|
|
|
// Trigger on `useLayoutEffect`
|
|
useLayoutUpdateEffect(() => {
|
|
callbackRef.current?.();
|
|
}, [count]);
|
|
|
|
// Trigger to update count
|
|
return () => {
|
|
if (effectRef.current !== count) {
|
|
return;
|
|
}
|
|
effectRef.current += 1;
|
|
setCount(effectRef.current);
|
|
};
|
|
}
|
|
export function useUpdateState(defaultState) {
|
|
const batchRef = useRef([]);
|
|
const [, forceUpdate] = useState({});
|
|
const state = useRef(typeof defaultState === 'function' ? defaultState() : defaultState);
|
|
const flushUpdate = useUpdate(() => {
|
|
let current = state.current;
|
|
batchRef.current.forEach(callback => {
|
|
current = callback(current);
|
|
});
|
|
batchRef.current = [];
|
|
state.current = current;
|
|
forceUpdate({});
|
|
});
|
|
function updater(callback) {
|
|
batchRef.current.push(callback);
|
|
flushUpdate();
|
|
}
|
|
return [state.current, updater];
|
|
} |